PG-1447 Assert there are no encrypted relations when using FILE_COPY

The FILE_COPY strategy of CREATE DATABASE does not use the SMGR so it
cannot properly copy and re-encrypt encrypted relations. So we check for
such relations and error out if any are found.

We look for encrypted relations simply by counting the number of keys in
the key map file of the database. This simple approach is fine even with
the risk of a left-over key from a crash or a bug due to the FILE_COPY
method being rare and that we therefore want to keep this code minimal.

contrib/pg_tde/expected/create_database.out

@@ -92,6 +92,14 @@ SELECT pg_tde_is_encrypted('test_plain_id_seq');
 (1 row)
 
 \c :regress_database
+CREATE DATABASE new_db_file_copy TEMPLATE template_db STRATEGY FILE_COPY;
+ERROR:  The FILE_COPY strategy cannot be used when there are encrypted objects in the template database: 3 objects found
+HINT:  Use the WAL_LOG strategy instead.
+\c template_db
+DROP TABLE test_enc;
+\c :regress_database
+CREATE DATABASE new_db_file_copy TEMPLATE template_db STRATEGY FILE_COPY;
+DROP DATABASE new_db_file_copy;
 DROP DATABASE new_db;
 DROP DATABASE template_db;
 DROP EXTENSION pg_tde;

contrib/pg_tde/sql/create_database.sql

@@ -45,6 +45,17 @@ SELECT pg_tde_is_encrypted('test_plain_id_seq');
 
 \c :regress_database
 
+CREATE DATABASE new_db_file_copy TEMPLATE template_db STRATEGY FILE_COPY;
+
+\c template_db
+
+DROP TABLE test_enc;
+
+\c :regress_database
+
+CREATE DATABASE new_db_file_copy TEMPLATE template_db STRATEGY FILE_COPY;
+
+DROP DATABASE new_db_file_copy;
 DROP DATABASE new_db;
 DROP DATABASE template_db;
 

contrib/pg_tde/src/access/pg_tde_tdemap.c

