Handle keyring_provider

src/access/pg_tde_tdemap.c

@@ -597,6 +597,7 @@ pg_tde_free_key_map_entry(const RelFileLocator *rlocator, off_t offset)
 			pg_tde_process_map_entry(NULL, db_map_path, &start, false) == -1)
 	{
 		pg_tde_delete_tde_files(rlocator->dbOid, rlocator->spcOid);
+		cleanup_key_provider_info(rlocator->dbOid, rlocator->spcOid);
 	}
 }
 
@@ -851,30 +852,61 @@ pg_tde_move_rel_key(const RelFileLocator *newrlocator, const RelFileLocator *old
 	RelKeyData 	*rel_key;
 	RelKeyData 	*enc_key;
 	TDEPrincipalKey *principal_key;
+	KeyringProvideRecord provider_rec;
+    GenericKeyring *keyring;
 	XLogRelKey	xlrec;
 	char		db_map_path[MAXPGPATH] = {0};
+	char		db_keydata_path[MAXPGPATH] = {0};
 	off_t		offset = 0;
+	int32		key_index = 0;
+
+	pg_tde_set_db_file_paths(oldrlocator, db_map_path, db_keydata_path);
 
 	LWLockAcquire(tde_lwlock_enc_keys(), LW_EXCLUSIVE);
-	
+
 	principal_key = GetPrincipalKey(oldrlocator->dbOid, oldrlocator->spcOid, LW_EXCLUSIVE);
-	rel_key = GetRelationKey(*oldrlocator);
-	Assert(rel_key);
-	enc_key = tde_encrypt_rel_key(principal_key, rel_key, newrlocator);
-	pg_tde_write_key_map_entry(newrlocator, enc_key, &principal_key->keyInfo);
-	pg_tde_put_key_into_cache(newrlocator->relNumber, rel_key);
-
-	pg_tde_free_key_map_entry(oldrlocator, offset);
+	Assert(principal_key);
 
-	LWLockRelease(tde_lwlock_enc_keys());
+	/* 
+	 * copy kering provider info
+	 */
+	keyring = GetKeyProviderByID(principal_key->keyInfo.keyringId, oldrlocator->dbOid, oldrlocator->spcOid);
+	Assert(keyring);
+	memcpy(provider_rec.provider_name, keyring->provider_name, sizeof(keyring->provider_name));
+	provider_rec.provider_type = keyring->type;
+	memcpy(provider_rec.options, keyring->options, sizeof(keyring->options));
+	copy_key_provider_info(&provider_rec, newrlocator->dbOid, newrlocator->spcOid, true);
 
+	principal_key->keyInfo.keyringId = provider_rec.provider_id;
+
+	key_index = pg_tde_process_map_entry(oldrlocator, db_map_path, &offset, false);
+	Assert(key_index != -1);
+	/* 
+	 * Re-encrypt relation key. We don't use internal_key cache to avoid locking
+	 * complications.
+	 */
+	enc_key = pg_tde_read_keydata(db_keydata_path, key_index, principal_key);
+	rel_key = tde_decrypt_rel_key(principal_key, enc_key, oldrlocator);
+	enc_key = tde_encrypt_rel_key(principal_key, rel_key, newrlocator);
 
 	xlrec.rlocator = *newrlocator;
 	xlrec.relKey = *enc_key;
+	xlrec.pkInfo = principal_key->keyInfo;
 	XLogBeginInsert();
 	XLogRegisterData((char *) &xlrec, sizeof(xlrec));
 	XLogInsert(RM_TDERMGR_ID, XLOG_TDE_ADD_RELATION_KEY);
 
+	pg_tde_write_key_map_entry(newrlocator, enc_key, &principal_key->keyInfo);
+	pg_tde_put_key_into_cache(newrlocator->relNumber, rel_key);
+
+	XLogBeginInsert();
+	XLogRegisterData((char *) oldrlocator, sizeof(RelFileLocator));
+	XLogInsert(RM_TDERMGR_ID, XLOG_TDE_FREE_MAP_ENTRY);
+
+	pg_tde_free_key_map_entry(oldrlocator, offset);
+
+	LWLockRelease(tde_lwlock_enc_keys());
+
 	pfree(enc_key);
 }
 
