revamp row-security tracking

From 7645f3a7215b052ab29be6d8e2b6e84954cb5caf Mon Sep 17 00:00:00 2001
From: Nathan Bossart <nathan@postgresql.org>
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)