From 83345fd5405039ed28b728ebe7f05ab84b9562f8 Mon Sep 17 00:00:00 2001
From: Anthony Leung <antholeu@amazon.com>
Date: Thu, 4 Apr 2024 23:33:13 +0000
Subject: [PATCH] Add tap test for pg_signal_autovacuum role

Add a tap test to validate the following:

1. Normal user cannot signal autovacuum worker backend
2. User with pg_signal_backend role cannot signal autovacuum worker
   backend
3. User with pg_signal_autovacuum role can signal autovacuum worker
   backend
---
 src/backend/postmaster/autovacuum.c         |  3 +
 src/test/signals/.gitignore                 |  2 +
 src/test/signals/Makefile                   | 27 +++++++
 src/test/signals/meson.build                | 15 ++++
 src/test/signals/t/001_signal_autovacuum.pl | 87 +++++++++++++++++++++
 5 files changed, 134 insertions(+)
 create mode 100644 src/test/signals/.gitignore
 create mode 100644 src/test/signals/Makefile
 create mode 100644 src/test/signals/meson.build
 create mode 100644 src/test/signals/t/001_signal_autovacuum.pl

diff --git a/src/backend/postmaster/autovacuum.c b/src/backend/postmaster/autovacuum.c
index c367ede6f8..241d8712c4 100644
--- a/src/backend/postmaster/autovacuum.c
+++ b/src/backend/postmaster/autovacuum.c
@@ -100,6 +100,7 @@
 #include "utils/fmgroids.h"
 #include "utils/fmgrprotos.h"
 #include "utils/guc_hooks.h"
+#include "utils/injection_point.h"
 #include "utils/lsyscache.h"
 #include "utils/memutils.h"
 #include "utils/ps_status.h"
@@ -1889,6 +1890,8 @@ do_autovacuum(void)
 	bool		found_concurrent_worker = false;
 	int			i;
 