@@ -1044,7 +1076,7 @@ pg_tde_process_map_entry(const RelFileLocator *rlocator, char *db_map_path, off_
 
 /*
  * Open the file and read the required key data from file and return encrypted key.
- * The caller should hold 
+ * The caller should hold a tde_lwlock_enc_keys lock
  */
 static RelKeyData *
 pg_tde_read_keydata(char *db_keydata_path, int32 key_index, TDEPrincipalKey *principal_key)

src/access/pg_tde_xlog.c

@@ -37,10 +37,14 @@ tdeheap_rmgr_redo(XLogReaderState *record)
 
 	if (info == XLOG_TDE_ADD_RELATION_KEY)
 	{
+		TDEPrincipalKeyInfo *pk = NULL;
 		XLogRelKey *xlrec = (XLogRelKey *) XLogRecGetData(record);
 
+		if (xlrec->pkInfo.databaseId != 0)
+			pk = &xlrec->pkInfo;
+
 		LWLockAcquire(tde_lwlock_enc_keys(), LW_EXCLUSIVE);
-		pg_tde_write_key_map_entry(&xlrec->rlocator, &xlrec->relKey, NULL);
+		pg_tde_write_key_map_entry(&xlrec->rlocator, &xlrec->relKey, pk);
 		LWLockRelease(tde_lwlock_enc_keys());
 	}
 	else if (info == XLOG_TDE_ADD_PRINCIPAL_KEY)
@@ -72,6 +76,16 @@ tdeheap_rmgr_redo(XLogReaderState *record)
 		xl_tde_perform_rotate_key(xlrec);
 		LWLockRelease(tde_lwlock_enc_keys());
 	}
