Fix: unprivileged copy from (#125)

(cherry picked from commit 3d8fbb6)

Author: aykut-bozkurt

src/parquet_copy_hook/copy_from.rs

@@ -18,9 +18,12 @@ use crate::{
     },
 };
 
-use super::copy_utils::{
-    copy_from_stmt_match_by, copy_stmt_attribute_list, copy_stmt_create_namespace_item,
-    copy_stmt_create_parse_state, create_filtered_tupledesc_for_relation,
+use super::{
+    copy_utils::{
+        copy_from_stmt_match_by, copy_stmt_attribute_list, copy_stmt_create_namespace_item,
+        copy_stmt_create_parse_state, create_filtered_tupledesc_for_relation,
+    },
+    pg_compat::check_copy_table_permission,
 };
 
 // stack to store parquet reader contexts for COPY FROM.
@@ -133,6 +136,8 @@ pub(crate) fn execute_copy_from(
         where_clause = copy_from_stmt_transform_where_clause(&p_state, &ns_item, where_clause);
     }
 
+    check_copy_table_permission(p_stmt, &p_state, &ns_item, &relation);
+
     let attribute_list = copy_stmt_attribute_list(p_stmt);
 
     let tupledesc = create_filtered_tupledesc_for_relation(p_stmt, &relation);

src/parquet_copy_hook/pg_compat.rs

@@ -1,6 +1,13 @@
 use std::ffi::{c_char, CStr};
 
-use pgrx::pg_sys::{AsPgCStr, List, Node, QueryEnvironment, RawStmt};
+use pgrx::{
+    pg_sys::{
+        bms_add_member, AsPgCStr, CopyGetAttnums, CopyStmt, FirstLowInvalidHeapAttributeNumber,
+        List, Node, ParseNamespaceItem, ParseState, PlannedStmt, QueryEnvironment, RawStmt,
+        ACL_INSERT, ACL_SELECT,
+    },
+    PgBox, PgList, PgRelation,
+};
 
 pub(crate) fn pg_analyze_and_rewrite(
     raw_stmt: *mut RawStmt,
@@ -65,3 +72,89 @@ pub(crate) fn MarkGUCPrefixReserved(guc_prefix: &str) {
         pgrx::pg_sys::MarkGUCPrefixReserved(guc_prefix.as_pg_cstr())
     }
 }
+
+/// check_copy_table_permission checks if the user has permission to copy from/to the table.
+/// This is taken from the original PostgreSQL DoCopy function.
+#[cfg(any(feature = "pg16", feature = "pg17"))]
+pub(crate) fn check_copy_table_permission(
+    p_stmt: &PgBox<PlannedStmt>,
+    p_state: &PgBox<ParseState>,
+    ns_item: &PgBox<ParseNamespaceItem>,
+    relation: &PgRelation,
+) {
+    let copy_stmt = unsafe { PgBox::<CopyStmt>::from_pg(p_stmt.utilityStmt as _) };
+
+    // init permissions
+    let mut perminfo =
+        unsafe { PgBox::<pgrx::pg_sys::RTEPermissionInfo>::from_pg(ns_item.p_perminfo) };
+
+    // set table access mode
+    perminfo.requiredPerms = if copy_stmt.is_from {
+        ACL_INSERT as _
+    } else {
+        ACL_SELECT as _
+    };
+
+    // set column access modes
+    let tup_desc = relation.tuple_desc();
+
+    let attnums =
+        unsafe { CopyGetAttnums(tup_desc.as_ptr(), relation.as_ptr(), copy_stmt.attlist) };
+    let attnums = unsafe { PgList::<i16>::from_pg(attnums) };
+
+    for attnum in attnums.iter_int() {
+        let attno = attnum - FirstLowInvalidHeapAttributeNumber;
+
+        if copy_stmt.is_from {
+            unsafe { perminfo.insertedCols = bms_add_member(perminfo.insertedCols, attno) };
+        } else {
+            unsafe { perminfo.selectedCols = bms_add_member(perminfo.selectedCols, attno) };
+        }
+    }
+
+    // check permissions
+    let mut perm_infos = PgList::<pgrx::pg_sys::RTEPermissionInfo>::new();
+    perm_infos.push(perminfo.as_ptr());
+
+    unsafe { pgrx::pg_sys::ExecCheckPermissions(p_state.p_rtable, perm_infos.as_ptr(), true) };
+}
+
+#[cfg(any(feature = "pg14", feature = "pg15"))]
+pub(crate) fn check_copy_table_permission(
+    p_stmt: &PgBox<PlannedStmt>,
+    p_state: &PgBox<ParseState>,
+    ns_item: &PgBox<ParseNamespaceItem>,
+    relation: &PgRelation,
+) {
+    let copy_stmt = unsafe { PgBox::<CopyStmt>::from_pg(p_stmt.utilityStmt as _) };
+
+    // init rte
+    let mut rte = unsafe { PgBox::<pgrx::pg_sys::RangeTblEntry>::from_pg(ns_item.p_rte) };
+
+    // set table access mode
+    rte.requiredPerms = if copy_stmt.is_from {
+        ACL_INSERT as _
+    } else {
+        ACL_SELECT as _
+    };
+
+    // set column access modes
+    let tup_desc = relation.tuple_desc();
+
+    let attnums =
+        unsafe { CopyGetAttnums(tup_desc.as_ptr(), relation.as_ptr(), copy_stmt.attlist) };
+    let attnums = unsafe { PgList::<i16>::from_pg(attnums) };
+
+    for attnum in attnums.iter_int() {
+        let attno = attnum - FirstLowInvalidHeapAttributeNumber;
+
+        if copy_stmt.is_from {
+            unsafe { rte.insertedCols = bms_add_member(rte.insertedCols, attno) };
+        } else {
+            unsafe { rte.selectedCols = bms_add_member(rte.selectedCols, attno) };
+        }
+    }
+
+    // check permissions
+    unsafe { pgrx::pg_sys::ExecCheckRTPerms(p_state.p_rtable, true) };
+}

src/pgrx_tests/copy_pg_rules.rs

@@ -483,7 +483,7 @@ mod tests {
 
     #[pg_test]
     #[should_panic(expected = "permission denied for table test_table")]
-    fn test_with_no_table_privilege() {
+    fn test_copy_to_with_no_table_privilege() {
         let create_table = "create table test_table(id int);";
         Spi::run(create_table).unwrap();
 
@@ -496,7 +496,29 @@ mod tests {
         let set_role = "set role test_role;";
         Spi::run(set_role).unwrap();
 
-        let copy_from_parquet = format!("copy test_table to '{}';", LOCAL_TEST_FILE_PATH);
+        let copy_to_parquet = format!("copy test_table to '{}';", LOCAL_TEST_FILE_PATH);
+        Spi::run(&copy_to_parquet).unwrap();
+    }
+
+    #[pg_test]
+    #[should_panic(expected = "permission denied for table test_table")]
+    fn test_copy_from_with_no_table_privilege() {
+        let copy_to_parquet = format!("copy (select 1) to '{LOCAL_TEST_FILE_PATH}';");
+        Spi::run(&copy_to_parquet).unwrap();
+
+        let create_table = "create table test_table(id int);";
+        Spi::run(create_table).unwrap();
+
+        let create_role = "create role test_role;";
+        Spi::run(create_role).unwrap();
+
+        let grant_role = "grant pg_read_server_files TO test_role;";
+        Spi::run(grant_role).unwrap();
+
+        let set_role = "set role test_role;";
+        Spi::run(set_role).unwrap();
+
+        let copy_from_parquet = format!("copy test_table from '{}';", LOCAL_TEST_FILE_PATH);
         Spi::run(&copy_from_parquet).unwrap();
     }
 }