fix: no permission check for copy stdin/out (#136)

Fixes #135

Author: aykut-bozkurt

src/arrow_parquet/uri_utils.rs

@@ -269,13 +269,18 @@ pub(crate) fn parquet_writer_from_uri(
         })
 }
 
-pub(crate) fn ensure_access_privilege_to_uri(uri: &Url, copy_from: bool) {
+pub(crate) fn ensure_access_privilege_to_uri(uri_info: &ParsedUriInfo, copy_from: bool) {
     if unsafe { superuser() } {
         return;
     }
 
+    // permission check is not needed for stdin/out
+    if uri_info.stdio_tmp_fd.is_some() {
+        return;
+    }
+
     let user_id = unsafe { GetUserId() };
-    let is_file = uri.scheme() == "file";
+    let is_file = uri_info.uri.scheme() == "file";
 
     let required_role_name = if is_file {
         if copy_from {

src/parquet_copy_hook/hook.rs

@@ -53,10 +53,8 @@ fn process_copy_to_parquet(
 ) -> u64 {
     let uri_info = copy_stmt_uri(p_stmt).unwrap_or_else(|e| panic!("{}", e));
 
-    let uri = uri_info.uri.clone();
-
     let copy_from = false;
-    ensure_access_privilege_to_uri(&uri, copy_from);
+    ensure_access_privilege_to_uri(&uri_info, copy_from);
 
     validate_copy_to_options(p_stmt, &uri_info);
 
@@ -68,7 +66,7 @@ fn process_copy_to_parquet(
     let compression_level = copy_to_stmt_compression_level(p_stmt, &uri_info);
 
     let parquet_split_dest = create_copy_to_parquet_split_dest_receiver(
-        uri_as_string(&uri).as_pg_cstr(),
+        uri_as_string(&uri_info.uri).as_pg_cstr(),
         uri_info.stdio_tmp_fd.is_some(),
         &file_size_bytes,
         field_ids,
@@ -103,10 +101,8 @@ fn process_copy_from_parquet(
 ) -> u64 {
     let uri_info = copy_stmt_uri(p_stmt).unwrap_or_else(|e| panic!("{}", e));
 
-    let uri = uri_info.uri.clone();
-
     let copy_from = true;
-    ensure_access_privilege_to_uri(&uri, copy_from);
+    ensure_access_privilege_to_uri(&uri_info, copy_from);
 
     validate_copy_from_options(p_stmt);
 

src/parquet_udfs/metadata.rs

@@ -45,7 +45,7 @@ mod parquet {
             panic!("{}", e.to_string());
         });
 
-        ensure_access_privilege_to_uri(&uri_info.uri, true);
+        ensure_access_privilege_to_uri(&uri_info, true);
         let parquet_metadata = parquet_metadata_from_uri(&uri_info);
 
         let mut rows = vec![];
@@ -148,7 +148,7 @@ mod parquet {
             panic!("{}", e.to_string());
         });
 
-        ensure_access_privilege_to_uri(&uri_info.uri, true);
+        ensure_access_privilege_to_uri(&uri_info, true);
         let parquet_metadata = parquet_metadata_from_uri(&uri_info);
 
         let created_by = parquet_metadata
@@ -188,7 +188,7 @@ mod parquet {
             panic!("{}", e.to_string());
         });
 
-        ensure_access_privilege_to_uri(&uri_info.uri, true);
+        ensure_access_privilege_to_uri(&uri_info, true);
         let parquet_metadata = parquet_metadata_from_uri(&uri_info);
 
         let kv_metadata = parquet_metadata.file_metadata().key_value_metadata();

src/parquet_udfs/schema.rs

@@ -36,7 +36,7 @@ mod parquet {
             panic!("{}", e.to_string());
         });
 
-        ensure_access_privilege_to_uri(&uri_info.uri, true);
+        ensure_access_privilege_to_uri(&uri_info, true);
         let parquet_schema = parquet_schema_from_uri(&uri_info);
 
         let root_type = parquet_schema.root_schema();

src/parquet_udfs/stats.rs

@@ -112,7 +112,7 @@ mod parquet {
             panic!("{}", e.to_string());
         });
 
-        ensure_access_privilege_to_uri(&uri_info.uri, true);
+        ensure_access_privilege_to_uri(&uri_info, true);
         let parquet_metadata = parquet_metadata_from_uri(&uri_info);
 
         let mut aggregated_column_stats = HashMap::new();

src/pgrx_tests/copy_stdin_out.rs

@@ -76,12 +76,33 @@ mod tests {
             "Failed to insert into test_expected table"
         );
 
-        // copy to stdout
+        // create a dummy role
+        let role = "dummy";
+        let output = Command::new("psql")
+            .arg("-p")
+            .arg(test_port.clone())
+            .arg("-h")
+            .arg("localhost")
+            .arg("-d")
+            .arg("pgrx_tests")
+            .arg("-c")
+            .arg(format!(
+                "CREATE ROLE {role} LOGIN;
+                 GRANT SELECT, INSERT, UPDATE, DELETE ON test_expected TO {role};
+                 GRANT SELECT, INSERT, UPDATE, DELETE ON test_result TO {role};"
+            ))
+            .output()
+            .expect("failed to execute process");
+        assert!(output.status.success(), "Failed to create role {role}");
+
+        // copy to stdout with dummy role
         let mut copy_to = Command::new("psql")
             .arg("-p")
             .arg(test_port.clone())
             .arg("-h")
             .arg("localhost")
+            .arg("-U")
+            .arg(role)
             .arg("-d")
             .arg("pgrx_tests")
             .arg("-c")
@@ -101,12 +122,14 @@ mod tests {
             assert!(status.success(), "psql COPY TO process did not succeed");
         }
 
-        // copy from stdin
+        // copy from stdin with dummy role
         let mut copy_from = Command::new("psql")
             .arg("-p")
             .arg(test_port.clone())
             .arg("-h")
             .arg("localhost")
+            .arg("-U")
+            .arg(role)
             .arg("-d")
             .arg("pgrx_tests")
             .arg("-c")
@@ -178,7 +201,7 @@ mod tests {
             ]
         );
 
-        // drop test_expected
+        // drop tables and role
         let output = Command::new("psql")
             .arg("-p")
             .arg(test_port.clone())
@@ -187,26 +210,11 @@ mod tests {
             .arg("-d")
             .arg("pgrx_tests")
             .arg("-c")
-            .arg("DROP TABLE test_expected;")
-            .output()
-            .expect("failed to execute process");
-        assert!(
-            output.status.success(),
-            "Failed to drop test_expected table"
-        );
-
-        // drop test_result
-        let output = Command::new("psql")
-            .arg("-p")
-            .arg(test_port.clone())
-            .arg("-h")
-            .arg("localhost")
-            .arg("-d")
-            .arg("pgrx_tests")
-            .arg("-c")
-            .arg("DROP TABLE test_result;")
+            .arg(format!(
+                "DROP TABLE test_expected, test_result; DROP ROLE {role};"
+            ))
             .output()
             .expect("failed to execute process");
-        assert!(output.status.success(), "Failed to drop test_result table");
+        assert!(output.status.success(), "Failed to clean up");
     }
 }