From 2c1da2d3db13603da4b2cded4a2dcf0d86614ca0 Mon Sep 17 00:00:00 2001
From: Nathan Bossart <nathan@postgresql.org>
Date: Wed, 24 Sep 2025 09:47:02 -0500
Subject: [PATCH v4 3/3] fix priv checks in pg_prewarm

---
 contrib/pg_prewarm/pg_prewarm.c   | 34 +++++++++++++++++++++++++++++--
 contrib/pg_prewarm/t/001_basic.pl | 29 +++++++++++++++++++++++++-
 2 files changed, 60 insertions(+), 3 deletions(-)

diff --git a/contrib/pg_prewarm/pg_prewarm.c b/contrib/pg_prewarm/pg_prewarm.c
index b968933ea8b..810a291204b 100644
--- a/contrib/pg_prewarm/pg_prewarm.c
+++ b/contrib/pg_prewarm/pg_prewarm.c
@@ -16,9 +16,11 @@
 #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/smgr.h"
 #include "utils/acl.h"
@@ -71,6 +73,8 @@ pg_prewarm(PG_FUNCTION_ARGS)
 	char	   *ttype;
 	PrewarmType ptype;
 	AclResult	aclresult;
+	char		relkind;
+	Oid			privOid;
 
 	/* Basic sanity checking. */
 	if (PG_ARGISNULL(0))
@@ -107,8 +111,31 @@ pg_prewarm(PG_FUNCTION_ARGS)
 	forkNumber = forkname_to_number(forkString);
 
 	/* Open relation and check privileges. */
+	relkind = get_rel_relkind(relOid);
+	if (relkind == RELKIND_INDEX ||
+		relkind == RELKIND_PARTITIONED_INDEX)
+	{
+		privOid = IndexGetRelation(relOid, false);
+		LockRelationOid(privOid, AccessShareLock);
+	}
+	else
+		privOid = relOid;
+
 	rel = relation_open(relOid, AccessShareLock);
-	aclresult = pg_class_aclcheck(relOid, GetUserId(), ACL_SELECT);
+
+	/*
+	 * If we did the IndexGetRelation() call above, it's barely possible that
+	 * a race against an index drop/recreation could have netted us the wrong
+	 * table.
+	 */
+	if (privOid != relOid &&
+		privOid != IndexGetRelation(relOid, true))
+		ereport(ERROR,
+				(errcode(ERRCODE_UNDEFINED_TABLE),
+				 errmsg("could not find parent table of index \"%s\"",
+						RelationGetRelationName(rel))));
+
+	aclresult = pg_class_aclcheck(privOid, GetUserId(), ACL_SELECT);
 	if (aclresult != ACLCHECK_OK)
 		aclcheck_error(aclresult, get_relkind_objtype(rel->rd_rel->relkind), get_rel_name(relOid));
 
@@ -233,8 +260,11 @@ pg_prewarm(PG_FUNCTION_ARGS)
 		read_stream_end(stream);
 	}
 
-	/* Close relation, release lock. */
+	/* Close relation, release locks. */
 	relation_close(rel, AccessShareLock);
 
+	if (privOid != relOid)
+		UnlockRelationOid(privOid, 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..a77ab67d29e 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,31 @@ 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');
+
 # 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)