From df9b604c0d8c3c90df2026406e6ec0302f22454c Mon Sep 17 00:00:00 2001 From: Nathan Bossart 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 #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)