COPY command with RLS bug

From ddf1247c550e7b30c9b8dc51cda438b7425bc7fa Mon Sep 17 00:00:00 2001
From: Stephen Frost <sfrost@snowman.net>
Date: Fri, 30 Sep 2016 13:31:22 -0400
Subject: [PATCH] Fix RLS with COPY (col1, col2) FROM tab

Attempting to COPY a subset of columns from a table with RLS enabled
would fail due to an invalid query being constructed (using a single
ColumnRef with the list of fields to exact in 'fields', but that's for
the different levels of an indirection for a single column, not for
specifying multiple columns).

Correct by building a ColumnRef and then RestTarget for each column
being requested and then adding those to the targetList for the select
query.  Include regression tests to hopefully catch if this is broken
again in the future.

Patch-By: Adam Brightwell
Reviewed-By: Michael Paquier
---
 src/backend/commands/copy.c         | 63 ++++++++++++++++++++++++------
 src/test/regress/expected/copy2.out | 78 +++++++++++++++++++++++++++++++++++++
 src/test/regress/sql/copy2.sql      | 63 ++++++++++++++++++++++++++++++
 3 files changed, 192 insertions(+), 12 deletions(-)

diff --git a/src/backend/commands/copy.c b/src/backend/commands/copy.c
index 432b0ca..e665f06 100644
--- a/src/backend/commands/copy.c
+++ b/src/backend/commands/copy.c
@@ -871,6 +871,7 @@ DoCopy(ParseState *pstate, const CopyStmt *stmt, uint64 *processed)
 			ColumnRef  *cr;
 			ResTarget  *target;
 			RangeVar   *from;
+			List	   *targetList = NIL;
 
 			if (is_from)
 				ereport(ERROR,
@@ -878,21 +879,59 @@ DoCopy(ParseState *pstate, const CopyStmt *stmt, uint64 *processed)
 				   errmsg("COPY FROM not supported with row-level security"),
 						 errhint("Use INSERT statements instead.")));
 
-			/* Build target list */
-			cr = makeNode(ColumnRef);
-
+			/*
+			 * Build target list
+			 *
+			 * If no columns are specified in the attribute list of the COPY
+			 * command, then the target list is 'all' columns. Therefore, '*'
+			 * should be used as the target list for the resulting SELECT
+			 * statement.
+			 *
+			 * In the case that columns are specified in the attribute list,
+			 * create a ColumnRef and ResTarget for each column and add them to
+			 * the target list for the resulting SELECT statement.
+			 */
 			if (!stmt->attlist)
+			{
+				cr = makeNode(ColumnRef);
 				cr->fields = list_make1(makeNode(A_Star));
-			else
-				cr->fields = stmt->attlist;
+				cr->location = -1;
+
+				target = makeNode(ResTarget);
+				target->name = NULL;
+				target->indirection = NIL;
+				target->val = (Node *) cr;
+				target->location = -1;
 
-			cr->location = 1;
+				targetList = list_make1(target);
+			}
+			else
+			{
+				ListCell   *lc;
 
-			target = makeNode(ResTarget);
-			target->name = NULL;
-			target->indirection = NIL;
-			target->val = (Node *) cr;
-			target->location = 1;
+				foreach(lc, stmt->attlist)
+				{
+					/*
+					 * Build the ColumnRef for each column.  The ColumnRef
+					 * 'fields' property is a String 'Value' node (see
+					 * nodes/value.h) that correspond to the column name
+					 * respectively.
+					 */
+					cr = makeNode(ColumnRef);
+					cr->fields = list_make1(lfirst(lc));
+					cr->location = -1;
+
+					/* Build the ResTarget and add the ColumnRef to it. */
+					target = makeNode(ResTarget);
+					target->name = NULL;
+					target->indirection = NIL;
+					target->val = (Node *) cr;
+					target->location = -1;
+
+					/* Add each column to the SELECT statement's target list */
+					targetList = lappend(targetList, target);
+				}
+			}
 
 			/*
 			 * Build RangeVar for from clause, fully qualified based on the
@@ -903,7 +942,7 @@ DoCopy(ParseState *pstate, const CopyStmt *stmt, uint64 *processed)
 
 			/* Build query */
 			select = makeNode(SelectStmt);
-			select->targetList = list_make1(target);
+			select->targetList = targetList;
 			select->fromClause = list_make1(from);
 
 			query = (Node *) select;
diff --git a/src/test/regress/expected/copy2.out b/src/test/regress/expected/copy2.out
index 5f6260a..b3c215c 100644
--- a/src/test/regress/expected/copy2.out
+++ b/src/test/regress/expected/copy2.out
@@ -460,9 +460,87 @@ select * from check_con_tbl;
    
 (2 rows)
 