+	INJECTION_POINT("autovacuum-start");
+
 	/*
 	 * StartTransactionCommand and CommitTransactionCommand will automatically
 	 * switch to other contexts.  We need this one to keep the list of
diff --git a/src/test/signals/.gitignore b/src/test/signals/.gitignore
new file mode 100644
index 0000000000..871e943d50
--- /dev/null
+++ b/src/test/signals/.gitignore
@@ -0,0 +1,2 @@
+# Generated by test suite
+/tmp_check/
diff --git a/src/test/signals/Makefile b/src/test/signals/Makefile
new file mode 100644
index 0000000000..18d8d2d0af
--- /dev/null
+++ b/src/test/signals/Makefile
@@ -0,0 +1,27 @@
+#-------------------------------------------------------------------------
+#
+# Makefile for src/test/recovery
+#
+# Portions Copyright (c) 1996-2024, PostgreSQL Global Development Group
+# Portions Copyright (c) 1994, Regents of the University of California
+#
+# src/test/recovery/Makefile
+#
+#-------------------------------------------------------------------------
+
+EXTRA_INSTALL=src/test/modules/injection_points
+
+subdir = src/test/signals
+top_builddir = ../../..
+include $(top_builddir)/src/Makefile.global
+
+export enable_injection_points enable_injection_points
+
+check:
+	$(prove_check)
+
+installcheck:
+	$(prove_installcheck)
+
+clean distclean maintainer-clean:
+	rm -rf tmp_check
diff --git a/src/test/signals/meson.build b/src/test/signals/meson.build
new file mode 100644
index 0000000000..3280d79ebc
--- /dev/null
+++ b/src/test/signals/meson.build
@@ -0,0 +1,15 @@
+# Copyright (c) 2022-2024, PostgreSQL Global Development Group
+
+tests += {
+  'name': 'signals',
+  'sd': meson.current_source_dir(),
+  'bd': meson.current_build_dir(),
+  'tap': {
+    'env': {
+       'enable_injection_points': get_option('injection_points') ? 'yes' : 'no',
+    },
+    'tests': [
+      't/001_signal_autovacuum.pl',
+    ],
+  },
+}
diff --git a/src/test/signals/t/001_signal_autovacuum.pl b/src/test/signals/t/001_signal_autovacuum.pl
new file mode 100644
index 0000000000..a219b3e2a1
--- /dev/null
+++ b/src/test/signals/t/001_signal_autovacuum.pl
@@ -0,0 +1,87 @@
+# Copyright (c) 2024-2024, PostgreSQL Global Development Group
+
+# Test signaling for pg_signal_autovacuum role.
+
+use strict;
+use warnings;
+use PostgreSQL::Test::Cluster;
+use PostgreSQL::Test::Utils;
+use Test::More;
+
+if ($ENV{enable_injection_points} ne 'yes')
+{
+	plan skip_all => 'Injection points not supported by this build';
+}
+
+# Initialize postgres
+my $psql_err = '';
+my $psql_out = '';
+my $node = PostgreSQL::Test::Cluster->new('node');
+$node->init();
+$node->append_conf(
+	'postgresql.conf', 'autovacuum_naptime = 1
+');
+$node->start;
+$node->safe_psql('postgres', 'CREATE EXTENSION injection_points;');
+
+# Regular user and user with pg_signal_backend role cannot signal autovacuum worker
+# User with pg_signal_backend can signal autovacuum worker 
+$node->safe_psql('postgres', qq(
+    CREATE ROLE regular_role;
+    CREATE ROLE signal_backend_role;
+    CREATE ROLE signal_autovacuum_role;
+    GRANT pg_signal_backend TO signal_backend_role;
+    GRANT pg_signal_autovacuum TO signal_autovacuum_role;
+));
+
+# From this point, autovacuum worker will wait before doing any vacuum.
+$node->safe_psql('postgres',
+	"SELECT injection_points_attach('autovacuum-start', 'wait');");
+
+# Create some content and set autovacuum setting such that it would be triggered.
+$node->safe_psql('postgres', qq(
+    CREATE TABLE tab_int(i int);
+    ALTER TABLE tab_int SET (autovacuum_vacuum_cost_limit = 1);
+    ALTER TABLE tab_int SET (autovacuum_vacuum_cost_delay = 100);
+));
+
+$node->safe_psql('postgres', qq(
+    INSERT INTO tab_int SELECT * FROM generate_series(1, 1000000);
+));
+
+# Wait until the autovacuum worker starts
+$node->wait_for_event('autovacuum worker', 'autovacuum-start');
+
+my $av_pid = $node->safe_psql('postgres', qq(
+    SELECT pid FROM pg_stat_activity WHERE backend_type = 'autovacuum worker';
+));
+
+# Regular user cannot terminate autovacuum worker
+my $terminate_with_no_pg_signal_av = $node->psql('postgres', qq( 
+    SET ROLE regular_role;
+    SELECT pg_terminate_backend($av_pid);
+), stdout => \$psql_out, stderr => \$psql_err);
+
+ok($terminate_with_no_pg_signal_av != 0, "Terminating autovacuum worker should not succeed without pg_signal_autovacuum role");
+like($psql_err, qr/ERROR:  permission denied to terminate autovacuum worker backend\nDETAIL:  Only roles with the SUPERUSER attribute or with privileges of the "pg_signal_autovacuum" role may terminate autovacuum worker backend/,
+    "passcheck errors gracefully when passcheck database is invalid");
+
+# User with signal_backend_role cannot terminate autovacuum worker
+my $terminate_with_signal_backend_role = $node->psql('postgres', qq( 
+    SET ROLE signal_backend_role;
+    SELECT pg_terminate_backend($av_pid);
+), stdout => \$psql_out, stderr => \$psql_err);
+
+ok($terminate_with_no_pg_signal_av != 0, "Terminating autovacuum worker should not succeed without pg_signal_autovacuum role");
+like($psql_err, qr/ERROR:  permission denied to terminate autovacuum worker backend\nDETAIL:  Only roles with the SUPERUSER attribute or with privileges of the "pg_signal_autovacuum" role may terminate autovacuum worker backend/,
+    "passcheck errors gracefully when passcheck database is invalid");
+
+# User with pg_signal_backend can terminate autovacuum worker 
+my $terminate_with_pg_signal_av = $node->psql('postgres', qq( 
+    SET ROLE signal_autovacuum_role;
+    SELECT pg_terminate_backend($av_pid);
+), stdout => \$psql_out, stderr => \$psql_err);
+
+ok($terminate_with_pg_signal_av == 0, "Terminating autovacuum worker should succeed with pg_signal_autovacuum role");
+
+done_testing();
-- 
2.40.1