@@ -115,7 +115,7 @@ static int	pg_tde_open_file_basic(const char *tde_filename, int fileFlags, bool
 static void pg_tde_file_header_read(const char *tde_filename, int fd, TDEFileHeader *fheader, off_t *bytes_read);
 static bool pg_tde_read_one_map_entry(int fd, TDEMapEntry *map_entry, off_t *offset);
 static void pg_tde_read_one_map_entry2(int keydata_fd, int32 key_index, TDEMapEntry *map_entry, Oid databaseId);
-static int	pg_tde_open_file_read(const char *tde_filename, off_t *curr_pos);
+static int	pg_tde_open_file_read(const char *tde_filename, bool ignore_missing, off_t *curr_pos);
 static InternalKey *pg_tde_get_key_from_cache(const RelFileLocator *rlocator, uint32 key_type);
 static WALKeyCacheRec *pg_tde_add_wal_key_to_cache(InternalKey *cached_key, XLogRecPtr start_lsn);
 static InternalKey *pg_tde_put_key_into_cache(const RelFileLocator *locator, InternalKey *key);
@@ -599,7 +599,7 @@ pg_tde_perform_rotate_key(TDEPrincipalKey *principal_key, TDEPrincipalKey *new_p
 
 	pg_tde_set_db_file_path(principal_key->keyInfo.databaseId, old_path);
 
-	old_fd = pg_tde_open_file_read(old_path, &old_curr_pos);
+	old_fd = pg_tde_open_file_read(old_path, false, &old_curr_pos);
 	new_fd = keyrotation_init_file(&new_signed_key_info, new_path, old_path, &new_curr_pos);
 
 	/* Read all entries until EOF */
@@ -831,7 +831,7 @@ pg_tde_find_map_entry(const RelFileLocator *rlocator, uint32 key_type, char *db_
 
 	Assert(rlocator != NULL);
 
-	map_fd = pg_tde_open_file_read(db_map_path, &curr_pos);
+	map_fd = pg_tde_open_file_read(db_map_path, false, &curr_pos);
 
 	while (pg_tde_read_one_map_entry(map_fd, map_entry, &curr_pos))
 	{
@@ -847,6 +847,47 @@ pg_tde_find_map_entry(const RelFileLocator *rlocator, uint32 key_type, char *db_
 	return found;
 }
 
+/*
+ * Counts number of encrypted objects in a database.
+ *
+ * Does not check if objects actually exist but just that they have keys in
+ * the map file. For the only current caller, checking if we can use
+ * FILE_COPY, this is good enough but for other workloads where a false
+ * positive is more harmful this might not be.
+ *
+ * Works even if the database has no map file.
+ */
+int
+pg_tde_count_relations(Oid dbOid)
+{
+	char		db_map_path[MAXPGPATH];
+	LWLock	   *lock_pk = tde_lwlock_enc_keys();
+	File		map_fd;
+	off_t		curr_pos = 0;
+	TDEMapEntry map_entry;
+	int			count = 0;
+
+	pg_tde_set_db_file_path(dbOid, db_map_path);
+
+	LWLockAcquire(lock_pk, LW_SHARED);
+
+	map_fd = pg_tde_open_file_read(db_map_path, true, &curr_pos);
+	if (map_fd < 0)
+		return count;
+
+	while (pg_tde_read_one_map_entry(map_fd, &map_entry, &curr_pos))
+	{
+		if (map_entry.flags & TDE_KEY_TYPE_SMGR)
+			count++;
+	}
+
+	close(map_fd);
+
+	LWLockRelease(lock_pk);
+
+	return count;
+}
+
 bool
 pg_tde_verify_principal_key_info(TDESignedPrincipalKeyInfo *signed_key_info, const TDEPrincipalKey *principal_key)
 {
@@ -883,7 +924,7 @@ tde_decrypt_rel_key(TDEPrincipalKey *principal_key, TDEMapEntry *map_entry)
  * is raised.
  */
 static int
-pg_tde_open_file_read(const char *tde_filename, off_t *curr_pos)
+pg_tde_open_file_read(const char *tde_filename, bool ignore_missing, off_t *curr_pos)
 {
 	int			fd;
 	TDEFileHeader fheader;
@@ -891,7 +932,9 @@ pg_tde_open_file_read(const char *tde_filename, off_t *curr_pos)
 
 	Assert(LWLockHeldByMeInMode(tde_lwlock_enc_keys(), LW_SHARED) || LWLockHeldByMeInMode(tde_lwlock_enc_keys(), LW_EXCLUSIVE));
 
-	fd = pg_tde_open_file_basic(tde_filename, O_RDONLY | PG_BINARY, false);
+	fd = pg_tde_open_file_basic(tde_filename, O_RDONLY | PG_BINARY, ignore_missing);
+	if (ignore_missing && fd < 0)
+		return fd;
 
 	pg_tde_file_header_read(tde_filename, fd, &fheader, &bytes_read);
 	*curr_pos = bytes_read;
@@ -1144,7 +1187,7 @@ pg_tde_read_last_wal_key(void)
 	}
 	pg_tde_set_db_file_path(rlocator.dbOid, db_map_path);
 
-	fd = pg_tde_open_file_read(db_map_path, &read_pos);
+	fd = pg_tde_open_file_read(db_map_path, false, &read_pos);
 	fsize = lseek(fd, 0, SEEK_END);
 	/* No keys */
 	if (fsize == TDE_FILE_HEADER_SIZE)
@@ -1187,7 +1230,7 @@ pg_tde_fetch_wal_keys(XLogRecPtr start_lsn)
 
 	pg_tde_set_db_file_path(rlocator.dbOid, db_map_path);
 
-	fd = pg_tde_open_file_read(db_map_path, &read_pos);
+	fd = pg_tde_open_file_read(db_map_path, false, &read_pos);
 
 	keys_count = (lseek(fd, 0, SEEK_END) - TDE_FILE_HEADER_SIZE) / MAP_ENTRY_SIZE;
 

contrib/pg_tde/src/include/access/pg_tde_tdemap.h

@@ -111,6 +111,7 @@ pg_tde_set_db_file_path(Oid dbOid, char *path)
 }
 
 extern InternalKey *GetSMGRRelationKey(RelFileLocatorBackend rel);
+extern int	pg_tde_count_relations(Oid dbOid);
 
 extern void pg_tde_delete_tde_files(Oid dbOid);
 

contrib/pg_tde/src/include/pg_tde_event_capture.h

@@ -26,8 +26,8 @@ typedef struct TdeCreateEvent
 										 * based on earlier encryption status. */
 } TdeCreateEvent;
 
+extern void TdeEventCaptureInit(void);
 extern TdeCreateEvent *GetCurrentTdeCreateEvent(void);
-
 extern void validateCurrentEventTriggerState(bool mightStartTransaction);
 
 #endif

contrib/pg_tde/src/pg_tde.c

@@ -33,6 +33,7 @@
 #include "utils/builtins.h"
 #include "smgr/pg_tde_smgr.h"
 #include "catalog/tde_global_space.h"
+#include "pg_tde_event_capture.h"
 #include "utils/percona.h"
 #include "pg_tde_guc.h"
 #include "access/tableam.h"
@@ -111,6 +112,7 @@ _PG_init(void)
 	check_percona_api_version();
 
 	TdeGucInit();
+	TdeEventCaptureInit();
 
 	InitializePrincipalKeyInfo();
 	InitializeKeyProviderInfo();

contrib/pg_tde/src/pg_tde_event_capture.c

@@ -16,6 +16,7 @@
 #include "utils/builtins.h"
 #include "utils/lsyscache.h"
 #include "catalog/pg_class.h"
+#include "catalog/pg_database.h"
 #include "commands/defrem.h"
 #include "commands/sequence.h"
 #include "access/table.h"
@@ -24,8 +25,13 @@
 #include "catalog/namespace.h"
 #include "commands/event_trigger.h"
 #include "common/pg_tde_utils.h"
+#include "storage/lmgr.h"
+#include "tcop/utility.h"
+#include "utils/fmgroids.h"
+#include "utils/syscache.h"
 #include "pg_tde_event_capture.h"
 #include "pg_tde_guc.h"
+#include "access/pg_tde_tdemap.h"
 #include "catalog/tde_principal_key.h"
 #include "miscadmin.h"
 #include "access/tableam.h"
@@ -34,6 +40,7 @@
 /* Global variable that gets set at ddl start and cleard out at ddl end*/
 static TdeCreateEvent tdeCurrentCreateEvent = {.tid = {.value = 0}};
 
+static Oid	get_db_oid(const char *name);
 static void reset_current_tde_create_event(void);
 static Oid	get_tde_table_am_oid(void);
 
@@ -354,3 +361,160 @@ get_tde_table_am_oid(void)
 {
 	return get_table_am_oid("tde_heap", false);
 }
+
+static ProcessUtility_hook_type next_ProcessUtility_hook = NULL;
+
+/*
+ * Handles utility commands which we cannot handle in the event trigger.
+ */
+static void
+pg_tde_proccess_utility(PlannedStmt *pstmt,
+						const char *queryString,
+						bool readOnlyTree,
+						ProcessUtilityContext context,
+						ParamListInfo params,
+						QueryEnvironment *queryEnv,
+						DestReceiver *dest,
+						QueryCompletion *qc)
+{
+	Node	   *parsetree = pstmt->utilityStmt;
+
+	switch (nodeTag(parsetree))
+	{
+		case T_CreatedbStmt:
+			{
+				CreatedbStmt *stmt = castNode(CreatedbStmt, parsetree);
+				ListCell   *option;
+				char	   *dbtemplate = "template1";
+				char	   *strategy = "wal_log";
+
+				foreach(option, stmt->options)
+				{
+					DefElem    *defel = (DefElem *) lfirst(option);
+
+					if (strcmp(defel->defname, "template") == 0)
+						dbtemplate = defGetString(defel);
+					else if (strcmp(defel->defname, "strategy") == 0)
+						strategy = defGetString(defel);
+				}
+
+				if (pg_strcasecmp(strategy, "file_copy") == 0)
+				{
+					Oid			dbOid = get_db_oid(dbtemplate);
+
+					if (dbOid != InvalidOid)
+					{
+						int			count = pg_tde_count_relations(dbOid);
+
+						if (count > 0)
+							ereport(ERROR,
+									errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
+									errmsg("The FILE_COPY strategy cannot be used when there are encrypted objects in the template database: %d objects found", count),
+									errhint("Use the WAL_LOG strategy instead."));
+					}
+				}
+			}
+			break;
+		default:
+			break;
+	}
+
+	if (next_ProcessUtility_hook)
+		(*next_ProcessUtility_hook) (pstmt, queryString, readOnlyTree,
+									 context, params, queryEnv,
+									 dest, qc);
+	else
+		standard_ProcessUtility(pstmt, queryString, readOnlyTree,
+								context, params, queryEnv,
+								dest, qc);
+}
+
+void
+TdeEventCaptureInit(void)
+{
+	next_ProcessUtility_hook = ProcessUtility_hook;
+	ProcessUtility_hook = pg_tde_proccess_utility;
+}
+
+/*
+ * A stripped down version of get_db_info() from src/backend/commands/dbcommands.c
+ */
+static Oid
+get_db_oid(const char *name)
+{
+	Oid			resDbOid = InvalidOid;
+	Relation	relation;
+
+	Assert(name);
+
+	relation = table_open(DatabaseRelationId, AccessShareLock);
+
+	/*
+	 * Loop covers the rare case where the database is renamed before we can
+	 * lock it.  We try again just in case we can find a new one of the same
+	 * name.
+	 */
+	for (;;)
+	{
+		ScanKeyData scanKey;
+		SysScanDesc scan;
+		HeapTuple	tuple;
+		Oid			dbOid;
+
+		/*
+		 * there's no syscache for database-indexed-by-name, so must do it the
+		 * hard way
+		 */
+		ScanKeyInit(&scanKey,
+					Anum_pg_database_datname,
+					BTEqualStrategyNumber, F_NAMEEQ,
+					CStringGetDatum(name));
+
+		scan = systable_beginscan(relation, DatabaseNameIndexId, true,
+								  NULL, 1, &scanKey);
+
+		tuple = systable_getnext(scan);
+
+		if (!HeapTupleIsValid(tuple))
+		{
+			/* definitely no database of that name */
+			systable_endscan(scan);
+			break;
+		}
+
+		dbOid = ((Form_pg_database) GETSTRUCT(tuple))->oid;
+
+		systable_endscan(scan);
+
+		/*
+		 * Now that we have a database OID, we can try to lock the DB.
+		 */
+		LockSharedObject(DatabaseRelationId, dbOid, 0, AccessExclusiveLock);
+
+		/*
+		 * And now, re-fetch the tuple by OID.  If it's still there and still
+		 * the same name, we win; else, drop the lock and loop back to try
+		 * again.
+		 */
+		tuple = SearchSysCache1(DATABASEOID, ObjectIdGetDatum(dbOid));
+		if (HeapTupleIsValid(tuple))
+		{
+			Form_pg_database dbform = (Form_pg_database) GETSTRUCT(tuple);
+
+			if (strcmp(name, NameStr(dbform->datname)) == 0)
+			{
+				resDbOid = dbOid;
+				ReleaseSysCache(tuple);
+				break;
+			}
+			/* can only get here if it was just renamed */
+			ReleaseSysCache(tuple);
+		}
+
+		UnlockSharedObject(DatabaseRelationId, dbOid, 0, AccessExclusiveLock);
+	}
+
+	table_close(relation, AccessShareLock);
+
+	return resDbOid;
+}