Subject: revamp row-security tracking

From: Nathan Bossart <[email-masked]>
Date: Thu, 21 Nov 2024 12:00:37 -0600

In light of CVE-2024-10976, which was fixed by commit cd7ab57, I'd like to
propose a bigger change to this area of the code that aims to future-proof
it a bit.  Instead of requiring hackers to carefully cart around whether a
query references a table with RLS enabled, I think we should instead
accumulate such information globally and require higher-level routines like
fireRIRrules() and inline_set_returning_function() to inspect it as needed.

The attached patch accomplishes this by establishing a global queue of
row-security "nest levels" that the aforementioned higher-level callers can
use.  Essentially, they will first call PushRowSecurityNestLevel(), then
any calls to get_row_security_policies() that encounter a table with
row-security enabled will mark the current nest level (and its parents).
Finally, PopRowSecurity() is used to retrieve whether row-security might
apply, and the query can be marked correctly.  I've also attempted to
handle resetting this queue after an ERROR or failed SRF inlining, but I'm
not yet positive that I've got all the details right.

With this patch applied, we should be able to revert some of the
row-security tracking code, especially the stuff added by commit cd7ab57.

Thoughts?

-- 
nathan
>From 7645f3a7215b052ab29be6d8e2b6e84954cb5caf Mon Sep 17 00:00:00 2001
From: Nathan Bossart <[email-masked]>
Date: Thu, 21 Nov 2024 11:57:43 -0600
Subject: [PATCH v1 1/1] revamp row-security tracking

---
 src/backend/access/transam/xact.c    |  6 ++
 src/backend/executor/functions.c     |  6 --
 src/backend/optimizer/util/clauses.c |  6 ++
 src/backend/rewrite/rewriteHandler.c | 77 ++++--------------------
 src/backend/rewrite/rowsecurity.c    | 88 +++++++++++++++++++++++++---
 src/include/rewrite/rowsecurity.h    |  6 +-
 src/tools/pgindent/typedefs.list     |  1 -
 7 files changed, 109 insertions(+), 81 deletions(-)

diff --git a/src/backend/access/transam/xact.c b/src/backend/access/transam/xact.c
index b7ebcc2a55..440bf34906 100644
--- a/src/backend/access/transam/xact.c
+++ b/src/backend/access/transam/xact.c
@@ -51,6 +51,7 @@
 #include "replication/origin.h"
 #include "replication/snapbuild.h"
 #include "replication/syncrep.h"
+#include "rewrite/rowsecurity.h"
 #include "storage/condition_variable.h"
 #include "storage/fd.h"
 #include "storage/lmgr.h"
@@ -2473,6 +2474,7 @@ CommitTransaction(void)
 	AtEOXact_Snapshot(true, false);
 	AtEOXact_ApplyLauncher(true);
 	AtEOXact_LogicalRepWorkers(true);
+	AtEOXact_RowSecurity(true);
 	pgstat_report_xact_timestamp(0);
 
 	ResourceOwnerDelete(TopTransactionResourceOwner);
@@ -2766,6 +2768,7 @@ PrepareTransaction(void)
 	/* we treat PREPARE as ROLLBACK so far as waking workers goes */
 	AtEOXact_ApplyLauncher(false);
 	AtEOXact_LogicalRepWorkers(false);
+	AtEOXact_RowSecurity(true);
 	pgstat_report_xact_timestamp(0);
 
 	CurrentResourceOwner = NULL;
@@ -2985,6 +2988,7 @@ AbortTransaction(void)
 		AtEOXact_PgStat(false, is_parallel_worker);
 		AtEOXact_ApplyLauncher(false);
 		AtEOXact_LogicalRepWorkers(false);
+		AtEOXact_RowSecurity(false);
 		pgstat_report_xact_timestamp(0);
 	}
 
@@ -5197,6 +5201,7 @@ CommitSubTransaction(void)
 	AtEOSubXact_HashTables(true, s->nestingLevel);
 	AtEOSubXact_PgStat(true, s->nestingLevel);
 	AtSubCommit_Snapshot(s->nestingLevel);
