Ensure decrypted internal keys never swapped

We have already had a locked memory for internal keys. But we read and
created keys into usually palloced memory before copying them to the
locked memory.

This commit changes the approach. Now, the user must allocate (acquire)
a locked memory first and use it as a buffer for decryption/creating
the key. This locked memory is kinda generic, so the user requests an
amount of bytes (rather than objects). Although currently, during the
key search, we assume that everything the memory contains only
`RelKeyCacheRec` objects. With two exceptions: 1) server creates its
own locked page to store the WAL write key;
2) `pg_tde_perform_rotate_key` cheekily allocates a space there for
the keys re-encryption but releases this memory when it's done.

In future, if we switch to the full _map files encryption, we can gulp
(mmap) the whole file into that memory.

Fixes PG-1445

Author: dAdAbird

contrib/pg_tde/src/access/pg_tde_tdemap.c

@@ -70,12 +70,7 @@ typedef struct EncryptionStateData
 
 static EncryptionStateData *EncryptionState = NULL;
 
-/* TODO: can be swapped out to the disk */
-static InternalKey EncryptionKey =
-{
-	.rel_type = MAP_ENTRY_EMPTY,
-	.start_lsn = InvalidXLogRecPtr,
-};
+static InternalKey *EncryptionKey = NULL;
 static void *EncryptionCryptCtx = NULL;
 
 static int	XLOGChooseNumBuffers(void);
@@ -165,48 +160,67 @@ TDEXLogWriteEncryptedPages(int fd, const void *buf, size_t count, off_t offset,
 						   TimeLineID tli, XLogSegNo segno)
 {
 	char		iv_prefix[16];
-	InternalKey *key = &EncryptionKey;
 	char	   *enc_buff = EncryptionState->segBuf;
 
 	Assert(count <= TDEXLogEncryptBuffSize());
 
 #ifdef TDE_XLOG_DEBUG
 	elog(DEBUG1, "write encrypted WAL, size: %lu, offset: %ld [%lX], seg: %X/%X, key_start_lsn: %X/%X",
-		 count, offset, offset, LSN_FORMAT_ARGS(segno), LSN_FORMAT_ARGS(key->start_lsn));
+		 count, offset, offset, LSN_FORMAT_ARGS(segno), LSN_FORMAT_ARGS(EncryptionKey->start_lsn));
 #endif
 
-	CalcXLogPageIVPrefix(tli, segno, key->base_iv, iv_prefix);
+	CalcXLogPageIVPrefix(tli, segno, EncryptionKey->base_iv, iv_prefix);
 	PG_TDE_ENCRYPT_DATA(iv_prefix, offset,
 						(char *) buf, count,
-						enc_buff, key, &EncryptionCryptCtx);
+						enc_buff, EncryptionKey, &EncryptionCryptCtx);
 
 	return pg_pwrite(fd, enc_buff, count, offset);
 }
 