+
+	else if (info == XLOG_TDE_FREE_MAP_ENTRY)
+	{
+		off_t		offset = 0;
+		RelFileLocator *xlrec = (RelFileLocator *) XLogRecGetData(record);
+
+		LWLockAcquire(tde_lwlock_enc_keys(), LW_EXCLUSIVE);
+		pg_tde_free_key_map_entry(xlrec, offset);
+		LWLockRelease(tde_lwlock_enc_keys());
+	}
 	else
 	{
 		elog(PANIC, "pg_tde_redo: unknown op code %u", info);

src/catalog/tde_global_space.c

@@ -114,7 +114,7 @@ init_default_keyring(void)
 		 * TODO: should we remove it automaticaly on
 		 * pg_tde_rotate_principal_key() ?
 		 */
-		save_new_key_provider_info(&provider, GLOBAL_DATA_TDE_OID, GLOBALTABLESPACE_OID, true);
+		save_new_key_provider_info(&provider, GLOBAL_DATA_TDE_OID, GLOBALTABLESPACE_OID, false);
 		elog(INFO,
 			 "default keyring has been created for the global tablespace (WAL)."
 			 " Change it with pg_tde_add_key_provider_* and run pg_tde_rotate_principal_key."

src/catalog/tde_keyring.c

@@ -78,7 +78,8 @@ static void key_provider_startup_cleanup(int tde_tbl_count, XLogExtensionInstall
 static const char *get_keyring_provider_typename(ProviderType p_type);
 static uint32 write_key_provider_info(KeyringProvideRecord *provider, 
 									Oid database_id, Oid tablespace_id,
-									off_t position, bool redo, bool recovery);
+									off_t position, bool error_if_exists, 
+									bool write_xlog);
 
 static Size initialize_shared_state(void *start_address);
 static Size required_shared_mem_size(void);
@@ -195,7 +196,7 @@ GetKeyProviderByName(const char *provider_name, Oid dbOid, Oid spcOid)
 
 static uint32
 write_key_provider_info(KeyringProvideRecord *provider, Oid database_id, 
-					Oid tablespace_id, off_t position, bool redo, bool recovery)
+					Oid tablespace_id, off_t position, bool error_if_exists, bool write_xlog)
 {
 	off_t bytes_written = 0;
 	off_t curr_pos = 0;
@@ -218,7 +219,7 @@ write_key_provider_info(KeyringProvideRecord *provider, Oid database_id,
 			(errcode_for_file_access(),
 				errmsg("could not open tde file \"%s\": %m", kp_info_path)));
 	}
-	if (!redo)
+	if (position == -1)
 	{
 		/* we also need to verify the name conflict and generate the next provider ID */
 		while (fetch_next_key_provider(fd, &curr_pos, &existing_provider))
@@ -227,9 +228,15 @@ write_key_provider_info(KeyringProvideRecord *provider, Oid database_id,
 			{
 				close(fd);
 				LWLockRelease(tde_provider_info_lock());
-				ereport(ERROR,
+				ereport(error_if_exists ? ERROR : DEBUG1,
 						(errcode(ERRCODE_DUPLICATE_OBJECT),
 						errmsg("key provider \"%s\" already exists", provider->provider_name)));
+
+				if (!error_if_exists)
+				{
+					provider->provider_id = existing_provider.provider_id;
+					return provider->provider_id;
+				}
 			}
 			if (max_provider_id < existing_provider.provider_id)
 				max_provider_id = existing_provider.provider_id;
@@ -240,7 +247,7 @@ write_key_provider_info(KeyringProvideRecord *provider, Oid database_id,
 		/* emit the xlog here. So that we can handle partial file write errors
 		 * but cannot make new WAL entries during recovery.
 		 */
-		if (!recovery)
+		if (write_xlog)
 		{
 			KeyringProviderXLRecord xlrec;
 
@@ -293,9 +300,18 @@ write_key_provider_info(KeyringProvideRecord *provider, Oid database_id,
  * Save the key provider info to the file
  */
 uint32
-save_new_key_provider_info(KeyringProvideRecord* provider, Oid databaseId, Oid tablespaceId, bool recovery)
+save_new_key_provider_info(KeyringProvideRecord* provider, Oid databaseId, Oid tablespaceId, bool write_xlog)
+{
+	return write_key_provider_info(provider, databaseId, tablespaceId, -1, true, write_xlog);
+}
+
+/*
+ * Save the key provider info to the file but don't fail if it is already exists.
+ */
+uint32
+copy_key_provider_info(KeyringProvideRecord* provider, Oid newdatabaseId, Oid newtablespaceId, bool write_xlog)
 {
-	return write_key_provider_info(provider, databaseId, tablespaceId, 0, false, recovery);
+	return write_key_provider_info(provider, newdatabaseId, newtablespaceId, -1, false, write_xlog);
 }
 
 uint32
@@ -334,7 +350,7 @@ pg_tde_add_key_provider_internal(PG_FUNCTION_ARGS)
 	strncpy(provider.options, options, sizeof(provider.options));
 	strncpy(provider.provider_name, provider_name, sizeof(provider.provider_name));
 	provider.provider_type = get_keyring_provider_from_typename(provider_type);
-	save_new_key_provider_info(&provider, dbOid, spcOid, false);
+	save_new_key_provider_info(&provider, dbOid, spcOid, true);
 
 	PG_RETURN_INT32(provider.provider_id);
 }

src/catalog/tde_principal_key.c

@@ -787,6 +787,12 @@ get_principal_key_from_keyring(Oid dbOid, Oid spcOid)
     if (spcOid != GLOBALTABLESPACE_OID)
     {
         push_principal_key_to_cache(principalKey);
+
+        /* If we do store key in cache we want to return a cache reference 
+         * rather then a palloc'ed copy.
+         */
+        pfree(principalKey);
+        principalKey = get_principal_key_from_cache(dbOid);
     }
 #endif
 

src/include/access/pg_tde_tdemap.h

@@ -31,6 +31,7 @@ typedef struct XLogRelKey
 {
 	RelFileLocator  rlocator;
 	RelKeyData      relKey;
+	TDEPrincipalKeyInfo pkInfo;
 } XLogRelKey;
 
 extern RelKeyData* pg_tde_create_key_map_entry(const RelFileLocator *newrlocator);

src/include/access/pg_tde_xlog.h

@@ -21,6 +21,7 @@
 #define XLOG_TDE_EXTENSION_INSTALL_KEY	0x20
 #define XLOG_TDE_ROTATE_KEY				0x30
 #define XLOG_TDE_ADD_KEY_PROVIDER_KEY 	0x40
+#define XLOG_TDE_FREE_MAP_ENTRY		 	0x50
 
 /* TODO: ID has to be registedred and changed: https://wiki.postgresql.org/wiki/CustomWALResourceManagers */
 #define RM_TDERMGR_ID	RM_EXPERIMENTAL_ID

src/include/catalog/tde_keyring.h

@@ -78,7 +78,12 @@ extern GenericKeyring *GetKeyProviderByID(int provider_id, Oid dbOid, Oid spcOid
 extern ProviderType get_keyring_provider_from_typename(char *provider_type);
 extern void cleanup_key_provider_info(Oid databaseId, Oid tablespaceId);
 extern void InitializeKeyProviderInfo(void);
-extern uint32 save_new_key_provider_info(KeyringProvideRecord *provider, Oid databaseId, Oid tablespaceId, bool recovery);
+extern uint32 save_new_key_provider_info(KeyringProvideRecord *provider, 
+											Oid databaseId, Oid tablespaceId, 
+											bool write_xlog);
+extern uint32 copy_key_provider_info(KeyringProvideRecord* provider, 
+										Oid newdatabaseId, Oid newtablespaceId, 
+										bool write_xlog);
 extern uint32 redo_key_provider_info(KeyringProviderXLRecord *xlrec);
 
 extern bool ParseKeyringJSONOptions(ProviderType provider_type, void *out_opts,

src/transam/pg_tde_xact_handler.c

@@ -157,8 +157,6 @@ reassign_pending_deletes_to_parent_xact(void)
         if (pending->nestLevel == nestLevel)
             pending->nestLevel--;
     }
-
-    LWLockRelease(tde_lwlock_enc_keys());
 }
 
 /*