Make sure that distributed functions always have the correct user

Objectives:

(a) both super user and regular user should have the correct owner for the function on the worker
(b) The transactional semantics would work fine for both super user and regular user
(c) non-super-user and non-function owner would get a reasonable error message if tries to distribute the function

Co-authored-by: @serprex

src/backend/distributed/commands/function.c

@@ -33,6 +33,7 @@
 #include "distributed/commands/utility_hook.h"
 #include "distributed/deparser.h"
 #include "distributed/maintenanced.h"
+#include "distributed/master_metadata_utility.h"
 #include "distributed/master_protocol.h"
 #include "distributed/metadata/distobject.h"
 #include "distributed/metadata/pg_dist_object.h"
@@ -54,6 +55,7 @@
 
 /* forward declaration for helper functions*/
 static char * GetFunctionDDLCommand(const RegProcedure funcOid);
+static char * GetFunctionAlterOwnerCommand(const RegProcedure funcOid);
 static int GetDistributionArgIndex(Oid functionOid, char *distributionArgumentName,
 								   Oid *distributionArgumentOid);
 static int GetFunctionColocationId(Oid functionOid, char *colocateWithName, Oid
@@ -144,6 +146,7 @@ create_distributed_function(PG_FUNCTION_ARGS)
 		colocateWithTableName = text_to_cstring(colocateWithText);
 	}
 
+	EnsureFunctionOwner(funcOid);
 
 	ObjectAddressSet(functionAddress, ProcedureRelationId, funcOid);
 	ErrorIfFunctionDependsOnExtension(&functionAddress);
@@ -157,7 +160,8 @@ create_distributed_function(PG_FUNCTION_ARGS)
 	EnsureDependenciesExistsOnAllNodes(&functionAddress);
 
 	ddlCommand = GetFunctionDDLCommand(funcOid);
-	SendCommandToWorkers(ALL_WORKERS, ddlCommand);
+
+	SendCommandToWorkersAsUser(ALL_WORKERS, CurrentUserName(), ddlCommand);
 
 	MarkObjectDistributed(&functionAddress);
 
@@ -530,14 +534,17 @@ UpdateFunctionDistributionInfo(const ObjectAddress *distAddress,
 
 /*
  * GetFunctionDDLCommand returns the complete "CREATE OR REPLACE FUNCTION ..." statement for
- * the specified function.
+ * the specified function followed by "ALTER FUNCTION .. SET OWNER ..".
  */
 static char *
 GetFunctionDDLCommand(const RegProcedure funcOid)
 {
+	StringInfo ddlCommand = makeStringInfo();
+
 	OverrideSearchPath *overridePath = NULL;
 	Datum sqlTextDatum = 0;
-	char *sql = NULL;
+	char *createFunctionSQL = NULL;
+	char *alterFunctionOwnerSQL = NULL;
 
 	/*
 	 * Set search_path to NIL so that all objects outside of pg_catalog will be
@@ -555,8 +562,84 @@ GetFunctionDDLCommand(const RegProcedure funcOid)
 	/* revert back to original search_path */
 	PopOverrideSearchPath();
 
-	sql = TextDatumGetCString(sqlTextDatum);
-	return sql;
+	createFunctionSQL = TextDatumGetCString(sqlTextDatum);
+	alterFunctionOwnerSQL = GetFunctionAlterOwnerCommand(funcOid);
+
+	appendStringInfo(ddlCommand, "%s;%s", createFunctionSQL, alterFunctionOwnerSQL);
+
+	return ddlCommand->data;
+}
+
+
+/*
+ * GetFunctionAlterOwnerCommand returns "ALTER FUNCTION .. SET OWNER .." statement for
+ * the specified function.
+ */
+static char *
+GetFunctionAlterOwnerCommand(const RegProcedure funcOid)
+{
+	HeapTuple proctup = SearchSysCache1(PROCOID, ObjectIdGetDatum(funcOid));
+	StringInfo alterCommand = makeStringInfo();
+	bool isProcedure = false;
+	Oid procOwner = InvalidOid;
+
+	char *functionSignature = NULL;
+	char *functionOwner = NULL;
+
+	OverrideSearchPath *overridePath = NULL;
+	Datum functionSignatureDatum = 0;
+
+	if (HeapTupleIsValid(proctup))
+	{
+		Form_pg_proc procform;
+
+		procform = (Form_pg_proc) GETSTRUCT(proctup);
+
+		procOwner = procform->proowner;
+
+#if (PG_VERSION_NUM >= 110000)
+		isProcedure = procform->prokind == PROKIND_PROCEDURE;
+#endif
+
+		ReleaseSysCache(proctup);
+	}
+	else if (!OidIsValid(funcOid) || !HeapTupleIsValid(proctup))
+	{
+		ereport(ERROR, (errmsg("cannot find function with oid: %d", funcOid)));
+	}
+
+	/*
+	 * Set search_path to NIL so that all objects outside of pg_catalog will be
+	 * schema-prefixed. pg_catalog will be added automatically when we call
+	 * PushOverrideSearchPath(), since we set addCatalog to true;
+	 */
+	overridePath = GetOverrideSearchPath(CurrentMemoryContext);
+	overridePath->schemas = NIL;
+	overridePath->addCatalog = true;
+
+	PushOverrideSearchPath(overridePath);
+
+	/*
+	 * If the function exists we want to use pg_get_function_identity_arguments to
+	 * serialize its canonical arguments
+	 */
+	functionSignatureDatum =
+		DirectFunctionCall1(regprocedureout, ObjectIdGetDatum(funcOid));
+
+	/* revert back to original search_path */
+	PopOverrideSearchPath();
+
+	/* regprocedureout returns cstring */
+	functionSignature = DatumGetCString(functionSignatureDatum);
+
+	functionOwner = GetUserNameFromId(procOwner, false);
+
+	appendStringInfo(alterCommand, "ALTER %s %s OWNER TO %s;",
+					 (isProcedure ? "PROCEDURE" : "FUNCTION"),
+					 functionSignature,
+					 quote_identifier(functionOwner));
+
+	return alterCommand->data;
 }
 
 

src/backend/distributed/master/master_metadata_utility.c

@@ -1412,6 +1412,22 @@ EnsureSequenceOwner(Oid sequenceOid)
 }
 
 