+	AtEOXact_RowSecurity(true);
 
 	/*
 	 * We need to restore the upper transaction's read-only state, in case the
@@ -5362,6 +5367,7 @@ AbortSubTransaction(void)
 		AtEOSubXact_HashTables(false, s->nestingLevel);
 		AtEOSubXact_PgStat(false, s->nestingLevel);
 		AtSubAbort_Snapshot(s->nestingLevel);
+		AtEOXact_RowSecurity(false);
 	}
 
 	/*
diff --git a/src/backend/executor/functions.c b/src/backend/executor/functions.c
index 3988066330..692854e2b3 100644
--- a/src/backend/executor/functions.c
+++ b/src/backend/executor/functions.c
@@ -1972,12 +1972,6 @@ tlist_coercion_finished:
 		rtr->rtindex = 1;
 		newquery->jointree = makeFromExpr(list_make1(rtr), NULL);
 
-		/*
-		 * Make sure the new query is marked as having row security if the
-		 * original one does.
-		 */
-		newquery->hasRowSecurity = parse->hasRowSecurity;
-
 		/* Replace original query in the correct element of the query list */
 		lfirst(parse_cell) = newquery;
 	}
diff --git a/src/backend/optimizer/util/clauses.c b/src/backend/optimizer/util/clauses.c
index 8e39795e24..775c79eca5 100644
--- a/src/backend/optimizer/util/clauses.c
+++ b/src/backend/optimizer/util/clauses.c
@@ -43,6 +43,7 @@
 #include "parser/parse_func.h"
 #include "rewrite/rewriteHandler.h"
 #include "rewrite/rewriteManip.h"
+#include "rewrite/rowsecurity.h"
 #include "tcop/tcopprot.h"
 #include "utils/acl.h"
 #include "utils/builtins.h"
@@ -5082,6 +5083,7 @@ inline_set_returning_function(PlannerInfo *root, RangeTblEntry *rte)
 	List	   *raw_parsetree_list;
 	List	   *querytree_list;
 	Query	   *querytree;
+	int			rowsec_level;
 
 	Assert(rte->rtekind == RTE_FUNCTION);
 
@@ -5169,6 +5171,8 @@ inline_set_returning_function(PlannerInfo *root, RangeTblEntry *rte)
 		return NULL;
 	}
 
+	rowsec_level = PushRowSecurityNestLevel();
+
 	/*
 	 * Make a temporary memory context, so that we don't leak all the stuff
 	 * that parsing might create.
@@ -5333,6 +5337,7 @@ inline_set_returning_function(PlannerInfo *root, RangeTblEntry *rte)
 	 * We must also notice if the inserted query adds a dependency on the
 	 * calling role due to RLS quals.
 	 */
+	querytree->hasRowSecurity = PopRowSecurityNestLevel(false, rowsec_level);
 	if (querytree->hasRowSecurity)
 		root->glob->dependsOnRole = true;
 
@@ -5340,6 +5345,7 @@ inline_set_returning_function(PlannerInfo *root, RangeTblEntry *rte)
 
 	/* Here if func is not inlinable: release temp memory and return NULL */
 fail:
+	PopRowSecurityNestLevel(true, rowsec_level);
 	MemoryContextSwitchTo(oldcxt);
 	MemoryContextDelete(mycxt);
 	error_context_stack = sqlerrcontext.previous;
diff --git a/src/backend/rewrite/rewriteHandler.c b/src/backend/rewrite/rewriteHandler.c
index 063afd4933..56c6ddf86d 100644
--- a/src/backend/rewrite/rewriteHandler.c
+++ b/src/backend/rewrite/rewriteHandler.c
@@ -58,12 +58,6 @@ typedef struct acquireLocksOnSubLinks_context
 	bool		for_execute;	/* AcquireRewriteLocks' forExecute param */
 } acquireLocksOnSubLinks_context;
 
