Repository: percona/postgres
Pull request #160: Add checksums to the key map entries
Author: login: jeltz
Created: 2025-03-22T12:45:21Z | Merged: 2025-03-31T13:48:15Z

This adds verification based on AES-128-GCM of the encrypted relation keys in the key map file to prevent use from issues where the wrong principal key is used to decrypt (and in the future encrypt) relation and WAL keys. As a bonus we get some protection against data corruption.

Another bonus is that I also, in a separate commit, added an IV to each key which is unique every time we write a key..

@dutow had a an alternative proposal for to do it with a checksum of the plain text relation key which is stored as part of the encrypted data but what appeals to me with this solution is that it uses standard cryptography tools but I am open for changing my mind and implementing Zsolt's solution would not be much additional work.

This work is based on top of my merging of the two key files.

What do people think?


Diff:
diff --git a/contrib/pg_tde/src/access/pg_tde_tdemap.c b/contrib/pg_tde/src/access/pg_tde_tdemap.c
index c38893e0ada07..5de3c481061ab 100644
--- a/contrib/pg_tde/src/access/pg_tde_tdemap.c
+++ b/contrib/pg_tde/src/access/pg_tde_tdemap.c
@@ -28,7 +28,6 @@
 #include "catalog/tde_global_space.h"
 #include "catalog/tde_principal_key.h"
 #include "encryption/enc_aes.h"
-#include "encryption/enc_tde.h"
 #include "keyring/keyring_api.h"
 #include "common/pg_tde_utils.h"
 
@@ -70,15 +69,6 @@ typedef struct TDEFileHeader
 	TDEPrincipalKeyInfo principal_key_info;
 } TDEFileHeader;
 
-/* We do not need the dbOid since the entries are stored in a file per db */
-typedef struct TDEMapEntry
-{
-	Oid			spcOid;
-	RelFileNumber relNumber;
-	uint32		flags;
-	InternalKey enc_key;
-} TDEMapEntry;
-
 typedef struct RelKeyCacheRec
 {
 	RelFileLocator locator;
@@ -121,7 +111,7 @@ static WALKeyCacheRec *tde_wal_key_last_rec = NULL;
 
 static InternalKey *pg_tde_get_key_from_file(const RelFileLocator *rlocator, uint32 key_type);
 static TDEMapEntry *pg_tde_find_map_entry(const RelFileLocator *rlocator, uint32 key_type, char *db_map_path);
-static InternalKey *tde_decrypt_rel_key(TDEPrincipalKey *principal_key, InternalKey *enc_rel_key_data, Oid dbOid);
+static InternalKey *tde_decrypt_rel_key(TDEPrincipalKey *principal_key, TDEMapEntry *map_entry);
 static int	pg_tde_open_file_basic(char *tde_filename, int fileFlags, bool ignore_missing);
 static void pg_tde_file_header_read(char *tde_filename, int fd, TDEFileHeader *fheader, off_t *bytes_read);
 static bool pg_tde_read_one_map_entry(int fd, const RelFileLocator *rlocator, int flags, TDEMapEntry *map_entry, off_t *offset);
@@ -136,10 +126,10 @@ static InternalKey *pg_tde_put_key_into_cache(const RelFileLocator *locator, Int
 static InternalKey *pg_tde_create_key_map_entry(const RelFileLocator *newrlocator, uint32 entry_type);
 static InternalKey *pg_tde_create_local_key(const RelFileLocator *newrlocator, uint32 entry_type);
 static void pg_tde_generate_internal_key(InternalKey *int_key, uint32 entry_type);
-static InternalKey *tde_encrypt_rel_key(TDEPrincipalKey *principal_key, InternalKey *rel_key_data, Oid dbOid);
 static int	pg_tde_file_header_write(char *tde_filename, int fd, TDEPrincipalKeyInfo *principal_key_info, off_t *bytes_written);
-static void pg_tde_write_map_entry(const RelFileLocator *rlocator, InternalKey *enc_rel_key_data, char *db_map_path, TDEPrincipalKeyInfo *principal_key_info);
-static off_t pg_tde_write_one_map_entry(int fd, const RelFileLocator *rlocator, InternalKey *enc_rel_key_data, off_t *offset, const char *db_map_path);
+static void pg_tde_initialize_map_entry(TDEMapEntry *map_entry, const TDEPrincipalKey *principal_key, const RelFileLocator *rlocator, const InternalKey *enc_rel_key_data);
+static off_t pg_tde_write_one_map_entry(int fd, const TDEMapEntry *map_entry, off_t *offset, const char *db_map_path);
+static void pg_tde_write_key_map_entry(const RelFileLocator *rlocator, InternalKey *rel_key_data, TDEPrincipalKey *principal_key, bool write_xlog);
 static bool pg_tde_delete_map_entry(const RelFileLocator *rlocator, uint32 key_type, char *db_map_path, off_t offset);
 static int	keyrotation_init_file(TDEPrincipalKeyInfo *new_principal_key_info, char *rotated_filename, char *filename, off_t *curr_pos);
 static void finalize_key_rotation(const char *path_old, const char *path_new);
@@ -161,9 +151,7 @@ static InternalKey *
 pg_tde_create_key_map_entry(const RelFileLocator *newrlocator, uint32 entry_type)
 {
 	InternalKey rel_key_data;
-	InternalKey *enc_rel_key_data;
 	TDEPrincipalKey *principal_key;
-	XLogRelKey	xlrec;
 	LWLock	   *lock_pk = tde_lwlock_enc_keys();
 
 	pg_tde_generate_internal_key(&rel_key_data, entry_type);
@@ -177,26 +165,11 @@ pg_tde_create_key_map_entry(const RelFileLocator *newrlocator, uint32 entry_type
 				 errhint("create one using pg_tde_set_principal_key before using encrypted tables")));
 	}
 
-	/* Encrypt the key */
-	enc_rel_key_data = tde_encrypt_rel_key(principal_key, &rel_key_data, newrlocator->dbOid);
-
-	/*
-	 * XLOG internal key
-	 */
-	xlrec.rlocator = *newrlocator;
-	xlrec.relKey = *enc_rel_key_data;
-	xlrec.pkInfo = principal_key->keyInfo;
-
-	XLogBeginInsert();
-	XLogRegisterData((char *) &xlrec, sizeof(xlrec));
-	XLogInsert(RM_TDERMGR_ID, XLOG_TDE_ADD_RELATION_KEY);
-
 	/*
 	 * Add the encrypted key to the key map data file structure.
 	 */
-	pg_tde_write_key_map_entry(newrlocator, enc_rel_key_data, &principal_key->keyInfo);
+	pg_tde_write_key_map_entry(newrlocator, &rel_key_data, principal_key, true);
 	LWLockRelease(lock_pk);
-	pfree(enc_rel_key_data);
 
 	return pg_tde_put_key_into_cache(newrlocator, &rel_key_data);
 }
