Allow lock_shard_resources to be called by the users with privileges

Before this commit, we required the user to be owner of the shard/table
in order to call lock_shard_resources.

However, that is too restrictive. We can have users with GRANTS
to the table who are not owners of the tables/shards.

With this commit, we allow such patterns.

src/backend/distributed/utils/resource_lock.c

@@ -160,6 +160,33 @@ EnsureShardOwner(uint64 shardId, bool missingOk)
 }
 
 
+/*
+ * EnsureCurrenUserCanModifyShard gets a shardId and throws error if the current
+ * user doesn't have privileges to modify the shard.
+ */
+void
+EnsureCurrenUserCanModifyShard(uint64 shardId, bool missingOk)
+{
+	Oid relationId = LookupShardRelationFromCatalog(shardId, missingOk);
+
+	if (!OidIsValid(relationId) && missingOk)
+	{
+		/*
+		 * This could happen in two ways. First, a malicious user is trying
+		 * to acquire locks on non-existing shards. Second, the metadata has
+		 * not been synced (or not yet visible) to this node. In the second
+		 * case, there is no point in locking the shards because no other
+		 * transaction can be accessing the table.
+		 */
+		return;
+	}
+
+	AclMode aclMask = ACL_INSERT | ACL_UPDATE | ACL_DELETE | ACL_TRUNCATE;
+
+	EnsureTablePermissions(relationId, aclMask);
+}
+
+
 /*
  * lock_shard_resources allows shard resources to be locked
  * remotely to serialise non-commutative writes on shards.
@@ -188,15 +215,12 @@ lock_shard_resources(PG_FUNCTION_ARGS)
 		int64 shardId = DatumGetInt64(shardIdArrayDatum[shardIdIndex]);
 
 		/*
-		 * We don't want random users to block writes. The callers of this
+		 * We don't want random users to block writes. If the current user
-		 * function either operates on all the colocated placements, such
+		 * has privileges to modify the shard, then the user can already
-		 * as shard moves, or requires superuser such as adding node.
+		 * acquire the lock. So, we allow.
-		 * In other words, the coordinator initiated operations has already
-		 * ensured table owner, we are preventing any malicious attempt to
-		 * use this function.
 		 */
 		bool missingOk = true;
-		EnsureShardOwner(shardId, missingOk);
+		EnsureCurrenUserCanModifyShard(shardId, missingOk);
 
 		LockShardResource(shardId, lockMode);
 	}

src/include/distributed/resource_lock.h

@@ -115,6 +115,7 @@ extern void LockShardDistributionMetadata(int64 shardId, LOCKMODE lockMode);
 extern void LockPlacementCleanup(void);
 extern bool TryLockPlacementCleanup(void);
 extern void EnsureShardOwner(uint64 shardId, bool missingOk);
+extern void EnsureCurrenUserCanModifyShard(uint64 shardId, bool missingOk);
 extern void LockShardListMetadataOnWorkers(LOCKMODE lockmode, List *shardIntervalList);
 extern void BlockWritesToShardList(List *shardList);
 

src/test/regress/expected/mx_regular_user.out

@@ -25,18 +25,95 @@ SELECT start_metadata_sync_to_node('localhost', :worker_2_port);
 SET client_min_messages TO ERROR;
 SET citus.enable_ddl_propagation TO OFF;
 CREATE USER regular_mx_user WITH LOGIN;
+SELECT 1 FROM run_command_on_workers($$CREATE USER regular_mx_user WITH LOGIN;$$);
+ ?column?
+---------------------------------------------------------------------
+        1
+        1
+(2 rows)
+
 GRANT ALL ON SCHEMA "Mx Regular User" TO regular_mx_user;
+-- create another table owned by the super user (e.g., current user of the session)
+-- and GRANT access to the user
+CREATE SCHEMA "Mx Super User";
+SELECT 1 FROM run_command_on_workers($$CREATE SCHEMA "Mx Super User";$$);
+ ?column?
+---------------------------------------------------------------------
+        1
+        1
+(2 rows)
+
+SET citus.next_shard_id TO 2980000;
+SET search_path TO "Mx Super User";
+CREATE TABLE super_user_owned_regular_user_granted (a int);
+SELECT create_reference_table ('"Mx Super User".super_user_owned_regular_user_granted');
+ create_reference_table
+---------------------------------------------------------------------
+
+(1 row)
+
+-- show that this table is owned by super user
+SELECT
+	rolsuper
+FROM
+	pg_roles
+		WHERE oid
+			IN
+		(SELECT relowner FROM pg_class WHERE oid = '"Mx Super User".super_user_owned_regular_user_granted'::regclass);
+ rolsuper
+---------------------------------------------------------------------
+ t
+(1 row)
+
+-- make sure that granting produce the same output for both community and enterprise
+SET client_min_messages TO ERROR;
+GRANT ALL ON SCHEMA "Mx Super User" TO regular_mx_user;
+GRANT ALL ON TABLE super_user_owned_regular_user_granted TO regular_mx_user;
+SELECT 1 FROM run_command_on_workers($$GRANT ALL ON SCHEMA "Mx Super User" TO regular_mx_user;$$);
+ ?column?
+---------------------------------------------------------------------
+        1
+        1
+(2 rows)
+
+SELECT 1 FROM run_command_on_workers($$GRANT ALL ON TABLE "Mx Super User".super_user_owned_regular_user_granted TO regular_mx_user;$$);
+ ?column?
+---------------------------------------------------------------------
+        1
+        1
+(2 rows)
+
+SELECT 1 FROM run_command_on_placements('super_user_owned_regular_user_granted', $$GRANT ALL ON TABLE %s TO regular_mx_user;$$);
+ ?column?
+---------------------------------------------------------------------
+        1
+        1
+        1
+(3 rows)
+
+-- now that the GRANT is given, the regular user should be able to
+-- modify the table
+\c - regular_mx_user - :master_port
+SET search_path TO "Mx Super User";
+COPY super_user_owned_regular_user_granted FROM STDIN;
+SELECT count(*) FROM super_user_owned_regular_user_granted;
+ count
+---------------------------------------------------------------------
+     2
+(1 row)
+
+\c - postgres - :master_port;
+SET client_min_messages TO ERROR;
+DROP SCHEMA "Mx Super User" CASCADE;
 \c - postgres - :worker_1_port;
 SET client_min_messages TO ERROR;
 SET citus.enable_ddl_propagation TO OFF;
 CREATE SCHEMA "Mx Regular User";
