Fix citus_add_node() to run SecLabel stmts, add tests&comments

src/backend/distributed/commands/role.c

@@ -23,6 +23,7 @@
 #include "catalog/pg_auth_members.h"
 #include "catalog/pg_authid.h"
 #include "catalog/pg_db_role_setting.h"
+#include "catalog/pg_shseclabel.h"
 #include "catalog/pg_type.h"
 #include "catalog/objectaddress.h"
 #include "commands/dbcommands.h"
@@ -65,6 +66,7 @@ static DefElem * makeDefElemBool(char *name, bool value);
 static List * GenerateRoleOptionsList(HeapTuple tuple);
 static List * GenerateGrantRoleStmtsFromOptions(RoleSpec *roleSpec, List *options);
 static List * GenerateGrantRoleStmtsOfRole(Oid roleid);
+static List * GenerateSecLabelOnRoleStmts(Oid roleid, char *rolename);
 static void EnsureSequentialModeForRoleDDL(void);
 
 static char * GetRoleNameFromDbRoleSetting(HeapTuple tuple,
@@ -515,13 +517,14 @@ GenerateCreateOrAlterRoleCommand(Oid roleOid)
 {
 	HeapTuple roleTuple = SearchSysCache1(AUTHOID, ObjectIdGetDatum(roleOid));
 	Form_pg_authid role = ((Form_pg_authid) GETSTRUCT(roleTuple));
+	char *rolename = pstrdup(NameStr(role->rolname));
 
 	CreateRoleStmt *createRoleStmt = NULL;
 	if (EnableCreateRolePropagation)
 	{
 		createRoleStmt = makeNode(CreateRoleStmt);
 		createRoleStmt->stmt_type = ROLESTMT_ROLE;
-		createRoleStmt->role = pstrdup(NameStr(role->rolname));
+		createRoleStmt->role = rolename;
 		createRoleStmt->options = GenerateRoleOptionsList(roleTuple);
 	}
 
@@ -532,7 +535,7 @@ GenerateCreateOrAlterRoleCommand(Oid roleOid)
 		alterRoleStmt->role = makeNode(RoleSpec);
 		alterRoleStmt->role->roletype = ROLESPEC_CSTRING;
 		alterRoleStmt->role->location = -1;
-		alterRoleStmt->role->rolename = pstrdup(NameStr(role->rolname));
+		alterRoleStmt->role->rolename = rolename;
 		alterRoleStmt->action = 1;
 		alterRoleStmt->options = GenerateRoleOptionsList(roleTuple);
 	}
@@ -544,7 +547,7 @@ GenerateCreateOrAlterRoleCommand(Oid roleOid)
 	{
 		/* add a worker_create_or_alter_role command if any of them are set */
 		char *createOrAlterRoleQuery = CreateCreateOrAlterRoleCommand(
-			pstrdup(NameStr(role->rolname)),
+			rolename,
 			createRoleStmt,
 			alterRoleStmt);
 
@@ -566,6 +569,14 @@ GenerateCreateOrAlterRoleCommand(Oid roleOid)
 		{
 			completeRoleList = lappend(completeRoleList, DeparseTreeNode(stmt));
 		}
+
+		/* append SECURITY LABEL ON ROLE commands fot this specific user */
+		List *secLabelOnRoleStmts = GenerateSecLabelOnRoleStmts(roleOid, rolename);
+		stmt = NULL;
+		foreach_ptr(stmt, secLabelOnRoleStmts)
+		{
+			completeRoleList = lappend(completeRoleList, DeparseTreeNode(stmt));
+		}
 	}
 
 	return completeRoleList;
@@ -895,6 +906,57 @@ GenerateGrantRoleStmtsOfRole(Oid roleid)
 }
 
 
