Move tde keys and keyring when chaging tablespace

When a relation moved to a new location it causes the change of relfilenode id
for it. Hence we must re-encrypt and store its internal key with the new id.
Also, we have to store the changed internal key in the new physical location,
and copy there principal key info and keyring data.

Fixes https://perconadev.atlassian.net/browse/PG-1038

expected/tablespace.out

@@ -1 +1,41 @@
-// TODO
+\set tde_am tde_heap
+\i sql/tablespace.inc
+CREATE EXTENSION pg_tde;
+SELECT pg_tde_add_key_provider_file('file-vault','/tmp/pg_tde_test_keyring.per');
+ pg_tde_add_key_provider_file 
+------------------------------
+                            1
+(1 row)
+
+SELECT pg_tde_set_principal_key('test-db-principal-key','file-vault');
+ pg_tde_set_principal_key 
+--------------------------
+ t
+(1 row)
+
+CREATE TABLE test(num1 bigint, num2 double precision, t text) USING :tde_am;
+INSERT INTO test(num1, num2, t)
+  SELECT round(random()*100), random(), 'text'
+  FROM generate_series(1, 10) s(i);
+CREATE INDEX test_idx ON test(num1);
+SET allow_in_place_tablespaces = true;
+CREATE TABLESPACE test_tblspace LOCATION '';
+ALTER TABLE test SET TABLESPACE test_tblspace;
+SELECT count(*) FROM test;
+ count 
+-------
+    10
+(1 row)
+
+ALTER TABLE test SET TABLESPACE pg_default;
+REINDEX (TABLESPACE test_tblspace, CONCURRENTLY) TABLE test;
+INSERT INTO test VALUES (110, 2);
+SELECT * FROM test WHERE num1=110;
+ num1 | num2 | t 
+------+------+---
+  110 |    2 | 
+(1 row)
+
+DROP TABLE test;
+DROP TABLESPACE test_tblspace;
+DROP EXTENSION pg_tde;

expected/tablespace_basic.out

@@ -0,0 +1,41 @@
+\set tde_am tde_heap_basic
+\i sql/tablespace.inc
+CREATE EXTENSION pg_tde;
+SELECT pg_tde_add_key_provider_file('file-vault','/tmp/pg_tde_test_keyring.per');
+ pg_tde_add_key_provider_file 
+------------------------------
+                            1
+(1 row)
+
+SELECT pg_tde_set_principal_key('test-db-principal-key','file-vault');
+ pg_tde_set_principal_key 
+--------------------------
+ t
+(1 row)
+
+CREATE TABLE test(num1 bigint, num2 double precision, t text) USING :tde_am;
+INSERT INTO test(num1, num2, t)
+  SELECT round(random()*100), random(), 'text'
+  FROM generate_series(1, 10) s(i);
+CREATE INDEX test_idx ON test(num1);
+SET allow_in_place_tablespaces = true;
+CREATE TABLESPACE test_tblspace LOCATION '';
+ALTER TABLE test SET TABLESPACE test_tblspace;
+SELECT count(*) FROM test;
+ count 
+-------
+    10
+(1 row)
+
+ALTER TABLE test SET TABLESPACE pg_default;
+REINDEX (TABLESPACE test_tblspace, CONCURRENTLY) TABLE test;
+INSERT INTO test VALUES (110, 2);
+SELECT * FROM test WHERE num1=110;
+ num1 | num2 | t 
+------+------+---
+  110 |    2 | 
+(1 row)
+
+DROP TABLE test;
+DROP TABLESPACE test_tblspace;
+DROP EXTENSION pg_tde;

meson.build

@@ -96,6 +96,7 @@ sql_tests = [
       'trigger_on_view_basic',
       'change_access_method_basic',
       'insert_update_delete_basic',
+      'tablespace_basic',
       'vault_v2_test_basic',
 ]
 

sql/tablespace.inc

@@ -0,0 +1,26 @@
+CREATE EXTENSION pg_tde;
+
+SELECT pg_tde_add_key_provider_file('file-vault','/tmp/pg_tde_test_keyring.per');
+SELECT pg_tde_set_principal_key('test-db-principal-key','file-vault');
+
+CREATE TABLE test(num1 bigint, num2 double precision, t text) USING :tde_am;
+INSERT INTO test(num1, num2, t)
+  SELECT round(random()*100), random(), 'text'
+  FROM generate_series(1, 10) s(i);
+CREATE INDEX test_idx ON test(num1);
+
+SET allow_in_place_tablespaces = true;
+CREATE TABLESPACE test_tblspace LOCATION '';
+
+ALTER TABLE test SET TABLESPACE test_tblspace;
+SELECT count(*) FROM test;
+ALTER TABLE test SET TABLESPACE pg_default;
+
+REINDEX (TABLESPACE test_tblspace, CONCURRENTLY) TABLE test;
+INSERT INTO test VALUES (110, 2);
+
+SELECT * FROM test WHERE num1=110;
+
+DROP TABLE test;
+DROP TABLESPACE test_tblspace;
+DROP EXTENSION pg_tde;

sql/tablespace.sql