+/*
+ * Check that the current user has owner rights to functionId, error out if
+ * not. Superusers are regarded as owners. Functions and procedures are
+ * treated equally.
+ */
+void
+EnsureFunctionOwner(Oid functionId)
+{
+	if (!pg_proc_ownercheck(functionId, GetUserId()))
+	{
+		aclcheck_error(ACLCHECK_NOT_OWNER, OBJECT_FUNCTION,
+					   get_func_name(functionId));
+	}
+}
+
+
 /*
  * EnsureSuperUser check that the current user is a superuser and errors out if not.
  */

src/backend/distributed/transaction/worker_transaction.c

@@ -46,7 +46,30 @@ SendCommandToWorker(char *nodeName, int32 nodePort, const char *command)
 
 
 /*
- * SendCommandToWorkerAsUSer sends a command to a particular worker as a particular user
+ * SendCommandToWorkersAsUser sends a command to targetWorkerSet as a particular user
+ * as part of the 2PC.
+ */
+void
+SendCommandToWorkersAsUser(TargetWorkerSet targetWorkerSet, const char *nodeUser,
+						   const char *command)
+{
+	List *workerNodeList = TargetWorkerSetNodeList(targetWorkerSet, ShareLock);
+	ListCell *workerNodeCell = NULL;
+
+	/* run commands serially */
+	foreach(workerNodeCell, workerNodeList)
+	{
+		WorkerNode *workerNode = (WorkerNode *) lfirst(workerNodeCell);
+		char *nodeName = workerNode->workerName;
+		int nodePort = workerNode->workerPort;
+
+		SendCommandToWorkerAsUser(nodeName, nodePort, nodeUser, command);
+	}
+}
+
+
+/*
+ * SendCommandToWorkerAsUser sends a command to a particular worker as a particular user
  * as part of the 2PC.
  */
 void

src/include/distributed/master_metadata_utility.h

@@ -143,6 +143,7 @@ extern void EnsureTablePermissions(Oid relationId, AclMode mode);
 extern void EnsureTableOwner(Oid relationId);
 extern void EnsureSchemaOwner(Oid schemaId);
 extern void EnsureSequenceOwner(Oid sequenceOid);
+extern void EnsureFunctionOwner(Oid functionId);
 extern void EnsureSuperUser(void);
 extern void EnsureReplicationSettings(Oid relationId, char replicationModel);
 extern bool RegularTable(Oid relationId);

src/include/distributed/worker_transaction.h

@@ -32,6 +32,9 @@ typedef enum TargetWorkerSet
 extern List * GetWorkerTransactions(void);
 extern List * TargetWorkerSetNodeList(TargetWorkerSet targetWorkerSet, LOCKMODE lockMode);
 extern void SendCommandToWorker(char *nodeName, int32 nodePort, const char *command);
