From df9b604c0d8c3c90df2026406e6ec0302f22454c Mon Sep 17 00:00:00 2001
From: Nathan Bossart <nathan@postgresql.org>
Date: Thu, 9 Oct 2025 15:49:18 -0500
Subject: [PATCH 1/1] pg_prewarm privilege test

---
 contrib/pg_prewarm/pg_prewarm.c   | 83 +++++++++++++++++++++++++++++--
 contrib/pg_prewarm/t/001_basic.pl | 36 +++++++++++++-
 2 files changed, 113 insertions(+), 6 deletions(-)

diff --git a/contrib/pg_prewarm/pg_prewarm.c b/contrib/pg_prewarm/pg_prewarm.c
index b968933ea8b..57b720cfbed 100644
--- a/contrib/pg_prewarm/pg_prewarm.c
+++ b/contrib/pg_prewarm/pg_prewarm.c
@@ -16,13 +16,17 @@
 #include <unistd.h>
 
 #include "access/relation.h"
+#include "catalog/index.h"
 #include "fmgr.h"
 #include "miscadmin.h"
 #include "storage/bufmgr.h"
+#include "storage/lmgr.h"
 #include "storage/read_stream.h"
+#include "storage/sinval.h"
 #include "storage/smgr.h"
 #include "utils/acl.h"
 #include "utils/builtins.h"
+#include "utils/inval.h"
 #include "utils/lsyscache.h"
 #include "utils/rel.h"
 
@@ -42,6 +46,74 @@ typedef enum
 
 static PGIOAlignedBlock blockbuffer;
 
+static void
+check_privs_and_lock_rel(Oid relid, Oid *heaprel)
+{
+	uint64		inval_count;
+	Oid			privOid = InvalidOid;
+
+	do
+	{
+		AclResult	aclresult;
+		bool		is_missing = false;
+		char		relkind;
+
+		/*
+		 * Remember this value so that, after locking the relation(s), we can
+		 * check whether any invalidation messages have been processed that
+		 * might require a do-over.
+		 */
+		inval_count = SharedInvalidMessageCounter;
+
+		/*
+		 * Unlock any previously-locked relations.  We could try to hold onto
+		 * these through iterations, but it's simpler to just start fresh each
+		 * time.
+		 */
+		if (OidIsValid(privOid))
+		{
+			if (privOid != relid)
+				UnlockRelationOid(privOid, AccessShareLock);
+			UnlockRelationOid(relid, AccessShareLock);
+		}
+
+		/*
+		 * For indexes, privilege checks must happen on the heap relation.
+		 */
+		relkind = get_rel_relkind(relid);
+		if (relkind == RELKIND_INDEX || relkind == RELKIND_PARTITIONED_INDEX)
+			privOid = IndexGetRelation(relid, false);
+		else
+			privOid = relid;
+
+		/*
+		 * Check privileges.  If the relation has gone missing, there's
+		 * nothing for us to do.  We'll either retry in the next loop
+		 * iteration, or the caller will fail when it tries to open the
+		 * relation.
+		 */
+		aclresult = pg_class_aclcheck_ext(privOid, GetUserId(),
+										  ACL_SELECT, &is_missing);
+		if (!is_missing && aclresult != ACLCHECK_OK)
+			aclcheck_error(aclresult,
+						   get_relkind_objtype(relkind),
+						   get_rel_name(relid));
+
+		/*
+		 * Lock the relation(s).  If relid is an index, make sure we lock its
+		 * heap first.  Note that we rely on LockRelationOid() to call
+		 * AcceptInvalidationMessages().
+		 */
+		if (privOid != relid)
+			LockRelationOid(privOid, AccessShareLock);
+		LockRelationOid(relid, AccessShareLock);
+
+	} while (inval_count != SharedInvalidMessageCounter);
+
+	if (privOid != relid)
+		*heaprel = privOid;
+}
+
 /*
  * pg_prewarm(regclass, mode text, fork text,
  *			  first_block int8, last_block int8)
@@ -70,7 +142,7 @@ pg_prewarm(PG_FUNCTION_ARGS)
 	char	   *forkString;
 	char	   *ttype;
 	PrewarmType ptype;
-	AclResult	aclresult;
+	Oid			heaprel = InvalidOid;
 
 	/* Basic sanity checking. */
 	if (PG_ARGISNULL(0))
