From 8f8606c007fdb58403792d36a62c6739341b0549 Mon Sep 17 00:00:00 2001
From: Ayush Vatsa <ayuvatsa@amazon.com>
Date: Mon, 17 Feb 2025 20:44:56 +0000
Subject: [PATCH v1] Improve ACL checks in pg_prewarm for indexes

When pg_prewarm is called on an index, perform the ACL check on its
underlying table instead of the index itself. Indexes do not have
independent access rights and depend on their associated table for
access control.
---
 contrib/pg_prewarm/pg_prewarm.c   | 29 ++++++++++++++++++++++++---
 contrib/pg_prewarm/t/001_basic.pl | 33 ++++++++++++++++++++++++++++++-
 2 files changed, 58 insertions(+), 4 deletions(-)

diff --git a/contrib/pg_prewarm/pg_prewarm.c b/contrib/pg_prewarm/pg_prewarm.c
index a2f0ac4af0..b2846cee2e 100644
--- a/contrib/pg_prewarm/pg_prewarm.c
+++ b/contrib/pg_prewarm/pg_prewarm.c
@@ -16,6 +16,7 @@
 #include <unistd.h>
 
 #include "access/relation.h"
+#include "catalog/index.h"
 #include "fmgr.h"
 #include "miscadmin.h"
 #include "storage/bufmgr.h"
@@ -55,6 +56,7 @@ Datum
 pg_prewarm(PG_FUNCTION_ARGS)
 {
 	Oid			relOid;
+	Oid			tableOid;
 	text	   *forkName;
 	text	   *type;
 	int64		first_block;
@@ -105,9 +107,30 @@ pg_prewarm(PG_FUNCTION_ARGS)
 
 	/* 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 access permissions for pg_prewarm. If the relation is an index,
+	 * perform the ACL check on its underlying table since indexes do not have
+	 * their own access rights.
+	 */
+	if (rel->rd_rel->relkind == RELKIND_INDEX)
+	{
+		/*
+		 * If it's an index, get the table OID and perform ACL check on the
+		 * table.
+		 */
+		tableOid = IndexGetRelation(relOid, false);
+		aclresult = pg_class_aclcheck(tableOid, GetUserId(), ACL_SELECT);
+		if (aclresult != ACLCHECK_OK)
+			aclcheck_error(aclresult, get_relkind_objtype(RELKIND_RELATION), get_rel_name(tableOid));
+	}
+	else
+	{
+		/* If it's not an index, perform ACL check on the relation itself. */
+		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 that the fork exists. */
 	if (!smgrexists(RelationGetSmgr(rel), forkNumber))
diff --git a/contrib/pg_prewarm/t/001_basic.pl b/contrib/pg_prewarm/t/001_basic.pl
index 0a8259d367..651f4ec371 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 WITH LOGIN;");
 
 # test read mode
 my $result =
@@ -42,6 +44,35 @@ ok( (        $stdout =~ qr/^[1-9][0-9]*$/
 		  or $stderr =~ qr/prefetch is not supported by this build/),
 	'prefetch mode succeeded');
 
+# grant SELECT permission on the table to test_user
+$node->safe_psql("postgres", "GRANT SELECT ON test TO test_user;");
+
+# test pg_prewarm on the table/index as test_user (should succeed)
+$result =
+  $node->safe_psql("postgres", "SELECT pg_prewarm('test', 'buffer');", extra_params => [ '--username' => "test_user" ]);
+like($result, qr/^[1-9][0-9]*$/, 'pg_prewarm on table succeeds with SELECT permission on table');
+
+$result =
+  $node->safe_psql("postgres", "SELECT pg_prewarm('test_idx', 'buffer');", extra_params => [ '--username' => "test_user" ]);
+like($result, qr/^[1-9][0-9]*$/, 'pg_prewarm on index succeeds with SELECT permission on table');
+
+# revoke SELECT permission on the table from test_user
+$node->safe_psql("postgres", "REVOKE SELECT ON test FROM test_user;");
+
+# test pg_prewarm on the table/index as test_user (should fail)
+($cmdret, $stdout, $stderr) =
+  $node->psql("postgres", "SELECT pg_prewarm('test', 'buffer');", extra_params => [ '--username' => "test_user" ]);
+ok( $stderr =~ /permission denied for table test/,
+	"error message indicates user doesn't have sufficient privileges for the parent table");
+
+($cmdret, $stdout, $stderr) =
+  $node->psql("postgres", "SELECT pg_prewarm('test_idx', 'buffer');", extra_params => [ '--username' => "test_user" ]);
+ok( $stderr =~ /permission denied for table test/,
+	"error message indicates user doesn't have sufficient privileges for the parent table");
+
+# clean up: drop the test_user role
+$node->safe_psql("postgres", "DROP ROLE test_user;");
+
 # 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.47.1