+-- test with RLS enabled.
+CREATE ROLE regress_rls_copy_user;
+CREATE ROLE regress_rls_copy_user_colperms;
+CREATE TABLE rls_t1 (a int, b int, c int);
+COPY rls_t1 (a, b, c) from stdin;
+CREATE POLICY p1 ON rls_t1 FOR SELECT USING (a % 2 = 0);
+ALTER TABLE rls_t1 ENABLE ROW LEVEL SECURITY;
+ALTER TABLE rls_t1 FORCE ROW LEVEL SECURITY;
+GRANT SELECT ON TABLE rls_t1 TO regress_rls_copy_user;
+GRANT SELECT (a, b) ON TABLE rls_t1 TO regress_rls_copy_user_colperms;
+-- all columns
+COPY rls_t1 TO stdout;
+1	4	1
+2	3	2
+3	2	3
+4	1	4
+COPY rls_t1 (a, b, c) TO stdout;
+1	4	1
+2	3	2
+3	2	3
+4	1	4
+-- subset of columns
+COPY rls_t1 (a) TO stdout;
+1
+2
+3
+4
+COPY rls_t1 (a, b) TO stdout;
+1	4
+2	3
+3	2
+4	1
+-- column reordering
+COPY rls_t1 (b, a) TO stdout;
+4	1
+3	2
+2	3
+1	4
+SET SESSION AUTHORIZATION regress_rls_copy_user;
+-- all columns
+COPY rls_t1 TO stdout;
+2	3	2
+4	1	4
+COPY rls_t1 (a, b, c) TO stdout;
+2	3	2
+4	1	4
+-- subset of columns
+COPY rls_t1 (a) TO stdout;
+2
+4
+COPY rls_t1 (a, b) TO stdout;
+2	3
+4	1
+-- column reordering
+COPY rls_t1 (b, a) TO stdout;
+3	2
+1	4
+RESET SESSION AUTHORIZATION;
+SET SESSION AUTHORIZATION regress_rls_copy_user_colperms;
+-- attempt all columns (should fail)
+COPY rls_t1 TO stdout;
+ERROR:  permission denied for relation rls_t1
+COPY rls_t1 (a, b, c) TO stdout;
+ERROR:  permission denied for relation rls_t1
+-- try to copy column with no privileges (should fail)
+COPY rls_t1 (c) TO stdout;
+ERROR:  permission denied for relation rls_t1
+-- subset of columns (should succeed)
+COPY rls_t1 (a) TO stdout;
+2
+4
+COPY rls_t1 (a, b) TO stdout;
+2	3
+4	1
+RESET SESSION AUTHORIZATION;
 DROP TABLE forcetest;
 DROP TABLE vistest;
 DROP FUNCTION truncate_in_subxact();
 DROP TABLE x, y;
+DROP TABLE rls_t1 CASCADE;
+DROP ROLE regress_rls_copy_user;
+DROP ROLE regress_rls_copy_user_colperms;
 DROP FUNCTION fn_x_before();
 DROP FUNCTION fn_x_after();
diff --git a/src/test/regress/sql/copy2.sql b/src/test/regress/sql/copy2.sql
index 39a9deb..89d0a39 100644
--- a/src/test/regress/sql/copy2.sql
+++ b/src/test/regress/sql/copy2.sql
@@ -327,9 +327,72 @@ copy check_con_tbl from stdin;
 \.
 select * from check_con_tbl;
 
+-- test with RLS enabled.
+CREATE ROLE regress_rls_copy_user;
+CREATE ROLE regress_rls_copy_user_colperms;
+CREATE TABLE rls_t1 (a int, b int, c int);
+
+COPY rls_t1 (a, b, c) from stdin;
+1	4	1
+2	3	2
+3	2	3
+4	1	4
+\.
+
+CREATE POLICY p1 ON rls_t1 FOR SELECT USING (a % 2 = 0);
+ALTER TABLE rls_t1 ENABLE ROW LEVEL SECURITY;
+ALTER TABLE rls_t1 FORCE ROW LEVEL SECURITY;
+
+GRANT SELECT ON TABLE rls_t1 TO regress_rls_copy_user;
+GRANT SELECT (a, b) ON TABLE rls_t1 TO regress_rls_copy_user_colperms;
+
+-- all columns
+COPY rls_t1 TO stdout;
+COPY rls_t1 (a, b, c) TO stdout;
+
+-- subset of columns
+COPY rls_t1 (a) TO stdout;
+COPY rls_t1 (a, b) TO stdout;
+
+-- column reordering
+COPY rls_t1 (b, a) TO stdout;
+
+SET SESSION AUTHORIZATION regress_rls_copy_user;
+
+-- all columns
+COPY rls_t1 TO stdout;
+COPY rls_t1 (a, b, c) TO stdout;
+
+-- subset of columns
+COPY rls_t1 (a) TO stdout;
+COPY rls_t1 (a, b) TO stdout;
+
+-- column reordering
+COPY rls_t1 (b, a) TO stdout;
+
+RESET SESSION AUTHORIZATION;
+
+SET SESSION AUTHORIZATION regress_rls_copy_user_colperms;
+
+-- attempt all columns (should fail)
+COPY rls_t1 TO stdout;
+COPY rls_t1 (a, b, c) TO stdout;
+
+-- try to copy column with no privileges (should fail)
+COPY rls_t1 (c) TO stdout;
+
+-- subset of columns (should succeed)
+COPY rls_t1 (a) TO stdout;
+COPY rls_t1 (a, b) TO stdout;
+
+RESET SESSION AUTHORIZATION;
+
 DROP TABLE forcetest;
 DROP TABLE vistest;
 DROP FUNCTION truncate_in_subxact();
 DROP TABLE x, y;
+DROP TABLE rls_t1 CASCADE;
+DROP ROLE regress_rls_copy_user;
+DROP ROLE regress_rls_copy_user_colperms;
 DROP FUNCTION fn_x_before();
 DROP FUNCTION fn_x_after();
-- 
2.7.4