[PG16] Add various hooks for SU functionality

src/backend/commands/event_trigger.c

@@ -74,6 +74,7 @@ typedef struct EventTriggerQueryState
 } EventTriggerQueryState;
 
 static EventTriggerQueryState *currentEventTriggerState = NULL;
+EventTrigger_hook_type EventTrigger_hook = NULL;
 
 /* Support for dropped objects */
 typedef struct SQLDropObject
@@ -921,7 +922,10 @@ EventTriggerInvoke(List *fn_oid_list, EventTriggerData *trigdata)
 		InitFunctionCallInfoData(*fcinfo, &flinfo, 0,
 								 InvalidOid, (Node *) trigdata, NULL);
 		pgstat_init_function_usage(fcinfo, &fcusage);
-		FunctionCallInvoke(fcinfo);
+		if (EventTrigger_hook)
+			EventTrigger_hook(fcinfo);
+		else
+			FunctionCallInvoke(fcinfo);
 		pgstat_end_function_usage(&fcusage, true);
 
 		/* Reclaim memory. */

src/backend/commands/trigger.c

@@ -72,6 +72,8 @@ int			SessionReplicationRole = SESSION_REPLICATION_ROLE_ORIGIN;
 /* How many levels deep into trigger execution are we? */
 static int	MyTriggerDepth = 0;
 