-typedef struct fireRIRonSubLink_context
-{
-	List	   *activeRIRs;
-	bool		hasRowSecurity;
-} fireRIRonSubLink_context;
-
 static bool acquireLocksOnSubLinks(Node *node,
 								   acquireLocksOnSubLinks_context *context);
 static Query *rewriteRuleAction(Query *parsetree,
@@ -1845,12 +1839,6 @@ ApplyRetrieveRule(Query *parsetree,
 	 */
 	rule_action = fireRIRrules(rule_action, activeRIRs);
 
-	/*
-	 * Make sure the query is marked as having row security if the view query
-	 * does.
-	 */
-	parsetree->hasRowSecurity |= rule_action->hasRowSecurity;
-
 	/*
 	 * Now, plug the view query in as a subselect, converting the relation's
 	 * original RTE to a subquery RTE.
@@ -1964,7 +1952,7 @@ markQueryForLocking(Query *qry, Node *jtnode,
  * the SubLink's subselect link with the possibly-rewritten subquery.
  */
 static bool
-fireRIRonSubLink(Node *node, fireRIRonSubLink_context *context)
+fireRIRonSubLink(Node *node, List *activeRIRs)
 {
 	if (node == NULL)
 		return false;
@@ -1974,12 +1962,7 @@ fireRIRonSubLink(Node *node, fireRIRonSubLink_context *context)
 
 		/* Do what we came for */
 		sub->subselect = (Node *) fireRIRrules((Query *) sub->subselect,
-											   context->activeRIRs);
-
-		/*
-		 * Remember if any of the sublinks have row security.
-		 */
-		context->hasRowSecurity |= ((Query *) sub->subselect)->hasRowSecurity;
+											   activeRIRs);
 
 		/* Fall through to process lefthand args of SubLink */
 	}
@@ -1989,7 +1972,7 @@ fireRIRonSubLink(Node *node, fireRIRonSubLink_context *context)
 	 * subselects of subselects for us.
 	 */
 	return expression_tree_walker(node, fireRIRonSubLink,
-								  (void *) context);
+								  (void *) activeRIRs);
 }

@@ -2006,6 +1989,7 @@ fireRIRrules(Query *parsetree, List *activeRIRs)
 	int			origResultRelation = parsetree->resultRelation;
 	int			rt_index;
 	ListCell   *lc;
+	int			rowsec_level = PushRowSecurityNestLevel();
 
 	/*
 	 * Expand SEARCH and CYCLE clauses in CTEs.
@@ -2050,13 +2034,6 @@ fireRIRrules(Query *parsetree, List *activeRIRs)
 		if (rte->rtekind == RTE_SUBQUERY)
 		{
 			rte->subquery = fireRIRrules(rte->subquery, activeRIRs);
-
-			/*
-			 * While we are here, make sure the query is marked as having row
-			 * security if any of its subqueries do.
-			 */
-			parsetree->hasRowSecurity |= rte->subquery->hasRowSecurity;
-
 			continue;
 		}
 
@@ -2170,12 +2147,6 @@ fireRIRrules(Query *parsetree, List *activeRIRs)
 
 		cte->ctequery = (Node *)
 			fireRIRrules((Query *) cte->ctequery, activeRIRs);
-
-		/*
-		 * While we are here, make sure the query is marked as having row
-		 * security if any of its CTEs do.
-		 */
-		parsetree->hasRowSecurity |= ((Query *) cte->ctequery)->hasRowSecurity;
 	}
 
 	/*
@@ -2183,22 +2154,9 @@ fireRIRrules(Query *parsetree, List *activeRIRs)
 	 * the rtable and cteList.
 	 */
 	if (parsetree->hasSubLinks)
-	{
-		fireRIRonSubLink_context context;
-
-		context.activeRIRs = activeRIRs;
-		context.hasRowSecurity = false;
-
-		query_tree_walker(parsetree, fireRIRonSubLink, (void *) &context,
+		query_tree_walker(parsetree, fireRIRonSubLink, (void *) activeRIRs,
 						  QTW_IGNORE_RC_SUBQUERIES);
 
-		/*
-		 * Make sure the query is marked as having row security if any of its
-		 * sublinks do.
-		 */
-		parsetree->hasRowSecurity |= context.hasRowSecurity;
-	}
-
 	/*
 	 * Apply any row-level security policies.  We do this last because it
 	 * requires special recursion detection if the new quals have sublink
@@ -2212,7 +2170,6 @@ fireRIRrules(Query *parsetree, List *activeRIRs)
 		Relation	rel;
 		List	   *securityQuals;
 		List	   *withCheckOptions;
-		bool		hasRowSecurity;
 		bool		hasSubLinks;
 
 		++rt_index;
@@ -2230,14 +2187,13 @@ fireRIRrules(Query *parsetree, List *activeRIRs)
 		 */
 		get_row_security_policies(parsetree, rte, rt_index,
 								  &securityQuals, &withCheckOptions,
-								  &hasRowSecurity, &hasSubLinks);
+								  &hasSubLinks);
 
 		if (securityQuals != NIL || withCheckOptions != NIL)
 		{
 			if (hasSubLinks)
 			{
 				acquireLocksOnSubLinks_context context;
-				fireRIRonSubLink_context fire_context;
 
 				/*
 				 * Recursively process the new quals, checking for infinite
@@ -2268,21 +2224,11 @@ fireRIRrules(Query *parsetree, List *activeRIRs)
 				 * Now that we have the locks on anything added by
 				 * get_row_security_policies, fire any RIR rules for them.
 				 */
-				fire_context.activeRIRs = activeRIRs;
-				fire_context.hasRowSecurity = false;
-
 				expression_tree_walker((Node *) securityQuals,
-									   fireRIRonSubLink, (void *) &fire_context);
+									   fireRIRonSubLink, (void *) activeRIRs);
 
 				expression_tree_walker((Node *) withCheckOptions,
-									   fireRIRonSubLink, (void *) &fire_context);
-
-				/*
-				 * We can ignore the value of fire_context.hasRowSecurity
-				 * since we only reach this code in cases where hasRowSecurity
-				 * is already true.
-				 */
-				Assert(hasRowSecurity);
+									   fireRIRonSubLink, (void *) activeRIRs);
 
 				activeRIRs = list_delete_last(activeRIRs);
 			}
