Check permissions in lock_relation_if_exists

src/backend/distributed/utils/resource_lock.c

@@ -19,6 +19,7 @@
 
 #include "access/xact.h"
 #include "catalog/namespace.h"
+#include "commands/tablecmds.h"
 #include "distributed/colocation_utils.h"
 #include "distributed/listutils.h"
 #include "distributed/master_metadata_utility.h"
@@ -36,6 +37,7 @@
 #include "distributed/version_compat.h"
 #include "storage/lmgr.h"
 #include "utils/builtins.h"
+#include "utils/lsyscache.h"
 #if (PG_VERSION_NUM >= 100000)
 #include "utils/varlena.h"
 #endif
@@ -73,6 +75,9 @@ static void LockShardListResources(List *shardIntervalList, LOCKMODE lockMode);
 static void LockShardListResourcesOnFirstWorker(LOCKMODE lockmode,
 												List *shardIntervalList);
 static bool IsFirstWorkerNode();
+static void CitusRangeVarCallbackForLockTable(const RangeVar *rangeVar, Oid relationId,
+											  Oid oldRelationId, void *arg);
+static AclResult CitusLockTableAclCheck(Oid relationId, LOCKMODE lockmode, Oid userId);
 
 
 /* exports for SQL callable functions */
@@ -742,16 +747,11 @@ lock_relation_if_exists(PG_FUNCTION_ARGS)
 	List *relationNameList = NIL;
 	RangeVar *relation = NULL;
 	LOCKMODE lockMode = NoLock;
+	bool relationExists = false;
 
     /* ensure that we're in a transaction block */
 	RequireTransactionBlock(true, "lock_relation_if_exists");
 
-	relationId = ResolveRelationId(relationName, true);
-	if (!OidIsValid(relationId))
-	{
-		PG_RETURN_BOOL(false);
-	}
-
     /* get the lock mode */
 	lockMode = LockModeTextToLockMode(lockModeCString);
 
@@ -760,7 +760,85 @@ lock_relation_if_exists(PG_FUNCTION_ARGS)
 	relation = makeRangeVarFromNameList(relationNameList);
 
     /* lock the relation with the lock mode */
-	RangeVarGetRelid(relation, lockMode, false);
+	relationId = RangeVarGetRelidInternal(relation, lockMode, RVR_MISSING_OK,
+										  CitusRangeVarCallbackForLockTable,
+										  (void *) &lockMode);
+	relationExists = OidIsValid(relationId);
 
-	PG_RETURN_BOOL(true);
+	PG_RETURN_BOOL(relationExists);
+}
+
+
+/*
+ * CitusRangeVarCallbackForLockTable is a callback for RangeVarGetRelidExtended used
+ * to check whether the user has permission to lock a table in a particular mode.
+ *
+ * This function is a copy of RangeVarCallbackForLockTable in lockcmds.c adapted to
+ * Citus code style.
+ */
+static void
+CitusRangeVarCallbackForLockTable(const RangeVar *rangeVar, Oid relationId,
+								  Oid oldRelationId, void *arg)
+{
+	LOCKMODE lockmode = *(LOCKMODE *) arg;
+	AclResult aclResult;
+
+	if (!OidIsValid(relationId))
+	{
+        /* table doesn't exist, so no permissions check */
+		return;
+	}
+
+    /* we only allow tables and views to be locked */
+	if (!RegularTable(relationId))
+	{
+		ereport(ERROR, (errcode(ERRCODE_WRONG_OBJECT_TYPE),
+						errmsg("\"%s\" is not a table", rangeVar->relname)));
+	}
+
+    /* check permissions */
+	aclResult = CitusLockTableAclCheck(relationId, lockmode, GetUserId());
+	if (aclResult != ACLCHECK_OK)
+	{
+#if (PG_VERSION_NUM >= 110000)
+		aclcheck_error(aclResult, get_relkind_objtype(get_rel_relkind(relationId)),
+					   rangeVar->relname);
+#else
+
+		aclcheck_error(aclResult, ACL_KIND_CLASS, rangeVar->relname);
+#endif
+	}
+}
+
+
+/*
+ * CitusLockTableAclCheck checks whether a user has permission to lock a relation
+ * in the given lock mode.
+ *
+ * This function is a copy of LockTableAclCheck in lockcmds.c adapted to Citus
+ * code style.
+ */
+static AclResult
+CitusLockTableAclCheck(Oid relationId, LOCKMODE lockmode, Oid userId)
+{
+	AclResult aclResult;
+	AclMode aclMask;
+
+    /* verify adequate privilege */
+	if (lockmode == AccessShareLock)
+	{
+		aclMask = ACL_SELECT;
+	}
+	else if (lockmode == RowExclusiveLock)
+	{
+		aclMask = ACL_INSERT | ACL_UPDATE | ACL_DELETE | ACL_TRUNCATE;
+	}
+	else
+	{
+		aclMask = ACL_UPDATE | ACL_DELETE | ACL_TRUNCATE;
+	}
+
+	aclResult = pg_class_aclcheck(relationId, userId, aclMask);
+
+	return aclResult;
 }