-CREATE USER regular_mx_user WITH LOGIN;
 GRANT ALL ON SCHEMA "Mx Regular User" TO regular_mx_user;
 \c - postgres - :worker_2_port;
 SET client_min_messages TO ERROR;
 SET citus.enable_ddl_propagation TO OFF;
 CREATE SCHEMA "Mx Regular User";
-CREATE USER regular_mx_user WITH LOGIN;
 GRANT ALL ON SCHEMA "Mx Regular User" TO regular_mx_user;
 -- now connect with that user
 \c - regular_mx_user - :master_port

src/test/regress/sql/mx_regular_user.sql

@@ -12,20 +12,60 @@ SELECT start_metadata_sync_to_node('localhost', :worker_2_port);
 SET client_min_messages TO ERROR;
 SET citus.enable_ddl_propagation TO OFF;
 CREATE USER regular_mx_user WITH LOGIN;
+SELECT 1 FROM run_command_on_workers($$CREATE USER regular_mx_user WITH LOGIN;$$);
 GRANT ALL ON SCHEMA "Mx Regular User" TO regular_mx_user;
 
+-- create another table owned by the super user (e.g., current user of the session)
+-- and GRANT access to the user
+CREATE SCHEMA "Mx Super User";
+SELECT 1 FROM run_command_on_workers($$CREATE SCHEMA "Mx Super User";$$);
+SET citus.next_shard_id TO 2980000;
+SET search_path TO "Mx Super User";
+CREATE TABLE super_user_owned_regular_user_granted (a int);
+SELECT create_reference_table ('"Mx Super User".super_user_owned_regular_user_granted');
+
+-- show that this table is owned by super user
+SELECT
+	rolsuper
+FROM
+	pg_roles
+		WHERE oid
+			IN
+		(SELECT relowner FROM pg_class WHERE oid = '"Mx Super User".super_user_owned_regular_user_granted'::regclass);
+
+-- make sure that granting produce the same output for both community and enterprise
+SET client_min_messages TO ERROR;
+GRANT ALL ON SCHEMA "Mx Super User" TO regular_mx_user;
+GRANT ALL ON TABLE super_user_owned_regular_user_granted TO regular_mx_user;
+
+SELECT 1 FROM run_command_on_workers($$GRANT ALL ON SCHEMA "Mx Super User" TO regular_mx_user;$$);
+SELECT 1 FROM run_command_on_workers($$GRANT ALL ON TABLE "Mx Super User".super_user_owned_regular_user_granted TO regular_mx_user;$$);
+SELECT 1 FROM run_command_on_placements('super_user_owned_regular_user_granted', $$GRANT ALL ON TABLE %s TO regular_mx_user;$$);
+
+-- now that the GRANT is given, the regular user should be able to
+-- modify the table
+\c - regular_mx_user - :master_port
+SET search_path TO "Mx Super User";
+COPY super_user_owned_regular_user_granted FROM STDIN;
+1
+2
+\.
+SELECT count(*) FROM super_user_owned_regular_user_granted;
+
+\c - postgres - :master_port;
+SET client_min_messages TO ERROR;
+DROP SCHEMA "Mx Super User" CASCADE;
+
 \c - postgres - :worker_1_port;
 SET client_min_messages TO ERROR;
 SET citus.enable_ddl_propagation TO OFF;
 CREATE SCHEMA "Mx Regular User";
-CREATE USER regular_mx_user WITH LOGIN;
 GRANT ALL ON SCHEMA "Mx Regular User" TO regular_mx_user;
 
 \c - postgres - :worker_2_port;
 SET client_min_messages TO ERROR;
 SET citus.enable_ddl_propagation TO OFF;
 CREATE SCHEMA "Mx Regular User";
-CREATE USER regular_mx_user WITH LOGIN;
 GRANT ALL ON SCHEMA "Mx Regular User" TO regular_mx_user;
 
 -- now connect with that user