From c602f6b1979dde088033cf353e4d9481ebd8d1ba Mon Sep 17 00:00:00 2001
From: David Christensen <david.christensen@crunchydata.com>
Date: Thu, 24 Oct 2024 12:52:21 -0500
Subject: [PATCH] Make permission checking conditional for DestReceiver (#56)

Perform access checking at the CopyHook or in the UDFs themselves explicitly,
instead of habing a lower-level perform this check.

This allows us to use the DestReceiver in other code while allowing our own
permission/hook checking there.
---
 src/arrow_parquet/uri_utils.rs | 52 +++++++++-------------------------
 src/lib.rs                     |  3 ++
 src/parquet_copy_hook/hook.rs  | 11 ++++++-
 src/parquet_udfs/metadata.rs   |  7 ++++-
 src/parquet_udfs/schema.rs     |  5 +++-
 5 files changed, 37 insertions(+), 41 deletions(-)

diff --git a/src/arrow_parquet/uri_utils.rs b/src/arrow_parquet/uri_utils.rs
index 745058f..e1dc475 100644
--- a/src/arrow_parquet/uri_utils.rs
+++ b/src/arrow_parquet/uri_utils.rs
@@ -58,8 +58,6 @@ fn parse_bucket_and_key(uri: &Url) -> (String, String) {
 
 fn object_store_with_location(uri: &Url, copy_from: bool) -> (Arc<dyn ObjectStore>, Path) {
     if uri.scheme() == "s3" {
-        ensure_object_store_access_privilege(copy_from);
-
         let (bucket_name, key) = parse_bucket_and_key(uri);
 
         let storage_container = PG_BACKEND_TOKIO_RUNTIME
@@ -71,8 +69,6 @@ fn object_store_with_location(uri: &Url, copy_from: bool) -> (Arc<dyn ObjectStor
     } else {
         debug_assert!(uri.scheme() == "file");
 
-        ensure_local_file_access_privilege(copy_from);
-
         let uri = uri_as_string(uri);
 
         if !copy_from {
@@ -251,14 +247,21 @@ pub(crate) fn parquet_writer_from_uri(
         .unwrap_or_else(|e| panic!("failed to create parquet writer for uri {}: {}", uri, e))
 }
 
-fn ensure_object_store_access_privilege(copy_from: bool) {
+pub(crate) fn ensure_access_privilege_to_uri(uri: &Url, copy_from: bool) {
     if unsafe { superuser() } {
         return;
     }
 
     let user_id = unsafe { GetUserId() };
+    let is_file = uri.scheme() == "file";
 
-    let required_role_name = if copy_from {
+    let required_role_name = if is_file {
+        if copy_from {
+            "pg_read_server_files"
+        } else {
+            "pg_write_server_files"
+        }
+    } else if copy_from {
         PARQUET_OBJECT_STORE_READ_ROLE
     } else {
         PARQUET_OBJECT_STORE_WRITE_ROLE
@@ -268,46 +271,19 @@ fn ensure_object_store_access_privilege(copy_from: bool) {
         unsafe { get_role_oid(required_role_name.to_string().as_pg_cstr(), false) };
 
     let operation_str = if copy_from { "from" } else { "to" };
+    let object_type = if is_file { "file" } else { "remote uri" };
 
     if !unsafe { has_privs_of_role(user_id, required_role_id) } {
         ereport!(
             pgrx::PgLogLevel::ERROR,
             pgrx::PgSqlErrorCode::ERRCODE_INSUFFICIENT_PRIVILEGE,
-            format!("permission denied to COPY {} a remote uri", operation_str),
             format!(
-                "Only roles with privileges of the \"{}\" role may COPY {} a remote uri.",
-                required_role_name, operation_str
+                "permission denied to COPY {} a {}",
+                operation_str, object_type
             ),
-        );
-    }
-}
-
-fn ensure_local_file_access_privilege(copy_from: bool) {
-    if unsafe { superuser() } {
-        return;
-    }
-
-    let user_id = unsafe { GetUserId() };
-
-    let required_role_name = if copy_from {
-        "pg_read_server_files"
-    } else {
-        "pg_write_server_files"
-    };
-
-    let required_role_id =
-        unsafe { get_role_oid(required_role_name.to_string().as_pg_cstr(), false) };
-
-    let operation_str = if copy_from { "from" } else { "to" };
-
-    if !unsafe { has_privs_of_role(user_id, required_role_id) } {
-        ereport!(
-            pgrx::PgLogLevel::ERROR,
-            pgrx::PgSqlErrorCode::ERRCODE_INSUFFICIENT_PRIVILEGE,
-            format!("permission denied to COPY {} a file", operation_str),
             format!(
-                "Only roles with privileges of the \"{}\" role may COPY {} a file.",
-                required_role_name, operation_str
+                "Only roles with privileges of the \"{}\" role may COPY {} a {}.",
+                required_role_name, operation_str, object_type
             ),
         );
     }
diff --git a/src/lib.rs b/src/lib.rs
index fdba58c..ce61540 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -2672,6 +2672,9 @@ mod tests {
         let create_role = "create role test_role;";
         Spi::run(create_role).unwrap();
 
+        let grant_role = "grant pg_write_server_files TO test_role;";
+        Spi::run(grant_role).unwrap();
+
         let set_role = "set role test_role;";
         Spi::run(set_role).unwrap();
 
diff --git a/src/parquet_copy_hook/hook.rs b/src/parquet_copy_hook/hook.rs
index 8d4cea0..1329817 100644
--- a/src/parquet_copy_hook/hook.rs
+++ b/src/parquet_copy_hook/hook.rs
@@ -7,7 +7,10 @@ use pg_sys::{
 use pgrx::{prelude::*, GucSetting};
 
 use crate::{
-    arrow_parquet::{compression::INVALID_COMPRESSION_LEVEL, uri_utils::uri_as_string},
+    arrow_parquet::{
+        compression::INVALID_COMPRESSION_LEVEL,
+        uri_utils::{ensure_access_privilege_to_uri, uri_as_string},
+    },
     parquet_copy_hook::{
         copy_to_dest_receiver::create_copy_to_parquet_dest_receiver,
         copy_utils::{
@@ -59,6 +62,9 @@ extern "C" fn parquet_copy_hook(
 
     if ENABLE_PARQUET_COPY_HOOK.get() && is_copy_to_parquet_stmt(&p_stmt) {
         let uri = copy_stmt_uri(&p_stmt).expect("uri is None");
+        let copy_from = false;
+
+        ensure_access_privilege_to_uri(&uri, copy_from);
 
         validate_copy_to_options(&p_stmt, &uri);
 
@@ -106,6 +112,9 @@ extern "C" fn parquet_copy_hook(
         return;
     } else if ENABLE_PARQUET_COPY_HOOK.get() && is_copy_from_parquet_stmt(&p_stmt) {
         let uri = copy_stmt_uri(&p_stmt).expect("uri is None");
+        let copy_from = true;
+
+        ensure_access_privilege_to_uri(&uri, copy_from);
 
         validate_copy_from_options(&p_stmt);
 
diff --git a/src/parquet_udfs/metadata.rs b/src/parquet_udfs/metadata.rs
index 4d167c6..946c3c4 100644
--- a/src/parquet_udfs/metadata.rs
+++ b/src/parquet_udfs/metadata.rs
@@ -1,7 +1,9 @@
 use ::parquet::file::statistics::Statistics;
 use pgrx::{iter::TableIterator, name, pg_extern, pg_schema};
 
-use crate::arrow_parquet::uri_utils::{parquet_metadata_from_uri, parse_uri, uri_as_string};
+use crate::arrow_parquet::uri_utils::{
+    ensure_access_privilege_to_uri, parquet_metadata_from_uri, parse_uri, uri_as_string,
+};
 
 #[pg_schema]
 mod parquet {
@@ -39,6 +41,7 @@ mod parquet {
     > {
         let uri = parse_uri(&uri);
 
+        ensure_access_privilege_to_uri(&uri, true);
         let parquet_metadata = parquet_metadata_from_uri(&uri);
 
         let mut rows = vec![];
@@ -137,6 +140,7 @@ mod parquet {
     > {
         let uri = parse_uri(&uri);
 
+        ensure_access_privilege_to_uri(&uri, true);
         let parquet_metadata = parquet_metadata_from_uri(&uri);
 
         let created_by = parquet_metadata
@@ -174,6 +178,7 @@ mod parquet {
     > {
         let uri = parse_uri(&uri);
 
+        ensure_access_privilege_to_uri(&uri, true);
         let parquet_metadata = parquet_metadata_from_uri(&uri);
         let kv_metadata = parquet_metadata.file_metadata().key_value_metadata();
 
diff --git a/src/parquet_udfs/schema.rs b/src/parquet_udfs/schema.rs
index 7e5c68f..c2fcf00 100644
--- a/src/parquet_udfs/schema.rs
+++ b/src/parquet_udfs/schema.rs
@@ -1,4 +1,6 @@
-use crate::arrow_parquet::uri_utils::{parquet_schema_from_uri, parse_uri, uri_as_string};
+use crate::arrow_parquet::uri_utils::{
+    ensure_access_privilege_to_uri, parquet_schema_from_uri, parse_uri, uri_as_string,
+};
 
 use ::parquet::{
     format::{ConvertedType, FieldRepetitionType, LogicalType, Type},
@@ -32,6 +34,7 @@ mod parquet {
     > {
         let uri = parse_uri(&uri);
 
+        ensure_access_privilege_to_uri(&uri, true);
         let parquet_schema = parquet_schema_from_uri(&uri);
 
         let root_type = parquet_schema.root_schema();