-#endif							/* !FRONTEND */
-
-void
-TDEXLogSmgrInit(void)
+static void
+tde_init_wal_encryption_key(void)
 {
-#ifndef FRONTEND
-	/* TODO: move to the separate func, it's not an SMGR init */
 	InternalKey *key = pg_tde_read_last_wal_key();
 
 	/* TDOO: clean-up this mess */
 	if ((!key && EncryptXLog) || (key &&
 								  ((key->rel_type & TDE_KEY_TYPE_WAL_ENCRYPTED && !EncryptXLog) ||
 								   (key->rel_type & TDE_KEY_TYPE_WAL_UNENCRYPTED && EncryptXLog))))
 	{
-		pg_tde_create_wal_key(&EncryptionKey, &GLOBAL_SPACE_RLOCATOR(XLOG_TDE_OID),
-							  (EncryptXLog ? TDE_KEY_TYPE_WAL_ENCRYPTED : TDE_KEY_TYPE_WAL_UNENCRYPTED));
+		key = pg_tde_create_wal_key(&GLOBAL_SPACE_RLOCATOR(XLOG_TDE_OID),
+									(EncryptXLog ? TDE_KEY_TYPE_WAL_ENCRYPTED : TDE_KEY_TYPE_WAL_UNENCRYPTED));
 	}
 	else if (key)
 	{
-		EncryptionKey = *key;
-		pfree(key);
-		pg_atomic_write_u64(&EncryptionState->enc_key_lsn, EncryptionKey.start_lsn);
+		pg_atomic_write_u64(&EncryptionState->enc_key_lsn, key->start_lsn);
 	}
 
+
+	/*
+	 * The relation cache might be moved to the new address space when growing
+	 * later on. It happens when there are a lot of tde_heap relations during
+	 * the recovery. So copy the key to the own locked mem. We could rewrite
+	 * the EncryptionKey pointer after the grow, but it'd be more complicated
+	 * and error-prone. We can sacrifice a couple of Kb.
+	 *
+	 * TODO: genralize and use `tde_alloc_safe_mem` instead?
+	 */
+	if (key)
+	{
+		long		pageSize = sysconf(_SC_PAGESIZE);
+
+		EncryptionKey = (InternalKey *) MemoryContextAllocAligned(TopMemoryContext, pageSize, pageSize, MCXT_ALLOC_ZERO);
+		memcpy(EncryptionKey, key, sizeof(InternalKey));
+	}
+}
+
+#endif							/* !FRONTEND */
+
+void
+TDEXLogInit(void)
+{
+#ifndef FRONTEND
+	tde_init_wal_encryption_key();
 	pg_tde_set_db_file_path(GLOBAL_SPACE_RLOCATOR(XLOG_TDE_OID).dbOid, EncryptionState->db_map_path);
 
 #endif
@@ -224,15 +238,14 @@ tdeheap_xlog_seg_write(int fd, const void *buf, size_t count, off_t offset,
 	 *
 	 * This func called with WALWriteLock held, so no need in any extra sync.
 	 */
-	if (EncryptionKey.rel_type & TDE_KEY_TYPE_GLOBAL &&
-		pg_atomic_read_u64(&EncryptionState->enc_key_lsn) == 0)
+	if (EncryptionKey && pg_atomic_read_u64(&EncryptionState->enc_key_lsn) == 0)
 	{
 		XLogRecPtr	lsn;
 
 		XLogSegNoOffsetToRecPtr(segno, offset, wal_segment_size, lsn);
 
 		pg_tde_wal_last_key_set_lsn(lsn, EncryptionState->db_map_path);
-		EncryptionKey.start_lsn = lsn;
+		EncryptionKey->start_lsn = lsn;
 		pg_atomic_write_u64(&EncryptionState->enc_key_lsn, lsn);
 	}
 

contrib/pg_tde/src/access/pg_tde_xlog_encrypt.c

@@ -100,7 +100,7 @@ extern WALKeyCacheRec *pg_tde_get_wal_cache_keys(void);
 extern void pg_tde_wal_last_key_set_lsn(XLogRecPtr lsn, const char *keyfile_path);
 
 extern InternalKey *pg_tde_create_smgr_key(const RelFileLocatorBackend *newrlocator);
-extern void pg_tde_create_wal_key(InternalKey *rel_key_data, const RelFileLocator *newrlocator, uint32 flags);
+extern InternalKey *pg_tde_create_wal_key(const RelFileLocator *newrlocator, uint32 flags);
 extern void pg_tde_free_key_map_entry(const RelFileLocator *rlocator, off_t offset);
 extern void pg_tde_write_key_map_entry_redo(const TDEMapEntry *write_map_entry, TDESignedPrincipalKeyInfo *signed_key_info);
 

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

@@ -13,6 +13,6 @@
 
 extern Size TDEXLogEncryptStateSize(void);
 extern void TDEXLogShmemInit(void);
-extern void TDEXLogSmgrInit(void);
+extern void TDEXLogInit(void);
 
 #endif							/* PG_TDE_XLOGENCRYPT_H */

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

@@ -92,7 +92,7 @@ tde_shmem_startup(void)
 	AesInit();
 
 	TDEXLogShmemInit();
-	TDEXLogSmgrInit();
+	TDEXLogInit();
 }
 
 void

contrib/pg_tde/src/pg_tde.c

@@ -21,7 +21,7 @@ typedef struct TDESMgrRelationData
 	struct _MdfdVec *md_seg_fds[MAX_FORKNUM + 1];
 
 	bool		encrypted_relation;
-	InternalKey relKey;
+	InternalKey *relKey;
 } TDESMgrRelationData;
 
 typedef TDESMgrRelationData *TDESMgrRelation;
