From 2fb2ae74f24cd14f3f08f612f1869517b145c61c Mon Sep 17 00:00:00 2001
From: Michael Paquier <michael@paquier.xyz>
Date: Wed, 4 Oct 2023 15:06:25 +0900
Subject: [PATCH v3 2/2] Allow background workers to bypass login check

This adds a new flag value (BGWORKER_BYPASS_ROLELOGINCHECK) to the flags
being used in BackgroundWorkerInitializeConnection() and
BackgroundWorkerInitializeConnectionByOid().  This flag allows the
background workers to bypass the login check done in
InitializeSessionUserId().
---
 src/include/miscadmin.h                       |  4 +-
 src/include/postmaster/bgworker.h             | 10 ++--
 src/backend/bootstrap/bootstrap.c             |  2 +-
 src/backend/postmaster/autovacuum.c           |  5 +-
 src/backend/postmaster/postmaster.c           |  6 ++-
 src/backend/tcop/postgres.c                   |  1 +
 src/backend/utils/init/miscinit.c             |  4 +-
 src/backend/utils/init/postinit.c             |  5 +-
 .../modules/worker_spi/t/001_worker_spi.pl    | 38 ++++++++++++++
 .../modules/worker_spi/worker_spi--1.0.sql    |  3 +-
 src/test/modules/worker_spi/worker_spi.c      | 49 +++++++++++++++++--
 doc/src/sgml/bgworker.sgml                    |  3 +-
 12 files changed, 109 insertions(+), 21 deletions(-)

diff --git a/src/include/miscadmin.h b/src/include/miscadmin.h
index 14bd574fc2..f1d03f2df8 100644
--- a/src/include/miscadmin.h
+++ b/src/include/miscadmin.h
@@ -364,7 +364,8 @@ extern bool InSecurityRestrictedOperation(void);
 extern bool InNoForceRLSOperation(void);
 extern void GetUserIdAndContext(Oid *userid, bool *sec_def_context);
 extern void SetUserIdAndContext(Oid userid, bool sec_def_context);
-extern void InitializeSessionUserId(const char *rolename, Oid roleid);
+extern void InitializeSessionUserId(const char *rolename, Oid roleid,
+									bool bypass_login_check);
 extern void InitializeSessionUserIdStandalone(void);
 extern void SetSessionAuthorization(Oid userid, bool is_superuser);
 extern Oid	GetCurrentRoleId(void);
@@ -469,6 +470,7 @@ extern void InitPostgres(const char *in_dbname, Oid dboid,
 						 const char *username, Oid useroid,
 						 bool load_session_libraries,
 						 bool override_allow_connections,
+						 bool bypass_login_check,
 						 char *out_dbname);
 extern void BaseInit(void);
 
diff --git a/src/include/postmaster/bgworker.h b/src/include/postmaster/bgworker.h
index 7815507e3d..0fdbfaf431 100644
--- a/src/include/postmaster/bgworker.h
+++ b/src/include/postmaster/bgworker.h
@@ -141,18 +141,20 @@ extern PGDLLIMPORT BackgroundWorker *MyBgworkerEntry;
  * If dbname is NULL, connection is made to no specific database;
  * only shared catalogs can be accessed.
  */
-extern void BackgroundWorkerInitializeConnection(const char *dbname, const char *username, uint32 flags);
+extern void BackgroundWorkerInitializeConnection(const char *dbname, const char *username, bits32 flags);
 
 /* Just like the above, but specifying database and user by OID. */
-extern void BackgroundWorkerInitializeConnectionByOid(Oid dboid, Oid useroid, uint32 flags);
+extern void BackgroundWorkerInitializeConnectionByOid(Oid dboid, Oid useroid, bits32 flags);
 
 /*
  * Flags to BackgroundWorkerInitializeConnection et al
  *
  *
- * Allow bypassing datallowconn restrictions when connecting to database
+ * Allow bypassing datallowconn restrictions and login check when connecting
+ * to database
  */
-#define BGWORKER_BYPASS_ALLOWCONN 1
+#define BGWORKER_BYPASS_ALLOWCONN 0x0001
+#define BGWORKER_BYPASS_ROLELOGINCHECK 0x0002
 
 
 /* Block/unblock signals in a background worker process */
