From 4972f29cef6dfbb948e843517ce5ff413628c2f1 Mon Sep 17 00:00:00 2001
From: Nathan Bossart <nathan@postgresql.org>
Date: Sat, 8 Jul 2023 21:35:03 -0700
Subject: [PATCH v4 1/1] move session auth permission check

---
 src/backend/commands/variable.c   | 23 +++++++++++++++++++++++
 src/backend/utils/init/miscinit.c | 29 ++++++++++-------------------
 src/include/miscadmin.h           |  1 +
 3 files changed, 34 insertions(+), 19 deletions(-)

diff --git a/src/backend/commands/variable.c b/src/backend/commands/variable.c
index f0f2e07655..f8e308f1d0 100644
--- a/src/backend/commands/variable.c
+++ b/src/backend/commands/variable.c
@@ -846,6 +846,29 @@ check_session_authorization(char **newval, void **extra, GucSource source)
 
 	ReleaseSysCache(roleTup);
 
+	/*
+	 * Only a superuser may set auth ID to something other than himself.  Note
+	 * that in case of multiple SETs in a single session, the original
+	 * userid's superuserness is what matters.  But we set the GUC variable
+	 * is_superuser to indicate whether the *current* session userid is a
+	 * superuser.
+	 */
+	if (roleid != GetAuthenticatedUserId() &&
+		!GetAuthenticatedUserIsSuperuser())
+	{
+		if (source == PGC_S_TEST)
+		{
+			ereport(NOTICE,
+					(errcode(ERRCODE_INSUFFICIENT_PRIVILEGE),
+					 errmsg("permission will be denied to set session authorization \"%s\"",
+							*newval)));
+			return true;
+		}
+		GUC_check_errcode(ERRCODE_INSUFFICIENT_PRIVILEGE);
+		GUC_check_errmsg("permission denied to set session authorization");
+		return false;
+	}
+
 	/* Set up "extra" struct for assign_session_authorization to use */
 	myextra = (role_auth_extra *) guc_malloc(LOG, sizeof(role_auth_extra));
 	if (!myextra)
diff --git a/src/backend/utils/init/miscinit.c b/src/backend/utils/init/miscinit.c
index a604432126..f5548a0f47 100644
--- a/src/backend/utils/init/miscinit.c
+++ b/src/backend/utils/init/miscinit.c
@@ -582,6 +582,16 @@ GetAuthenticatedUserId(void)
 	return AuthenticatedUserId;
 }
 
+/*
+ * Return whether the authenticated user was superuser at connection start.
+ */
+bool
+GetAuthenticatedUserIsSuperuser(void)
+{
+	Assert(OidIsValid(AuthenticatedUserId));
+	return AuthenticatedUserIsSuperuser;
+}
+
 
 /*
  * GetUserIdAndSecContext/SetUserIdAndSecContext - get/set the current user ID
@@ -888,29 +898,10 @@ system_user(PG_FUNCTION_ARGS)
 
 /*
  * Change session auth ID while running
- *
- * Only a superuser may set auth ID to something other than himself.  Note
- * that in case of multiple SETs in a single session, the original userid's
- * superuserness is what matters.  But we set the GUC variable is_superuser
- * to indicate whether the *current* session userid is a superuser.
- *
- * Note: this is not an especially clean place to do the permission check.
- * It's OK because the check does not require catalog access and can't
- * fail during an end-of-transaction GUC reversion, but we may someday
- * have to push it up into assign_session_authorization.
  */
 void
 SetSessionAuthorization(Oid userid, bool is_superuser)
 {
-	/* Must have authenticated already, else can't make permission check */
-	Assert(OidIsValid(AuthenticatedUserId));
-
-	if (userid != AuthenticatedUserId &&
-		!AuthenticatedUserIsSuperuser)
-		ereport(ERROR,
-				(errcode(ERRCODE_INSUFFICIENT_PRIVILEGE),
-				 errmsg("permission denied to set session authorization")));
-
 	SetSessionUserId(userid, is_superuser);
 
 	SetConfigOption("is_superuser",
diff --git a/src/include/miscadmin.h b/src/include/miscadmin.h
index 14bd574fc2..11d6e6869d 100644
--- a/src/include/miscadmin.h
+++ b/src/include/miscadmin.h
@@ -357,6 +357,7 @@ extern Oid	GetUserId(void);
 extern Oid	GetOuterUserId(void);
 extern Oid	GetSessionUserId(void);
 extern Oid	GetAuthenticatedUserId(void);
+extern bool GetAuthenticatedUserIsSuperuser(void);
 extern void GetUserIdAndSecContext(Oid *userid, int *sec_context);
 extern void SetUserIdAndSecContext(Oid userid, int sec_context);
 extern bool InLocalUserIdChange(void);
-- 
2.25.1