+extern void SendCommandToWorkersAsUser(TargetWorkerSet targetWorkerSet, const
+									   char *nodeUser,
+									   const char *command);
 extern void SendCommandToWorkerAsUser(char *nodeName, int32 nodePort,
 									  const char *nodeUser, const char *command);
 extern void SendCommandToWorkers(TargetWorkerSet targetWorkerSet, const char *command);

src/test/regress/expected/multi_multiuser.out

@@ -14,6 +14,18 @@ SELECT substring(:'server_version', '\d+')::int > 10 AS version_above_ten;
 SET citus.next_shard_id TO 1420000;
 ALTER SEQUENCE pg_catalog.pg_dist_jobid_seq RESTART 1420000;
 SET citus.shard_replication_factor TO 1;
+ALTER SYSTEM SET citus.metadata_sync_interval TO 3000;
+ALTER SYSTEM SET citus.metadata_sync_retry_interval TO 500;
+SELECT pg_reload_conf();
+ pg_reload_conf 
+----------------
+ t
+(1 row)
+
+CREATE FUNCTION wait_until_metadata_sync(timeout INTEGER DEFAULT 15000)
+    RETURNS void
+    LANGUAGE C STRICT
+    AS 'citus';
 CREATE TABLE test (id integer, val integer);
 SELECT create_distributed_table('test', 'id');
  create_distributed_table 
@@ -428,7 +440,117 @@ INSERT INTO full_access_user_schema.t1 VALUES (1),(2),(3);
 -- not allowed to create a table
 SELECT create_distributed_table('full_access_user_schema.t1', 'id');
 ERROR:  permission denied for schema full_access_user_schema
-CONTEXT:  while executing command on localhost:57638
+CONTEXT:  while executing command on localhost:57637
+RESET ROLE;
+SET ROLE usage_access;
+CREATE TYPE usage_access_type AS ENUM ('a', 'b');
+CREATE FUNCTION usage_access_func(x usage_access_type, variadic v int[]) RETURNS int[]
+    LANGUAGE plpgsql AS 'begin return v; end;';
+SET ROLE no_access;
+SELECT create_distributed_function('usage_access_func(usage_access_type,int[])');
+ERROR:  must be owner of function usage_access_func
+SET ROLE usage_access;
+SELECT create_distributed_function('usage_access_func(usage_access_type,int[])');
+ create_distributed_function 
+-----------------------------
+ 
+(1 row)
+
+SELECT typowner::regrole FROM pg_type WHERE typname = 'usage_access_type';
+   typowner   
+--------------
+ usage_access
+(1 row)
+
+SELECT proowner::regrole FROM pg_proc WHERE proname = 'usage_access_func';
+   proowner   
+--------------
+ usage_access
+(1 row)
+
+SELECT run_command_on_workers($$SELECT typowner::regrole FROM pg_type WHERE typname = 'usage_access_type'$$);
+      run_command_on_workers      
+----------------------------------
+ (localhost,57637,t,usage_access)
+ (localhost,57638,t,usage_access)
+(2 rows)
+
+SELECT run_command_on_workers($$SELECT proowner::regrole FROM pg_proc WHERE proname = 'usage_access_func'$$);
+      run_command_on_workers      
+----------------------------------
+ (localhost,57637,t,usage_access)
+ (localhost,57638,t,usage_access)
+(2 rows)
+
+SELECT wait_until_metadata_sync();
+ wait_until_metadata_sync 
+--------------------------
+ 
+(1 row)
+
+-- now, make sure that the user can use the function 
+-- created in the transaction
+BEGIN;
+CREATE FUNCTION usage_access_func_second(key int, variadic v int[]) RETURNS text
+    LANGUAGE plpgsql AS 'begin return current_user; end;';
+SELECT create_distributed_function('usage_access_func_second(int,int[])', '$1');
+ create_distributed_function 
+-----------------------------
+ 
+(1 row)
+
+SELECT usage_access_func_second(1, 2,3,4,5) FROM full_access_user_schema.t1 LIMIT 1; 
+ usage_access_func_second 
+--------------------------
+ usage_access
+(1 row)
+
+ROLLBACK;
+CREATE FUNCTION usage_access_func_third(key int, variadic v int[]) RETURNS text
+    LANGUAGE plpgsql AS 'begin return current_user; end;';
+-- connect back as super user
+\c - - - :master_port
+-- show that the current user is a super user
+SELECT usesuper FROM pg_user where usename IN (SELECT current_user);
+ usesuper 
+----------
+ t
+(1 row)
+
+-- superuser creates the distributed function that is owned by a regular user
+SELECT create_distributed_function('usage_access_func_third(int,int[])', '$1');
+ create_distributed_function 
+-----------------------------
+ 
+(1 row)
+
+SELECT proowner::regrole FROM pg_proc WHERE proname = 'usage_access_func_third';
+   proowner   
+--------------
+ usage_access
+(1 row)
+
+SELECT run_command_on_workers($$SELECT proowner::regrole FROM pg_proc WHERE proname = 'usage_access_func_third'$$);
+      run_command_on_workers      
+----------------------------------
+ (localhost,57637,t,usage_access)
+ (localhost,57638,t,usage_access)
+(2 rows)
+
+-- we don't want other tests to have metadata synced
+-- that might change the test outputs, so we're just trying to be careful
+SELECT stop_metadata_sync_to_node('localhost', :worker_1_port);
+ stop_metadata_sync_to_node 
+----------------------------
+ 
+(1 row)
+
+SELECT stop_metadata_sync_to_node('localhost', :worker_2_port);
+ stop_metadata_sync_to_node 
+----------------------------
+ 
+(1 row)
+
 RESET ROLE;
 -- now we distribute the table as super user
 SELECT create_distributed_table('full_access_user_schema.t1', 'id');