@@ -2301,17 +2247,16 @@ fireRIRrules(Query *parsetree, List *activeRIRs)
 		}
 
 		/*
-		 * Make sure the query is marked correctly if row-level security
-		 * applies, or if the new quals had sublinks.
+		 * Make sure the query is marked correctly if the new quals had
+		 * sublinks.
 		 */
-		if (hasRowSecurity)
-			parsetree->hasRowSecurity = true;
 		if (hasSubLinks)
 			parsetree->hasSubLinks = true;
 
 		table_close(rel, NoLock);
 	}
 
+	parsetree->hasRowSecurity = PopRowSecurityNestLevel(false, rowsec_level);
 	return parsetree;
 }
 
diff --git a/src/backend/rewrite/rowsecurity.c b/src/backend/rewrite/rowsecurity.c
index 59fd305dd7..de3fb9325f 100644
--- a/src/backend/rewrite/rowsecurity.c
+++ b/src/backend/rewrite/rowsecurity.c
@@ -45,6 +45,7 @@
 #include "rewrite/rewriteManip.h"
 #include "rewrite/rowsecurity.h"
 #include "utils/acl.h"
+#include "utils/memutils.h"
 #include "utils/rel.h"
 #include "utils/rls.h"
 
@@ -74,6 +75,11 @@ static void add_with_check_options(Relation rel,
 
 static bool check_role_for_policy(ArrayType *policy_roles, Oid user_id);
 
+static void SetRowSecurityForCurrentNestLevel(void);
+
+static Bitmapset *row_sec_bms;
+static int	row_sec_nest_level = -1;
+
 /*
  * hooks to allow extensions to add their own security policies
  *
@@ -90,14 +96,14 @@ row_security_policy_hook_type row_security_policy_hook_restrictive = NULL;
  * Get any row security quals and WithCheckOption checks that should be
  * applied to the specified RTE.
  *
- * In addition, hasRowSecurity is set to true if row-level security is enabled
+ * In addition, we mark the current nest level if row security is enabled
  * (even if this RTE doesn't have any row security quals), and hasSubLinks is
  * set to true if any of the quals returned contain sublinks.
  */
 void
 get_row_security_policies(Query *root, RangeTblEntry *rte, int rt_index,
 						  List **securityQuals, List **withCheckOptions,
-						  bool *hasRowSecurity, bool *hasSubLinks)
+						  bool *hasSubLinks)
 {
 	Oid			user_id;
 	int			rls_status;
@@ -110,7 +116,6 @@ get_row_security_policies(Query *root, RangeTblEntry *rte, int rt_index,
 	/* Defaults for the return values */
 	*securityQuals = NIL;
 	*withCheckOptions = NIL;
-	*hasRowSecurity = false;
 	*hasSubLinks = false;
 
 	Assert(rte->rtekind == RTE_RELATION);
@@ -135,8 +140,8 @@ get_row_security_policies(Query *root, RangeTblEntry *rte, int rt_index,
 
 	/*
 	 * RLS_NONE_ENV means we are not doing any RLS now, but that may change
-	 * with changes to the environment, so we mark it as hasRowSecurity to
-	 * force a re-plan when the environment changes.
+	 * with changes to the environment, so we mark it as having row-security
+	 * to force a re-plan when the environment changes.
 	 */
 	if (rls_status == RLS_NONE_ENV)
 	{
@@ -145,7 +150,7 @@ get_row_security_policies(Query *root, RangeTblEntry *rte, int rt_index,
 		 * replanned if the environment changes (GUCs, role), but we are not
 		 * adding anything here.
 		 */
-		*hasRowSecurity = true;
+		SetRowSecurityForCurrentNestLevel();
 
 		return;
 	}
@@ -526,7 +531,7 @@ get_row_security_policies(Query *root, RangeTblEntry *rte, int rt_index,
 	 * Mark this query as having row security, so plancache can invalidate it
 	 * when necessary (eg: role changes)
 	 */
-	*hasRowSecurity = true;
+	SetRowSecurityForCurrentNestLevel();
 }
 
 /*
@@ -930,3 +935,72 @@ check_role_for_policy(ArrayType *policy_roles, Oid user_id)
 
 	return false;
 }
+
+/*
+ * The following functions are used to track whether any calls to
+ * get_row_security_policies() set hasRowSecurity between two points of code.
+ * This can be used to ensure that a Query is correctly marked as having
+ * row-level security dependencies without needing callers to preserve it in
+ * all cases.
+ *
+ * Before any potential calls to get_row_security_policies() that you care
+ * about, call PushRowSecurityNestLevel() to register your current level of
+ * recursion.  Then, once any potential calls to get_row_security_policies()
+ * are done, call PopRowSecurityNestLevel() and save its result in your Query
+ * object's hasRowSecurity flag.
+ *
+ * The value returned by PushRowSecurityNestLevel() should be passed to the
+ * matching call to PopRowSecurityNestLevel().  This is used to verify that we
+ * aren't stranding any nest levels or popping extra ones inadvertently.
+ */
+
+int
+PushRowSecurityNestLevel(void)
+{
+	return ++row_sec_nest_level;
+}
+
+bool
+PopRowSecurityNestLevel(bool discard, int nest_level)
+{
+	bool		ret = bms_is_member(row_sec_nest_level, row_sec_bms);
+
+	if (unlikely(nest_level != row_sec_nest_level))
+		elog(ERROR, "row security nest level mismatch");
+
+	if (ret)
+	{
+		MemoryContext oldcontext = MemoryContextSwitchTo(CurTransactionContext);
+
+		if (!discard && row_sec_nest_level > 0)
+			row_sec_bms = bms_add_member(row_sec_bms, row_sec_nest_level - 1);
+		row_sec_bms = bms_del_member(row_sec_bms, row_sec_nest_level);
+
+		MemoryContextSwitchTo(oldcontext);
+	}
+
+	row_sec_nest_level--;
+	return ret;
+}
+
+void
+AtEOXact_RowSecurity(bool isCommit)
+{
+	Assert(!isCommit || row_sec_nest_level == -1);
+
+	row_sec_bms = NULL;
+	row_sec_nest_level = -1;
+}
+
+static void
+SetRowSecurityForCurrentNestLevel(void)
+{
+	MemoryContext oldcontext = MemoryContextSwitchTo(CurTransactionContext);
+
+	if (unlikely(row_sec_nest_level < 0))
+		elog(ERROR, "no row security nest level was pushed");
+
+	row_sec_bms = bms_add_member(row_sec_bms, row_sec_nest_level);
+
+	MemoryContextSwitchTo(oldcontext);
+}
diff --git a/src/include/rewrite/rowsecurity.h b/src/include/rewrite/rowsecurity.h
index 1717ed4e5e..023f7f4091 100644
--- a/src/include/rewrite/rowsecurity.h
+++ b/src/include/rewrite/rowsecurity.h
@@ -44,6 +44,10 @@ extern PGDLLIMPORT row_security_policy_hook_type row_security_policy_hook_restri
 extern void get_row_security_policies(Query *root,
 									  RangeTblEntry *rte, int rt_index,
 									  List **securityQuals, List **withCheckOptions,
-									  bool *hasRowSecurity, bool *hasSubLinks);
+									  bool *hasSubLinks);
+
+extern int	PushRowSecurityNestLevel(void);
+extern bool PopRowSecurityNestLevel(bool discard, int nest_level);
+extern void AtEOXact_RowSecurity(bool isCommit);
 
 #endif							/* ROWSECURITY_H */
diff --git a/src/tools/pgindent/typedefs.list b/src/tools/pgindent/typedefs.list
index 08521d51a9..d7fd8507ba 100644
--- a/src/tools/pgindent/typedefs.list
+++ b/src/tools/pgindent/typedefs.list
@@ -3475,7 +3475,6 @@ fill_string_relopt
 finalize_primnode_context
 find_dependent_phvs_context
 find_expr_references_context
-fireRIRonSubLink_context
 fix_join_expr_context
 fix_scan_expr_context
 fix_upper_expr_context
-- 
2.39.5 (Apple Git-154)