diff --git a/src/backend/bootstrap/bootstrap.c b/src/backend/bootstrap/bootstrap.c
index 5810f8825e..bdd005b66a 100644
--- a/src/backend/bootstrap/bootstrap.c
+++ b/src/backend/bootstrap/bootstrap.c
@@ -345,7 +345,7 @@ BootstrapModeMain(int argc, char *argv[], bool check_only)
 	if (pg_link_canary_is_frontend())
 		elog(ERROR, "backend is incorrectly linked to frontend functions");
 
-	InitPostgres(NULL, InvalidOid, NULL, InvalidOid, false, false, NULL);
+	InitPostgres(NULL, InvalidOid, NULL, InvalidOid, false, false, false, NULL);
 
 	/* Initialize stuff for bootstrap-file processing */
 	for (i = 0; i < MAXATTR; i++)
diff --git a/src/backend/postmaster/autovacuum.c b/src/backend/postmaster/autovacuum.c
index ae9be9b911..4686b81f68 100644
--- a/src/backend/postmaster/autovacuum.c
+++ b/src/backend/postmaster/autovacuum.c
@@ -488,7 +488,7 @@ AutoVacLauncherMain(int argc, char *argv[])
 	/* Early initialization */
 	BaseInit();
 
-	InitPostgres(NULL, InvalidOid, NULL, InvalidOid, false, false, NULL);
+	InitPostgres(NULL, InvalidOid, NULL, InvalidOid, false, false, false, NULL);
 
 	SetProcessingMode(NormalProcessing);
 
@@ -1706,8 +1706,7 @@ AutoVacWorkerMain(int argc, char *argv[])
 		 * Note: if we have selected a just-deleted database (due to using
 		 * stale stats info), we'll fail and exit here.
 		 */
