tmp DO NOT MERGE

supabase · Feb 7, 2024 · 66a6eb8 · 66a6eb8
1 parent c9cf5bb
commit 66a6eb8
Show file tree

Hide file tree

Showing 5 changed files with 169 additions and 30 deletions.
diff --git a/src/supautils.c b/src/supautils.c
@@ -283,6 +283,8 @@ supautils_hook(PROCESS_UTILITY_PARAMS)
 			{
 				AlterRoleStmt *stmt = (AlterRoleStmt *)utility_stmt;
 				DefElem *bypassrls_option = NULL;
+				DefElem *replication_option = NULL;
+				List *deferred_options = NIL;
 				ListCell *option_cell = NULL;
 
 				if (!IsTransactionState()) {
@@ -298,7 +300,8 @@ supautils_hook(PROCESS_UTILITY_PARAMS)
 					break;
 				}
 
-				/* Check to see if there are any descriptions related to bypassrls. */
+				/* Check to see if there are any descriptions related to
+				 * bypassrls or replication. */
 				foreach(option_cell, stmt->options)
 				{
 					DefElem *defel = (DefElem *) lfirst(option_cell);
@@ -311,17 +314,32 @@ supautils_hook(PROCESS_UTILITY_PARAMS)
 						}
 						bypassrls_option = defel;
 					}
+
+					if (strcmp(defel->defname, "isreplication") == 0) {
+						if (replication_option != NULL) {
+							ereport(ERROR,
+									(errcode(ERRCODE_SYNTAX_ERROR),
+									 errmsg("conflicting or redundant options")));
+						}
+						replication_option = defel;
+					}
 				}
 
-				// Defer setting bypassrls attribute if using `privileged_role`.
+				// Defer setting bypassrls or replication attributes if using
+				// `privileged_role`.
 				if (bypassrls_option != NULL) {
 					stmt->options = list_delete_ptr(stmt->options, bypassrls_option);
+					deferred_options = lappend(deferred_options, bypassrls_option);
+				}
+				if (replication_option != NULL) {
+					stmt->options = list_delete_ptr(stmt->options, replication_option);
+					deferred_options = lappend(deferred_options, replication_option);
 				}
 
 				run_process_utility_hook(prev_hook);
 