@@ -107,10 +179,8 @@ pg_prewarm(PG_FUNCTION_ARGS)
 	forkNumber = forkname_to_number(forkString);
 
 	/* Open relation and check privileges. */
-	rel = relation_open(relOid, AccessShareLock);
-	aclresult = pg_class_aclcheck(relOid, GetUserId(), ACL_SELECT);
-	if (aclresult != ACLCHECK_OK)
-		aclcheck_error(aclresult, get_relkind_objtype(rel->rd_rel->relkind), get_rel_name(relOid));
+	check_privs_and_lock_rel(relOid, &heaprel);
+	rel = relation_open(relOid, NoLock);
 
 	/* Check that the relation has storage. */
 	if (!RELKIND_HAS_STORAGE(rel->rd_rel->relkind))
@@ -236,5 +306,8 @@ pg_prewarm(PG_FUNCTION_ARGS)
 	/* Close relation, release lock. */
 	relation_close(rel, AccessShareLock);
 
+	if (OidIsValid(heaprel))
+		UnlockRelationOid(heaprel, AccessShareLock);
+
 	PG_RETURN_INT64(blocks_done);
 }
diff --git a/contrib/pg_prewarm/t/001_basic.pl b/contrib/pg_prewarm/t/001_basic.pl
index 0a8259d3678..5d7010eb44e 100644
--- a/contrib/pg_prewarm/t/001_basic.pl
+++ b/contrib/pg_prewarm/t/001_basic.pl
@@ -23,7 +23,9 @@ $node->start;
 $node->safe_psql("postgres",
 		"CREATE EXTENSION pg_prewarm;\n"
 	  . "CREATE TABLE test(c1 int);\n"
-	  . "INSERT INTO test SELECT generate_series(1, 100);");
+	  . "INSERT INTO test SELECT generate_series(1, 100);\n"
+	  . "CREATE INDEX test_idx ON test(c1);\n"
+	  . "CREATE ROLE test_user LOGIN;");
 
 # test read mode
 my $result =
@@ -42,6 +44,38 @@ ok( (        $stdout =~ qr/^[1-9][0-9]*$/
 		  or $stderr =~ qr/prefetch is not supported by this build/),
 	'prefetch mode succeeded');
 
+# test_user should be unable to prewarm table/index without privileges
+($cmdret, $stdout, $stderr) =
+  $node->psql(
+    "postgres", "SELECT pg_prewarm('test');",
+    extra_params => [ '--username' => 'test_user' ]);
+ok($stderr =~ /permission denied for table test/, 'pg_prewarm failed as expected');
+($cmdret, $stdout, $stderr) =
+  $node->psql(
+    "postgres", "SELECT pg_prewarm('test_idx');",
+    extra_params => [ '--username' => 'test_user' ]);
+ok($stderr =~ /permission denied for index test_idx/, 'pg_prewarm failed as expected');
+
+# test_user should be able to prewarm table/index with privileges
+$node->safe_psql("postgres", "GRANT SELECT ON test TO test_user;");
+$result =
+  $node->safe_psql(
+    "postgres", "SELECT pg_prewarm('test');",
+    extra_params => [ '--username' => 'test_user' ]);
+like($result, qr/^[1-9][0-9]*$/, 'pg_prewarm succeeded as expected');
+$result =
+  $node->safe_psql(
+    "postgres", "SELECT pg_prewarm('test_idx');",
+    extra_params => [ '--username' => 'test_user' ]);
+like($result, qr/^[1-9][0-9]*$/, 'pg_prewarm succeeded as expected');
+
+# prewarm fails for nonexistent table
+($cmdret, $stdout, $stderr) =
+  $node->psql(
+    "postgres", "SELECT pg_prewarm(0);",
+    extra_params => [ '--username' => 'test_user' ]);
+ok($stderr =~ /could not open relation with OID 0/, 'pg_prewarm failed as expected');
+
 # test autoprewarm_dump_now()
 $result = $node->safe_psql("postgres", "SELECT autoprewarm_dump_now();");
 like($result, qr/^[1-9][0-9]*$/, 'autoprewarm_dump_now succeeded');
-- 
2.39.5 (Apple Git-154)