@@ -248,7 +221,6 @@ tde_sprint_key(InternalKey *k)
 void
 pg_tde_create_wal_key(InternalKey *rel_key_data, const RelFileLocator *newrlocator, uint32 entry_type)
 {
-	InternalKey *enc_rel_key_data;
 	TDEPrincipalKey *principal_key;
 
 	principal_key = get_principal_key_from_keyring(newrlocator->dbOid);
@@ -261,27 +233,11 @@ pg_tde_create_wal_key(InternalKey *rel_key_data, const RelFileLocator *newrlocat
 
 	/* TODO: no need in generating key if TDE_KEY_TYPE_WAL_UNENCRYPTED */
 	pg_tde_generate_internal_key(rel_key_data, TDE_KEY_TYPE_GLOBAL | entry_type);
-	enc_rel_key_data = tde_encrypt_rel_key(principal_key, rel_key_data, newrlocator->dbOid);
 
 	/*
 	 * Add the encrypted key to the key map data file structure.
 	 */
-	pg_tde_write_key_map_entry(newrlocator, enc_rel_key_data, &principal_key->keyInfo);
-	pfree(enc_rel_key_data);
-}
-
-/*
- * Encrypts a given key and returns the encrypted one.
- */
-static InternalKey *
-tde_encrypt_rel_key(TDEPrincipalKey *principal_key, InternalKey *rel_key_data, Oid dbOid)
-{
-	InternalKey *enc_rel_key_data;
-	size_t		enc_key_bytes;
-
-	AesEncryptKey(principal_key, dbOid, rel_key_data, &enc_rel_key_data, &enc_key_bytes);
-
-	return enc_rel_key_data;
+	pg_tde_write_key_map_entry(newrlocator, rel_key_data, principal_key, false);
 }
 
 /*
@@ -363,24 +319,79 @@ pg_tde_file_header_write(char *tde_filename, int fd, TDEPrincipalKeyInfo *princi
 	return fd;
 }
 
+static void
+pg_tde_initialize_map_entry(TDEMapEntry *map_entry, const TDEPrincipalKey *principal_key, const RelFileLocator *rlocator, const InternalKey *rel_key_data)
+{
+	map_entry->spcOid = rlocator->spcOid;
+	map_entry->relNumber = rlocator->relNumber;
+	map_entry->flags = rel_key_data->rel_type;
+	map_entry->enc_key = *rel_key_data;
+
+	if (!RAND_bytes(map_entry->entry_iv, MAP_ENTRY_EMPTY_IV_SIZE))
+		ereport(ERROR,
+				(errcode(ERRCODE_INTERNAL_ERROR),
+				 errmsg("could not generate iv for key map: %s", ERR_error_string(ERR_get_error(), NULL))));
+
+	AesGcmEncrypt(principal_key->keyData, map_entry->entry_iv, (unsigned char *) map_entry, offsetof(TDEMapEntry, enc_key), rel_key_data->key, INTERNAL_KEY_LEN, map_entry->enc_key.key, map_entry->aead_tag);
+}
+
+/*
+ * Based on the given arguments,write the entry into the key map file.
+ */
+static off_t
+pg_tde_write_one_map_entry(int fd, const TDEMapEntry *map_entry, off_t *offset, const char *db_map_path)
+{
+	int			bytes_written = 0;
+
+	/* TODO: pgstat_report_wait_start / pgstat_report_wait_end */
+	bytes_written = pg_pwrite(fd, map_entry, MAP_ENTRY_SIZE, *offset);
+
+	/* Add the entry to the file */
+	if (bytes_written != MAP_ENTRY_SIZE)
+	{
+		ereport(ERROR,
+				(errcode_for_file_access(),
+				 errmsg("could not write tde map file \"%s\": %m",
+						db_map_path)));
+	}
+	if (pg_fsync(fd) != 0)
+	{
+		ereport(data_sync_elevel(ERROR),
+				(errcode_for_file_access(),
+				 errmsg("could not fsync file \"%s\": %m", db_map_path)));
+	}
+
+	return (*offset + bytes_written);
+}
+
 /*
+ * Calls the create map entry function to get an index into the keydata. This
+ * The keydata function will then write the encrypted key on the desired
+ * location.
+ *
  * Key Map Table [pg_tde.map]:
  * 		header: {Format Version, Principal Key Name}
  * 		data: {OID, Flag, index of key in pg_tde.dat}...
  *
- * Returns the index of the key to be written in the key data file.
  * The caller must hold an exclusive lock on the map file to avoid
  * concurrent in place updates leading to data conflicts.
  */
-static void
-pg_tde_write_map_entry(const RelFileLocator *rlocator, InternalKey *enc_rel_key_data, char *db_map_path, TDEPrincipalKeyInfo *principal_key_info)
+void
+pg_tde_write_key_map_entry(const RelFileLocator *rlocator, InternalKey *rel_key_data, TDEPrincipalKey *principal_key, bool write_xlog)
 {
+	char		db_map_path[MAXPGPATH] = {0};
 	int			map_fd = -1;
 	off_t		curr_pos = 0;
 	off_t		prev_pos = 0;
+	TDEMapEntry write_map_entry;
+
+	Assert(rlocator);
+
+	/* Set the file paths */
+	pg_tde_set_db_file_path(rlocator->dbOid, db_map_path);
 
 	/* Open and validate file for basic correctness. */
-	map_fd = pg_tde_open_file(db_map_path, principal_key_info, O_RDWR | O_CREAT, &curr_pos);
+	map_fd = pg_tde_open_file(db_map_path, &principal_key->keyInfo, O_RDWR | O_CREAT, &curr_pos);
 	prev_pos = curr_pos;
 
 	/*
@@ -390,11 +401,11 @@ pg_tde_write_map_entry(const RelFileLocator *rlocator, InternalKey *enc_rel_key_
 	 */
 	while (1)
 	{
-		TDEMapEntry map_entry;
+		TDEMapEntry read_map_entry;
 		bool		found;
 
 		prev_pos = curr_pos;
-		found = pg_tde_read_one_map_entry(map_fd, NULL, MAP_ENTRY_EMPTY, &map_entry, &curr_pos);
+		found = pg_tde_read_one_map_entry(map_fd, NULL, MAP_ENTRY_EMPTY, &read_map_entry, &curr_pos);
 
 		/*
 		 * We either reach EOF or found an empty slot in the middle of the
@@ -404,12 +415,27 @@ pg_tde_write_map_entry(const RelFileLocator *rlocator, InternalKey *enc_rel_key_
 			break;
 	}
 
+	/* Initialize map entry and encrypt key */
+	pg_tde_initialize_map_entry(&write_map_entry, principal_key, rlocator, rel_key_data);
+
+	if (write_xlog)
+	{
+		XLogRelKey	xlrec;
+
+		xlrec.mapEntry = write_map_entry;
+		xlrec.pkInfo = principal_key->keyInfo;
+
+		XLogBeginInsert();
+		XLogRegisterData((char *) &xlrec, sizeof(xlrec));
+		XLogInsert(RM_TDERMGR_ID, XLOG_TDE_ADD_RELATION_KEY);
+	}
+
 	/*
 	 * Write the given entry at the location pointed by prev_pos; i.e. the
 	 * free entry
 	 */
 	curr_pos = prev_pos;
-	pg_tde_write_one_map_entry(map_fd, rlocator, enc_rel_key_data, &prev_pos, db_map_path);
+	pg_tde_write_one_map_entry(map_fd, &write_map_entry, &prev_pos, db_map_path);
 
 	/* Let's close the file. */
 	close(map_fd);
@@ -419,61 +445,56 @@ pg_tde_write_map_entry(const RelFileLocator *rlocator, InternalKey *enc_rel_key_
 }
 
 /*
- * Based on the given arguments, creates and write the entry into the key
- * map file.
+ * Write and already encrypted entry to the key map.
+ *
+ * The caller must hold an exclusive lock on the map file to avoid
+ * concurrent in place updates leading to data conflicts.
  */
-static off_t
-pg_tde_write_one_map_entry(int fd, const RelFileLocator *rlocator, InternalKey *enc_rel_key_data, off_t *offset, const char *db_map_path)
+void
+pg_tde_write_key_map_entry_redo(const TDEMapEntry *write_map_entry, TDEPrincipalKeyInfo *principal_key_info)
 {
-	int			bytes_written = 0;
-	TDEMapEntry map_entry;
+	char		db_map_path[MAXPGPATH] = {0};
+	int			map_fd = -1;
+	off_t		curr_pos = 0;
+	off_t		prev_pos = 0;
 
-	/* Fill in the map entry structure */
-	map_entry.spcOid = (rlocator == NULL) ? 0 : rlocator->spcOid;
-	map_entry.relNumber = (rlocator == NULL) ? 0 : rlocator->relNumber;
-	map_entry.flags = enc_rel_key_data->rel_type;
-	map_entry.enc_key = *enc_rel_key_data;
+	/* Set the file paths */
+	pg_tde_set_db_file_path(principal_key_info->databaseId, db_map_path);
 
-	/* TODO: pgstat_report_wait_start / pgstat_report_wait_end */
-	bytes_written = pg_pwrite(fd, &map_entry, MAP_ENTRY_SIZE, *offset);
+	/* Open and validate file for basic correctness. */
+	map_fd = pg_tde_open_file(db_map_path, principal_key_info, O_RDWR | O_CREAT, &curr_pos);
+	prev_pos = curr_pos;
 
-	/* Add the entry to the file */
-	if (bytes_written != MAP_ENTRY_SIZE)
-	{
-		ereport(ERROR,
-				(errcode_for_file_access(),
-				 errmsg("could not write tde map file \"%s\": %m",
-						db_map_path)));
-	}
-	if (pg_fsync(fd) != 0)
+	/*
+	 * Read until we find an empty slot. Otherwise, read until end. This seems
+	 * to be less frequent than vacuum. So let's keep this function here
+	 * rather than overloading the vacuum process.
+	 */
+	while (1)
 	{
-		ereport(data_sync_elevel(ERROR),
-				(errcode_for_file_access(),
-				 errmsg("could not fsync file \"%s\": %m", db_map_path)));
-	}
-
-	return (*offset + bytes_written);
-}
+		TDEMapEntry read_map_entry;
+		bool		found;
 
-/*
- * Calls the create map entry function to get an index into the keydata. This
- * The keydata function will then write the encrypted key on the desired
- * location.
- *
- * The caller must hold an exclusive lock tde_lwlock_enc_keys.
- */
-void
-pg_tde_write_key_map_entry(const RelFileLocator *rlocator, InternalKey *enc_rel_key_data, TDEPrincipalKeyInfo *principal_key_info)
-{
-	char		db_map_path[MAXPGPATH] = {0};
+		prev_pos = curr_pos;
+		found = pg_tde_read_one_map_entry(map_fd, NULL, MAP_ENTRY_EMPTY, &read_map_entry, &curr_pos);
 
-	Assert(rlocator);
+		/*
+		 * We either reach EOF or found an empty slot in the middle of the
+		 * file
+		 */
+		if (prev_pos == curr_pos || found)
+			break;
+	}
 
-	/* Set the file paths */
-	pg_tde_set_db_file_path(rlocator->dbOid, db_map_path);
+	/*
+	 * Write the given entry at the location pointed by prev_pos; i.e. the
+	 * free entry
+	 */
+	curr_pos = prev_pos;
+	pg_tde_write_one_map_entry(map_fd, write_map_entry, &prev_pos, db_map_path);
 
-	/* Create the map entry and add the encrypted key to the data file */
-	pg_tde_write_map_entry(rlocator, enc_rel_key_data, db_map_path, principal_key_info);
+	/* Let's close the file. */
+	close(map_fd);
 }
 
 static bool
@@ -514,10 +535,10 @@ pg_tde_delete_map_entry(const RelFileLocator *rlocator, uint32 key_type, char *d
 	 */
 	while (1)
 	{
-		TDEMapEntry map_entry;
+		TDEMapEntry read_map_entry;
 		off_t		prev_pos = curr_pos;
 
-		found = pg_tde_read_one_map_entry(map_fd, rlocator, key_type, &map_entry, &curr_pos);
+		found = pg_tde_read_one_map_entry(map_fd, rlocator, key_type, &read_map_entry, &curr_pos);
 
 		/* We've reached EOF */
 		if (curr_pos == prev_pos)
@@ -526,9 +547,14 @@ pg_tde_delete_map_entry(const RelFileLocator *rlocator, uint32 key_type, char *d
 		/* We found a valid entry for the relation */
 		if (found)
 		{
-			InternalKey dummy_key = {.rel_type = MAP_ENTRY_EMPTY};
-
-			pg_tde_write_one_map_entry(map_fd, NULL, &dummy_key, &prev_pos, db_map_path);
+			TDEMapEntry empty_map_entry = {
+				.flags = MAP_ENTRY_EMPTY,
+				.enc_key = {
+					.rel_type = MAP_ENTRY_EMPTY,
+				},
+			};
+
+			pg_tde_write_one_map_entry(map_fd, &empty_map_entry, &prev_pos, db_map_path);
 			break;
 		}
 	}
@@ -636,9 +662,9 @@ pg_tde_perform_rotate_key(TDEPrincipalKey *principal_key, TDEPrincipalKey *new_p
 	while (1)
 	{
 		InternalKey *rel_key_data;
-		InternalKey *new_enc_rel_key_data;
 		off_t		prev_pos[PRINCIPAL_KEY_COUNT];
-		TDEMapEntry read_map_entry;
+		TDEMapEntry read_map_entry,
+					write_map_entry;
 		RelFileLocator rloc;
 		bool		found;
 
@@ -657,16 +683,15 @@ pg_tde_perform_rotate_key(TDEPrincipalKey *principal_key, TDEPrincipalKey *new_p
 		rloc.dbOid = principal_key->keyInfo.databaseId;
 		rloc.relNumber = read_map_entry.relNumber;
 
-		/* Decrypt and re-encrypt keys */
-		rel_key_data = tde_decrypt_rel_key(principal_key, &read_map_entry.enc_key, principal_key->keyInfo.databaseId);
-		new_enc_rel_key_data = tde_encrypt_rel_key(new_principal_key, rel_key_data, principal_key->keyInfo.databaseId);
+		/* Decrypt and re-encrypt key */
+		rel_key_data = tde_decrypt_rel_key(principal_key, &read_map_entry);
+		pg_tde_initialize_map_entry(&write_map_entry, new_principal_key, &rloc, rel_key_data);
 
 		/* Write the given entry at the location pointed by prev_pos */
 		prev_pos[NEW_PRINCIPAL_KEY] = curr_pos[NEW_PRINCIPAL_KEY];
-		curr_pos[NEW_PRINCIPAL_KEY] = pg_tde_write_one_map_entry(fd[NEW_PRINCIPAL_KEY], &rloc, new_enc_rel_key_data, &prev_pos[NEW_PRINCIPAL_KEY], path[NEW_PRINCIPAL_KEY]);
+		curr_pos[NEW_PRINCIPAL_KEY] = pg_tde_write_one_map_entry(fd[NEW_PRINCIPAL_KEY], &write_map_entry, &prev_pos[NEW_PRINCIPAL_KEY], path[NEW_PRINCIPAL_KEY]);
 
 		pfree(rel_key_data);
-		pfree(new_enc_rel_key_data);
 	}
 
 	close(fd[OLD_PRINCIPAL_KEY]);
@@ -882,7 +907,7 @@ pg_tde_get_key_from_file(const RelFileLocator *rlocator, uint32 key_type)
 				 errhint("create one using pg_tde_set_principal_key before using encrypted tables")));
 	LWLockRelease(lock_pk);
 
-	return tde_decrypt_rel_key(principal_key, &map_entry->enc_key, rlocator->dbOid);
+	return tde_decrypt_rel_key(principal_key, map_entry);
 }
 
 /*
@@ -935,12 +960,20 @@ pg_tde_find_map_entry(const RelFileLocator *rlocator, uint32 key_type, char *db_
  * Decrypts a given key and returns the decrypted one.
  */
 static InternalKey *
-tde_decrypt_rel_key(TDEPrincipalKey *principal_key, InternalKey *enc_rel_key_data, Oid dbOid)
+tde_decrypt_rel_key(TDEPrincipalKey *principal_key, TDEMapEntry *map_entry)
 {
-	InternalKey *rel_key_data = NULL;
-	size_t		key_bytes;
+	InternalKey *rel_key_data = palloc_object(InternalKey);
+
+	/* Ensure we are getting a valid pointer here */
+	Assert(principal_key);
+
+	/* Fill in the structure */
+	*rel_key_data = map_entry->enc_key;
+
+	if (!AesGcmDecrypt(principal_key->keyData, map_entry->entry_iv, (unsigned char *) map_entry, offsetof(TDEMapEntry, enc_key), map_entry->enc_key.key, INTERNAL_KEY_LEN, rel_key_data->key, map_entry->aead_tag))
+		ereport(ERROR,
+				(errmsg("Failed to decrypt key, incorrect principal key or corrupted key file")));
 
-	AesDecryptKey(principal_key, dbOid, &rel_key_data, enc_rel_key_data, &key_bytes);
 
 	return rel_key_data;
 }
@@ -1295,7 +1328,7 @@ pg_tde_read_last_wal_key(void)
 		return NULL;
 	}
 
-	rel_key_data = tde_decrypt_rel_key(principal_key, &map_entry->enc_key, rlocator.dbOid);
+	rel_key_data = tde_decrypt_rel_key(principal_key, map_entry);
 	LWLockRelease(lock_pk);
 	close(fd);
 	pfree(map_entry);
@@ -1363,7 +1396,7 @@ pg_tde_fetch_wal_keys(XLogRecPtr start_lsn)
 			WALKeyIsValid(&map_entry->enc_key) &&
 			map_entry->enc_key.start_lsn >= start_lsn)
 		{
-			InternalKey *rel_key_data = tde_decrypt_rel_key(principal_key, &map_entry->enc_key, rlocator.dbOid);
+			InternalKey *rel_key_data = tde_decrypt_rel_key(principal_key, map_entry);
 			InternalKey *cached_key = pg_tde_put_key_into_cache(&rlocator, rel_key_data);
 			WALKeyCacheRec *wal_rec;
 
diff --git a/contrib/pg_tde/src/access/pg_tde_xlog.c b/contrib/pg_tde/src/access/pg_tde_xlog.c
index bd2e4519023b6..93ca9cb21d542 100644
--- a/contrib/pg_tde/src/access/pg_tde_xlog.c
+++ b/contrib/pg_tde/src/access/pg_tde_xlog.c
@@ -58,7 +58,7 @@ tdeheap_rmgr_redo(XLogReaderState *record)
 		XLogRelKey *xlrec = (XLogRelKey *) XLogRecGetData(record);
 
 		LWLockAcquire(tde_lwlock_enc_keys(), LW_EXCLUSIVE);
-		pg_tde_write_key_map_entry(&xlrec->rlocator, &xlrec->relKey, &xlrec->pkInfo);
+		pg_tde_write_key_map_entry_redo(&xlrec->mapEntry, &xlrec->pkInfo);
 		LWLockRelease(tde_lwlock_enc_keys());
 	}
 	else if (info == XLOG_TDE_ADD_PRINCIPAL_KEY)
@@ -116,7 +116,7 @@ tdeheap_rmgr_desc(StringInfo buf, XLogReaderState *record)
 	{
 		XLogRelKey *xlrec = (XLogRelKey *) XLogRecGetData(record);
 
-		appendStringInfo(buf, "add tde internal key for relation %u/%u", xlrec->rlocator.dbOid, xlrec->rlocator.relNumber);
+		appendStringInfo(buf, "add tde internal key for relation %u/%u", xlrec->pkInfo.databaseId, xlrec->mapEntry.relNumber);
 	}
 	if (info == XLOG_TDE_ADD_PRINCIPAL_KEY)
 	{
diff --git a/contrib/pg_tde/src/encryption/enc_aes.c b/contrib/pg_tde/src/encryption/enc_aes.c
index a223b0d121655..b0ec21489a086 100644
--- a/contrib/pg_tde/src/encryption/enc_aes.c
+++ b/contrib/pg_tde/src/encryption/enc_aes.c
@@ -42,6 +42,7 @@
  */
 
 static const EVP_CIPHER *cipher_cbc;
+static const EVP_CIPHER *cipher_gcm;
 static const EVP_CIPHER *cipher_ctr_ecb;
 
 void
@@ -55,6 +56,7 @@ AesInit(void)
 		ERR_load_crypto_strings();
 
 		cipher_cbc = EVP_aes_128_cbc();
+		cipher_gcm = EVP_aes_128_gcm();
 		cipher_ctr_ecb = EVP_aes_128_ecb();
 
 		initialized = 1;
@@ -63,8 +65,10 @@ AesInit(void)
 
 /*  TODO: a few things could be optimized in this. It's good enough for a prototype. */
 static void
-AesRunCtr(EVP_CIPHER_CTX **ctxPtr, int enc, const unsigned char *key, const unsigned char *iv, const unsigned char *in, int in_len, unsigned char *out, int *out_len)
+AesRunCtr(EVP_CIPHER_CTX **ctxPtr, int enc, const unsigned char *key, const unsigned char *iv, const unsigned char *in, int in_len, unsigned char *out)
 {
+	int			out_len;
+
 	if (*ctxPtr == NULL)
 	{
 		*ctxPtr = EVP_CIPHER_CTX_new();
@@ -77,15 +81,18 @@ AesRunCtr(EVP_CIPHER_CTX **ctxPtr, int enc, const unsigned char *key, const unsi
 		EVP_CIPHER_CTX_set_padding(*ctxPtr, 0);
 	}
 
-	if (EVP_CipherUpdate(*ctxPtr, out, out_len, in, in_len) == 0)
+	if (EVP_CipherUpdate(*ctxPtr, out, &out_len, in, in_len) == 0)
 		ereport(ERROR,
 				(errmsg("EVP_CipherUpdate failed. OpenSSL error: %s", ERR_error_string(ERR_get_error(), NULL))));
+
+	Assert(out_len == in_len);
 }
 
 static void
-AesRunCbc(int enc, const unsigned char *key, const unsigned char *iv, const unsigned char *in, int in_len, unsigned char *out, int *out_len)
+AesRunCbc(int enc, const unsigned char *key, const unsigned char *iv, const unsigned char *in, int in_len, unsigned char *out)
 {
-	int			out_len_final = 0;
+	int			out_len;
+	int			out_len_final;
 	EVP_CIPHER_CTX *ctx = NULL;
 
 	Assert(in_len % EVP_CIPHER_block_size(cipher_cbc) == 0);
@@ -99,11 +106,11 @@ AesRunCbc(int enc, const unsigned char *key, const unsigned char *iv, const unsi
 
 	EVP_CIPHER_CTX_set_padding(ctx, 0);
 
-	if (EVP_CipherUpdate(ctx, out, out_len, in, in_len) == 0)
+	if (EVP_CipherUpdate(ctx, out, &out_len, in, in_len) == 0)
 		ereport(ERROR,
 				(errmsg("EVP_CipherUpdate failed. OpenSSL error: %s", ERR_error_string(ERR_get_error(), NULL))));
 
-	if (EVP_CipherFinal_ex(ctx, out + *out_len, &out_len_final) == 0)
+	if (EVP_CipherFinal_ex(ctx, out + out_len, &out_len_final) == 0)
 		ereport(ERROR,
 				(errmsg("EVP_CipherFinal_ex failed. OpenSSL error: %s", ERR_error_string(ERR_get_error(), NULL))));
 
@@ -111,23 +118,138 @@ AesRunCbc(int enc, const unsigned char *key, const unsigned char *iv, const unsi
 	 * We encrypt one block (16 bytes) Our expectation is that the result
 	 * should also be 16 bytes, without any additional padding
 	 */
-	*out_len += out_len_final;
-	Assert(in_len == *out_len);
+	out_len += out_len_final;
+	Assert(in_len == out_len);
 
 	EVP_CIPHER_CTX_cleanup(ctx);
 	EVP_CIPHER_CTX_free(ctx);
 }
 
 void
-AesEncrypt(const unsigned char *key, const unsigned char *iv, const unsigned char *in, int in_len, unsigned char *out, int *out_len)
+AesEncrypt(const unsigned char *key, const unsigned char *iv, const unsigned char *in, int in_len, unsigned char *out)
 {
-	AesRunCbc(1, key, iv, in, in_len, out, out_len);
+	AesRunCbc(1, key, iv, in, in_len, out);
 }
 
 void
-AesDecrypt(const unsigned char *key, const unsigned char *iv, const unsigned char *in, int in_len, unsigned char *out, int *out_len)
+AesDecrypt(const unsigned char *key, const unsigned char *iv, const unsigned char *in, int in_len, unsigned char *out)
+{
+	AesRunCbc(0, key, iv, in, in_len, out);
+}
+
+void
+AesGcmEncrypt(const unsigned char *key, const unsigned char *iv, const unsigned char *aad, int aad_len, const unsigned char *in, int in_len, unsigned char *out, unsigned char *tag)
+{
+	int			out_len;
+	int			out_len_final;
+	EVP_CIPHER_CTX *ctx;
+
+	Assert(in_len % EVP_CIPHER_block_size(cipher_gcm) == 0);
+
+	ctx = EVP_CIPHER_CTX_new();
+	EVP_CIPHER_CTX_init(ctx);
+
+	if (EVP_EncryptInit_ex(ctx, cipher_gcm, NULL, NULL, NULL) == 0)
+		ereport(ERROR,
+				(errmsg("EVP_EncryptInit_ex failed. OpenSSL error: %s", ERR_error_string(ERR_get_error(), NULL))));
+
+	if (EVP_CIPHER_CTX_set_padding(ctx, 0) == 0)
+		ereport(ERROR,
+				(errmsg("EVP_CIPHER_CTX_set_padding failed. OpenSSL error: %s", ERR_error_string(ERR_get_error(), NULL))));
+
+	if (EVP_CIPHER_CTX_ctrl(ctx, EVP_CTRL_GCM_SET_IVLEN, 16, NULL) == 0)
+		ereport(ERROR,
+				(errmsg("EVP_CTRL_GCM_SET_IVLEN failed. OpenSSL error: %s", ERR_error_string(ERR_get_error(), NULL))));
+
+	if (EVP_EncryptInit_ex(ctx, NULL, NULL, key, iv) == 0)
+		ereport(ERROR,
+				(errmsg("EVP_EncryptInit_ex failed. OpenSSL error: %s", ERR_error_string(ERR_get_error(), NULL))));
+
+	if (EVP_EncryptUpdate(ctx, NULL, &out_len, (unsigned char *) aad, aad_len) == 0)
+		ereport(ERROR,
+				(errmsg("EVP_CipherUpdate failed. OpenSSL error: %s", ERR_error_string(ERR_get_error(), NULL))));
+
+	if (EVP_EncryptUpdate(ctx, out, &out_len, in, in_len) == 0)
+		ereport(ERROR,
+				(errmsg("EVP_CipherUpdate failed. OpenSSL error: %s", ERR_error_string(ERR_get_error(), NULL))));
+
+	if (EVP_EncryptFinal_ex(ctx, out + out_len, &out_len_final) == 0)
+		ereport(ERROR,
+				(errmsg("EVP_CipherFinal_ex failed. OpenSSL error: %s", ERR_error_string(ERR_get_error(), NULL))));
+
+	if (EVP_CIPHER_CTX_ctrl(ctx, EVP_CTRL_GCM_GET_TAG, 16, tag) == 0)
+		ereport(ERROR,
+				(errmsg("EVP_CTRL_GCM_GET_TAG failed. OpenSSL error: %s", ERR_error_string(ERR_get_error(), NULL))));
+
+	/*
+	 * We encrypt one block (16 bytes) Our expectation is that the result
+	 * should also be 16 bytes, without any additional padding
+	 */
+	out_len += out_len_final;
+	Assert(in_len == out_len);
+
+	EVP_CIPHER_CTX_cleanup(ctx);
+	EVP_CIPHER_CTX_free(ctx);
+}
+
+bool
+AesGcmDecrypt(const unsigned char *key, const unsigned char *iv, const unsigned char *aad, int aad_len, const unsigned char *in, int in_len, unsigned char *out, unsigned char *tag)
 {
-	AesRunCbc(0, key, iv, in, in_len, out, out_len);
+	int			out_len;
+	int			out_len_final;
+	EVP_CIPHER_CTX *ctx;
+
+	Assert(in_len % EVP_CIPHER_block_size(cipher_gcm) == 0);
+
+	ctx = EVP_CIPHER_CTX_new();
+	EVP_CIPHER_CTX_init(ctx);
+
+	if (EVP_DecryptInit_ex(ctx, cipher_gcm, NULL, NULL, NULL) == 0)
+		ereport(ERROR,
+				(errmsg("EVP_EncryptInit_ex failed. OpenSSL error: %s", ERR_error_string(ERR_get_error(), NULL))));
+
+	if (EVP_CIPHER_CTX_set_padding(ctx, 0) == 0)
+		ereport(ERROR,
+				(errmsg("EVP_CIPHER_CTX_set_padding failed. OpenSSL error: %s", ERR_error_string(ERR_get_error(), NULL))));
+
+	if (EVP_CIPHER_CTX_ctrl(ctx, EVP_CTRL_GCM_SET_IVLEN, 16, NULL) == 0)
+		ereport(ERROR,
+				(errmsg("EVP_CTRL_GCM_SET_IVLEN failed. OpenSSL error: %s", ERR_error_string(ERR_get_error(), NULL))));
+
+	if (EVP_DecryptInit_ex(ctx, NULL, NULL, key, iv) == 0)
+		ereport(ERROR,
+				(errmsg("EVP_EncryptInit_ex failed. OpenSSL error: %s", ERR_error_string(ERR_get_error(), NULL))));
+
+	if (EVP_CIPHER_CTX_ctrl(ctx, EVP_CTRL_GCM_SET_TAG, 16, tag) == 0)
+		ereport(ERROR,
+				(errmsg("EVP_CTRL_GCM_SET_TAG failed. OpenSSL error: %s", ERR_error_string(ERR_get_error(), NULL))));
+
+	if (EVP_DecryptUpdate(ctx, NULL, &out_len, aad, aad_len) == 0)
+		ereport(ERROR,
+				(errmsg("EVP_CipherUpdate failed. OpenSSL error: %s", ERR_error_string(ERR_get_error(), NULL))));
+
+	if (EVP_DecryptUpdate(ctx, out, &out_len, in, in_len) == 0)
+		ereport(ERROR,
+				(errmsg("EVP_CipherUpdate failed. OpenSSL error: %s", ERR_error_string(ERR_get_error(), NULL))));
+
+	if (EVP_DecryptFinal_ex(ctx, out + out_len, &out_len_final) == 0)
+	{
+		EVP_CIPHER_CTX_cleanup(ctx);
+		EVP_CIPHER_CTX_free(ctx);
+		return false;
+	}
+
+	/*
+	 * We encrypt one block (16 bytes) Our expectation is that the result
+	 * should also be 16 bytes, without any additional padding
+	 */
+	out_len += out_len_final;
+	Assert(in_len == out_len);
+
+	EVP_CIPHER_CTX_cleanup(ctx);
+	EVP_CIPHER_CTX_free(ctx);
+
+	return true;
 }
 
 /* This function assumes that the out buffer is big enough: at least (blockNumber2 - blockNumber1) * 16 bytes
@@ -137,7 +259,6 @@ Aes128EncryptedZeroBlocks(void *ctxPtr, const unsigned char *key, const char *iv
 {
 	const unsigned char iv[16] = {0,};
 	unsigned char *p;
-	int			outLen;
 
 	Assert(blockNumber2 >= blockNumber1);
 
@@ -157,6 +278,5 @@ Aes128EncryptedZeroBlocks(void *ctxPtr, const unsigned char *key, const char *iv
 		p += sizeof(j);
 	}
 
-	AesRunCtr(ctxPtr, 1, key, iv, out, p - out, out, &outLen);
-	Assert(outLen == p - out);
+	AesRunCtr(ctxPtr, 1, key, iv, out, p - out, out);
 }
diff --git a/contrib/pg_tde/src/encryption/enc_tde.c b/contrib/pg_tde/src/encryption/enc_tde.c
index 312ee42668fd0..f4779037a207d 100644
--- a/contrib/pg_tde/src/encryption/enc_tde.c
+++ b/contrib/pg_tde/src/encryption/enc_tde.c
@@ -6,7 +6,6 @@
 #include "encryption/enc_tde.h"
 #include "encryption/enc_aes.h"
 #include "storage/bufmgr.h"
-#include "keyring/keyring_api.h"
 
 #define AES_BLOCK_SIZE 		        16
 #define NUM_AES_BLOCKS_IN_BATCH     200
@@ -152,54 +151,3 @@ pg_tde_crypt(const char *iv_prefix, uint32 start_offset, const char *data, uint3
 		pg_tde_crypt_simple(iv_prefix, start_offset, data, data_len, out, key, ctxPtr, context);
 	}
 }
-
-#ifndef FRONTEND
-/*
- * Provide a simple interface to encrypt a given key.
- *
- * The function pallocs and updates the p_enc_rel_key_data along with key bytes. The memory
- * is allocated in the current memory context as this key should be ephemeral with a very
- * short lifespan until it is written to disk.
- */
-void
-AesEncryptKey(const TDEPrincipalKey *principal_key, Oid dbOid, InternalKey *rel_key_data, InternalKey **p_enc_rel_key_data, size_t *enc_key_bytes)
-{
-	unsigned char iv[16] = {0,};
-
-	/* Ensure we are getting a valid pointer here */
-	Assert(principal_key);
-
-	memcpy(iv, &dbOid, sizeof(Oid));
-
-	*p_enc_rel_key_data = palloc_object(InternalKey);
-	**p_enc_rel_key_data = *rel_key_data;
-
-	AesEncrypt(principal_key->keyData, iv, (unsigned char *) rel_key_data, INTERNAL_KEY_LEN, (unsigned char *) *p_enc_rel_key_data, (int *) enc_key_bytes);
-}
-
-#endif							/* FRONTEND */
-
-/*
- * Provide a simple interface to decrypt a given key.
- *
- * The function pallocs and updates the p_rel_key_data along with key bytes. It's important
- * to note that memory is allocated in the TopMemoryContext so we expect this to be added
- * to our key cache.
- */
-void
-AesDecryptKey(const TDEPrincipalKey *principal_key, Oid dbOid, InternalKey **p_rel_key_data, InternalKey *enc_rel_key_data, size_t *key_bytes)
-{
-	unsigned char iv[16] = {0,};
-
-	/* Ensure we are getting a valid pointer here */
-	Assert(principal_key);
-
-	memcpy(iv, &dbOid, sizeof(Oid));
-
-	*p_rel_key_data = palloc_object(InternalKey);
-
-	/* Fill in the structure */
-	**p_rel_key_data = *enc_rel_key_data;
-
-	AesDecrypt(principal_key->keyData, iv, (unsigned char *) enc_rel_key_data, INTERNAL_KEY_LEN, (unsigned char *) *p_rel_key_data, (int *) key_bytes);
-}
diff --git a/contrib/pg_tde/src/include/access/pg_tde_tdemap.h b/contrib/pg_tde/src/include/access/pg_tde_tdemap.h
index c1783647c06e1..b43f97a375f79 100644
--- a/contrib/pg_tde/src/include/access/pg_tde_tdemap.h
+++ b/contrib/pg_tde/src/include/access/pg_tde_tdemap.h
@@ -39,10 +39,24 @@ typedef struct InternalKey
 	(((key)->rel_type & TDE_KEY_TYPE_WAL_UNENCRYPTED) != 0 || \
 	((key)->rel_type & TDE_KEY_TYPE_WAL_ENCRYPTED) != 0)
 
+#define MAP_ENTRY_EMPTY_IV_SIZE 16
+#define MAP_ENTRY_EMPTY_AEAD_TAG_SIZE 16
+
+/* We do not need the dbOid since the entries are stored in a file per db */
+typedef struct TDEMapEntry
+{
+	Oid			spcOid;
+	RelFileNumber relNumber;
+	uint32		flags;
+	InternalKey enc_key;
+	/* IV and tag used when encrypting the key itself */
+	unsigned char entry_iv[MAP_ENTRY_EMPTY_IV_SIZE];
+	unsigned char aead_tag[MAP_ENTRY_EMPTY_AEAD_TAG_SIZE];
+} TDEMapEntry;
+
 typedef struct XLogRelKey
 {
-	RelFileLocator rlocator;
-	InternalKey relKey;
+	TDEMapEntry mapEntry;
 	TDEPrincipalKeyInfo pkInfo;
 } XLogRelKey;
 
@@ -76,8 +90,8 @@ 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 void pg_tde_write_key_map_entry(const RelFileLocator *rlocator, InternalKey *enc_rel_key_data, TDEPrincipalKeyInfo *principal_key_info);
 extern void pg_tde_free_key_map_entry(const RelFileLocator *rlocator, uint32 key_type, off_t offset);
+extern void pg_tde_write_key_map_entry_redo(const TDEMapEntry *write_map_entry, TDEPrincipalKeyInfo *principal_key_info);
 
 #define PG_TDE_MAP_FILENAME			"pg_tde_%d_map"
 
diff --git a/contrib/pg_tde/src/include/encryption/enc_aes.h b/contrib/pg_tde/src/include/encryption/enc_aes.h
index a8b4cb71c5153..1cbece5768d55 100644
--- a/contrib/pg_tde/src/include/encryption/enc_aes.h
+++ b/contrib/pg_tde/src/include/encryption/enc_aes.h
@@ -15,8 +15,9 @@
 extern void AesInit(void);
 extern void Aes128EncryptedZeroBlocks(void *ctxPtr, const unsigned char *key, const char *iv_prefix, uint64_t blockNumber1, uint64_t blockNumber2, unsigned char *out);
 
-/* Only used for testing */
-extern void AesEncrypt(const unsigned char *key, const unsigned char *iv, const unsigned char *in, int in_len, unsigned char *out, int *out_len);
-extern void AesDecrypt(const unsigned char *key, const unsigned char *iv, const unsigned char *in, int in_len, unsigned char *out, int *out_len);
+extern void AesEncrypt(const unsigned char *key, const unsigned char *iv, const unsigned char *in, int in_len, unsigned char *out);
+extern void AesDecrypt(const unsigned char *key, const unsigned char *iv, const unsigned char *in, int in_len, unsigned char *out);
+extern void AesGcmEncrypt(const unsigned char *key, const unsigned char *iv, const unsigned char *aad, int aad_len, const unsigned char *in, int in_len, unsigned char *out, unsigned char *tag);
+extern bool AesGcmDecrypt(const unsigned char *key, const unsigned char *iv, const unsigned char *aad, int aad_len, const unsigned char *in, int in_len, unsigned char *out, unsigned char *tag);
 
 #endif							/* ENC_AES_H */
diff --git a/contrib/pg_tde/src/include/encryption/enc_tde.h b/contrib/pg_tde/src/include/encryption/enc_tde.h
index 4ec8716891715..092cb5df4f5b0 100644
--- a/contrib/pg_tde/src/include/encryption/enc_tde.h
+++ b/contrib/pg_tde/src/include/encryption/enc_tde.h
@@ -22,7 +22,4 @@ extern void pg_tde_crypt(const char *iv_prefix, uint32 start_offset, const char
 #define PG_TDE_DECRYPT_DATA(_iv_prefix, _start_offset, _data, _data_len, _out, _key, _ctxptr) \
 	pg_tde_crypt(_iv_prefix, _start_offset, _data, _data_len, _out, _key, _ctxptr, "DECRYPT")
 
-extern void AesEncryptKey(const TDEPrincipalKey *principal_key, Oid dbOid, InternalKey *rel_key_data, InternalKey **p_enc_rel_key_data, size_t *enc_key_bytes);
-extern void AesDecryptKey(const TDEPrincipalKey *principal_key, Oid dbOid, InternalKey **p_rel_key_data, InternalKey *enc_rel_key_data, size_t *key_bytes);
-
 #endif							/* ENC_TDE_H */
diff --git a/contrib/pg_tde/src/smgr/pg_tde_smgr.c b/contrib/pg_tde/src/smgr/pg_tde_smgr.c
index 85f7504f67a52..328ed3a421b6e 100644
--- a/contrib/pg_tde/src/smgr/pg_tde_smgr.c
+++ b/contrib/pg_tde/src/smgr/pg_tde_smgr.c
@@ -95,7 +95,6 @@ tde_mdwritev(SMgrRelation reln, ForkNumber forknum, BlockNumber blocknum,
 
 		for (int i = 0; i < nblocks; ++i)
 		{
-			int			out_len = BLCKSZ;
 			BlockNumber bn = blocknum + i;
 			unsigned char iv[16] = {0,};
 
@@ -104,7 +103,7 @@ tde_mdwritev(SMgrRelation reln, ForkNumber forknum, BlockNumber blocknum,
 
 			memcpy(iv + 4, &bn, sizeof(BlockNumber));
 
-			AesEncrypt(int_key->key, iv, ((unsigned char **) buffers)[i], BLCKSZ, local_buffers[i], &out_len);
+			AesEncrypt(int_key->key, iv, ((unsigned char **) buffers)[i], BLCKSZ, local_buffers[i]);
 		}
 
 		mdwritev(reln, forknum, blocknum,
@@ -130,7 +129,6 @@ tde_mdextend(SMgrRelation reln, ForkNumber forknum, BlockNumber blocknum,
 	{
 		unsigned char *local_blocks = palloc(BLCKSZ * (1 + 1));
 		unsigned char *local_blocks_aligned = (unsigned char *) TYPEALIGN(PG_IO_ALIGN_SIZE, local_blocks);
-		int			out_len = BLCKSZ;
 		unsigned char iv[16] = {
 			0,
 		};
@@ -138,7 +136,7 @@ tde_mdextend(SMgrRelation reln, ForkNumber forknum, BlockNumber blocknum,
 		AesInit();
 		memcpy(iv + 4, &blocknum, sizeof(BlockNumber));
 
-		AesEncrypt(int_key->key, iv, ((unsigned char *) buffer), BLCKSZ, local_blocks_aligned, &out_len);
+		AesEncrypt(int_key->key, iv, ((unsigned char *) buffer), BLCKSZ, local_blocks_aligned);
 
 		mdextend(reln, forknum, blocknum, local_blocks_aligned, skipFsync);
 
@@ -150,7 +148,6 @@ static void
 tde_mdreadv(SMgrRelation reln, ForkNumber forknum, BlockNumber blocknum,
 			void **buffers, BlockNumber nblocks)
 {
-	int			out_len = BLCKSZ;
 	TDESMgrRelation tdereln = (TDESMgrRelation) reln;
 	InternalKey *int_key = &tdereln->relKey;
 
@@ -194,7 +191,7 @@ tde_mdreadv(SMgrRelation reln, ForkNumber forknum, BlockNumber blocknum,
 
 		memcpy(iv + 4, &bn, sizeof(BlockNumber));
 
-		AesDecrypt(int_key->key, iv, ((unsigned char **) buffers)[i], BLCKSZ, ((unsigned char **) buffers)[i], &out_len);
+		AesDecrypt(int_key->key, iv, ((unsigned char **) buffers)[i], BLCKSZ, ((unsigned char **) buffers)[i]);
 	}
 }
 
diff --git a/contrib/pg_tde/t/010_change_key_provider.pl b/contrib/pg_tde/t/010_change_key_provider.pl
index 94ac25da537e7..c7add5bb551bf 100644
--- a/contrib/pg_tde/t/010_change_key_provider.pl
+++ b/contrib/pg_tde/t/010_change_key_provider.pl
@@ -24,6 +24,7 @@
 unlink('/tmp/change_key_provider_1.per');
 unlink('/tmp/change_key_provider_2.per');
 unlink('/tmp/change_key_provider_3.per');
+unlink('/tmp/change_key_provider_4.per');
 
 # Start server
 my $rt_value = $node->start;
@@ -131,6 +132,70 @@
 (undef, $stdout, $stderr) = $node->psql('postgres', 'DROP EXTENSION pg_tde CASCADE;', extra_params => ['-a']);
 PGTDE::append_to_file($stdout);
 PGTDE::append_to_file($stderr);
+
+# CREATE EXTENSION
+($cmdret, $stdout, $stderr) = $node->psql('postgres', 'CREATE EXTENSION IF NOT EXISTS pg_tde;', extra_params => ['-a']);
+ok($cmdret == 0, "CREATE PGTDE EXTENSION");
+PGTDE::append_to_file($stdout);
+
+# Change provider and generate a new principal key
+$stdout = $node->safe_psql('postgres', "SELECT pg_tde_add_key_provider_file('file-vault', '/tmp/change_key_provider_4.per');", extra_params => ['-a']);
+PGTDE::append_to_file($stdout);
+$stdout = $node->psql('postgres', "SELECT pg_tde_set_principal_key('test-key', 'file-vault');", extra_params => ['-a']);
+PGTDE::append_to_file($stdout);
+
+$stdout = $node->safe_psql('postgres', 'CREATE TABLE test_enc (id serial, k integer, PRIMARY KEY (id)) USING tde_heap;', extra_params => ['-a']);
+PGTDE::append_to_file($stdout);
+$stdout = $node->safe_psql('postgres', 'INSERT INTO test_enc (k) VALUES (5), (6);', extra_params => ['-a']);
+PGTDE::append_to_file($stdout);
+
+$stdout = $node->safe_psql('postgres', "SELECT pg_tde_verify_principal_key();", extra_params => ['-a']);
+PGTDE::append_to_file($stdout);
+$stdout = $node->safe_psql('postgres', "SELECT pg_tde_is_encrypted('test_enc');", extra_params => ['-a']);
+PGTDE::append_to_file($stdout);
+$stdout = $node->safe_psql('postgres', 'SELECT * FROM test_enc ORDER BY id;', extra_params => ['-a']);
+PGTDE::append_to_file($stdout);
+
+$stdout = $node->safe_psql('postgres', "SELECT pg_tde_change_key_provider_file('file-vault', '/tmp/change_key_provider_3.per');", extra_params => ['-a']);
+PGTDE::append_to_file($stdout);
+
+# Restart the server
+PGTDE::append_to_file("-- server restart");
+$rt_value = $node->stop();
+$rt_value = $node->start();
+
+# Verify
+(undef, $stdout, $stderr) = $node->psql('postgres', "SELECT pg_tde_verify_principal_key();", extra_params => ['-a']);
+PGTDE::append_to_file($stdout);
+PGTDE::append_to_file($stderr);
+(undef, $stdout, $stderr) = $node->psql('postgres', "SELECT pg_tde_is_encrypted('test_enc');", extra_params => ['-a']);
+PGTDE::append_to_file($stdout);
+PGTDE::append_to_file($stderr);
+(undef, $stdout, $stderr) = $node->psql('postgres', 'SELECT * FROM test_enc ORDER BY id;', extra_params => ['-a']);
+PGTDE::append_to_file($stdout);
+PGTDE::append_to_file($stderr);
+
+$stdout = $node->safe_psql('postgres', "SELECT pg_tde_change_key_provider_file('file-vault', '/tmp/change_key_provider_4.per');", extra_params => ['-a']);
+PGTDE::append_to_file($stdout);
+
+# Restart the server
+PGTDE::append_to_file("-- server restart");
+$rt_value = $node->stop();
+$rt_value = $node->start();
+
+# Verify
+$stdout = $node->safe_psql('postgres', "SELECT pg_tde_verify_principal_key();", extra_params => ['-a']);
+PGTDE::append_to_file($stdout);
+$stdout = $node->safe_psql('postgres', "SELECT pg_tde_is_encrypted('test_enc');", extra_params => ['-a']);
+PGTDE::append_to_file($stdout);
+$stdout = $node->safe_psql('postgres', 'SELECT * FROM test_enc ORDER BY id;', extra_params => ['-a']);
+PGTDE::append_to_file($stdout);
+
+# DROP EXTENSION
+(undef, $stdout, $stderr) = $node->psql('postgres', 'DROP EXTENSION pg_tde CASCADE;', extra_params => ['-a']);
+PGTDE::append_to_file($stdout);
+PGTDE::append_to_file($stderr);
+
 # Stop the server
 $node->stop();
 
diff --git a/contrib/pg_tde/t/expected/010_change_key_provider.out b/contrib/pg_tde/t/expected/010_change_key_provider.out
index 81459c8d2578c..147a7048efc88 100644
--- a/contrib/pg_tde/t/expected/010_change_key_provider.out
+++ b/contrib/pg_tde/t/expected/010_change_key_provider.out
@@ -63,3 +63,38 @@ SELECT * FROM test_enc ORDER BY id;
 2|6
 DROP EXTENSION pg_tde CASCADE;
 psql:<stdin>:1: NOTICE:  drop cascades to table test_enc
+CREATE EXTENSION IF NOT EXISTS pg_tde;
+SELECT pg_tde_add_key_provider_file('file-vault', '/tmp/change_key_provider_4.per');
+1
+0
+CREATE TABLE test_enc (id serial, k integer, PRIMARY KEY (id)) USING tde_heap;
+INSERT INTO test_enc (k) VALUES (5), (6);
+SELECT pg_tde_verify_principal_key();
+
+SELECT pg_tde_is_encrypted('test_enc');
+t
+SELECT * FROM test_enc ORDER BY id;
+1|5
+2|6
+SELECT pg_tde_change_key_provider_file('file-vault', '/tmp/change_key_provider_3.per');
+1
+-- server restart
+SELECT pg_tde_verify_principal_key();
+
+
+SELECT pg_tde_is_encrypted('test_enc');
+psql:<stdin>:1: ERROR:  Failed to decrypt key, incorrect principal key or corrupted key file
+SELECT * FROM test_enc ORDER BY id;
+psql:<stdin>:1: ERROR:  Failed to decrypt key, incorrect principal key or corrupted key file
+SELECT pg_tde_change_key_provider_file('file-vault', '/tmp/change_key_provider_4.per');
+1
+-- server restart
+SELECT pg_tde_verify_principal_key();
+
+SELECT pg_tde_is_encrypted('test_enc');
+t
+SELECT * FROM test_enc ORDER BY id;
+1|5
+2|6
+DROP EXTENSION pg_tde CASCADE;
+psql:<stdin>:1: NOTICE:  drop cascades to table test_enc


Reviews:
user: login: dutow
body: Forgot to finish this comment, will add as a separate one
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-25T10:12:28Z
user: login: dAdAbird
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-26T19:13:34Z
user: login: jeltz
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-26T21:09:40Z
user: login: jeltz
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-26T21:11:13Z
user: login: jeltz
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-26T21:12:03Z
user: login: jeltz
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-27T13:33:26Z
user: login: dAdAbird
body: Yeah, errors are a bit confusing. It sucks that we don't have an easy way to bump errors up wrap into caller context. We could probably use PG_TRY but it seems like overkill here.
What if make messages something like "Decryption/Encryption failed on EVP_CipherFinal_ex. OpenSSL error...."?

LGTM other than that
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-27T15:49:00Z
user: login: dAdAbird
state: APPROVED
author_association: MEMBER
submitted_at: 2025-03-31T13:05:06Z

Review comments:
diff_hunk: @@ -41,10 +41,20 @@ typedef struct InternalKey
 	(((key)->rel_type & TDE_KEY_TYPE_WAL_UNENCRYPTED) != 0 || \
 	((key)->rel_type & TDE_KEY_TYPE_WAL_ENCRYPTED) != 0)
 
+/* We do not need the dbOid since the entries are stored in a file per db */
+typedef struct TDEMapEntry
path: contrib/pg_tde/src/include/access/pg_tde_tdemap.h
user: login: dutow
body: Since we are changing the format, we could add an algorithm and key length field so we can support more in the future without changing the disk format again.
created_at: 2025-03-24T20:38:03Z
author_association: COLLABORATOR
diff_hunk: @@ -360,24 +316,79 @@ pg_tde_file_header_write(char *tde_filename, int fd, TDEPrincipalKeyInfo *princi
 	return fd;
 }
 
+static void
+pg_tde_initialize_map_entry(TDEMapEntry *map_entry, const TDEPrincipalKey *principal_key, const RelFileLocator *rlocator, const InternalKey *rel_key_data)
+{
+	map_entry->spcOid = rlocator->spcOid;
+	map_entry->relNumber = rlocator->relNumber;
+	map_entry->flags = rel_key_data->rel_type;
+	map_entry->enc_key = *rel_key_data;
+
+	if (!RAND_bytes(map_entry->iv, 16))
+		ereport(ERROR,
+				(errcode(ERRCODE_INTERNAL_ERROR),
+				 errmsg("could not generate iv for key map: %s", ERR_error_string(ERR_get_error(), NULL))));
+
+	AesGcmEncrypt(principal_key->keyData, map_entry->iv, (unsigned char *) map_entry, offsetof(TDEMapEntry, enc_key), rel_key_data->key, INTERNAL_KEY_LEN, map_entry->enc_key.key, map_entry->aead_tag);
path: contrib/pg_tde/src/access/pg_tde_tdemap.c
user: login: dutow
body: I'm not 100% sure about this. We are using the IV meant for use with the rel_key for encrypting the rel_key.
created_at: 2025-03-24T20:54:27Z
author_association: COLLABORATOR
diff_hunk: @@ -401,12 +412,27 @@ pg_tde_write_map_entry(const RelFileLocator *rlocator, InternalKey *enc_rel_key_
 			break;
 	}
 
+	/* Initialize map entry and encrypt key */
+	pg_tde_initialize_map_entry(&write_map_entry, principal_key, rlocator, rel_key_data);
+
+	if (write_xlog)
path: contrib/pg_tde/src/access/pg_tde_tdemap.c
user: login: dutow
body: Do we still xlog internal keys?
created_at: 2025-03-24T21:01:01Z
author_association: COLLABORATOR
diff_hunk: @@ -401,12 +412,27 @@ pg_tde_write_map_entry(const RelFileLocator *rlocator, InternalKey *enc_rel_key_
 			break;
 	}
 
+	/* Initialize map entry and encrypt key */
+	pg_tde_initialize_map_entry(&write_map_entry, principal_key, rlocator, rel_key_data);
+
+	if (write_xlog)
path: contrib/pg_tde/src/access/pg_tde_tdemap.c
user: login: dAdAbird
body: We don't [XLog the copy of the default key](https://github.com/percona/postgres/commit/d1eace54279550509d81dff64e245f2b1a2ba68c) anymore. But XLog of internal keys is still there. We should get rid of that but I reckon with the separate PR
created_at: 2025-03-26T19:13:33Z
author_association: MEMBER
diff_hunk: @@ -41,10 +41,20 @@ typedef struct InternalKey
 	(((key)->rel_type & TDE_KEY_TYPE_WAL_UNENCRYPTED) != 0 || \
 	((key)->rel_type & TDE_KEY_TYPE_WAL_ENCRYPTED) != 0)
 
+/* We do not need the dbOid since the entries are stored in a file per db */
+typedef struct TDEMapEntry
path: contrib/pg_tde/src/include/access/pg_tde_tdemap.h
user: login: jeltz
body: Agreed, that is a good thing to have before GA.
created_at: 2025-03-26T21:09:40Z
author_association: COLLABORATOR
diff_hunk: @@ -360,24 +316,79 @@ pg_tde_file_header_write(char *tde_filename, int fd, TDEPrincipalKeyInfo *princi
 	return fd;
 }
 
+static void
+pg_tde_initialize_map_entry(TDEMapEntry *map_entry, const TDEPrincipalKey *principal_key, const RelFileLocator *rlocator, const InternalKey *rel_key_data)
+{
+	map_entry->spcOid = rlocator->spcOid;
+	map_entry->relNumber = rlocator->relNumber;
+	map_entry->flags = rel_key_data->rel_type;
+	map_entry->enc_key = *rel_key_data;
+
+	if (!RAND_bytes(map_entry->iv, 16))
+		ereport(ERROR,
+				(errcode(ERRCODE_INTERNAL_ERROR),
+				 errmsg("could not generate iv for key map: %s", ERR_error_string(ERR_get_error(), NULL))));
+
+	AesGcmEncrypt(principal_key->keyData, map_entry->iv, (unsigned char *) map_entry, offsetof(TDEMapEntry, enc_key), rel_key_data->key, INTERNAL_KEY_LEN, map_entry->enc_key.key, map_entry->aead_tag);
path: contrib/pg_tde/src/access/pg_tde_tdemap.c
user: login: jeltz
body: No, my idea was to have two IVs (one for encrypting the key and one for encrypting the relation) but I never got around changing the name from just "IV" because I could not think of a good name.
created_at: 2025-03-26T21:11:12Z
author_association: COLLABORATOR
diff_hunk: @@ -401,12 +412,27 @@ pg_tde_write_map_entry(const RelFileLocator *rlocator, InternalKey *enc_rel_key_
 			break;
 	}
 
+	/* Initialize map entry and encrypt key */
+	pg_tde_initialize_map_entry(&write_map_entry, principal_key, rlocator, rel_key_data);
+
+	if (write_xlog)
path: contrib/pg_tde/src/access/pg_tde_tdemap.c
user: login: jeltz
body: Yeah, this PR is big enough as-is but maybe we should priotize doing that first so we then can make this PR smaller?
created_at: 2025-03-26T21:12:03Z
author_association: COLLABORATOR
diff_hunk: @@ -360,24 +316,79 @@ pg_tde_file_header_write(char *tde_filename, int fd, TDEPrincipalKeyInfo *princi
 	return fd;
 }
 
+static void
+pg_tde_initialize_map_entry(TDEMapEntry *map_entry, const TDEPrincipalKey *principal_key, const RelFileLocator *rlocator, const InternalKey *rel_key_data)
+{
+	map_entry->spcOid = rlocator->spcOid;
+	map_entry->relNumber = rlocator->relNumber;
+	map_entry->flags = rel_key_data->rel_type;
+	map_entry->enc_key = *rel_key_data;
+
+	if (!RAND_bytes(map_entry->iv, 16))
+		ereport(ERROR,
+				(errcode(ERRCODE_INTERNAL_ERROR),
+				 errmsg("could not generate iv for key map: %s", ERR_error_string(ERR_get_error(), NULL))));
+
+	AesGcmEncrypt(principal_key->keyData, map_entry->iv, (unsigned char *) map_entry, offsetof(TDEMapEntry, enc_key), rel_key_data->key, INTERNAL_KEY_LEN, map_entry->enc_key.key, map_entry->aead_tag);
path: contrib/pg_tde/src/access/pg_tde_tdemap.c
user: login: jeltz
body: Cleared this up
created_at: 2025-03-27T13:33:25Z
author_association: COLLABORATOR

PR comments:
user: login: dutow
created_at: 2025-03-25T12:08:10Z
body: My original proposal was about verifying that the file itself is 100% consistent - this is a bit different, as it verifies things per encrypted file instead. This could be both a good and a bad thing.

The previous "encrypt only the key" was a quick implementation mainly because of our two file architecture, it was easy to implement with it. With that gone, we could possibly encrypt simply the entire file.

The current solution is probably better for using a single or few tables per connection, while the other could be better in the opposite case, with tons of small encrypted files.
author_association: COLLABORATOR
user: login: jeltz
created_at: 2025-03-26T21:14:28Z
body: Yeah, I also thought about encrypting and decrypting the whole file but got stuck on our current random access writes which for example update the LSN. Let me think a bit more on this and let's discuss this in a meeting or on Slack.
author_association: COLLABORATOR
user: login: dAdAbird
created_at: 2025-03-27T08:03:29Z
body: Not only writing LSN but also deleting keys. Also it can be thousands of tables per database.
Plus let's consider that we should read (decrypt) internal keys right into the locked memory which might be limited. That's not the case now but I'm planning to do in PG-1445
author_association: MEMBER
user: login: jeltz
created_at: 2025-03-27T13:54:23Z
body: Makring this ready for review. Maybe we should have better error messages though. What do you think? You can see in the test output what they look like.
author_association: COLLABORATOR
user: login: dutow
created_at: 2025-03-27T19:32:48Z
body: The openssl error part doesn't seem to be useful, looks like openssl doesn't report anything meaningful there - maybe we could skip that part and only say that the call failed, without the openssl details part.
author_association: COLLABORATOR
user: login: jeltz
created_at: 2025-03-31T12:57:47Z
body: Added an improved error message.
author_association: COLLABORATOR