Handle ALTER EXTENSION ADD/DROP with pg_init_privs

In commit 6c268df, pg_init_privs was added to track the initial
privileges of catalog objects and extensions.  Unfortunately, that
commit didn't include understanding of ALTER EXTENSION ADD/DROP, which
allows the objects associated with an extension to be changed after the
initial CREATE EXTENSION script has been run.

The result of this meant that ACLs for objects added through
ALTER EXTENSION ADD were not recorded into pg_init_privs and we would
end up including those ACLs in pg_dump when we shouldn't have.

This commit corrects that by making sure to have pg_init_privs updated
when ALTER EXTENSION ADD/DROP is run, recording the permissions as they
are at ALTER EXTENSION ADD time, and removing any if/when ALTER
EXTENSION DROP is called.

This issue was pointed out by Moshe Jacobson as commentary on bug #14456
(which was actually a bug about versions prior to 9.6 not handling
custom ACLs on extensions correctly, an issue now addressed with
pg_init_privs in 9.6).

Back-patch to 9.6 where pg_init_privs was introduced.

2 files changed, 435 insertions, 22 deletions

diff --git a/src/backend/catalog/aclchk.c b/src/backend/catalog/aclchk.c
index a585c3ad19a..066b7fdb740 100644
--- a/src/backend/catalog/aclchk.c
+++ b/src/backend/catalog/aclchk.c
@@ -27,7 +27,10 @@
 #include "catalog/dependency.h"
 #include "catalog/indexing.h"
 #include "catalog/objectaccess.h"
+#include "catalog/pg_aggregate.h"
+#include "catalog/pg_am.h"
 #include "catalog/pg_authid.h"
+#include "catalog/pg_cast.h"
 #include "catalog/pg_collation.h"
 #include "catalog/pg_conversion.h"
 #include "catalog/pg_database.h"
@@ -49,6 +52,9 @@
 #include "catalog/pg_type.h"
 #include "catalog/pg_ts_config.h"
 #include "catalog/pg_ts_dict.h"
+#include "catalog/pg_ts_parser.h"
+#include "catalog/pg_ts_template.h"
+#include "catalog/pg_transform.h"
 #include "commands/dbcommands.h"
 #include "commands/event_trigger.h"
 #include "commands/extension.h"
@@ -130,6 +136,8 @@ static AclMode pg_aclmask(AclObjectKind objkind, Oid table_oid, AttrNumber attnu
 		   Oid roleid, AclMode mask, AclMaskHow how);
 static void recordExtensionInitPriv(Oid objoid, Oid classoid, int objsubid,
 						Acl *new_acl);
+static void recordExtensionInitPrivWorker(Oid objoid, Oid classoid, int objsubid,
+							  Acl *new_acl);
 
 
 #ifdef ACLDEBUG