+/*
+ * GenerateSecLabelOnRoleStmts generates the SecLabelStmts for the role
+ * whose oid is roleid.
+ */
+static List *
+GenerateSecLabelOnRoleStmts(Oid roleid, char *rolename)
+{
+	ScanKeyData skey[1];
+	HeapTuple tuple = NULL;
+	List *secLabelStmts = NIL;
+
+	/*
+	 * Note that roles are shared database objects, therefore their
+	 * security labels are stored in pg_shseclabel instead of pg_seclabel.
+	 */
+	Relation pg_shseclabel = table_open(SharedSecLabelRelationId, AccessShareLock);
+	ScanKeyInit(&skey[0], Anum_pg_shseclabel_objoid, BTEqualStrategyNumber, F_OIDEQ,
+				ObjectIdGetDatum(roleid));
+	SysScanDesc scan = systable_beginscan(pg_shseclabel, SharedSecLabelObjectIndexId,
+										  true, NULL, 1, &skey[0]);
+
+	while (HeapTupleIsValid(tuple = systable_getnext(scan)))
+	{
+		SecLabelStmt *secLabelStmt = makeNode(SecLabelStmt);
+		secLabelStmt->objtype = OBJECT_ROLE;
+		secLabelStmt->object = (Node *) makeString(rolename);
+
+		Datum datumArray[Natts_pg_shseclabel];
+		bool isNullArray[Natts_pg_shseclabel];
+
+		heap_deform_tuple(tuple, RelationGetDescr(pg_shseclabel), datumArray,
+						  isNullArray);
+
+		secLabelStmt->provider = TextDatumGetCString(
+			datumArray[Anum_pg_shseclabel_provider - 1]);
+		if (!isNullArray[Anum_pg_shseclabel_label - 1])
+		{
+			secLabelStmt->label = TextDatumGetCString(
+				datumArray[Anum_pg_shseclabel_label - 1]);
+		}
+
+		secLabelStmts = lappend(secLabelStmts, secLabelStmt);
+	}
+
+	systable_endscan(scan);
+	table_close(pg_shseclabel, AccessShareLock);
+
+	return secLabelStmts;
+}
+
+
 /*
  * PreprocessCreateRoleStmt creates a worker_create_or_alter_role query for the
  * role that is being created. With that query we can create the role in the

src/backend/distributed/commands/seclabel.c

@@ -23,6 +23,11 @@
 
 PG_FUNCTION_INFO_V1(citus_test_register_label_provider);
 
+
+/*
+ * citus_test_register_label_provider registers a dummy label provider
+ * named 'citus_tests_label_provider'. This is aimed to be used for testing.
+ */
 Datum
 citus_test_register_label_provider(PG_FUNCTION_ARGS)
 {
@@ -84,7 +89,9 @@ PreprocessSecLabelStmt(Node *node, const char *queryString,
 
 
 /*
- * PostprocessSecLabelStmt
+ * PostprocessSecLabelStmt ensures that all object dependencies exist on all
+ * nodes for the object in the SecLabelStmt. Currently, we only support SecLabelStmts
+ * operating on a ROLE object.
  */
 List *
 PostprocessSecLabelStmt(Node *node, const char *queryString)
@@ -112,17 +119,25 @@ PostprocessSecLabelStmt(Node *node, const char *queryString)
 
 
 /*
- * SecLabelStmtObjectAddress
+ * SecLabelStmtObjectAddress returns the object address of the object on
+ * which this statement operates (secLabelStmt->object). Note that it has no limitation
+ * on the object type being OBJECT_ROLE. This is intentionally implemented like this
+ * since it is fairly simple to implement and we might extend SECURITY LABEL propagation
+ * in the future to include more object types.
  */
 List *
 SecLabelStmtObjectAddress(Node *node, bool missing_ok, bool isPostprocess)
 {
 	SecLabelStmt *secLabelStmt = castNode(SecLabelStmt, node);
 
-	Relation rel = NULL; /* not used, but required to pass to get_object_address */
+	Relation rel = NULL;
 	ObjectAddress address = get_object_address(secLabelStmt->objtype,
 											   secLabelStmt->object, &rel,
 											   AccessShareLock, missing_ok);
+	if (rel != NULL)
+	{
+		relation_close(rel, AccessShareLock);
+	}
 
 	ObjectAddress *addressPtr = palloc0(sizeof(ObjectAddress));
 	*addressPtr = address;
@@ -131,7 +146,8 @@ SecLabelStmtObjectAddress(Node *node, bool missing_ok, bool isPostprocess)
 
 
 /*
- * citus_test_object_relabel
+ * citus_test_object_relabel is a dummy function for check_object_relabel_type hook.
+ * It is meant to be used in tests combined with citus_test_register_label_provider
  */
 void
 citus_test_object_relabel(const ObjectAddress *object, const char *seclabel)

src/backend/distributed/deparser/deparse_seclabel_stmts.c

@@ -57,6 +57,7 @@ AppendSecLabelStmt(StringInfo buf, SecLabelStmt *stmt)
 			break;
 		}
 
+		/* normally, we shouldn't reach this */
 		default:
 		{
 			ereport(ERROR, (errmsg("unsupported security label statement for"

src/test/regress/expected/seclabel.out

@@ -1,9 +1,17 @@
+--
+-- SECLABEL
+--
+-- Test suite for SECURITY LABEL ON ROLE statements
+--
+-- first we remove one of the worker nodes to be able to test
+-- citus_add_node later
 SELECT citus_remove_node('localhost', :worker_2_port);
  citus_remove_node
 ---------------------------------------------------------------------
 
 (1 row)
 
+-- now we register a label provider
 CREATE FUNCTION citus_test_register_label_provider()
   RETURNS void
   LANGUAGE C
@@ -15,6 +23,48 @@ SELECT citus_test_register_label_provider();
 (1 row)
 
 CREATE ROLE user1;
+-- check an invalid label for our current dummy hook citus_test_object_relabel
+SECURITY LABEL ON ROLE user1 IS 'invalid_label';
+ERROR:  'invalid_label' is not a valid security label for Citus tests.
+-- if we disable metadata_sync, the command will not be propagated
+SET citus.enable_metadata_sync TO off;
+SECURITY LABEL ON ROLE user1 IS 'citus_unclassified';
+SELECT objtype, objname, provider, label FROM pg_seclabels;
+ objtype | objname |          provider          |       label
+---------------------------------------------------------------------
+ role    | user1   | citus_tests_label_provider | citus_unclassified
+(1 row)
+
+\c - - - :worker_1_port
+SELECT objtype, objname, provider, label FROM pg_seclabels;
+ objtype | objname | provider | label
+---------------------------------------------------------------------
+(0 rows)
+
+\c - - - :master_port
+SELECT citus_test_register_label_provider();
+ citus_test_register_label_provider
+---------------------------------------------------------------------
+
+(1 row)
+
+-- check that we only support propagating for roles
+SET citus.shard_replication_factor to 1;
+CREATE TABLE a (a int);
+SELECT create_distributed_table('a', 'a');
+ create_distributed_table
+---------------------------------------------------------------------
+
+(1 row)
+
+SECURITY LABEL ON TABLE a IS 'citus_classified';
+NOTICE:  not propagating SECURITY LABEL commands whose object type is not role
+HINT:  Connect to worker nodes directly to manually run the same SECURITY LABEL command after disabling DDL propagation.
+DROP TABLE a;
+-- the registered label provider is per session only
+-- this means that we need to maintain the same connection to the worker node
+-- in order for the label provider to be visible there
+-- hence here we create the necessary session_level_connection_to_node functions
 SET citus.enable_metadata_sync TO off;
 CREATE OR REPLACE FUNCTION start_session_level_connection_to_node(text, integer)
         RETURNS void
@@ -42,13 +92,14 @@ CREATE OR REPLACE FUNCTION stop_session_level_connection_to_node()
         LANGUAGE C STRICT VOLATILE
         AS 'citus', $$stop_session_level_connection_to_node$$;
 RESET citus.enable_metadata_sync;
-ALTER SYSTEM SET citus.max_cached_conns_per_worker TO 2;
+-- now we establish a connection to the worker node
 SELECT start_session_level_connection_to_node('localhost', :worker_1_port);
  start_session_level_connection_to_node
 ---------------------------------------------------------------------
 
 (1 row)
 
+-- with that same connection, we register the label provider in the worker node
 SELECT run_commands_on_session_level_connection_to_node('SELECT citus_test_register_label_provider()');
  run_commands_on_session_level_connection_to_node
 ---------------------------------------------------------------------
@@ -56,16 +107,17 @@ SELECT run_commands_on_session_level_connection_to_node('SELECT citus_test_regis
 (1 row)
 
 SET citus.log_remote_commands TO on;
+SET citus.grep_remote_commands = '%SECURITY LABEL%';
+-- then we run a security label statement which will use the same connection to the worker node
+-- it should finish successfully
 SECURITY LABEL for citus_tests_label_provider ON ROLE user1 IS 'citus_classified';
-NOTICE:  issuing BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;SELECT assign_distributed_transaction_id(xx, xx, 'xxxxxxx');
-DETAIL:  on server postgres@localhost:xxxxx connectionId: xxxxxxx
-NOTICE:  issuing SET citus.enable_ddl_propagation TO 'off'
-DETAIL:  on server postgres@localhost:xxxxx connectionId: xxxxxxx
 NOTICE:  issuing SECURITY LABEL FOR citus_tests_label_provider ON ROLE user1 IS 'citus_classified'
 DETAIL:  on server postgres@localhost:xxxxx connectionId: xxxxxxx
-NOTICE:  issuing SET citus.enable_ddl_propagation TO 'on'
+SECURITY LABEL for citus_tests_label_provider ON ROLE user1 IS NULL;
+NOTICE:  issuing SECURITY LABEL FOR citus_tests_label_provider ON ROLE user1 IS NULL
 DETAIL:  on server postgres@localhost:xxxxx connectionId: xxxxxxx
-NOTICE:  issuing COMMIT
+SECURITY LABEL for citus_tests_label_provider ON ROLE user1 IS 'citus_unclassified';
+NOTICE:  issuing SECURITY LABEL FOR citus_tests_label_provider ON ROLE user1 IS 'citus_unclassified'
 DETAIL:  on server postgres@localhost:xxxxx connectionId: xxxxxxx
 RESET citus.log_remote_commands;
 SELECT stop_session_level_connection_to_node();
@@ -74,9 +126,35 @@ SELECT stop_session_level_connection_to_node();
 
 (1 row)
 
-ALTER SYSTEM RESET citus.max_cached_conns_per_worker;
+\c - - - :worker_1_port
+SELECT objtype, objname, provider, label FROM pg_seclabels;
+ objtype | objname |          provider          |       label
+---------------------------------------------------------------------
+ role    | user1   | citus_tests_label_provider | citus_unclassified
+(1 row)
+
+\c - - - :master_port
+-- adding a new node will fail because the label provider is not there
+-- however, this is enough for testing as we can see that the SECURITY LABEL commands
+-- will be propagated when adding a new node
+SET citus.log_remote_commands TO on;
+SET citus.grep_remote_commands = '%SECURITY LABEL%';
+SELECT 1 FROM citus_add_node('localhost', :worker_2_port);
+NOTICE:  issuing SELECT worker_create_or_alter_role('user1', 'CREATE ROLE user1 NOSUPERUSER NOCREATEDB NOCREATEROLE INHERIT NOLOGIN NOREPLICATION NOBYPASSRLS CONNECTION LIMIT -1 PASSWORD NULL VALID UNTIL ''infinity''', 'ALTER ROLE user1 NOSUPERUSER NOCREATEDB NOCREATEROLE INHERIT NOLOGIN NOREPLICATION NOBYPASSRLS CONNECTION LIMIT -1 PASSWORD NULL VALID UNTIL ''infinity''');SECURITY LABEL FOR citus_tests_label_provider ON ROLE user1 IS 'citus_unclassified'
+DETAIL:  on server postgres@localhost:xxxxx connectionId: xxxxxxx
+WARNING:  security label provider "citus_tests_label_provider" is not loaded
+CONTEXT:  while executing command on localhost:xxxxx
+ERROR:  failure on connection marked as essential: localhost:xxxxx
+-- cleanup
+RESET citus.log_remote_commands;
 DROP FUNCTION stop_session_level_connection_to_node, run_commands_on_session_level_connection_to_node,
               override_backend_data_gpid, start_session_level_connection_to_node;
+SELECT run_command_on_workers($$ DROP FUNCTION override_backend_data_gpid $$);
+       run_command_on_workers
+---------------------------------------------------------------------
+ (localhost,57637,t,"DROP FUNCTION")
+(1 row)
+
 DROP FUNCTION citus_test_register_label_provider;
 DROP ROLE user1;
 SELECT 1 FROM citus_add_node('localhost', :worker_2_port);

src/test/regress/sql/seclabel.sql

@@ -1,63 +1,105 @@
+--
+-- SECLABEL
+--
+-- Test suite for SECURITY LABEL ON ROLE statements
+--
+
+-- first we remove one of the worker nodes to be able to test
+-- citus_add_node later
 SELECT citus_remove_node('localhost', :worker_2_port);
 
+-- now we register a label provider
 CREATE FUNCTION citus_test_register_label_provider()
   RETURNS void
   LANGUAGE C
   AS 'citus', $$citus_test_register_label_provider$$;
-
 SELECT citus_test_register_label_provider();
+
 CREATE ROLE user1;
 
-SET citus.enable_metadata_sync TO off;
+-- check an invalid label for our current dummy hook citus_test_object_relabel
+SECURITY LABEL ON ROLE user1 IS 'invalid_label';
 
+-- if we disable metadata_sync, the command will not be propagated
+SET citus.enable_metadata_sync TO off;
+SECURITY LABEL ON ROLE user1 IS 'citus_unclassified';
+SELECT objtype, objname, provider, label FROM pg_seclabels;
+
+\c - - - :worker_1_port
+SELECT objtype, objname, provider, label FROM pg_seclabels;
+
+\c - - - :master_port
+SELECT citus_test_register_label_provider();
+
+-- check that we only support propagating for roles
+SET citus.shard_replication_factor to 1;
+CREATE TABLE a (a int);
+SELECT create_distributed_table('a', 'a');
+SECURITY LABEL ON TABLE a IS 'citus_classified';
+DROP TABLE a;
+
+-- the registered label provider is per session only
+-- this means that we need to maintain the same connection to the worker node
+-- in order for the label provider to be visible there
+-- hence here we create the necessary session_level_connection_to_node functions
+SET citus.enable_metadata_sync TO off;
 CREATE OR REPLACE FUNCTION start_session_level_connection_to_node(text, integer)
         RETURNS void
         LANGUAGE C STRICT VOLATILE
         AS 'citus', $$start_session_level_connection_to_node$$;
-
 CREATE OR REPLACE FUNCTION override_backend_data_gpid(bigint)
         RETURNS void
         LANGUAGE C STRICT IMMUTABLE
         AS 'citus', $$override_backend_data_gpid$$;
-
 SELECT run_command_on_workers($$SET citus.enable_metadata_sync TO off;CREATE OR REPLACE FUNCTION override_backend_data_gpid(bigint)
         RETURNS void
         LANGUAGE C STRICT IMMUTABLE
         AS 'citus'$$);
-
 CREATE OR REPLACE FUNCTION run_commands_on_session_level_connection_to_node(text)
         RETURNS void
         LANGUAGE C STRICT VOLATILE
         AS 'citus', $$run_commands_on_session_level_connection_to_node$$;
-
 CREATE OR REPLACE FUNCTION stop_session_level_connection_to_node()
         RETURNS void
         LANGUAGE C STRICT VOLATILE
         AS 'citus', $$stop_session_level_connection_to_node$$;
-
 RESET citus.enable_metadata_sync;
 
-ALTER SYSTEM SET citus.max_cached_conns_per_worker TO 2;
+-- now we establish a connection to the worker node
-
 SELECT start_session_level_connection_to_node('localhost', :worker_1_port);
 
+-- with that same connection, we register the label provider in the worker node
 SELECT run_commands_on_session_level_connection_to_node('SELECT citus_test_register_label_provider()');
 
 SET citus.log_remote_commands TO on;
+SET citus.grep_remote_commands = '%SECURITY LABEL%';
 
+-- then we run a security label statement which will use the same connection to the worker node
+-- it should finish successfully
 SECURITY LABEL for citus_tests_label_provider ON ROLE user1 IS 'citus_classified';
+SECURITY LABEL for citus_tests_label_provider ON ROLE user1 IS NULL;
+SECURITY LABEL for citus_tests_label_provider ON ROLE user1 IS 'citus_unclassified';
 
 RESET citus.log_remote_commands;
-
 SELECT stop_session_level_connection_to_node();
 
-ALTER SYSTEM RESET citus.max_cached_conns_per_worker;
+\c - - - :worker_1_port
+SELECT objtype, objname, provider, label FROM pg_seclabels;
 
+\c - - - :master_port
+
+-- adding a new node will fail because the label provider is not there
+-- however, this is enough for testing as we can see that the SECURITY LABEL commands
+-- will be propagated when adding a new node
+SET citus.log_remote_commands TO on;
+SET citus.grep_remote_commands = '%SECURITY LABEL%';
+SELECT 1 FROM citus_add_node('localhost', :worker_2_port);
+
+-- cleanup
+RESET citus.log_remote_commands;
 DROP FUNCTION stop_session_level_connection_to_node, run_commands_on_session_level_connection_to_node,
               override_backend_data_gpid, start_session_level_connection_to_node;
-
+SELECT run_command_on_workers($$ DROP FUNCTION override_backend_data_gpid $$);
 DROP FUNCTION citus_test_register_label_provider;
-
 DROP ROLE user1;
-
 SELECT 1 FROM citus_add_node('localhost', :worker_2_port);