-		InitPostgres(NULL, dbid, NULL, InvalidOid, false, false,
-					 dbname);
+		InitPostgres(NULL, dbid, NULL, InvalidOid, false, false, false, dbname);
 		SetProcessingMode(NormalProcessing);
 		set_ps_display(dbname);
 		ereport(DEBUG1,
diff --git a/src/backend/postmaster/postmaster.c b/src/backend/postmaster/postmaster.c
index 54e9bfb8c4..37f5b94469 100644
--- a/src/backend/postmaster/postmaster.c
+++ b/src/backend/postmaster/postmaster.c
@@ -5575,7 +5575,7 @@ MaxLivePostmasterChildren(void)
  * Connect background worker to a database.
  */
 void
-BackgroundWorkerInitializeConnection(const char *dbname, const char *username, uint32 flags)
+BackgroundWorkerInitializeConnection(const char *dbname, const char *username, bits32 flags)
 {
 	BackgroundWorker *worker = MyBgworkerEntry;
 
@@ -5589,6 +5589,7 @@ BackgroundWorkerInitializeConnection(const char *dbname, const char *username, u
 				 username, InvalidOid,	/* role to connect as */
 				 false,			/* never honor session_preload_libraries */
 				 (flags & BGWORKER_BYPASS_ALLOWCONN) != 0,	/* ignore datallowconn? */
+				 (flags & BGWORKER_BYPASS_ROLELOGINCHECK) != 0, /* bypass login check? */
 				 NULL);			/* no out_dbname */
 
 	/* it had better not gotten out of "init" mode yet */
@@ -5602,7 +5603,7 @@ BackgroundWorkerInitializeConnection(const char *dbname, const char *username, u
  * Connect background worker to a database using OIDs.
  */
 void
-BackgroundWorkerInitializeConnectionByOid(Oid dboid, Oid useroid, uint32 flags)
+BackgroundWorkerInitializeConnectionByOid(Oid dboid, Oid useroid, bits32 flags)
 {
 	BackgroundWorker *worker = MyBgworkerEntry;
 
@@ -5616,6 +5617,7 @@ BackgroundWorkerInitializeConnectionByOid(Oid dboid, Oid useroid, uint32 flags)
 				 NULL, useroid, /* role to connect as */
 				 false,			/* never honor session_preload_libraries */
 				 (flags & BGWORKER_BYPASS_ALLOWCONN) != 0,	/* ignore datallowconn? */
+				 (flags & BGWORKER_BYPASS_ROLELOGINCHECK) != 0, /* bypass login check? */
 				 NULL);			/* no out_dbname */
 
 	/* it had better not gotten out of "init" mode yet */
diff --git a/src/backend/tcop/postgres.c b/src/backend/tcop/postgres.c
index 21b9763183..5c1e21b99d 100644
--- a/src/backend/tcop/postgres.c
+++ b/src/backend/tcop/postgres.c
@@ -4211,6 +4211,7 @@ PostgresMain(const char *dbname, const char *username)
 				 username, InvalidOid,	/* role to connect as */
 				 !am_walsender, /* honor session_preload_libraries? */
 				 false,			/* don't ignore datallowconn */
+				 false,			/* don't bypass login check */
 				 NULL);			/* no out_dbname */
 
 	/*
diff --git a/src/backend/utils/init/miscinit.c b/src/backend/utils/init/miscinit.c
index 1e671c560c..182d666852 100644
--- a/src/backend/utils/init/miscinit.c
+++ b/src/backend/utils/init/miscinit.c
@@ -725,7 +725,7 @@ has_rolreplication(Oid roleid)
  * Initialize user identity during normal backend startup
  */
 void
-InitializeSessionUserId(const char *rolename, Oid roleid)
+InitializeSessionUserId(const char *rolename, Oid roleid, bool bypass_login_check)
 {
 	HeapTuple	roleTup;
 	Form_pg_authid rform;
@@ -789,7 +789,7 @@ InitializeSessionUserId(const char *rolename, Oid roleid)
 		/*
 		 * Is role allowed to login at all?
 		 */
-		if (!rform->rolcanlogin)
+		if (!bypass_login_check && !rform->rolcanlogin)
 			ereport(FATAL,
 					(errcode(ERRCODE_INVALID_AUTHORIZATION_SPECIFICATION),
 					 errmsg("role \"%s\" is not permitted to log in",
diff --git a/src/backend/utils/init/postinit.c b/src/backend/utils/init/postinit.c
index df4d15a50f..dcfeb30832 100644
--- a/src/backend/utils/init/postinit.c
+++ b/src/backend/utils/init/postinit.c
@@ -719,6 +719,7 @@ InitPostgres(const char *in_dbname, Oid dboid,
 			 const char *username, Oid useroid,
 			 bool load_session_libraries,
 			 bool override_allow_connections,
+			 bool bypass_login_check,
 			 char *out_dbname)
 {
 	bool		bootstrap = IsBootstrapProcessingMode();
@@ -901,7 +902,7 @@ InitPostgres(const char *in_dbname, Oid dboid,
 		}
 		else
 		{
-			InitializeSessionUserId(username, useroid);
+			InitializeSessionUserId(username, useroid, bypass_login_check);
 			am_superuser = superuser();
 		}
 	}
@@ -910,7 +911,7 @@ InitPostgres(const char *in_dbname, Oid dboid,
 		/* normal multiuser case */
 		Assert(MyProcPort != NULL);
 		PerformAuthentication(MyProcPort);
-		InitializeSessionUserId(username, useroid);
+		InitializeSessionUserId(username, useroid, false);
 		/* ensure that auth_method is actually valid, aka authn_id is not NULL */
 		if (MyClientConnectionInfo.authn_id)
 			InitializeSystemUser(MyClientConnectionInfo.authn_id,
diff --git a/src/test/modules/worker_spi/t/001_worker_spi.pl b/src/test/modules/worker_spi/t/001_worker_spi.pl
index a026042c65..81f47e09d9 100644
--- a/src/test/modules/worker_spi/t/001_worker_spi.pl
+++ b/src/test/modules/worker_spi/t/001_worker_spi.pl
@@ -98,4 +98,42 @@ ok( $node->poll_query_until(
 	'dynamic bgworkers all launched'
 ) or die "Timed out while waiting for dynamic bgworkers to be launched";
 
+# Check BGWORKER_BYPASS_ROLELOGINCHECK.
+# First create a role without login access.
+$node->safe_psql(
+	'postgres', qq[
+  CREATE ROLE nologrole with nologin;
+  ALTER ROLE nologrole with superuser;
+]);
+my $nologrole_id = $node->safe_psql('mydb', "SELECT oid FROM pg_roles where rolname = 'nologrole';");
+
+my $log_start = -s $node->logfile;
+
+# Ask the background workers to connect with this role, failing.
+$node->safe_psql('mydb', "SELECT worker_spi_launch(12, $mydb_id, $nologrole_id);");
+
+# An error message should be issued.
+ok( $node->log_contains(
+		"role \"nologrole\" is not permitted to log in", $log_start),
+	"nologrole not allowed to connect if ROLELOGINCHECK is not set"
+);
+
+$log_start = -s $node->logfile;
+
+my $worker3_pid = $node->safe_psql('mydb', qq(SELECT worker_spi_launch(13, $mydb_id, $nologrole_id, '{"ROLELOGINCHECK"}');));
+
+ok( $node->poll_query_until(
+		'mydb',
+		qq[SELECT datname, usename, wait_event FROM pg_stat_activity
+            WHERE backend_type = 'worker_spi dynamic' AND
+            pid = $worker3_pid;],
+		'mydb|nologrole|worker_spi_main'),
+	'dynamic bgworkers '
+) or die "Timed out while waiting for dynamic bgworkers to be launched";
+
+# An error message should not be issued.
+ok( !$node->log_contains(
+		"role \"nologrole\" is not permitted to log in", $log_start),
+	"nologrole allowed to connect if ROLELOGINCHECK is set");
+
 done_testing();
diff --git a/src/test/modules/worker_spi/worker_spi--1.0.sql b/src/test/modules/worker_spi/worker_spi--1.0.sql
index f5b91128ff..59c0c686f3 100644
--- a/src/test/modules/worker_spi/worker_spi--1.0.sql
+++ b/src/test/modules/worker_spi/worker_spi--1.0.sql
@@ -3,7 +3,8 @@
 -- complain if script is sourced in psql, rather than via CREATE EXTENSION
 \echo Use "CREATE EXTENSION worker_spi" to load this file. \quit
 
-CREATE FUNCTION worker_spi_launch(pg_catalog.int4, pg_catalog.oid, pg_catalog.oid)
+CREATE FUNCTION worker_spi_launch(pg_catalog.int4, pg_catalog.oid,
+				  pg_catalog.oid, pg_catalog.text[] DEFAULT '{}')
 RETURNS pg_catalog.int4 STRICT
 AS 'MODULE_PATHNAME'
 LANGUAGE C;
diff --git a/src/test/modules/worker_spi/worker_spi.c b/src/test/modules/worker_spi/worker_spi.c
index f35897ddac..07d8bab721 100644
--- a/src/test/modules/worker_spi/worker_spi.c
+++ b/src/test/modules/worker_spi/worker_spi.c
@@ -144,6 +144,7 @@ worker_spi_main(Datum main_arg)
 	Oid			dboid;
 	Oid			roleoid;
 	char	   *p;
+	bits32		flags = 0;
 
 	table = palloc(sizeof(worktable));
 	sprintf(name, "schema%d", index);
@@ -155,6 +156,8 @@ worker_spi_main(Datum main_arg)
 	memcpy(&dboid, p, sizeof(Oid));
 	p += sizeof(Oid);
 	memcpy(&roleoid, p, sizeof(Oid));
+	p += sizeof(bits32);
+	memcpy(&flags, p, sizeof(bits32));
 
 	/* Establish signal handlers before unblocking signals. */
 	pqsignal(SIGHUP, SignalHandlerForConfigReload);
@@ -165,9 +168,10 @@ worker_spi_main(Datum main_arg)
 
 	/* Connect to our database */
 	if (OidIsValid(dboid))
-		BackgroundWorkerInitializeConnectionByOid(dboid, roleoid, 0);
+		BackgroundWorkerInitializeConnectionByOid(dboid, roleoid, flags);
 	else
-		BackgroundWorkerInitializeConnection(worker_spi_database, worker_spi_role, 0);
+		BackgroundWorkerInitializeConnection(worker_spi_database,
+											 worker_spi_role, flags);
 
 	elog(LOG, "%s initialized with %s.%s",
 		 MyBgworkerEntry->bgw_name, table->schema, table->name);
@@ -371,8 +375,9 @@ _PG_init(void)
 	/*
 	 * Now fill in worker-specific data, and do the actual registrations.
 	 *
-	 * bgw_extra can optionally include a dabatase OID and a role OID.  This
-	 * is left empty here to fallback to the related GUCs at startup.
+	 * bgw_extra can optionally include a dabatase OID, a role OID and a set
+	 * of flags.  This is left empty here to fallback to the related GUCs at
+	 * startup (or 0 for the flags).
 	 */
 	for (int i = 1; i <= worker_spi_total_workers; i++)
 	{
@@ -398,6 +403,11 @@ worker_spi_launch(PG_FUNCTION_ARGS)
 	BgwHandleStatus status;
 	pid_t		pid;
 	char	   *p;
+	bits32		flags = 0;
+	ArrayType  *arr = PG_GETARG_ARRAYTYPE_P(3);
+	Size		ndim;
+	int			nelems;
+	Datum	   *datum_flags;
 
 	memset(&worker, 0, sizeof(worker));
 	worker.bgw_flags = BGWORKER_SHMEM_ACCESS |
@@ -412,6 +422,35 @@ worker_spi_launch(PG_FUNCTION_ARGS)
 	/* set bgw_notify_pid so that we can use WaitForBackgroundWorkerStartup */
 	worker.bgw_notify_pid = MyProcPid;
 
+	/* extract flags, if any */
+	ndim = ARR_NDIM(arr);
+	if (ndim > 1)
+		ereport(ERROR,
+				(errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
+				 errmsg("flags array must be one-dimensional")));
+
+	if (array_contains_nulls(arr))
+		ereport(ERROR,
+				(errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
+				 errmsg("flags array must not contain nulls")));
+
+	Assert(ARR_ELEMTYPE(arr) == TEXTOID);
+	deconstruct_array_builtin(arr, TEXTOID, &datum_flags, NULL, &nelems);
+
+	for (i = 0; i < nelems; i++)
+	{
+		char	   *optname = TextDatumGetCString(datum_flags[i]);
+
+		if (strcmp(optname, "ROLELOGINCHECK") == 0)
+			flags |= BGWORKER_BYPASS_ROLELOGINCHECK;
+		else if (strcmp(optname, "ALLOWCONN") == 0)
+			flags |= BGWORKER_BYPASS_ALLOWCONN;
+		else
+			ereport(ERROR,
+					(errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
+					 errmsg("incorrect flag value found in array")));
+	}
+
 	/*
 	 * Register database and role to use for the worker started in bgw_extra.
 	 * If none have been provided, fallback to the GUCs.
@@ -433,6 +472,8 @@ worker_spi_launch(PG_FUNCTION_ARGS)
 	memcpy(p, &dboid, sizeof(Oid));
 	p += sizeof(Oid);
 	memcpy(p, &roleoid, sizeof(Oid));
+	p += sizeof(bits32);
+	memcpy(p, &flags, sizeof(bits32));
 
 	if (!RegisterDynamicBackgroundWorker(&worker, &handle))
 		PG_RETURN_NULL();
diff --git a/doc/src/sgml/bgworker.sgml b/doc/src/sgml/bgworker.sgml
index 9ad1146ba0..7335c02ed8 100644
--- a/doc/src/sgml/bgworker.sgml
+++ b/doc/src/sgml/bgworker.sgml
@@ -200,7 +200,8 @@ typedef struct BackgroundWorker
    <literal>InvalidOid</literal>, the process will run as the superuser created
    during <command>initdb</command>. If <literal>BGWORKER_BYPASS_ALLOWCONN</literal>
    is specified as <varname>flags</varname> it is possible to bypass the restriction
-   to connect to databases not allowing user connections.
+   to connect to databases not allowing user connections. If <literal>BGWORKER_BYPASS_ROLELOGINCHECK</literal>
+   is specified as <varname>flags</varname> it is possible to bypass the login check to connect to databases.
    A background worker can only call one of these two functions, and only
    once.  It is not possible to switch databases.
   </para>
-- 
2.42.0