+DataTrigger_hook_type DataTrigger_hook = NULL;
+
 /* Local function prototypes */
 static void renametrig_internal(Relation tgrel, Relation targetrel,
 								HeapTuple trigtup, const char *newname,
@@ -2366,7 +2368,10 @@ ExecCallTriggerFunc(TriggerData *trigdata,
 	MyTriggerDepth++;
 	PG_TRY();
 	{
-		result = FunctionCallInvoke(fcinfo);
+		if (DataTrigger_hook)
+			result = DataTrigger_hook(fcinfo);
+		else
+			result = FunctionCallInvoke(fcinfo);
 	}
 	PG_FINALLY();
 	{

src/backend/replication/libpqwalreceiver/libpqwalreceiver.c

@@ -163,15 +163,11 @@ libpqrcv_connect(const char *conninfo, bool logical, bool must_use_password,
 	 */
 	if (pg_strcasecmp(appname, "walreceiver") == 0)
 	{
-		if (neon_storage_token[0] != '\0')
+		if (neon_storage_token && neon_storage_token[0] != '\0')
 		{
 			keys[++i] = "password";
 			vals[i] = neon_storage_token;
 		}
-		else
-		{
-			elog(LOG, "no storage token set");
-		}
 	}
 /* END_NEON */
 

src/backend/replication/walreceiver.c

@@ -1396,22 +1396,6 @@ WalRcvGetStateString(WalRcvState state)
 	return "UNKNOWN";
 }
 
-/*
- * We currently grant the privileged role pg_monitor, which implies
- * pg_read_all_settings. Until we fix that, let's just redact the content unless
- * the user requesting the value is a superuser.
- *
- * See: https://databricks.atlassian.net/browse/LKB-7128
- */
-const char *
-show_neon_storage_token(void)
-{
-	if (superuser())
-		return neon_storage_token;
-
-	return "**********";
-}
-
 /*
  * Returns activity of WAL receiver, including pid, state and xlog locations
  * received from the WAL sender of another server.

src/backend/utils/adt/acl.c

@@ -129,6 +129,24 @@ static void RoleMembershipCacheCallback(Datum arg, int cacheid, uint32 hashvalue
  */
 char *privileged_role_name = NULL;
 
+static bool
+is_privileged_role_arg_super(Oid roleid, bool nosuper)
+{
+	Oid privileged_role_oid;
+
+	if (privileged_role_name == NULL)
+		return false;
+
+	privileged_role_oid = get_role_oid(privileged_role_name, true /* missing_ok */);
+
+	if (privileged_role_oid == InvalidOid)
+		return false;
+
+	if (nosuper)
+		return has_privs_of_role_nosuper(roleid, privileged_role_oid);
+	return has_privs_of_role(roleid, privileged_role_oid);
+}
+
 bool
 is_privileged_role(void)
 {
@@ -138,14 +156,12 @@ is_privileged_role(void)
 bool
 is_privileged_role_arg(Oid roleid)
 {
-	Oid privileged_role_oid;
-
-	if (privileged_role_name == NULL)
-		return false;
-
-	privileged_role_oid = get_role_oid(privileged_role_name, true /* missing_ok */);
+	return is_privileged_role_arg_super(roleid, false);
+}
 
-	return privileged_role_oid != InvalidOid && has_privs_of_role(roleid, privileged_role_oid);
+bool is_privileged_role_nosuper(void)
+{
+	return is_privileged_role_arg_super(GetUserId(), true);
 }
 
 /*
@@ -5018,6 +5034,24 @@ has_privs_of_role(Oid member, Oid role)
 						   role);
 }
 
+/*
+ * Same as has_privs_of_role, but ignores checking superuser.
+ */
+ bool
+ has_privs_of_role_nosuper(Oid member, Oid role)
+ {
+	 /* Fast path for simple case */
+	 if (member == role)
+		 return true;
+	 /*
+	  * Find all the roles that member has the privileges of, including
+	  * multi-level recursion, then see if target role is any one of them.
+	  */
+	 return list_member_oid(roles_is_member_of(member, ROLERECURSE_PRIVS,
+											   InvalidOid, NULL),
+							role);
+ }
+
 /*
  * Can member use SET ROLE to this role?
  *

src/backend/utils/misc/guc_tables.c

@@ -4640,18 +4640,6 @@ struct config_string ConfigureNamesString[] =
 		check_restrict_nonsystem_relation_kind, assign_restrict_nonsystem_relation_kind, NULL
 	},
 
-
-	{
-		{"neon_storage_token", PGC_SUSET, REPLICATION_STANDBY,
-			"Authentication token for Neon storage",
-			NULL,
-			GUC_NO_SHOW_ALL | GUC_NO_RESET | GUC_NO_RESET_ALL | GUC_NOT_IN_SAMPLE | GUC_SUPERUSER_ONLY
-		},
-		&neon_storage_token,
-		"",
-		NULL, NULL, show_neon_storage_token
-	},
-
 	/* End-of-list marker */
 	{
 		{NULL, 0, 0, NULL, NULL}, NULL, NULL, NULL, NULL, NULL

src/backend/utils/misc/superuser.c

@@ -35,6 +35,7 @@
 static Oid	last_roleid = InvalidOid;	/* InvalidOid == cache not valid */
 static bool last_roleid_is_super = false;
 static bool roleid_callback_registered = false;
+SUForUser_hook_type SUForUser_hook = NULL;
 
 static void RoleidCallback(Datum arg, int cacheid, uint32 hashvalue);
 
@@ -72,6 +73,11 @@ superuser_arg(Oid roleid)
 	{
 		result = ((Form_pg_authid) GETSTRUCT(rtup))->rolsuper;
 		ReleaseSysCache(rtup);
+
+		if (!result && SUForUser_hook != NULL)
+		{
+			result = SUForUser_hook(roleid);
+		}
 	}
 	else
 	{

src/include/commands/event_trigger.h

@@ -16,6 +16,7 @@
 #include "catalog/dependency.h"
 #include "catalog/objectaddress.h"
 #include "catalog/pg_event_trigger.h"
+#include "fmgr.h"
 #include "nodes/parsenodes.h"
 #include "tcop/cmdtag.h"
 #include "tcop/deparse_utility.h"
@@ -47,6 +48,9 @@ typedef struct EventTriggerData
 #define CALLED_AS_EVENT_TRIGGER(fcinfo) \
 	((fcinfo)->context != NULL && IsA((fcinfo)->context, EventTriggerData))
 
+typedef void (*EventTrigger_hook_type)(FunctionCallInfo fcinfo);
+extern PGDLLEXPORT EventTrigger_hook_type EventTrigger_hook;
+
 extern Oid	CreateEventTrigger(CreateEventTrigStmt *stmt);
 extern Oid	get_event_trigger_oid(const char *trigname, bool missing_ok);
 

src/include/commands/trigger.h

@@ -151,6 +151,9 @@ extern PGDLLIMPORT int SessionReplicationRole;
 #define TRIGGER_FIRES_ON_REPLICA			'R'
 #define TRIGGER_DISABLED					'D'
 
+typedef Datum (*DataTrigger_hook_type)(FunctionCallInfo fcinfo);
+extern PGDLLEXPORT DataTrigger_hook_type DataTrigger_hook;
+
 extern ObjectAddress CreateTrigger(CreateTrigStmt *stmt, const char *queryString,
 								   Oid relOid, Oid refRelOid, Oid constraintOid, Oid indexOid,
 								   Oid funcoid, Oid parentTriggerOid, Node *whenClause,

src/include/miscadmin.h

@@ -387,9 +387,13 @@ extern const char *GetSystemUser(void);
 extern bool superuser(void);	/* current user is superuser */
 extern bool superuser_arg(Oid roleid);	/* given user is superuser */
 
+typedef bool (*SUForUser_hook_type) (Oid roleid);
+extern SUForUser_hook_type SUForUser_hook;
+
 /* in utils/adt/acl.c */
 extern PGDLLIMPORT char *privileged_role_name;
 extern bool is_privileged_role(void); /* current user is a privileged role */
+extern bool is_privileged_role_nosuper(void); /* current user is a privileged role */
 extern bool is_privileged_role_arg(Oid roleid); /* given user is a privileged role */
 
 /*****************************************************************************

src/include/replication/walreceiver.h

@@ -462,8 +462,6 @@ extern void WalReceiverMain(void) pg_attribute_noreturn();
 extern void ProcessWalRcvInterrupts(void);
 extern void WalRcvForceReply(void);
 
-extern const char *show_neon_storage_token(void);
-
 /* prototypes for functions in walreceiverfuncs.c */
 extern Size WalRcvShmemSize(void);
 extern void WalRcvShmemInit(void);

src/include/utils/acl.h

@@ -209,6 +209,7 @@ extern AclMode aclmask(const Acl *acl, Oid roleid, Oid ownerId,
 extern int	aclmembers(const Acl *acl, Oid **roleids);
 
 extern bool has_privs_of_role(Oid member, Oid role);
+extern bool has_privs_of_role_nosuper(Oid member, Oid role);
 extern bool member_can_set_role(Oid member, Oid role);
 extern void check_can_set_role(Oid member, Oid role);
 extern bool is_member_of_role(Oid member, Oid role);