@@ -5222,10 +5230,367 @@ get_user_default_acl(GrantObjectType objtype, Oid ownerId, Oid nsp_oid)
 }
 
 /*
- * Record initial ACL for an extension object
+ * Record initial privileges for the top-level object passed in.
  *
- * This will perform a wholesale replacement of the entire ACL for the object
- * passed in, therefore be sure to pass in the complete new ACL to use.
+ * For the object passed in, this will record its ACL (if any) and the ACLs of
+ * any sub-objects (eg: columns) into pg_init_privs.
+ *
+ * Any new kinds of objects which have ACLs associated with them and can be
+ * added to an extension should be added to the if-else tree below.
+ */
+void
+recordExtObjInitPriv(Oid objoid, Oid classoid)
+{
+	/*
+	 * pg_class / pg_attribute
+	 *
+	 * If this is a relation then we need to see if there are any sub-objects
+	 * (eg: columns) for it and, if so, be sure to call
+	 * recordExtensionInitPrivWorker() for each one.
+	 */
+	if (classoid == RelationRelationId)
+	{
+		Form_pg_class pg_class_tuple;
+		Datum		aclDatum;
+		bool		isNull;
+		HeapTuple	tuple;
+
+		tuple = SearchSysCache1(RELOID, ObjectIdGetDatum(objoid));
+		if (!HeapTupleIsValid(tuple))
+			elog(ERROR, "cache lookup failed for relation %u", objoid);
+		pg_class_tuple = (Form_pg_class) GETSTRUCT(tuple);
+
+		/* Indexes don't have permissions */
+		if (pg_class_tuple->relkind == RELKIND_INDEX)
+			return;
+
+		/* Composite types don't have permissions either */
+		if (pg_class_tuple->relkind == RELKIND_COMPOSITE_TYPE)
+			return;
+
+		/*
+		 * If this isn't a sequence, index, or composite type then it's
+		 * possibly going to have columns associated with it that might have
+		 * ACLs.
+		 */
+		if (pg_class_tuple->relkind != RELKIND_SEQUENCE)
+		{
+			AttrNumber	curr_att;
+			AttrNumber	nattrs = pg_class_tuple->relnatts;
+
+			for (curr_att = 1; curr_att <= nattrs; curr_att++)
+			{
+				HeapTuple	attTuple;
+				Datum		attaclDatum;
+
+				attTuple = SearchSysCache2(ATTNUM,
+										   ObjectIdGetDatum(objoid),
+										   Int16GetDatum(curr_att));
+
+				if (!HeapTupleIsValid(attTuple))
+					continue;
+
+				/* ignore dropped columns */
+				if (((Form_pg_attribute) GETSTRUCT(attTuple))->attisdropped)
+				{
+					ReleaseSysCache(attTuple);
+					continue;
+				}
+
+				attaclDatum = SysCacheGetAttr(ATTNUM, attTuple,
+											  Anum_pg_attribute_attacl,
+											  &isNull);
+
+				/* no need to do anything for a NULL ACL */
+				if (isNull)
+				{
+					ReleaseSysCache(attTuple);
+					continue;
+				}
+
+				recordExtensionInitPrivWorker(objoid, classoid, curr_att,
+											  DatumGetAclP(attaclDatum));
+
+				ReleaseSysCache(attTuple);
+			}
+		}
+
+		aclDatum = SysCacheGetAttr(RELOID, tuple, Anum_pg_class_relacl,
+								   &isNull);
+
+		/* Add the record, if any, for the top-level object */
+		if (!isNull)
+			recordExtensionInitPrivWorker(objoid, classoid, 0,
+										  DatumGetAclP(aclDatum));
+
+		ReleaseSysCache(tuple);
+	}
+	/* pg_foreign_data_wrapper */
+	else if (classoid == ForeignDataWrapperRelationId)
+	{
+		Datum		aclDatum;
+		bool		isNull;
+		HeapTuple	tuple;
+
+		tuple = SearchSysCache1(FOREIGNDATAWRAPPEROID,
+								ObjectIdGetDatum(objoid));
+		if (!HeapTupleIsValid(tuple))
+			elog(ERROR, "cache lookup failed for foreign data wrapper %u",
+				 objoid);
+
+		aclDatum = SysCacheGetAttr(FOREIGNDATAWRAPPEROID, tuple,
+								   Anum_pg_foreign_data_wrapper_fdwacl,
+								   &isNull);
+
+		/* Add the record, if any, for the top-level object */
+		if (!isNull)
+			recordExtensionInitPrivWorker(objoid, classoid, 0,
+										  DatumGetAclP(aclDatum));
+
+		ReleaseSysCache(tuple);
+	}
+	/* pg_foreign_server */
+	else if (classoid == ForeignServerRelationId)
+	{
+		Datum		aclDatum;
+		bool		isNull;
+		HeapTuple	tuple;
+
+		tuple = SearchSysCache1(FOREIGNSERVEROID, ObjectIdGetDatum(objoid));
+		if (!HeapTupleIsValid(tuple))
+			elog(ERROR, "cache lookup failed for foreign data wrapper %u",
+				 objoid);
+
+		aclDatum = SysCacheGetAttr(FOREIGNSERVEROID, tuple,
+								   Anum_pg_foreign_server_srvacl,
+								   &isNull);
+
+		/* Add the record, if any, for the top-level object */
+		if (!isNull)
+			recordExtensionInitPrivWorker(objoid, classoid, 0,
+										  DatumGetAclP(aclDatum));
+
+		ReleaseSysCache(tuple);
+	}
+	/* pg_language */
+	else if (classoid == LanguageRelationId)
+	{
+		Datum		aclDatum;
+		bool		isNull;
+		HeapTuple	tuple;
+
+		tuple = SearchSysCache1(LANGOID, ObjectIdGetDatum(objoid));
+		if (!HeapTupleIsValid(tuple))
+			elog(ERROR, "cache lookup failed for language %u", objoid);
+
+		aclDatum = SysCacheGetAttr(LANGOID, tuple, Anum_pg_language_lanacl,
+								   &isNull);
+
+		/* Add the record, if any, for the top-level object */
+		if (!isNull)
+			recordExtensionInitPrivWorker(objoid, classoid, 0,
+										  DatumGetAclP(aclDatum));
+
+		ReleaseSysCache(tuple);
+	}
+	/* pg_largeobject_metadata */
+	else if (classoid == LargeObjectMetadataRelationId)
+	{
+		Datum		aclDatum;
+		bool		isNull;
+		HeapTuple	tuple;
+		ScanKeyData entry[1];
+		SysScanDesc scan;
+		Relation	relation;
+
+		relation = heap_open(LargeObjectMetadataRelationId, RowExclusiveLock);
+
+		/* There's no syscache for pg_largeobject_metadata */
+		ScanKeyInit(&entry[0],
+					ObjectIdAttributeNumber,
+					BTEqualStrategyNumber, F_OIDEQ,
+					ObjectIdGetDatum(objoid));
+
+		scan = systable_beginscan(relation,
+								  LargeObjectMetadataOidIndexId, true,
+								  NULL, 1, entry);
+
+		tuple = systable_getnext(scan);
+		if (!HeapTupleIsValid(tuple))
+			elog(ERROR, "cache lookup failed for large object %u", objoid);
+
+		aclDatum = heap_getattr(tuple,
+								Anum_pg_largeobject_metadata_lomacl,
+								RelationGetDescr(relation), &isNull);
+
+		/* Add the record, if any, for the top-level object */
+		if (!isNull)
+			recordExtensionInitPrivWorker(objoid, classoid, 0,
+										  DatumGetAclP(aclDatum));
+
+		systable_endscan(scan);
+	}
+	/* pg_namespace */
+	else if (classoid == NamespaceRelationId)
+	{
+		Datum		aclDatum;
+		bool		isNull;
+		HeapTuple	tuple;
+
+		tuple = SearchSysCache1(NAMESPACEOID, ObjectIdGetDatum(objoid));
+		if (!HeapTupleIsValid(tuple))
+			elog(ERROR, "cache lookup failed for function %u", objoid);
+
+		aclDatum = SysCacheGetAttr(NAMESPACEOID, tuple,
+								   Anum_pg_namespace_nspacl, &isNull);
+
+		/* Add the record, if any, for the top-level object */
+		if (!isNull)
+			recordExtensionInitPrivWorker(objoid, classoid, 0,
+										  DatumGetAclP(aclDatum));
+
+		ReleaseSysCache(tuple);
+	}
+	/* pg_proc */
+	else if (classoid == ProcedureRelationId)
+	{
+		Datum		aclDatum;
+		bool		isNull;
+		HeapTuple	tuple;
+
+		tuple = SearchSysCache1(PROCOID, ObjectIdGetDatum(objoid));
+		if (!HeapTupleIsValid(tuple))
+			elog(ERROR, "cache lookup failed for function %u", objoid);
+
+		aclDatum = SysCacheGetAttr(PROCOID, tuple, Anum_pg_proc_proacl,
+								   &isNull);
+
+		/* Add the record, if any, for the top-level object */
+		if (!isNull)
+			recordExtensionInitPrivWorker(objoid, classoid, 0,
+										  DatumGetAclP(aclDatum));
+
+		ReleaseSysCache(tuple);
+	}
+	/* pg_type */
+	else if (classoid == TypeRelationId)
+	{
+		Datum		aclDatum;
+		bool		isNull;
+		HeapTuple	tuple;
+
+		tuple = SearchSysCache1(TYPEOID, ObjectIdGetDatum(objoid));
+		if (!HeapTupleIsValid(tuple))
+			elog(ERROR, "cache lookup failed for function %u", objoid);
+
+		aclDatum = SysCacheGetAttr(TYPEOID, tuple, Anum_pg_type_typacl,
+								   &isNull);
+
+		/* Add the record, if any, for the top-level object */
+		if (!isNull)
+			recordExtensionInitPrivWorker(objoid, classoid, 0,
+										  DatumGetAclP(aclDatum));
+
+		ReleaseSysCache(tuple);
+	}
+	else if (classoid == AccessMethodRelationId ||
+			 classoid == AggregateRelationId ||
+			 classoid == CastRelationId ||
+			 classoid == CollationRelationId ||
+			 classoid == ConversionRelationId ||
+			 classoid == EventTriggerRelationId ||
+			 classoid == OperatorRelationId ||
+			 classoid == OperatorClassRelationId ||
+			 classoid == OperatorFamilyRelationId ||
+			 classoid == NamespaceRelationId ||
+			 classoid == TSConfigRelationId ||
+			 classoid == TSDictionaryRelationId ||
+			 classoid == TSParserRelationId ||
+			 classoid == TSTemplateRelationId ||
+			 classoid == TransformRelationId
+		)
+	{
+		/* no ACL for these object types, so do nothing. */
+	}
+
+	/*
+	 * complain if we are given a class OID for a class that extensions don't
+	 * support or that we don't recognize.
+	 */
+	else
+	{
+		elog(ERROR, "unrecognized or unsupported class OID: %u", classoid);
+	}
+}
+
+/*
+ * For the object passed in, remove its ACL and the ACLs of any object subIds
+ * from pg_init_privs (via recordExtensionInitPrivWorker()).
+ */
+void
+removeExtObjInitPriv(Oid objoid, Oid classoid)
+{
+	/*
+	 * If this is a relation then we need to see if there are any sub-objects
+	 * (eg: columns) for it and, if so, be sure to call
+	 * recordExtensionInitPrivWorker() for each one.
+	 */
+	if (classoid == RelationRelationId)
+	{
+		Form_pg_class pg_class_tuple;
+		HeapTuple	tuple;
+
+		tuple = SearchSysCache1(RELOID, ObjectIdGetDatum(objoid));
+		if (!HeapTupleIsValid(tuple))
+			elog(ERROR, "cache lookup failed for relation %u", objoid);
+		pg_class_tuple = (Form_pg_class) GETSTRUCT(tuple);
+
+		/* Indexes don't have permissions */
+		if (pg_class_tuple->relkind == RELKIND_INDEX)
+			return;
+
+		/* Composite types don't have permissions either */
+		if (pg_class_tuple->relkind == RELKIND_COMPOSITE_TYPE)
+			return;
+
+		/*
+		 * If this isn't a sequence, index, or composite type then it's
+		 * possibly going to have columns associated with it that might have
+		 * ACLs.
+		 */
+		if (pg_class_tuple->relkind != RELKIND_SEQUENCE)
+		{
+			AttrNumber	curr_att;
+			AttrNumber	nattrs = pg_class_tuple->relnatts;
+
+			for (curr_att = 1; curr_att <= nattrs; curr_att++)
+			{
+				HeapTuple	attTuple;
+
+				attTuple = SearchSysCache2(ATTNUM,
+										   ObjectIdGetDatum(objoid),
+										   Int16GetDatum(curr_att));
+
+				if (!HeapTupleIsValid(attTuple))
+					continue;
+
+				/* when removing, remove all entires, even dropped columns */
+
+				recordExtensionInitPrivWorker(objoid, classoid, curr_att, NULL);
+
+				ReleaseSysCache(attTuple);
+			}
+		}
+
+		ReleaseSysCache(tuple);
+	}
+
+	/* Remove the record, if any, for the top-level object */
+	recordExtensionInitPrivWorker(objoid, classoid, 0, NULL);
+}
+
+/*
+ * Record initial ACL for an extension object
  *
  * Can be called at any time, we check if 'creating_extension' is set and, if
  * not, exit immediately.
@@ -5244,12 +5609,6 @@ get_user_default_acl(GrantObjectType objtype, Oid ownerId, Oid nsp_oid)
 static void
 recordExtensionInitPriv(Oid objoid, Oid classoid, int objsubid, Acl *new_acl)
 {
-	Relation	relation;
-	ScanKeyData key[3];
-	SysScanDesc scan;
-	HeapTuple	tuple;
-	HeapTuple	oldtuple;
-
 	/*
 	 * Generally, we only record the initial privileges when an extension is
 	 * being created, but because we don't actually use CREATE EXTENSION
@@ -5261,6 +5620,30 @@ recordExtensionInitPriv(Oid objoid, Oid classoid, int objsubid, Acl *new_acl)
 	if (!creating_extension && !binary_upgrade_record_init_privs)
 		return;
 
+	recordExtensionInitPrivWorker(objoid, classoid, objsubid, new_acl);
+}
+
+/*
+ * Record initial ACL for an extension object, worker.
+ *
+ * This will perform a wholesale replacement of the entire ACL for the object
+ * passed in, therefore be sure to pass in the complete new ACL to use.
+ *
+ * Generally speaking, do *not* use this function directly but instead use
+ * recordExtensionInitPriv(), which checks if 'creating_extension' is set.
+ * This function does *not* check if 'creating_extension' is set as it is also
+ * used when an object is added to or removed from an extension via ALTER
+ * EXTENSION ... ADD/DROP.
+ */
+static void
+recordExtensionInitPrivWorker(Oid objoid, Oid classoid, int objsubid, Acl *new_acl)
+{
+	Relation	relation;
+	ScanKeyData key[3];
+	SysScanDesc scan;
+	HeapTuple	tuple;
+	HeapTuple	oldtuple;
+
 	relation = heap_open(InitPrivsRelationId, RowExclusiveLock);
 
 	ScanKeyInit(&key[0],
@@ -5313,28 +5696,37 @@ recordExtensionInitPriv(Oid objoid, Oid classoid, int objsubid, Acl *new_acl)
 	}
 	else
 	{
-		/* No entry found, so add it. */
 		Datum		values[Natts_pg_init_privs];
 		bool		nulls[Natts_pg_init_privs];
 
-		MemSet(nulls, false, sizeof(nulls));
+		/*
+		 * Only add a new entry if the new ACL is non-NULL.
+		 *
+		 * If we are passed in a NULL ACL and no entry exists, we can just
+		 * fall through and do nothing.
+		 */
+		if (new_acl)
+		{
+			/* No entry found, so add it. */
+			MemSet(nulls, false, sizeof(nulls));
 
-		values[Anum_pg_init_privs_objoid - 1] = ObjectIdGetDatum(objoid);
-		values[Anum_pg_init_privs_classoid - 1] = ObjectIdGetDatum(classoid);
-		values[Anum_pg_init_privs_objsubid - 1] = Int32GetDatum(objsubid);
+			values[Anum_pg_init_privs_objoid - 1] = ObjectIdGetDatum(objoid);
+			values[Anum_pg_init_privs_classoid - 1] = ObjectIdGetDatum(classoid);
+			values[Anum_pg_init_privs_objsubid - 1] = Int32GetDatum(objsubid);
 
-		/* This function only handles initial privileges of extensions */
-		values[Anum_pg_init_privs_privtype - 1] =
-			CharGetDatum(INITPRIVS_EXTENSION);
+			/* This function only handles initial privileges of extensions */
+			values[Anum_pg_init_privs_privtype - 1] =
+				CharGetDatum(INITPRIVS_EXTENSION);
 
-		values[Anum_pg_init_privs_privs - 1] = PointerGetDatum(new_acl);
+			values[Anum_pg_init_privs_privs - 1] = PointerGetDatum(new_acl);
 
-		tuple = heap_form_tuple(RelationGetDescr(relation), values, nulls);
+			tuple = heap_form_tuple(RelationGetDescr(relation), values, nulls);
 
-		simple_heap_insert(relation, tuple);
+			simple_heap_insert(relation, tuple);
 
-		/* keep the catalog indexes up to date */
-		CatalogUpdateIndexes(relation, tuple);
+			/* keep the catalog indexes up to date */
+			CatalogUpdateIndexes(relation, tuple);
+		}
 	}
 
 	systable_endscan(scan);
diff --git a/src/backend/commands/extension.c b/src/backend/commands/extension.c
index 518fefcf09b..33b0de0a765 100644
--- a/src/backend/commands/extension.c
+++ b/src/backend/commands/extension.c
@@ -52,6 +52,7 @@
 #include "nodes/makefuncs.h"
 #include "storage/fd.h"
 #include "tcop/utility.h"
+#include "utils/acl.h"
 #include "utils/builtins.h"
 #include "utils/fmgroids.h"
 #include "utils/lsyscache.h"
@@ -3046,6 +3047,16 @@ ExecAlterExtensionContentsStmt(AlterExtensionContentsStmt *stmt,
 		 * OK, add the dependency.
 		 */
 		recordDependencyOn(&object, &extension, DEPENDENCY_EXTENSION);
+
+		/*
+		 * Also record the initial ACL on the object, if any.
+		 *
+		 * Note that this will handle the object's ACLs, as well as any ACLs
+		 * on object subIds.  (In other words, when the object is a table,
+		 * this will record the table's ACL and the ACLs for the columns on
+		 * the table, if any).
+		 */
+		recordExtObjInitPriv(object.objectId, object.classId);
 	}
 	else
 	{
@@ -3073,6 +3084,16 @@ ExecAlterExtensionContentsStmt(AlterExtensionContentsStmt *stmt,
 		 */
 		if (object.classId == RelationRelationId)
 			extension_config_remove(extension.objectId, object.objectId);
+
+		/*
+		 * Remove all the initial ACLs, if any.
+		 *
+		 * Note that this will remove the object's ACLs, as well as any ACLs
+		 * on object subIds.  (In other words, when the object is a table,
+		 * this will remove the table's ACL and the ACLs for the columns on
+		 * the table, if any).
+		 */
+		removeExtObjInitPriv(object.objectId, object.classId);
 	}
 
 	InvokeObjectPostAlterHook(ExtensionRelationId, extension.objectId, 0);