v4-0001-Introduce-pg_signal_autovacuum-role-to-allow-non-.patch

From 591fffa188eaecf59280ff32fe377fc782dd6457 Mon Sep 17 00:00:00 2001
From: Anthony Leung <antholeu@amazon.com>
Date: Thu, 4 Apr 2024 20:24:33 +0000
Subject: [PATCH v4 1/2] Introduce pg_signal_autovacuum role to allow
 non-superuser to signal autovacuum worker

Non-superusers are currently not allowed to signal autovacuum
workers which can create friction in some scenarios when performing
maintenance. This commit introduce a new pre-defined role
'pg_signal_autovacuum'. Non-superusers granted with this role are
allowed to signal backends that are running autovacuum.
---
 doc/src/sgml/user-manag.sgml                |  4 +++
 src/backend/storage/ipc/signalfuncs.c       | 40 +++++++++++++++++----
 src/backend/utils/activity/backend_status.c | 19 ++++++++++
 src/include/catalog/pg_authid.dat           |  5 +++
 src/include/utils/backend_status.h          |  2 +-
 5 files changed, 63 insertions(+), 7 deletions(-)

diff --git a/doc/src/sgml/user-manag.sgml b/doc/src/sgml/user-manag.sgml
index 07a16247d7..2c5b569d4d 100644
--- a/doc/src/sgml/user-manag.sgml
+++ b/doc/src/sgml/user-manag.sgml
@@ -705,6 +705,10 @@ DROP ROLE doomed_role;
        database to issue
        <link linkend="sql-createsubscription"><command>CREATE SUBSCRIPTION</command></link>.</entry>
       </row>
+      <row>
+       <entry>pg_signal_autovacuum_worker</entry>
+       <entry>Allow signaling autovacuum worker backend to cancel or terminate</entry>
+      </row>
      </tbody>
     </tgroup>
    </table>
diff --git a/src/backend/storage/ipc/signalfuncs.c b/src/backend/storage/ipc/signalfuncs.c
index 88e9bf8125..12d93ba83f 100644
--- a/src/backend/storage/ipc/signalfuncs.c
+++ b/src/backend/storage/ipc/signalfuncs.c
@@ -45,6 +45,7 @@
 #define SIGNAL_BACKEND_ERROR 1
 #define SIGNAL_BACKEND_NOPERMISSION 2
 #define SIGNAL_BACKEND_NOSUPERUSER 3
+#define SIGNAL_BACKEND_NOAUTOVACUUM 4
 static int
 pg_signal_backend(int pid, int sig)
 {
@@ -74,19 +75,32 @@ pg_signal_backend(int pid, int sig)
 		return SIGNAL_BACKEND_ERROR;
 	}
 
+	/*
+	 * If the backend is autovacuum worker, allow user with the privileges of
+	 * pg_signal_autovacuum_worker role to signal the backend.
+	 */
+	if ((!OidIsValid(proc->roleId) || superuser_arg(proc->roleId))
+		&& pgstat_get_backend_type(GetNumberFromPGProc(proc)) == B_AUTOVAC_WORKER)
+	{
+		if (!has_privs_of_role(GetUserId(), ROLE_PG_SIGNAL_AUTOVACUUM_WORKER))
+			return SIGNAL_BACKEND_NOAUTOVACUUM;
+	}
 	/*
 	 * Only allow superusers to signal superuser-owned backends.  Any process
 	 * not advertising a role might have the importance of a superuser-owned
 	 * backend, so treat it that way.
-	 */
-	if ((!OidIsValid(proc->roleId) || superuser_arg(proc->roleId)) &&
-		!superuser())
+	*/
+	else if ((!OidIsValid(proc->roleId) || superuser_arg(proc->roleId)) &&
+			!superuser())
+	{
 		return SIGNAL_BACKEND_NOSUPERUSER;
-
+	}
 	/* Users can signal backends they have role membership in. */