-				if (bypassrls_option != NULL) {
-					alter_role_with_bypassrls_option_as_superuser(stmt->role->rolename, bypassrls_option, privileged_extensions_superuser);
+				if (deferred_options != NIL) {
+					alter_role_with_options_as_superuser(stmt->role->rolename, deferred_options, privileged_extensions_superuser);
 				}
 
 				return;
@@ -387,8 +405,10 @@ supautils_hook(PROCESS_UTILITY_PARAMS)
 					const char *created_role = stmt->role;
 					List *addroleto = NIL;	/* roles to make this a member of */
 					bool hasrolemembers = false;	/* has roles to be members of this role */
+					List *deferred_options = NIL;
 					bool stmt_has_bypassrls = false;
-					bool bypassrls_is_allowed = true;
+					bool stmt_has_replication = false;
+					bool role_is_privileged = is_privileged_role();
 					ListCell *option_cell;
 
 					/* if role already exists, bypass the hook to let it fail with the usual error */
@@ -398,11 +418,6 @@ supautils_hook(PROCESS_UTILITY_PARAMS)
 					/* CREATE ROLE <reserved_role> */
 					confirm_reserved_roles(created_role, false);
 
-					// Allow bypassrls attribute if using `privileged_role`.
-					if (!is_privileged_role()){
-						bypassrls_is_allowed = false;
-					}
-
 					/* Check to see if there are any descriptions related to membership and bypassrls. */
 					foreach(option_cell, stmt->options)
 					{
@@ -414,15 +429,20 @@ supautils_hook(PROCESS_UTILITY_PARAMS)
 							strcmp(defel->defname, "adminmembers") == 0)
 							hasrolemembers = true;
 
-						// Defer setting bypassrls attribute if using `privileged_role`.
+						// Defer setting bypassrls & replication attributes if
+						// using `privileged_role`.
 						//
 						// Duplicate/conflicting attributes will be caught by
 						// the standard process utility hook, so we can assume
 						// there's at most one bypassrls DefElem.
-						if (bypassrls_is_allowed && strcmp(defel->defname, "bypassrls") == 0) {
+						if (role_is_privileged && strcmp(defel->defname, "bypassrls") == 0) {
 							stmt_has_bypassrls = intVal(defel->arg) != 0;
 							intVal(defel->arg) = 0;
 						}
+						if (role_is_privileged && strcmp(defel->defname, "isreplication") == 0) {
+							stmt_has_replication = intVal(defel->arg) != 0;
+							intVal(defel->arg) = 0;
+						}
 					}
 
 					/* CREATE ROLE <any_role> IN ROLE/GROUP <role_with_reserved_membership> */
@@ -445,22 +465,32 @@ supautils_hook(PROCESS_UTILITY_PARAMS)
 					if (hasrolemembers)
 						confirm_reserved_memberships(created_role);
 
-					// CREATE ROLE <any_role> BYPASSRLS
+					// CREATE ROLE <any_role> < BYPASSRLS | REPLICATION >
 					//
 					// Allow `privileged_role` (in addition to superusers) to
-					// set bypassrls attribute. The setting of the attribute is
-					// deferred in the original CREATE ROLE - the actual setting
-					// is done here.
-					if (bypassrls_is_allowed && stmt_has_bypassrls) {
+					// set bypassrls and replication attributes. The setting of
+					// the attributes is deferred in the original CREATE ROLE -
+					// the actual setting is done here.
+					if (role_is_privileged) {
 						Node *true_node = (Node *)makeInteger(true);
 						DefElem *bypassrls_option = makeDefElem("bypassrls", true_node, -1);
+						DefElem *replication_option = makeDefElem("isreplication", true_node, -1);
+
+						if (stmt_has_bypassrls) {
+							deferred_options = lappend(deferred_options, bypassrls_option);
+						}
+						if (stmt_has_replication) {
+							deferred_options = lappend(deferred_options, replication_option);
+						}
 
 						run_process_utility_hook(prev_hook);
 
-						alter_role_with_bypassrls_option_as_superuser(stmt->role, bypassrls_option, privileged_extensions_superuser);
+						alter_role_with_options_as_superuser(stmt->role, deferred_options, privileged_extensions_superuser);
 
 						pfree(true_node);
 						pfree(bypassrls_option);
+						pfree(replication_option);
+						list_free(deferred_options);
 
 						return;
 					}

diff --git a/src/utils.c b/src/utils.c
@@ -12,9 +12,8 @@ static bool is_switched_to_superuser = false;
 
 static bool strstarts(const char *, const char *);
 
-void alter_role_with_bypassrls_option_as_superuser(const char *role_name,
-                                                   DefElem *bypassrls_option,
-                                                   const char *superuser_name) {
+void alter_role_with_options_as_superuser(const char *role_name, List *options,
+                                          const char *superuser_name) {
     RoleSpec *role = makeNode(RoleSpec);
     AlterRoleStmt *bypassrls_stmt = makeNode(AlterRoleStmt);
     bool already_switched_to_superuser = false;
@@ -24,7 +23,7 @@ void alter_role_with_bypassrls_option_as_superuser(const char *role_name,
     role->location = -1;
 
     bypassrls_stmt->role = role;
-    bypassrls_stmt->options = list_make1(bypassrls_option);
+    bypassrls_stmt->options = options;
 
     switch_to_superuser(superuser_name, &already_switched_to_superuser);
 
@@ -40,7 +39,6 @@ void alter_role_with_bypassrls_option_as_superuser(const char *role_name,
 
     pfree(role->rolename);
     pfree(role);
-    list_free(bypassrls_stmt->options);
     pfree(bypassrls_stmt);
 
     return;

diff --git a/src/utils.h b/src/utils.h
@@ -62,10 +62,9 @@
 #define SUPAUTILS_GUC_CONTEXT_SIGHUP PGC_SIGHUP
 #endif
 
-extern void
-alter_role_with_bypassrls_option_as_superuser(const char *role_name,
-                                              DefElem *bypassrls_option,
-                                              const char *superuser_name);
+extern void alter_role_with_options_as_superuser(const char *role_name,
+                                                 List *options,
+                                                 const char *superuser_name);
 
 /**
  * Switch to a superuser and save the original role. Caller is responsible for

diff --git a/test/expected/privileged_role.out b/test/expected/privileged_role.out
@@ -64,6 +64,33 @@ select rolbypassrls from pg_roles where rolname = 'r';
 drop role r;
 \echo
 
+-- can manage replication role attribute
+create role r replication;
+select rolreplication from pg_roles where rolname = 'r';
+ rolreplication 
+----------------
+ t
+(1 row)
+
+alter role r noreplication;
+ERROR:  must be superuser to alter replication users
+select rolreplication from pg_roles where rolname = 'r';
+ rolreplication 
+----------------
+ t
+(1 row)
+
+alter role r replication;
+ERROR:  must be superuser to alter replication users
+select rolreplication from pg_roles where rolname = 'r';
+ rolreplication 
+----------------
+ t
+(1 row)
+
+drop role r;
+\echo
+
 -- can manage foreign data wrappers
 create extension postgres_fdw;
 create foreign data wrapper new_fdw
@@ -75,8 +102,6 @@ NOTICE:  drop cascades to foreign-data wrapper new_fdw
 
 -- non-superuser non-privileged role cannot manage bypassrls role attribute
 set role rolecreator;
-\echo
-
 -- the error message changed in PG14
 do $$
 begin
@@ -93,6 +118,24 @@ drop role r;
 set role privileged_role;
 \echo
 
+-- non-superuser non-privileged role cannot manage replication role attribute
+set role rolecreator;
+-- the error message changed in PG14
+do $$
+begin
+  create role r replication;
+  exception when insufficient_privilege then null;
+end;
+$$ language plpgsql;
+create role r;
+alter role r noreplication;
+ERROR:  must be superuser to alter replication users
+alter role r replication;
+ERROR:  must be superuser to alter replication users
+drop role r;
+set role privileged_role;
+\echo
+
 -- superuser can manage bypassrls role attribute
 set role postgres;
 create role r bypassrls;
@@ -120,6 +163,33 @@ drop role r;
 set role privileged_role;
 \echo
 
+-- superuser can manage replication role attribute
+set role postgres;
+create role r replication;
+select rolreplication from pg_roles where rolname = 'r';
+ rolreplication 
+----------------
+ t
+(1 row)
+
+alter role r noreplication;
+select rolreplication from pg_roles where rolname = 'r';
+ rolreplication 
+----------------
+ f
+(1 row)
+
+alter role r replication;
+select rolreplication from pg_roles where rolname = 'r';
+ rolreplication 
+----------------
+ t
+(1 row)
+
+drop role r;
+set role privileged_role;
+\echo
+
 -- regression: https://github.com/supabase/supautils/issues/34
 create role tmp;
 alter role tmp;

diff --git a/test/sql/privileged_role.sql b/test/sql/privileged_role.sql
@@ -50,6 +50,16 @@ select rolbypassrls from pg_roles where rolname = 'r';
 drop role r;
 \echo
 
+-- can manage replication role attribute
+create role r replication;
+select rolreplication from pg_roles where rolname = 'r';
+alter role r noreplication;
+select rolreplication from pg_roles where rolname = 'r';
+alter role r replication;
+select rolreplication from pg_roles where rolname = 'r';
+drop role r;
+\echo
+
 -- can manage foreign data wrappers
 create extension postgres_fdw;
 create foreign data wrapper new_fdw
@@ -61,7 +71,6 @@ drop extension postgres_fdw cascade;
 
 -- non-superuser non-privileged role cannot manage bypassrls role attribute
 set role rolecreator;
-\echo
 
 -- the error message changed in PG14
 do $$
@@ -79,6 +88,25 @@ drop role r;
 set role privileged_role;
 \echo
 
+-- non-superuser non-privileged role cannot manage replication role attribute
+set role rolecreator;
+
+-- the error message changed in PG14
+do $$
+begin
+  create role r replication;
+  exception when insufficient_privilege then null;
+end;
+$$ language plpgsql;
+
+create role r;
+alter role r noreplication;
+alter role r replication;
+drop role r;
+
+set role privileged_role;
+\echo
+
 -- superuser can manage bypassrls role attribute
 set role postgres;
 
@@ -93,6 +121,20 @@ drop role r;
 set role privileged_role;
 \echo
 
+-- superuser can manage replication role attribute
+set role postgres;
+
+create role r replication;
+select rolreplication from pg_roles where rolname = 'r';
+alter role r noreplication;
+select rolreplication from pg_roles where rolname = 'r';
+alter role r replication;
+select rolreplication from pg_roles where rolname = 'r';
+drop role r;
+
+set role privileged_role;
+\echo
+
 -- regression: https://github.com/supabase/supautils/issues/34
 create role tmp;
 alter role tmp;