@@ -1,26 +1,2 @@
-CREATE EXTENSION pg_tde;
-
-SELECT  * FROM pg_tde_principal_key_info();
-
-SELECT pg_tde_add_key_provider_file('file-vault','/tmp/pg_tde_test_keyring.per');
-SELECT pg_tde_set_principal_key('test-db-principal-key','file-vault');
-
-
-CREATE TABLE test(num1 bigint, num2 double precision, t text);
-INSERT INTO test(num1, num2, t)
-  SELECT round(random()*100), random(), 'text'
-  FROM generate_series(1, 10) s(i);
-CREATE INDEX test_idx ON test(num1);
-
-SET allow_in_place_tablespaces = true;
-CREATE TABLESPACE test_tblspace LOCATION '';
-
-ALTER TABLE test SET TABLESPACE test_tblspace;
-ALTER TABLE test SET TABLESPACE pg_default;
-
-REINDEX (TABLESPACE test_tblspace, CONCURRENTLY) TABLE test;
-INSERT INTO test VALUES (10, 2);
-
-DROP TABLE test;
-DROP TABLESPACE test_tblspace;
-DROP EXTENSION pg_tde;
+\set tde_am tde_heap
+\i sql/tablespace.inc

sql/tablespace_basic.sql

@@ -0,0 +1,2 @@
+\set tde_am tde_heap_basic
+\i sql/tablespace.inc

src/access/pg_tde_tdemap.c

@@ -555,24 +555,23 @@ pg_tde_delete_key_map_entry(const RelFileLocator *rlocator)
  *
  * The offset allows us to simply seek to the desired location and mark the entry
  * as MAP_ENTRY_FREE without needing any further processing.
+ * 
+ * A caller should hold an EXCLUSIVE tde_lwlock_enc_keys lock.
  */
 void
 pg_tde_free_key_map_entry(const RelFileLocator *rlocator, off_t offset)
 {
 	int32	key_index = 0;
-	LWLock	*lock_files = tde_lwlock_enc_keys();
 	char	db_map_path[MAXPGPATH] = {0};
-	char	db_keydata_path[MAXPGPATH] = {0};
+	off_t	start = 0;
 
 	Assert(rlocator);
 
 	/* Get the file paths */
-	pg_tde_set_db_file_paths(rlocator->dbOid, rlocator->spcOid, db_map_path, db_keydata_path);
+	pg_tde_set_db_file_paths(rlocator->dbOid, rlocator->spcOid, db_map_path, NULL);
 
 	/* Remove the map entry if found */
-	LWLockAcquire(lock_files, LW_EXCLUSIVE);
 	key_index = pg_tde_process_map_entry(rlocator, db_map_path, &offset, true);
-	LWLockRelease(lock_files);
 
 	if (key_index == -1)
 	{
@@ -582,7 +581,17 @@ pg_tde_free_key_map_entry(const RelFileLocator *rlocator, off_t offset)
 						rlocator->relNumber,
 						db_map_path)));
 
-		return;
+	}
+	/* 
+	 * Remove TDE files it was the last TDE relation in a custom tablespace.
+	 * DROP TABLESPACE needs an empty dir.
+	 */
+	if (rlocator->spcOid != GLOBALTABLESPACE_OID && 
+			rlocator->spcOid != DEFAULTTABLESPACE_OID &&
+			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);
 	}
 }
 
@@ -822,6 +831,87 @@ pg_tde_write_map_keydata_files(off_t map_size, char *m_file_data, off_t keydata_
 	return !is_err;
 }
 
+/* 
+ * Move relation's key to the new physical location and cache it with the new
+ * relfilenode. It recreates *.map and *.dat files with the old principal key
+ * and re-encrypted with the new relfilenode internal key. And copies the
+ * old keyring to the new location.
+ * Needed by ALTER TABLE SET TABLESPACE for example.
+ */
+bool
+pg_tde_move_rel_key(const RelFileLocator *newrlocator, const RelFileLocator *oldrlocator)
+{
+	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->dbOid, oldrlocator->spcOid, db_map_path, db_keydata_path);
+
+	LWLockAcquire(tde_lwlock_enc_keys(), LW_EXCLUSIVE);
+
+	principal_key = GetPrincipalKey(oldrlocator->dbOid, oldrlocator->spcOid, LW_EXCLUSIVE);
+	Assert(principal_key);
+
+	/* 
+	 * Copy kering provider info.
+	 *
+	 * TODO: we can potentially avoid moving keyring and key tde files keeping 
+	 * these files always in dbOid+MyDatabaseTableSpace path. But the
+	 * background writer isn't aware of MyDatabaseTableSpace hence it won't
+	 * work with SMGR -> tde_heap. Revisit this after chages in SMGR (mdcreate)
+	 * interface.
+	 */
+	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);
+
+	/* Clean-up map/dat entries. It will also remove physical files (*.map,
+	 * *.dat and keyring) if it was the last tde_heap_basic relation in the old
+	 * locator AND it was a custom tablespace.
+	 */
+	pg_tde_free_key_map_entry(oldrlocator, offset);
+
+	LWLockRelease(tde_lwlock_enc_keys());
+
+	pfree(enc_key);
+}
+
 #endif		/* !FRONTEND */
 
 /*
@@ -988,7 +1078,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)
@@ -1167,7 +1257,6 @@ pg_tde_read_one_map_entry(File map_file, const RelFileLocator *rlocator, int fla
 	return found;
 }
 
-
 /*
  * Reads a single keydata from the file.
  */

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

@@ -113,7 +113,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."