src/test/regress/sql/multi_multiuser.sql

@@ -13,6 +13,15 @@ ALTER SEQUENCE pg_catalog.pg_dist_jobid_seq RESTART 1420000;
 
 SET citus.shard_replication_factor TO 1;
 
+ALTER SYSTEM SET citus.metadata_sync_interval TO 3000;
+ALTER SYSTEM SET citus.metadata_sync_retry_interval TO 500;
+SELECT pg_reload_conf();
+CREATE FUNCTION wait_until_metadata_sync(timeout INTEGER DEFAULT 15000)
+    RETURNS void
+    LANGUAGE C STRICT
+    AS 'citus';
+
+
 CREATE TABLE test (id integer, val integer);
 SELECT create_distributed_table('test', 'id');
 
@@ -272,6 +281,58 @@ INSERT INTO full_access_user_schema.t1 VALUES (1),(2),(3);
 SELECT create_distributed_table('full_access_user_schema.t1', 'id');
 RESET ROLE;
 
+SET ROLE usage_access;
+
+CREATE TYPE usage_access_type AS ENUM ('a', 'b');
+CREATE FUNCTION usage_access_func(x usage_access_type, variadic v int[]) RETURNS int[]
+    LANGUAGE plpgsql AS 'begin return v; end;';
+
+SET ROLE no_access;
+SELECT create_distributed_function('usage_access_func(usage_access_type,int[])');
+
+
+SET ROLE usage_access;
+SELECT create_distributed_function('usage_access_func(usage_access_type,int[])');
+
+SELECT typowner::regrole FROM pg_type WHERE typname = 'usage_access_type';
+SELECT proowner::regrole FROM pg_proc WHERE proname = 'usage_access_func';
+SELECT run_command_on_workers($$SELECT typowner::regrole FROM pg_type WHERE typname = 'usage_access_type'$$);
+SELECT run_command_on_workers($$SELECT proowner::regrole FROM pg_proc WHERE proname = 'usage_access_func'$$);
+
+SELECT wait_until_metadata_sync();
+
+-- now, make sure that the user can use the function 
+-- created in the transaction
+BEGIN;
+CREATE FUNCTION usage_access_func_second(key int, variadic v int[]) RETURNS text
+    LANGUAGE plpgsql AS 'begin return current_user; end;';
+SELECT create_distributed_function('usage_access_func_second(int,int[])', '$1');
+
+SELECT usage_access_func_second(1, 2,3,4,5) FROM full_access_user_schema.t1 LIMIT 1; 
+
+ROLLBACK;
+
+CREATE FUNCTION usage_access_func_third(key int, variadic v int[]) RETURNS text
+    LANGUAGE plpgsql AS 'begin return current_user; end;';
+
+-- connect back as super user
+\c - - - :master_port
+
+-- show that the current user is a super user
+SELECT usesuper FROM pg_user where usename IN (SELECT current_user);
+
+-- superuser creates the distributed function that is owned by a regular user
+SELECT create_distributed_function('usage_access_func_third(int,int[])', '$1');
+
+SELECT proowner::regrole FROM pg_proc WHERE proname = 'usage_access_func_third';
+SELECT run_command_on_workers($$SELECT proowner::regrole FROM pg_proc WHERE proname = 'usage_access_func_third'$$);
+
+-- we don't want other tests to have metadata synced
+-- that might change the test outputs, so we're just trying to be careful
+SELECT stop_metadata_sync_to_node('localhost', :worker_1_port);
+SELECT stop_metadata_sync_to_node('localhost', :worker_2_port);
+
+RESET ROLE;
 -- now we distribute the table as super user
 SELECT create_distributed_table('full_access_user_schema.t1', 'id');