src/test/regress/expected/multi_multiuser.out

@@ -165,6 +165,17 @@ SELECT count(*) FROM test a JOIN test b ON (a.val = b.val) WHERE a.id = 1 AND b.
 COPY "postgresql.conf" TO STDOUT WITH (format transmit);
 ERROR:  operation is not allowed
 HINT:  Run the command with a superuser.
+-- should not be allowed to take aggressive locks on table
+BEGIN;
+SELECT lock_relation_if_exists('test', 'ACCESS SHARE');
+ lock_relation_if_exists 
+-------------------------
+ t
+(1 row)
+
+SELECT lock_relation_if_exists('test', 'EXCLUSIVE');
+ERROR:  permission denied for table test
+ABORT;
 SET citus.task_executor_type TO 'real-time';
 -- check no permission
 SET ROLE no_access;

src/test/regress/expected/multi_multiuser_0.out

@@ -21,6 +21,21 @@ SELECT create_distributed_table('test', 'id');
  
 (1 row)
 
+CREATE TABLE test_coloc (id integer, val integer);
+SELECT create_distributed_table('test_coloc', 'id', colocate_with := 'none');
+ create_distributed_table 
+--------------------------
+ 
+(1 row)
+
+SET citus.shard_count TO 1;
+CREATE TABLE singleshard (id integer, val integer);
+SELECT create_distributed_table('singleshard', 'id');
+ create_distributed_table 
+--------------------------
+ 
+(1 row)
+
 -- turn off propagation to avoid Enterprise processing the following section
 SET citus.enable_ddl_propagation TO off;
 CREATE USER full_access;
@@ -157,6 +172,17 @@ SELECT count(*) FROM test a JOIN test b ON (a.val = b.val) WHERE a.id = 1 AND b.
 COPY "postgresql.conf" TO STDOUT WITH (format transmit);
 ERROR:  operation is not allowed
 HINT:  Run the command with a superuser.
+-- should not be allowed to take aggressive locks on table
+BEGIN;
+SELECT lock_relation_if_exists('test', 'ACCESS SHARE');
+ lock_relation_if_exists 
+-------------------------
+ t
+(1 row)
+
+SELECT lock_relation_if_exists('test', 'EXCLUSIVE');
+ERROR:  permission denied for relation test
+ABORT;
 SET citus.task_executor_type TO 'real-time';
 -- check no permission
 SET ROLE no_access;
@@ -212,6 +238,12 @@ ERROR:  permission denied for relation test
 ABORT;
 SELECT * FROM citus_stat_statements_reset();
 ERROR:  permission denied for function citus_stat_statements_reset
+-- should not be allowed to upgrade to reference table
+SELECT upgrade_to_reference_table('singleshard');
+ERROR:  must be owner of relation singleshard
+-- should not be allowed to co-located tables
+SELECT mark_tables_colocated('test', ARRAY['test_coloc'::regclass]);
+ERROR:  must be owner of relation test
 -- table owner should be the same on the shards, even when distributing the table as superuser
 SET ROLE full_access;
 CREATE TABLE my_table (id integer, val integer);
@@ -229,8 +261,7 @@ SELECT result FROM run_command_on_workers($$SELECT tableowner FROM pg_tables WHE
  full_access
 (2 rows)
 
-DROP TABLE my_table;
-DROP TABLE test;
+DROP TABLE my_table, singleshard, test, test_coloc;
 DROP USER full_access;
 DROP USER read_access;
 DROP USER no_access;

src/test/regress/sql/multi_multiuser.sql

@@ -106,6 +106,12 @@ SELECT count(*) FROM test a JOIN test b ON (a.val = b.val) WHERE a.id = 1 AND b.
 -- should not be able to transmit directly
 COPY "postgresql.conf" TO STDOUT WITH (format transmit);
 
+-- should not be allowed to take aggressive locks on table
+BEGIN;
+SELECT lock_relation_if_exists('test', 'ACCESS SHARE');
+SELECT lock_relation_if_exists('test', 'EXCLUSIVE');
+ABORT;
+
 SET citus.task_executor_type TO 'real-time';
 
 -- check no permission