@@ -81,7 +81,7 @@ tde_mdwritev(SMgrRelation reln, ForkNumber forknum, BlockNumber blocknum,
 			 const void **buffers, BlockNumber nblocks, bool skipFsync)
 {
 	TDESMgrRelation tdereln = (TDESMgrRelation) reln;
-	InternalKey *int_key = &tdereln->relKey;
+	InternalKey *int_key = tdereln->relKey;
 
 	if (!tdereln->encrypted_relation)
 	{
@@ -120,7 +120,7 @@ tde_mdextend(SMgrRelation reln, ForkNumber forknum, BlockNumber blocknum,
 			 const void *buffer, bool skipFsync)
 {
 	TDESMgrRelation tdereln = (TDESMgrRelation) reln;
-	InternalKey *int_key = &tdereln->relKey;
+	InternalKey *int_key = tdereln->relKey;
 
 	if (!tdereln->encrypted_relation)
 	{
@@ -149,7 +149,7 @@ tde_mdreadv(SMgrRelation reln, ForkNumber forknum, BlockNumber blocknum,
 			void **buffers, BlockNumber nblocks)
 {
 	TDESMgrRelation tdereln = (TDESMgrRelation) reln;
-	InternalKey *int_key = &tdereln->relKey;
+	InternalKey *int_key = tdereln->relKey;
 
 	mdreadv(reln, forknum, blocknum, buffers, nblocks);
 
@@ -236,7 +236,7 @@ tde_mdcreate(RelFileLocator relold, SMgrRelation reln, ForkNumber forknum, bool
 		if (key)
 		{
 			tdereln->encrypted_relation = true;
-			tdereln->relKey = *key;
+			tdereln->relKey = key;
 		}
 		else
 		{
@@ -257,7 +257,7 @@ tde_mdopen(SMgrRelation reln)
 	if (key)
 	{
 		tdereln->encrypted_relation = true;
-		tdereln->relKey = *key;
+		tdereln->relKey = key;
 	}
 	else
 	{

contrib/pg_tde/src/smgr/pg_tde_smgr.c

@@ -91,9 +91,28 @@
 
 $stdout = $node->safe_psql('postgres', "SELECT pg_drop_replication_slot('tde_slot');", extra_params => ['-a']);
 PGTDE::append_to_file($stdout);
+PGTDE::append_to_file("-- server restart with wal encryption and recovery with a lot of tde_heap relations");
+$node->safe_psql('postgres',
+	q{
+SELECT pg_tde_add_key_provider_file('file-keyring-010-2','/tmp/pg_tde_test_keyring010_2.per');
+SELECT pg_tde_set_principal_key('test-db-principal-key','file-keyring-010-2');
+
+do $$
+    DECLARE idx integer;
+begin
+    for idx in 0..700 loop
+        EXECUTE format('CREATE TABLE t%s (c1 int) USING tde_heap', idx);
+    end loop;
+end; $$;
+});
+
+$rt_value = $node->restart();
+ok($rt_value == 1, "Restart Server");
+
+# TODO: add WAL content testing after the wal rework
 
 # DROP EXTENSION
-$stdout = $node->safe_psql('postgres', 'DROP EXTENSION pg_tde;', extra_params => ['-a']);
+$stdout = $node->safe_psql('postgres', 'DROP EXTENSION pg_tde CASCADE;', extra_params => ['-a']);
 PGTDE::append_to_file($stdout);
 # Stop the server
 $node->stop();

contrib/pg_tde/t/009_wal_encrypt.pl

@@ -38,4 +38,5 @@ table public.test_wal: INSERT: id[integer]:6 k[integer]:6
 COMMIT 742
 SELECT pg_drop_replication_slot('tde_slot');
 
-DROP EXTENSION pg_tde;
+-- server restart with wal encryption and recovery with a lot of tde_heap relations
+DROP EXTENSION pg_tde CASCADE;

contrib/pg_tde/t/expected/009_wal_encrypt.out

@@ -1156,7 +1156,7 @@ main(int argc, char **argv)
 	if (kringdir != NULL)
 	{
 		pg_tde_fe_init(kringdir);
-		TDEXLogSmgrInit();
+		TDEXLogInit();
 	}
 #endif