-	if (!has_privs_of_role(GetUserId(), proc->roleId) &&
-		!has_privs_of_role(GetUserId(), ROLE_PG_SIGNAL_BACKEND))
+	else if (!has_privs_of_role(GetUserId(), proc->roleId) &&
+			!has_privs_of_role(GetUserId(), ROLE_PG_SIGNAL_BACKEND))
+	{
 		return SIGNAL_BACKEND_NOPERMISSION;
+	}
 
 	/*
 	 * Can the process we just validated above end, followed by the pid being
@@ -137,6 +151,13 @@ pg_cancel_backend(PG_FUNCTION_ARGS)
 				 errdetail("Only roles with privileges of the role whose query is being canceled or with privileges of the \"%s\" role may cancel this query.",
 						   "pg_signal_backend")));
 
+	if (r == SIGNAL_BACKEND_NOAUTOVACUUM)
+		ereport(ERROR,
+				(errcode(ERRCODE_INSUFFICIENT_PRIVILEGE),
+				 errmsg("permission denied to cancel autovacuum worker backend"),
+				 errdetail("Only roles with the %s attribute or with privileges of the \"%s\" role may cancel autovacuum worker backend",
+						   "SUPERUSER", "pg_signal_autovacuum_worker")));
+
 	PG_RETURN_BOOL(r == SIGNAL_BACKEND_SUCCESS);
 }
 
@@ -243,6 +264,13 @@ pg_terminate_backend(PG_FUNCTION_ARGS)
 				 errdetail("Only roles with privileges of the role whose process is being terminated or with privileges of the \"%s\" role may terminate this process.",
 						   "pg_signal_backend")));
 
+	if (r == SIGNAL_BACKEND_NOAUTOVACUUM)
+		ereport(ERROR,
+				(errcode(ERRCODE_INSUFFICIENT_PRIVILEGE),
+				 errmsg("permission denied to terminate autovacuum worker backend"),
+				 errdetail("Only roles with the %s attribute or with privileges of the \"%s\" role may terminate autovacuum worker backend",
+						   "SUPERUSER", "pg_signal_autovacuum_worker")));
+
 	/* Wait only on success and if actually requested */
 	if (r == SIGNAL_BACKEND_SUCCESS && timeout > 0)
 		PG_RETURN_BOOL(pg_wait_until_termination(pid, timeout));
diff --git a/src/backend/utils/activity/backend_status.c b/src/backend/utils/activity/backend_status.c
index 1ccf4c6d83..a6c7899318 100644
--- a/src/backend/utils/activity/backend_status.c
+++ b/src/backend/utils/activity/backend_status.c
@@ -1038,6 +1038,25 @@ pgstat_get_my_query_id(void)
 	return MyBEEntry->st_query_id;
 }
 
+/* ----------
+ * pgstat_get_backend_type() -
+ *
+ * Return the backend type of the backend for the given proc number.
+ * ----------
+ */
+BackendType
+pgstat_get_backend_type(ProcNumber procNumber)
+{
+	PgBackendStatus *ret;
+
+	ret = pgstat_get_beentry_by_proc_number(procNumber);
+
+	if (!ret)
+		return B_INVALID;
+
+	return ret->st_backendType;
+}
+
 /* ----------
  * cmp_lbestatus
  *
diff --git a/src/include/catalog/pg_authid.dat b/src/include/catalog/pg_authid.dat
index 55cabdda6f..6156fc55dd 100644
--- a/src/include/catalog/pg_authid.dat
+++ b/src/include/catalog/pg_authid.dat
@@ -99,5 +99,10 @@
   rolcreaterole => 'f', rolcreatedb => 'f', rolcanlogin => 'f',
   rolreplication => 'f', rolbypassrls => 'f', rolconnlimit => '-1',
   rolpassword => '_null_', rolvaliduntil => '_null_' },
+{ oid => '8916', oid_symbol => 'ROLE_PG_SIGNAL_AUTOVACUUM_WORKER',
+  rolname => 'pg_signal_autovacuum_worker', rolsuper => 'f', rolinherit => 't',
+  rolcreaterole => 'f', rolcreatedb => 'f', rolcanlogin => 'f',
+  rolreplication => 'f', rolbypassrls => 'f', rolconnlimit => '-1',
+  rolpassword => '_null_', rolvaliduntil => '_null_' },
 
 ]
diff --git a/src/include/utils/backend_status.h b/src/include/utils/backend_status.h
index 7b7f6f59d0..7fd9cd7167 100644
--- a/src/include/utils/backend_status.h
+++ b/src/include/utils/backend_status.h
@@ -323,7 +323,7 @@ extern const char *pgstat_get_backend_current_activity(int pid, bool checkUser);
 extern const char *pgstat_get_crashed_backend_activity(int pid, char *buffer,
 													   int buflen);
 extern uint64 pgstat_get_my_query_id(void);
-
+extern BackendType pgstat_get_backend_type(ProcNumber procNumber);
 
 /* ----------
  * Support functions for the SQL-callable functions to
-- 
2.25.1