From 8e93fe58704ea1f6b4e83d1c54b263c2dbdd3b80 Mon Sep 17 00:00:00 2001
From: Marco Slot <marco.slot@gmail.com>
Date: Thu, 22 Nov 2018 05:08:12 +0100
Subject: [PATCH] Check schema owner in task_tracker_assign_task

---
 .../master/master_metadata_utility.c          | 15 ++++++++++
 .../worker/task_tracker_protocol.c            | 28 +++++++++++++++++--
 .../distributed/master_metadata_utility.h     |  1 +
 3 files changed, 42 insertions(+), 2 deletions(-)

diff --git a/src/backend/distributed/master/master_metadata_utility.c b/src/backend/distributed/master/master_metadata_utility.c
index 134d40eb2..b8798611f 100644
--- a/src/backend/distributed/master/master_metadata_utility.c
+++ b/src/backend/distributed/master/master_metadata_utility.c
@@ -1346,6 +1346,21 @@ EnsureTableOwner(Oid relationId)
 }
 
 
+/*
+ * Check that the current user has owner rights to the schema, error out if
+ * not. Superusers are regarded as owners.
+ */
+void
+EnsureSchemaOwner(Oid schemaId)
+{
+	if (!pg_namespace_ownercheck(schemaId, GetUserId()))
+	{
+		aclcheck_error(ACLCHECK_NOT_OWNER, ACLCHECK_OBJECT_TABLE,
+					   get_namespace_name(schemaId));
+	}
+}
+
+
 /*
  * Check that the current user has owner rights to sequenceRelationId, error out if
  * not. Superusers are regarded as owners.
diff --git a/src/backend/distributed/worker/task_tracker_protocol.c b/src/backend/distributed/worker/task_tracker_protocol.c
index 2cd19c4d2..c33df2a01 100644
--- a/src/backend/distributed/worker/task_tracker_protocol.c
+++ b/src/backend/distributed/worker/task_tracker_protocol.c
@@ -19,7 +19,10 @@
 
 #include <time.h>
 
+#include "access/htup_details.h"
 #include "access/xact.h"
+#include "catalog/pg_namespace.h"
+#include "catalog/namespace.h"
 #include "commands/dbcommands.h"
 #include "commands/schemacmds.h"
 #include "commands/trigger.h"
@@ -33,6 +36,8 @@
 #include "storage/lwlock.h"
 #include "storage/pmsignal.h"
 #include "utils/builtins.h"
+#include "utils/syscache.h"
+#include "utils/lsyscache.h"
 
 
 /* Local functions forward declarations */
@@ -105,6 +110,10 @@ task_tracker_assign_task(PG_FUNCTION_ARGS)
 	}
 	else
 	{
+		Oid schemaId = get_namespace_oid(jobSchemaName->data, false);
+
+		EnsureSchemaOwner(schemaId);
+
 		UnlockJobResource(jobId, AccessExclusiveLock);
 	}
 
@@ -179,6 +188,7 @@ task_tracker_cleanup_job(PG_FUNCTION_ARGS)
 {
 	uint64 jobId = PG_GETARG_INT64(0);
 
+	bool schemaExists = false;
 	HASH_SEQ_STATUS status;
 	WorkerTask *currentTask = NULL;
 	StringInfo jobDirectoryName = NULL;
@@ -186,6 +196,22 @@ task_tracker_cleanup_job(PG_FUNCTION_ARGS)
 
 	CheckCitusVersion(ERROR);
 
+	jobSchemaName = JobSchemaName(jobId);
+
+	/*
+	 * We'll keep this lock for a while, but that's ok because nothing
+	 * else should be happening on this job.
+	 */
+	LockJobResource(jobId, AccessExclusiveLock);
+
+	schemaExists = JobSchemaExists(jobSchemaName);
+	if (schemaExists)
+	{
+		Oid schemaId = get_namespace_oid(jobSchemaName->data, false);
+
+		EnsureSchemaOwner(schemaId);
+	}
+
 	/*
 	 * We first clean up any open connections, and remove tasks belonging to
 	 * this job from the shared hash.
@@ -216,8 +242,6 @@ task_tracker_cleanup_job(PG_FUNCTION_ARGS)
 	jobDirectoryName = JobDirectoryName(jobId);
 	CitusRemoveDirectory(jobDirectoryName);
 
-	LockJobResource(jobId, AccessExclusiveLock);
-	jobSchemaName = JobSchemaName(jobId);
 	RemoveJobSchema(jobSchemaName);
 	UnlockJobResource(jobId, AccessExclusiveLock);
 
diff --git a/src/include/distributed/master_metadata_utility.h b/src/include/distributed/master_metadata_utility.h
index 1461fd473..125f17cb4 100644
--- a/src/include/distributed/master_metadata_utility.h
+++ b/src/include/distributed/master_metadata_utility.h
@@ -155,6 +155,7 @@ extern void CreateTruncateTrigger(Oid relationId);
 extern char * TableOwner(Oid relationId);
 extern void EnsureTablePermissions(Oid relationId, AclMode mode);
 extern void EnsureTableOwner(Oid relationId);
+extern void EnsureSchemaOwner(Oid schemaId);
 extern void EnsureSequenceOwner(Oid sequenceOid);
 extern void EnsureSuperUser(void);
 extern void EnsureReplicationSettings(Oid relationId, char replicationModel);