diff --git a/contrib/pg_tde/src/access/pg_tde_tdemap.c b/contrib/pg_tde/src/access/pg_tde_tdemap.c
index 70b6722a2ebc4..77159f535bf6e 100644
--- a/contrib/pg_tde/src/access/pg_tde_tdemap.c
+++ b/contrib/pg_tde/src/access/pg_tde_tdemap.c
@@ -39,8 +39,11 @@
 }
 #endif
 
-#define PG_TDE_FILEMAGIC			0x03454454	/* version ID value = TDE 03 */
-#define PG_TDE_MAP_FILENAME			"%d_keys"
+#define PG_TDE_FILEMAGIC_OLD			0x03454454	/* old version ID value = TDE 03 */
+#define PG_TDE_MAP_OLD_FNAME_SUFFIX			"_keys"
+
+#define PG_TDE_FILEMAGIC			0x04454454	/* version ID value = TDE 04 */
+#define PG_TDE_MAP_FILENAME			"%d_keys_v2" /* TODO: should be _v2 of _v4 ? */
 
 typedef enum
 {
@@ -63,14 +66,14 @@ typedef struct TDEFileHeader
  * encrypting/decrypting existing keys from the key files, so any changes here
  * might break existing clusters.
  */
-typedef struct TDEMapEntry
+typedef struct TDEMapEntryOld
 {
 	Oid			spcOid;			/* Part of AAD */
 	RelFileNumber relNumber;	/* Part of AAD */
 	uint32		type;			/* Part of AAD */
 	uint32		_unused3;		/* Part of AAD */
 
-	uint8		encrypted_key_data[INTERNAL_KEY_LEN];
+	uint8		encrypted_key_data[INTERNAL_KEY_OLD_LEN];
 	uint8		key_base_iv[INTERNAL_KEY_IV_LEN];
 
 	uint32		_unused1;		/* Will be 1 in existing files entries. */
@@ -80,6 +83,21 @@ typedef struct TDEMapEntry
 	/* IV and tag used when encrypting the key itself */
 	unsigned char entry_iv[MAP_ENTRY_IV_SIZE];
 	unsigned char aead_tag[MAP_ENTRY_AEAD_TAG_SIZE];
+} TDEMapEntryOld;
+
+typedef struct TDEMapEntry
+{
+	uint32		key_len;		/* Part of AAD */
+	Oid			spcOid;			/* Part of AAD */
+	RelFileNumber relNumber;	/* Part of AAD */
+	uint32		type;			/* Part of AAD */
+
+	/* IV and tag used when encrypting the key itself */
+	unsigned char entry_iv[MAP_ENTRY_IV_SIZE];
+	unsigned char aead_tag[MAP_ENTRY_AEAD_TAG_SIZE];
+
+	uint8		key_base_iv[INTERNAL_KEY_IV_LEN];
+	uint8		encrypted_key_data[INTERNAL_KEY_MAX_LEN];
 } TDEMapEntry;
 
 static void pg_tde_set_db_file_path(Oid dbOid, char *path);
@@ -380,7 +398,7 @@ pg_tde_sign_principal_key_info(TDESignedPrincipalKeyInfo *signed_key_info, const
 				errcode(ERRCODE_INTERNAL_ERROR),
 				errmsg("could not generate iv for key map: %s", ERR_error_string(ERR_get_error(), NULL)));
 
-	AesGcmEncrypt(principal_key->keyData,
+	AesGcmEncrypt(principal_key->keyData, principal_key->keyLength,
 				  signed_key_info->sign_iv, MAP_ENTRY_IV_SIZE,
 				  (unsigned char *) &signed_key_info->data, sizeof(signed_key_info->data),
 				  NULL, 0,
@@ -397,22 +415,17 @@ pg_tde_initialize_map_entry(TDEMapEntry *map_entry, const TDEPrincipalKey *princ
 	map_entry->type = MAP_ENTRY_TYPE_KEY;
 	memcpy(map_entry->key_base_iv, rel_key_data->base_iv, INTERNAL_KEY_IV_LEN);
 
-	/*
-	 * We set these fields here so that existing file entries will be
-	 * consistent and future use of these fields easier.
-	 */
-	map_entry->_unused1 = 1;
-	map_entry->_unused2 = 0;
+	map_entry->key_len = rel_key_data->key_len;
 
 	if (!RAND_bytes(map_entry->entry_iv, MAP_ENTRY_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,
+	AesGcmEncrypt(principal_key->keyData, principal_key->keyLength,
 				  map_entry->entry_iv, MAP_ENTRY_IV_SIZE,
-				  (unsigned char *) map_entry, offsetof(TDEMapEntry, encrypted_key_data),
-				  rel_key_data->key, INTERNAL_KEY_LEN,
+				  (unsigned char *) map_entry, offsetof(TDEMapEntry, entry_iv),
+				  rel_key_data->key, rel_key_data->key_len,
 				  map_entry->encrypted_key_data,
 				  map_entry->aead_tag, MAP_ENTRY_AEAD_TAG_SIZE);
 }
@@ -573,7 +586,7 @@ pg_tde_count_encryption_keys(Oid dbOid)
 bool
 pg_tde_verify_principal_key_info(TDESignedPrincipalKeyInfo *signed_key_info, const KeyData *principal_key_data)
 {
-	return AesGcmDecrypt(principal_key_data->data,
+	return AesGcmDecrypt(principal_key_data->data, principal_key_data->len,
 						 signed_key_info->sign_iv, MAP_ENTRY_IV_SIZE,
 						 (unsigned char *) &signed_key_info->data, sizeof(signed_key_info->data),
 						 NULL, 0,
@@ -588,16 +601,17 @@ tde_decrypt_rel_key(const TDEPrincipalKey *principal_key, TDEMapEntry *map_entry
 
 	Assert(principal_key);
 
-	if (!AesGcmDecrypt(principal_key->keyData,
+	if (!AesGcmDecrypt(principal_key->keyData, principal_key->keyLength,
 					   map_entry->entry_iv, MAP_ENTRY_IV_SIZE,
-					   (unsigned char *) map_entry, offsetof(TDEMapEntry, encrypted_key_data),
-					   map_entry->encrypted_key_data, INTERNAL_KEY_LEN,
+					   (unsigned char *) map_entry, offsetof(TDEMapEntry, entry_iv),
+					   map_entry->encrypted_key_data, map_entry->key_len,
 					   key->key,
 					   map_entry->aead_tag, MAP_ENTRY_AEAD_TAG_SIZE))
 		ereport(ERROR,
 				errmsg("Failed to decrypt key, incorrect principal key or corrupted key file"));
 
 	memcpy(key->base_iv, map_entry->key_base_iv, INTERNAL_KEY_IV_LEN);
+	key->key_len = map_entry->key_len;
 
 	return key;
 }
@@ -785,10 +799,36 @@ pg_tde_get_principal_key_info(Oid dbOid)
 	fd = pg_tde_open_file_basic(db_map_path, O_RDONLY, true);
 
 	/* The file does not exist. */
-	if (fd < 0)
-		return NULL;
+	if (fd >= 0)
+		pg_tde_file_header_read(db_map_path, fd, &fheader, &bytes_read);
+	else
+	{
+		/* 
+		 * TODO: An ugly hack for now, we need to get a key info when rewriting 
+		 * an old file...
+		 */
+		char	old_key_file_path[MAXPGPATH] = {0};
+		char	*fname = psprintf("%d"PG_TDE_MAP_OLD_FNAME_SUFFIX, dbOid);
+
+		join_path_components(old_key_file_path, pg_tde_get_data_dir(), fname);
+		pfree(fname);
+
+		fd = pg_tde_open_file_basic(old_key_file_path, O_RDONLY, true);
+
+		/* The file does not exist */
+		if (fd < 0)
+			return NULL;
 
-	pg_tde_file_header_read(db_map_path, fd, &fheader, &bytes_read);
+		bytes_read = pg_pread(fd, &fheader, sizeof(TDEFileHeader), 0);
+
+		if (bytes_read > 0 && (bytes_read != sizeof(TDEFileHeader)
+			|| fheader.file_version != PG_TDE_FILEMAGIC_OLD))
+		{
+			ereport(FATAL,
+					errcode_for_file_access(),
+					errmsg("old smgr key file \"%s\" is corrupted: %m", old_key_file_path));
+		}
+	}
 
 	CloseTransientFile(fd);
 
@@ -881,3 +921,189 @@ pg_tde_get_smgr_key(RelFileLocator rel)
 
 	return rel_key;
 }
+
+#ifndef FRONTEND
+
+/* 
+ * Functions for rewriting old smgr _key into a new format file.
+ *
+ * TODO: The old format should be deprecated. And this code should be removed
+ * eventually.
+ */
+
+static int
+pg_tde_open_old_file_read(const char *tde_filename, bool ignore_missing, off_t *curr_pos)
+{
+	int			fd;
+	TDEFileHeader fheader;
+	off_t		bytes_read = 0;
+
+	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, ignore_missing);
+	if (ignore_missing && fd < 0)
+		return fd;
+
+	bytes_read = pg_pread(fd, &fheader, sizeof(TDEFileHeader), 0);
+
+	/* File is empty */
+	if (bytes_read == 0)
+		return fd;
+
+	if (bytes_read != sizeof(TDEFileHeader)
+		|| fheader.file_version != PG_TDE_FILEMAGIC_OLD)
+	{
+		ereport(FATAL,
+				errcode_for_file_access(),
+				errmsg("old TDE map file \"%s\" is corrupted: %m", tde_filename));
+	}
+
+	*curr_pos = bytes_read;
+
+	return fd;
+}
+
+static bool
+pg_tde_read_one_old_map_entry(int map_file, TDEMapEntryOld *map_entry, off_t *offset)
+{
+	off_t		bytes_read = 0;
+
+	Assert(map_entry);
+	Assert(offset);
+
+	bytes_read = pg_pread(map_file, map_entry, sizeof(TDEMapEntryOld), *offset);
+
+	/* We've reached the end of the file. */
+	if (bytes_read != sizeof(TDEMapEntryOld))
+		return false;
+
+	*offset += bytes_read;
+
+	return true;
+}
+
+static InternalKey *
+tde_decrypt_old_rel_key(const TDEPrincipalKey *principal_key, TDEMapEntryOld *map_entry)
+{
+	InternalKey *key = palloc_object(InternalKey);
+
+	Assert(principal_key);
+
+	if (!AesGcmDecrypt(principal_key->keyData, principal_key->keyLength,
+					   map_entry->entry_iv, MAP_ENTRY_IV_SIZE,
+					   (unsigned char *) map_entry, offsetof(TDEMapEntryOld, encrypted_key_data),
+					   map_entry->encrypted_key_data, INTERNAL_KEY_OLD_LEN,
+					   key->key,
+					   map_entry->aead_tag, MAP_ENTRY_AEAD_TAG_SIZE))
+		ereport(ERROR,
+				errmsg("Failed to decrypt key, incorrect principal key or corrupted key file"));
+
+	memcpy(key->base_iv, map_entry->key_base_iv, INTERNAL_KEY_IV_LEN);
+	key->key_len = INTERNAL_KEY_OLD_LEN;
+
+	return key;
+}
+
+void
+pg_tde_migrate_smgr_keys_file(void)
+{
+	DIR			*dir;
+	LWLock	   *lock_pk = tde_lwlock_enc_keys();
+	struct dirent *file;
+	TDEPrincipalKey *principal_key = NULL;
+	TDESignedPrincipalKeyInfo signed_key_info;
+	int old_suffix_len = strlen(PG_TDE_MAP_OLD_FNAME_SUFFIX);
+
+	/* 
+	 * No real need in lock here as the func should be called only on the server
+	 * start, but GetPrincipalKey() expects lock.
+	*/
+	LWLockAcquire(lock_pk, LW_EXCLUSIVE);
+
+	dir = opendir(pg_tde_get_data_dir());
+	if (dir == NULL && errno != ENOENT)
+		elog(ERROR, "could not open directory \"%s\": %m",
+						pg_tde_get_data_dir());
+
+	/* 
+	 * TODO: create a "current_version_v4" file after the successfull migration
+	 * or during pg_tde dir creatation, so no need to scan the dir all the time,
+	 * but just test the file?  
+	 */
+	while (errno = 0, (file = readdir(dir)) != NULL)
+	{
+		if (strlen(file->d_name) > old_suffix_len &&
+			strncmp(file->d_name + strlen(file->d_name) - old_suffix_len, PG_TDE_MAP_OLD_FNAME_SUFFIX, old_suffix_len) == 0)
+		{
+			char	old_db_map_path[MAXPGPATH] = {0};
+			char	db_map_path[MAXPGPATH] = {0};
+			off_t	read_pos,
+					write_pos;
+			int		old_fd,
+					new_fd;
+			Oid		dbOid;
+			char 	*suffix;
+			char	*old_fname;
+
+
+			dbOid = strtoul(file->d_name, &suffix, 10);
+			if (strcmp(suffix, PG_TDE_MAP_OLD_FNAME_SUFFIX) != 0)
+				continue;
+
+			old_fname = psprintf("%d"PG_TDE_MAP_OLD_FNAME_SUFFIX, dbOid);
+			join_path_components(old_db_map_path, pg_tde_get_data_dir(), old_fname);
+			pfree(old_fname);
+
+			old_fd = pg_tde_open_old_file_read(old_db_map_path, false, &read_pos);
+
+			pg_tde_set_db_file_path(dbOid, db_map_path);
+			/* 
+			 * The old file exists and it's not empty, hece a principal key
+			 * should exist as well.
+			 */
+			if (principal_key == NULL)
+			{
+				principal_key = GetPrincipalKey(dbOid, LW_EXCLUSIVE);
+				if (principal_key == NULL)
+				{
+					ereport(ERROR,
+							errmsg("could not get server principal key"),
+							errdetail("Failed to migrate the keys file of %u database.", dbOid));
+				}
+				pg_tde_sign_principal_key_info(&signed_key_info, principal_key);
+			}
+			new_fd = pg_tde_open_file_write(db_map_path, &signed_key_info, true, &write_pos);
+
+			while (1)
+			{
+				TDEMapEntryOld	old_entry;
+				TDEMapEntry		new_entry;
+				InternalKey *rel_key_data;
+				RelFileLocator rloc;
+
+				if (!pg_tde_read_one_old_map_entry(old_fd, &old_entry, &read_pos))
+					break;
+
+				rloc.spcOid = old_entry.spcOid;
+				rloc.dbOid = dbOid;
+				rloc.relNumber = old_entry.relNumber;
+
+				rel_key_data = tde_decrypt_old_rel_key(principal_key, &old_entry);
+				pg_tde_initialize_map_entry(&new_entry, principal_key, &rloc, rel_key_data);
+				pg_tde_write_one_map_entry(new_fd, &new_entry, &write_pos, db_map_path);
+
+				pfree(rel_key_data);
+			}
+
+			CloseTransientFile(old_fd);
+			CloseTransientFile(new_fd);
+			durable_unlink(old_db_map_path, ERROR);
+		}
+	}
+
+	closedir(dir);
+	LWLockRelease(lock_pk);
+}
+
+#endif
\ No newline at end of file
diff --git a/contrib/pg_tde/src/access/pg_tde_xlog_keys.c b/contrib/pg_tde/src/access/pg_tde_xlog_keys.c
index 06253573fc9ac..b31538b5a19d3 100644
--- a/contrib/pg_tde/src/access/pg_tde_xlog_keys.c
+++ b/contrib/pg_tde/src/access/pg_tde_xlog_keys.c
@@ -22,8 +22,11 @@
 #include "pg_tde_fe.h"
 #endif
 
-#define PG_TDE_WAL_KEY_FILE_MAGIC 0x014B4557	/* version ID value = WEK 01 */
-#define PG_TDE_WAL_KEY_FILE_NAME "wal_keys"
+#define PG_TDE_WAL_KEY_FILE_MAGIC_OLD 0x014B4557	/* old version ID value = WEK 01 */
+#define PG_TDE_WAL_KEY_FILE_NAME_OLD "wal_keys"
+
+#define PG_TDE_WAL_KEY_FILE_MAGIC 0x024B4557	/* version ID value = WEK 02 */
+#define PG_TDE_WAL_KEY_FILE_NAME "wal_keys_v2"
 
 typedef struct WalKeyFileHeader
 {
@@ -40,12 +43,12 @@ typedef struct WalKeyFileHeader
  * encrypting/decrypting existing keys from the key files, so any changes here
  * might break existing clusters.
  */
-typedef struct WalKeyFileEntry
+typedef struct WalKeyFileEntryOld
 {
 	uint32		_unused1;		/* Part of AAD, is 1 or 2 in existing entries */
 	uint32		_unused2;		/* Part of AAD */
 
-	uint8		encrypted_key_data[INTERNAL_KEY_LEN];
+	uint8		encrypted_key_data[INTERNAL_KEY_OLD_LEN];
 	uint8		key_base_iv[INTERNAL_KEY_IV_LEN];
 
 	uint32		range_type;		/* WalEncryptionRangeType */
@@ -55,6 +58,20 @@ typedef struct WalKeyFileEntry
 	/* IV and tag used when encrypting the key itself */
 	unsigned char entry_iv[MAP_ENTRY_IV_SIZE];
 	unsigned char aead_tag[MAP_ENTRY_AEAD_TAG_SIZE];
+} WalKeyFileEntryOld;
+
+typedef struct WalKeyFileEntry
+{
+	uint32		key_len;
+	uint32		range_type;		/* WalEncryptionRangeType */
+	WalLocation range_start;
+
+	/* IV and tag used when encrypting the key itself */
+	unsigned char entry_iv[MAP_ENTRY_IV_SIZE];
+	unsigned char aead_tag[MAP_ENTRY_AEAD_TAG_SIZE];
+
+	uint8		key_base_iv[INTERNAL_KEY_IV_LEN];
+	uint8		encrypted_key_data[INTERNAL_KEY_MAX_LEN];
 } WalKeyFileEntry;
 
 static WALKeyCacheRec *tde_wal_key_cache = NULL;
@@ -178,7 +195,7 @@ pg_tde_wal_last_range_set_location(WalLocation loc)
  * with the actual lsn by the first WAL write.
  */
 void
-pg_tde_create_wal_range(WalEncryptionRange *range, WalEncryptionRangeType type)
+pg_tde_create_wal_range(WalEncryptionRange *range, WalEncryptionRangeType type, int key_len)
 {
 	TDEPrincipalKey *principal_key;
 
@@ -199,7 +216,7 @@ pg_tde_create_wal_range(WalEncryptionRange *range, WalEncryptionRangeType type)
 	range->end.lsn = MaxXLogRecPtr;
 	range->end.tli = MaxTimeLineID;
 
-	pg_tde_generate_internal_key(&range->key);
+	pg_tde_generate_internal_key(&range->key, key_len);
 
 	pg_tde_write_wal_key_file_entry(range, principal_key);
 
@@ -615,12 +632,13 @@ pg_tde_wal_range_from_entry(const TDEPrincipalKey *principal_key, WalKeyFileEntr
 	range->start = entry->range_start;
 	range->end.tli = MaxTimeLineID;
 	range->end.lsn = MaxXLogRecPtr;
+	range->key.key_len = entry->key_len;
 
 	memcpy(range->key.base_iv, entry->key_base_iv, INTERNAL_KEY_IV_LEN);
-	if (!AesGcmDecrypt(principal_key->keyData,
+	if (!AesGcmDecrypt(principal_key->keyData, principal_key->keyLength,
 					   entry->entry_iv, MAP_ENTRY_IV_SIZE,
-					   (unsigned char *) entry, offsetof(WalKeyFileEntry, encrypted_key_data),
-					   entry->encrypted_key_data, INTERNAL_KEY_LEN,
+					   (unsigned char *) entry, offsetof(WalKeyFileEntry, range_start),
+					   entry->encrypted_key_data, entry->key_len,
 					   range->key.key,
 					   entry->aead_tag, MAP_ENTRY_AEAD_TAG_SIZE))
 		ereport(ERROR,
@@ -664,13 +682,7 @@ pg_tde_initialize_wal_key_file_entry(WalKeyFileEntry *entry,
 
 	memset(entry, 0, sizeof(WalKeyFileEntry));
 
-	/*
-	 * We set this field here so that existing file entries will be consistent
-	 * and future use of this field easier. Some existing entries will have 2
-	 * here.
-	 */
-	entry->_unused1 = 1;
-
+	entry->key_len = range->key.key_len;
 	entry->range_type = range->type;
 	entry->range_start = range->start;
 	memcpy(entry->key_base_iv, range->key.base_iv, INTERNAL_KEY_IV_LEN);
@@ -680,10 +692,10 @@ pg_tde_initialize_wal_key_file_entry(WalKeyFileEntry *entry,
 				errcode(ERRCODE_INTERNAL_ERROR),
 				errmsg("could not generate iv for wal key file entry: %s", ERR_error_string(ERR_get_error(), NULL)));
 
-	AesGcmEncrypt(principal_key->keyData,
+	AesGcmEncrypt(principal_key->keyData, principal_key->keyLength,
 				  entry->entry_iv, MAP_ENTRY_IV_SIZE,
-				  (unsigned char *) entry, offsetof(WalKeyFileEntry, encrypted_key_data),
-				  range->key.key, INTERNAL_KEY_LEN,
+				  (unsigned char *) entry, offsetof(WalKeyFileEntry, range_start),
+				  range->key.key, range->key.key_len,
 				  entry->encrypted_key_data,
 				  entry->aead_tag, MAP_ENTRY_AEAD_TAG_SIZE);
 }
@@ -834,11 +846,34 @@ pg_tde_get_server_key_info(void)
 	 */
 	fd = pg_tde_open_wal_key_file_basic(get_wal_key_file_path(), O_RDONLY, true);
 
-	/* The file does not exist. */
-	if (fd < 0)
-		return NULL;
+	if (fd >= 0)
+		pg_tde_wal_key_file_header_read(get_wal_key_file_path(), fd, &fheader, &bytes_read);
+	else
+	{
+		/* 
+		 * TODO: An ugly hack for now, we need to get a key info when rewriting 
+		 * an old file...
+		 */
+		char	old_wal_key_file_path[MAXPGPATH] = {0};
+
+		join_path_components(old_wal_key_file_path, pg_tde_get_data_dir(), PG_TDE_WAL_KEY_FILE_NAME_OLD);
 
-	pg_tde_wal_key_file_header_read(get_wal_key_file_path(), fd, &fheader, &bytes_read);
+		fd = pg_tde_open_wal_key_file_basic(old_wal_key_file_path, O_RDONLY, true);
+
+		/* The file does not exist */
+		if (fd < 0)
+			return NULL;
+
+		bytes_read = pg_pread(fd, &fheader, sizeof(WalKeyFileHeader), 0);
+
+		if (bytes_read > 0 && (bytes_read != sizeof(WalKeyFileHeader)
+			|| fheader.file_version != PG_TDE_WAL_KEY_FILE_MAGIC_OLD))
+		{
+			ereport(FATAL,
+					errcode_for_file_access(),
+					errmsg("old WAL key file \"%s\" is corrupted: %m", old_wal_key_file_path));
+		}
+	}
 
 	CloseTransientFile(fd);
 
@@ -894,4 +929,157 @@ pg_tde_delete_server_key(void)
 	/* Remove whole key map file */
 	durable_unlink(get_wal_key_file_path(), ERROR);
 }
+
+/* 
+ * Functions for rewriting old wal_keys into a new format file.
+ *
+ * TODO: The old format should be deprecated. And this code should be removed
+ * eventually.
+ */
+static int
+pg_tde_open_old_wal_key_file_read(const char *filename,
+							  bool ignore_missing,
+							  off_t *curr_pos)
+{
+	int			fd;
+	WalKeyFileHeader fheader;
+	off_t		bytes_read = 0;
+
+	Assert(LWLockHeldByMeInMode(tde_lwlock_enc_keys(), LW_SHARED) ||
+		   LWLockHeldByMeInMode(tde_lwlock_enc_keys(), LW_EXCLUSIVE));
+
+	fd = pg_tde_open_wal_key_file_basic(filename, O_RDONLY | PG_BINARY, ignore_missing);
+	if (ignore_missing && fd < 0)
+		return fd;
+
+	bytes_read = pg_pread(fd, &fheader, sizeof(WalKeyFileHeader), 0);
+
+	/* File is empty */
+	if (bytes_read == 0)
+		return fd;
+
+	if (bytes_read != sizeof(WalKeyFileHeader)
+		|| fheader.file_version != PG_TDE_WAL_KEY_FILE_MAGIC_OLD)
+	{
+		ereport(FATAL,
+				errcode_for_file_access(),
+				errmsg("old WAL key file \"%s\" is corrupted: %m", filename));
+	}
+	*curr_pos = bytes_read;
+
+	return fd;
+}
+
+static bool
+pg_tde_read_one_wal_key_file_old_entry(int fd,
+								   WalKeyFileEntryOld *entry,
+								   off_t *offset)
+{
+	off_t		bytes_read = 0;
+
+	Assert(entry);
+	Assert(offset);
+
+	bytes_read = pg_pread(fd, entry, sizeof(WalKeyFileEntryOld), *offset);
+
+	/* We've reached the end of the file. */
+	if (bytes_read != sizeof(WalKeyFileEntryOld))
+		return false;
+
+	*offset += bytes_read;
+
+	return true;
+}
+
+static WalEncryptionRange *
+pg_tde_wal_range_from_old_entry(const TDEPrincipalKey *principal_key, WalKeyFileEntryOld *entry)
+{
+	WalEncryptionRange *range = tde_wal_prealloc_range == NULL ? palloc0_object(WalEncryptionRange) : tde_wal_prealloc_range;
+
+	tde_wal_prealloc_range = NULL;
+
+	Assert(principal_key);
+
+	range->type = entry->range_type;
+	range->start = entry->range_start;
+	range->end.tli = MaxTimeLineID;
+	range->end.lsn = MaxXLogRecPtr;
+	range->key.key_len = INTERNAL_KEY_OLD_LEN;
+
+	memcpy(range->key.base_iv, entry->key_base_iv, INTERNAL_KEY_IV_LEN);
+	if (!AesGcmDecrypt(principal_key->keyData, principal_key->keyLength,
+					   entry->entry_iv, MAP_ENTRY_IV_SIZE,
+					   (unsigned char *) entry, offsetof(WalKeyFileEntryOld, encrypted_key_data),
+					   entry->encrypted_key_data, INTERNAL_KEY_OLD_LEN,
+					   range->key.key,
+					   entry->aead_tag, MAP_ENTRY_AEAD_TAG_SIZE))
+		ereport(ERROR,
+				errmsg("Failed to decrypt key, incorrect principal key or corrupted key file %u", principal_key->keyLength));
+
+	return range;
+}
+
+void
+pg_tde_update_wal_keys_file(void)
+{
+	LWLock	   *lock_pk = tde_lwlock_enc_keys();
+	TDEPrincipalKey *principal_key;
+	TDESignedPrincipalKeyInfo signed_key_info;
+	char	wal_key_file_path[MAXPGPATH] = {0};
+	off_t	read_pos,
+			write_pos;
+	int		old_fd,
+			new_fd;
+
+	join_path_components(wal_key_file_path, pg_tde_get_data_dir(), PG_TDE_WAL_KEY_FILE_NAME_OLD);
+
+	/* Check if there is anything to do */
+	if (access(wal_key_file_path, F_OK) != 0)
+	 	return;
+
+	/* 
+	 * No real need in lock here as the func should be called only on the server
+	 * start, but GetPrincipalKey() expects one. 
+	 */
+	LWLockAcquire(lock_pk, LW_EXCLUSIVE);
+
+
+	old_fd = pg_tde_open_old_wal_key_file_read(wal_key_file_path, false, &read_pos);
+
+	/* 
+	 * The old file exists and it's not empty, hece a principal key
+	 * should exist as well.
+	 */
+	principal_key = GetPrincipalKey(GLOBAL_DATA_TDE_OID, LW_EXCLUSIVE);
+	if (principal_key == NULL)
+	{
+		ereport(ERROR,
+				errmsg("could not get server principal key"),
+				errdetail("Failed to updated format of WAL keys."));
+	}
+	pg_tde_sign_principal_key_info(&signed_key_info, principal_key);
+
+	new_fd = pg_tde_open_wal_key_file_write(get_wal_key_file_path(), &signed_key_info, true, &write_pos);
+
+	while (1)
+	{
+		WalKeyFileEntryOld	old_entry;
+		WalKeyFileEntry		new_entry;
+		WalEncryptionRange	*range;
+
+		if (!pg_tde_read_one_wal_key_file_old_entry(old_fd, &old_entry, &read_pos))
+			break;
+
+		range = pg_tde_wal_range_from_old_entry(principal_key, &old_entry);
+		pg_tde_initialize_wal_key_file_entry(&new_entry, principal_key, range);
+		pg_tde_write_one_wal_key_file_entry(new_fd, &new_entry, &write_pos, get_wal_key_file_path());
+		pfree(range);
+	}
+
+	CloseTransientFile(old_fd);
+	CloseTransientFile(new_fd);
+	durable_unlink(wal_key_file_path, ERROR);
+
+	LWLockRelease(lock_pk);
+}
 #endif
diff --git a/contrib/pg_tde/src/access/pg_tde_xlog_smgr.c b/contrib/pg_tde/src/access/pg_tde_xlog_smgr.c
index c8e62a1edc144..a2117763c3fd0 100644
--- a/contrib/pg_tde/src/access/pg_tde_xlog_smgr.c
+++ b/contrib/pg_tde/src/access/pg_tde_xlog_smgr.c
@@ -221,7 +221,7 @@ TDEXLogSmgrInit()
 }
 
 void
-TDEXLogSmgrInitWrite(bool encrypt_xlog)
+TDEXLogSmgrInitWrite(bool encrypt_xlog, int key_len)
 {
 	WalEncryptionRange *range;
 	WALKeyCacheRec *keys;
@@ -242,11 +242,11 @@ TDEXLogSmgrInitWrite(bool encrypt_xlog)
 	 */
 	if (encrypt_xlog)
 	{
-		pg_tde_create_wal_range(&CurrentWalEncryptionRange, WAL_ENCRYPTION_RANGE_ENCRYPTED);
+		pg_tde_create_wal_range(&CurrentWalEncryptionRange, WAL_ENCRYPTION_RANGE_ENCRYPTED, key_len);
 	}
 	else if (range && range->type == WAL_ENCRYPTION_RANGE_ENCRYPTED)
 	{
-		pg_tde_create_wal_range(&CurrentWalEncryptionRange, WAL_ENCRYPTION_RANGE_UNENCRYPTED);
+		pg_tde_create_wal_range(&CurrentWalEncryptionRange, WAL_ENCRYPTION_RANGE_UNENCRYPTED, key_len);
 	}
 	else if (range)
 	{
@@ -353,6 +353,7 @@ TDEXLogWriteEncryptedPages(int fd, const void *buf, size_t count, off_t offset,
 						count,
 						enc_buff,
 						range->key.key,
+						range->key.key_len,
 						&EncryptionCryptCtx);
 
 	return pg_pwrite(fd, enc_buff, count, offset);
@@ -584,6 +585,7 @@ TDEXLogCryptBuffer(const void *buf, void *out_buf, size_t count, off_t offset,
 									dec_sz,
 									o_buf,
 									curr_key->range.key.key,
+									curr_key->range.key.key_len,
 									&curr_key->crypt_ctx);
 			}
 		}
diff --git a/contrib/pg_tde/src/catalog/tde_principal_key.c b/contrib/pg_tde/src/catalog/tde_principal_key.c
index 97ae2a612babe..d0a6c8753fde7 100644
--- a/contrib/pg_tde/src/catalog/tde_principal_key.c
+++ b/contrib/pg_tde/src/catalog/tde_principal_key.c
@@ -62,9 +62,6 @@ typedef struct TdePrincipalKeylocalState
 	dshash_table *sharedHash;
 } TdePrincipalKeylocalState;
 
-/* Length of newly generated principal keys */
-#define PRINCIPAL_KEY_LEN 16
-
 /* Parameters for the principal key info shared hash */
 static dshash_parameters principal_key_dsh_params = {
 	.key_size = sizeof(Oid),
@@ -534,7 +531,7 @@ pg_tde_create_principal_key_internal(Oid providerOid,
 				errcode(ERRCODE_INVALID_PARAMETER_VALUE),
 				errmsg("cannot create key \"%s\" because it already exists", key_name));
 
-	key_info = KeyringGenerateNewKeyAndStore(provider, key_name, PRINCIPAL_KEY_LEN);
+	key_info = KeyringGenerateNewKeyAndStore(provider, key_name, TdeKeyLength);
 
 	pfree(key_info);
 	pfree(provider);
diff --git a/contrib/pg_tde/src/encryption/enc_aes.c b/contrib/pg_tde/src/encryption/enc_aes.c
index 8c98a3750741a..d5ae680a92322 100644
--- a/contrib/pg_tde/src/encryption/enc_aes.c
+++ b/contrib/pg_tde/src/encryption/enc_aes.c
@@ -36,6 +36,10 @@ static const EVP_CIPHER *cipher_cbc = NULL;
 static const EVP_CIPHER *cipher_gcm = NULL;
 static const EVP_CIPHER *cipher_ctr_ecb = NULL;
 
+static const EVP_CIPHER *cipher_cbc_256 = NULL;
+static const EVP_CIPHER *cipher_gcm_256 = NULL;
+static const EVP_CIPHER *cipher_ctr_ecb_256 = NULL;
+
 void
 AesInit(void)
 {
@@ -45,21 +49,26 @@ AesInit(void)
 	cipher_cbc = EVP_aes_128_cbc();
 	cipher_gcm = EVP_aes_128_gcm();
 	cipher_ctr_ecb = EVP_aes_128_ecb();
+
+	cipher_cbc_256 = EVP_aes_256_cbc();
+	cipher_gcm_256 = EVP_aes_256_gcm();
+	cipher_ctr_ecb_256 = EVP_aes_256_ecb();
 }
 
 static void
-AesEcbEncrypt(EVP_CIPHER_CTX **ctxPtr, const unsigned char *key, const unsigned char *in, int in_len, unsigned char *out)
+AesEcbEncrypt(EVP_CIPHER_CTX **ctxPtr, const unsigned char *key, int key_len, const unsigned char *in, int in_len, unsigned char *out)
 {
 	int			out_len;
+	const EVP_CIPHER	*cipher = key_len == 32 ? cipher_ctr_ecb_256 : cipher_ctr_ecb;
 
 	if (*ctxPtr == NULL)
 	{
-		Assert(cipher_ctr_ecb != NULL);
+		Assert(cipher != NULL);
 
 		*ctxPtr = EVP_CIPHER_CTX_new();
 		EVP_CIPHER_CTX_init(*ctxPtr);
 
-		if (EVP_CipherInit_ex(*ctxPtr, cipher_ctr_ecb, NULL, key, NULL, 1) == 0)
+		if (EVP_CipherInit_ex(*ctxPtr, cipher, NULL, key, NULL, 1) == 0)
 			ereport(ERROR,
 					errmsg("EVP_CipherInit_ex failed. OpenSSL error: %s", ERR_error_string(ERR_get_error(), NULL)));
 
@@ -74,19 +83,20 @@ AesEcbEncrypt(EVP_CIPHER_CTX **ctxPtr, const unsigned char *key, const unsigned
 }
 
 static void
-AesRunCbc(int enc, const unsigned char *key, const unsigned char *iv, const unsigned char *in, int in_len, unsigned char *out)
+AesRunCbc(int enc, const unsigned char *key, int key_len, const unsigned char *iv, const unsigned char *in, int in_len, unsigned char *out)
 {
 	int			out_len;
 	int			out_len_final;
 	EVP_CIPHER_CTX *ctx = NULL;
+	const EVP_CIPHER	*cipher = key_len == 32 ? cipher_cbc_256 : cipher_cbc;
 
-	Assert(cipher_cbc != NULL);
-	Assert(in_len % EVP_CIPHER_block_size(cipher_cbc) == 0);
+	Assert(cipher != NULL);
+	Assert(in_len % EVP_CIPHER_block_size(cipher) == 0);
 
 	ctx = EVP_CIPHER_CTX_new();
 	EVP_CIPHER_CTX_init(ctx);
 
-	if (EVP_CipherInit_ex(ctx, cipher_cbc, NULL, key, iv, enc) == 0)
+	if (EVP_CipherInit_ex(ctx, cipher, NULL, key, iv, enc) == 0)
 		ereport(ERROR,
 				errmsg("EVP_CipherInit_ex failed. OpenSSL error: %s", ERR_error_string(ERR_get_error(), NULL)));
 
@@ -112,31 +122,32 @@ AesRunCbc(int enc, const unsigned char *key, const unsigned char *iv, const unsi
 }
 
 void
-AesEncrypt(const unsigned char *key, const unsigned char *iv, const unsigned char *in, int in_len, unsigned char *out)
+AesEncrypt(const unsigned char *key, int key_len, const unsigned char *iv, const unsigned char *in, int in_len, unsigned char *out)
 {
-	AesRunCbc(1, key, iv, in, in_len, out);
+	AesRunCbc(1, key, key_len, 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)
+AesDecrypt(const unsigned char *key, int key_len, const unsigned char *iv, const unsigned char *in, int in_len, unsigned char *out)
 {
-	AesRunCbc(0, key, iv, in, in_len, out);
+	AesRunCbc(0, key, key_len, iv, in, in_len, out);
 }
 
 void
-AesGcmEncrypt(const unsigned char *key, const unsigned char *iv, int iv_len, const unsigned char *aad, int aad_len, const unsigned char *in, int in_len, unsigned char *out, unsigned char *tag, int tag_len)
+AesGcmEncrypt(const unsigned char *key, int key_len, const unsigned char *iv, int iv_len, const unsigned char *aad, int aad_len, const unsigned char *in, int in_len, unsigned char *out, unsigned char *tag, int tag_len)
 {
 	int			out_len;
 	int			out_len_final;
 	EVP_CIPHER_CTX *ctx;
+	const EVP_CIPHER	*cipher = key_len == 32 ? cipher_gcm_256 : cipher_gcm;
 
-	Assert(cipher_gcm != NULL);
-	Assert(in_len % EVP_CIPHER_block_size(cipher_gcm) == 0);
+	Assert(cipher != NULL);
+	Assert(in_len % EVP_CIPHER_block_size(cipher) == 0);
 
 	ctx = EVP_CIPHER_CTX_new();
 	EVP_CIPHER_CTX_init(ctx);
 
-	if (EVP_EncryptInit_ex(ctx, cipher_gcm, NULL, NULL, NULL) == 0)
+	if (EVP_EncryptInit_ex(ctx, cipher, NULL, NULL, NULL) == 0)
 		ereport(ERROR,
 				errmsg("EVP_EncryptInit_ex failed. OpenSSL error: %s", ERR_error_string(ERR_get_error(), NULL)));
 
@@ -180,18 +191,19 @@ AesGcmEncrypt(const unsigned char *key, const unsigned char *iv, int iv_len, con
 }
 
 bool
-AesGcmDecrypt(const unsigned char *key, const unsigned char *iv, int iv_len, const unsigned char *aad, int aad_len, const unsigned char *in, int in_len, unsigned char *out, unsigned char *tag, int tag_len)
+AesGcmDecrypt(const unsigned char *key, int key_len, const unsigned char *iv, int iv_len, const unsigned char *aad, int aad_len, const unsigned char *in, int in_len, unsigned char *out, unsigned char *tag, int tag_len)
 {
 	int			out_len;
 	int			out_len_final;
 	EVP_CIPHER_CTX *ctx;
+	const EVP_CIPHER	*cipher = key_len == 32 ? cipher_gcm_256 : cipher_gcm;
 
-	Assert(in_len % EVP_CIPHER_block_size(cipher_gcm) == 0);
+	Assert(in_len % EVP_CIPHER_block_size(cipher) == 0);
 
 	ctx = EVP_CIPHER_CTX_new();
 	EVP_CIPHER_CTX_init(ctx);
 
-	if (EVP_DecryptInit_ex(ctx, cipher_gcm, NULL, NULL, NULL) == 0)
+	if (EVP_DecryptInit_ex(ctx, cipher, NULL, NULL, NULL) == 0)
 		ereport(ERROR,
 				errmsg("EVP_EncryptInit_ex failed. OpenSSL error: %s", ERR_error_string(ERR_get_error(), NULL)));
 
@@ -243,7 +255,7 @@ AesGcmDecrypt(const unsigned char *key, const unsigned char *iv, int iv_len, con
  * This function assumes that the out buffer is big enough: at least (blockNumber2 - blockNumber1) * 16 bytes
  */
 void
-AesCtrEncryptedZeroBlocks(void *ctxPtr, const unsigned char *key, const char *iv_prefix, uint64_t blockNumber1, uint64_t blockNumber2, unsigned char *out)
+AesCtrEncryptedZeroBlocks(void *ctxPtr, const unsigned char *key, int key_len, const char *iv_prefix, uint64_t blockNumber1, uint64_t blockNumber2, unsigned char *out)
 {
 	unsigned char *p;
 
@@ -265,5 +277,5 @@ AesCtrEncryptedZeroBlocks(void *ctxPtr, const unsigned char *key, const char *iv
 		p += sizeof(j);
 	}
 
-	AesEcbEncrypt(ctxPtr, key, out, p - out, out);
+	AesEcbEncrypt(ctxPtr, key, key_len, 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 bceaaf7cf8b52..f5efcfc0ebff9 100644
--- a/contrib/pg_tde/src/encryption/enc_tde.c
+++ b/contrib/pg_tde/src/encryption/enc_tde.c
@@ -27,9 +27,9 @@ iv_prefix_debug(const char *iv_prefix, char *out_hex)
 #endif
 
 void
-pg_tde_generate_internal_key(InternalKey *int_key)
+pg_tde_generate_internal_key(InternalKey *int_key, int key_len)
 {
-	if (!RAND_bytes(int_key->key, INTERNAL_KEY_LEN))
+	if (!RAND_bytes(int_key->key, key_len))
 		ereport(ERROR,
 				errcode(ERRCODE_INTERNAL_ERROR),
 				errmsg("could not generate internal key: %s",
@@ -39,6 +39,8 @@ pg_tde_generate_internal_key(InternalKey *int_key)
 				errcode(ERRCODE_INTERNAL_ERROR),
 				errmsg("could not generate IV: %s",
 					   ERR_error_string(ERR_get_error(), NULL)));
+
+	int_key->key_len = key_len;
 }
 
 /*
@@ -53,6 +55,7 @@ pg_tde_stream_crypt(const char *iv_prefix,
 					uint32 data_len,
 					char *out,
 					const uint8 *key,
+					int key_len,
 					void **ctxPtr)
 {
 	const uint64 aes_start_block = start_offset / AES_BLOCK_SIZE;
@@ -68,7 +71,7 @@ pg_tde_stream_crypt(const char *iv_prefix,
 		uint32		current_batch_bytes;
 		uint64		batch_end_block = Min(batch_start_block + NUM_AES_BLOCKS_IN_BATCH, aes_end_block);
 
-		AesCtrEncryptedZeroBlocks(ctxPtr, key, iv_prefix, batch_start_block, batch_end_block, enc_key);
+		AesCtrEncryptedZeroBlocks(ctxPtr, key, key_len, iv_prefix, batch_start_block, batch_end_block, enc_key);
 
 #ifdef ENCRYPTION_DEBUG
 		{
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 40ee9eb9cbc34..263a5ed5b489a 100644
--- a/contrib/pg_tde/src/include/access/pg_tde_tdemap.h
+++ b/contrib/pg_tde/src/include/access/pg_tde_tdemap.h
@@ -27,4 +27,6 @@ extern void pg_tde_sign_principal_key_info(TDESignedPrincipalKeyInfo *signed_key
 
 const char *tde_sprint_key(InternalKey *k);
 
+extern void pg_tde_migrate_smgr_keys_file(void);
+
 #endif							/* PG_TDE_MAP_H */
diff --git a/contrib/pg_tde/src/include/access/pg_tde_xlog_keys.h b/contrib/pg_tde/src/include/access/pg_tde_xlog_keys.h
index 047d76414e962..d80dd8299a0ec 100644
--- a/contrib/pg_tde/src/include/access/pg_tde_xlog_keys.h
+++ b/contrib/pg_tde/src/include/access/pg_tde_xlog_keys.h
@@ -75,7 +75,7 @@ typedef struct WALKeyCacheRec
 } WALKeyCacheRec;
 
 extern int	pg_tde_count_wal_ranges_in_file(void);
-extern void pg_tde_create_wal_range(WalEncryptionRange *range, WalEncryptionRangeType type);
+extern void pg_tde_create_wal_range(WalEncryptionRange *range, WalEncryptionRangeType type, int key_len);
 extern void pg_tde_delete_server_key(void);
 extern WALKeyCacheRec *pg_tde_fetch_wal_keys(WalLocation start);
 extern void pg_tde_free_wal_key_cache(void);
@@ -88,5 +88,6 @@ extern void pg_tde_save_server_key(const TDEPrincipalKey *principal_key, bool wr
 extern void pg_tde_save_server_key_redo(const TDESignedPrincipalKeyInfo *signed_key_info);
 extern void pg_tde_wal_last_range_set_location(WalLocation loc);
 extern void pg_tde_wal_cache_extra_palloc(void);
+extern void pg_tde_update_wal_keys_file(void);
 
 #endif							/* PG_TDE_XLOG_KEYS_H */
diff --git a/contrib/pg_tde/src/include/access/pg_tde_xlog_smgr.h b/contrib/pg_tde/src/include/access/pg_tde_xlog_smgr.h
index 8d5ec4bc08b42..3ff29d0957811 100644
--- a/contrib/pg_tde/src/include/access/pg_tde_xlog_smgr.h
+++ b/contrib/pg_tde/src/include/access/pg_tde_xlog_smgr.h
@@ -12,7 +12,7 @@
 extern Size TDEXLogEncryptStateSize(void);
 extern void TDEXLogShmemInit(void);
 extern void TDEXLogSmgrInit(void);
-extern void TDEXLogSmgrInitWrite(bool encrypt_xlog);
+extern void TDEXLogSmgrInitWrite(bool encrypt_xlog, int key_len);
 extern void TDEXLogSmgrInitWriteOldKeys(void);
 
 extern void TDEXLogCryptBuffer(const void *buf, void *out_buf, size_t count, off_t offset,
diff --git a/contrib/pg_tde/src/include/encryption/enc_aes.h b/contrib/pg_tde/src/include/encryption/enc_aes.h
index 68b9416972fd6..ecaf4b82b90b8 100644
--- a/contrib/pg_tde/src/include/encryption/enc_aes.h
+++ b/contrib/pg_tde/src/include/encryption/enc_aes.h
@@ -8,10 +8,10 @@
 #include <stdint.h>
 
 extern void AesInit(void);
-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, int iv_len, const unsigned char *aad, int aad_len, const unsigned char *in, int in_len, unsigned char *out, unsigned char *tag, int tag_len);
-extern bool AesGcmDecrypt(const unsigned char *key, const unsigned char *iv, int iv_len, const unsigned char *aad, int aad_len, const unsigned char *in, int in_len, unsigned char *out, unsigned char *tag, int tag_len);
-extern void AesCtrEncryptedZeroBlocks(void *ctxPtr, const unsigned char *key, const char *iv_prefix, uint64_t blockNumber1, uint64_t blockNumber2, unsigned char *out);
+extern void AesEncrypt(const unsigned char *key, int key_len, const unsigned char *iv, const unsigned char *in, int in_len, unsigned char *out);
+extern void AesDecrypt(const unsigned char *key, int key_len, const unsigned char *iv, const unsigned char *in, int in_len, unsigned char *out);
+extern void AesGcmEncrypt(const unsigned char *key, int key_len, const unsigned char *iv, int iv_len, const unsigned char *aad, int aad_len, const unsigned char *in, int in_len, unsigned char *out, unsigned char *tag, int tag_len);
+extern bool AesGcmDecrypt(const unsigned char *key, int key_len, const unsigned char *iv, int iv_len, const unsigned char *aad, int aad_len, const unsigned char *in, int in_len, unsigned char *out, unsigned char *tag, int tag_len);
+extern void AesCtrEncryptedZeroBlocks(void *ctxPtr, const unsigned char *key, int key_len, const char *iv_prefix, uint64_t blockNumber1, uint64_t blockNumber2, unsigned char *out);
 
 #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 64299c85e33f5..2cfc271e704b5 100644
--- a/contrib/pg_tde/src/include/encryption/enc_tde.h
+++ b/contrib/pg_tde/src/include/encryption/enc_tde.h
@@ -5,22 +5,25 @@
 #ifndef ENC_TDE_H
 #define ENC_TDE_H
 
-#define INTERNAL_KEY_LEN 16
+#define INTERNAL_KEY_OLD_LEN 16 /* Old max size of an Internal Key - for compatibility */
+#define INTERNAL_KEY_MAX_LEN 32 /* Max size of an Internal Key */
 #define INTERNAL_KEY_IV_LEN 16
 
 typedef struct InternalKey
 {
-	uint8		key[INTERNAL_KEY_LEN];
+	uint32		key_len;
 	uint8		base_iv[INTERNAL_KEY_IV_LEN];
+	uint8		key[INTERNAL_KEY_MAX_LEN];
 } InternalKey;
 
-extern void pg_tde_generate_internal_key(InternalKey *int_key);
+extern void pg_tde_generate_internal_key(InternalKey *int_key, int key_len);
 extern void pg_tde_stream_crypt(const char *iv_prefix,
 								uint32 start_offset,
 								const char *data,
 								uint32 data_len,
 								char *out,
 								const uint8 *key,
+								int key_len,
 								void **ctxPtr);
 
 #endif							/* ENC_TDE_H */
diff --git a/contrib/pg_tde/src/include/pg_tde_guc.h b/contrib/pg_tde/src/include/pg_tde_guc.h
index c4ce064402d6d..837f4b1767e44 100644
--- a/contrib/pg_tde/src/include/pg_tde_guc.h
+++ b/contrib/pg_tde/src/include/pg_tde_guc.h
@@ -10,6 +10,14 @@
 extern bool AllowInheritGlobalProviders;
 extern bool EncryptXLog;
 extern bool EnforceEncryption;
+extern int Cipher;
+extern int TdeKeyLength;
+
+typedef enum CipherOption
+{
+	TDE_CIPHER_AES_128,
+	TDE_CIPHER_AES_256,
+} CipherOption;
 
 extern void TdeGucInit(void);
 
diff --git a/contrib/pg_tde/src/keyring/keyring_api.c b/contrib/pg_tde/src/keyring/keyring_api.c
index 11da816681f4d..6ef6d5fd2e712 100644
--- a/contrib/pg_tde/src/keyring/keyring_api.c
+++ b/contrib/pg_tde/src/keyring/keyring_api.c
@@ -142,8 +142,8 @@ ValidateKey(KeyInfo *key)
 		return false;
 	}
 
-	/* For now we only support 128-bit keys */
-	if (key->data.len != KEY_DATA_SIZE_128)
+	/* For now we only support 128 and 256-bit keys */
+	if (key->data.len != KEY_DATA_SIZE_128 && key->data.len != KEY_DATA_SIZE_256)
 	{
 		ereport(WARNING,
 				errmsg("invalid key: unsupported key length \"%u\"", key->data.len));
diff --git a/contrib/pg_tde/src/pg_tde.c b/contrib/pg_tde/src/pg_tde.c
index e6f32f099bbe0..5388133bff108 100644
--- a/contrib/pg_tde/src/pg_tde.c
+++ b/contrib/pg_tde/src/pg_tde.c
@@ -70,7 +70,9 @@ tde_shmem_startup(void)
 	PrincipalKeyShmemInit();
 	TDEXLogShmemInit();
 	TDEXLogSmgrInit();
-	TDEXLogSmgrInitWrite(EncryptXLog);
+	pg_tde_update_wal_keys_file();
+	pg_tde_migrate_smgr_keys_file();
+	TDEXLogSmgrInitWrite(EncryptXLog, TdeKeyLength);
 
 	LWLockRelease(AddinShmemInitLock);
 }
diff --git a/contrib/pg_tde/src/pg_tde_guc.c b/contrib/pg_tde/src/pg_tde_guc.c
index 912385cc0df75..1c8392ef75b4b 100644
--- a/contrib/pg_tde/src/pg_tde_guc.c
+++ b/contrib/pg_tde/src/pg_tde_guc.c
@@ -6,11 +6,36 @@
 
 #include "utils/guc.h"
 
+#include "keyring/keyring_api.h"
 #include "pg_tde_guc.h"
 
 bool		AllowInheritGlobalProviders = true;
 bool		EncryptXLog = false;
 bool		EnforceEncryption = false;
+int			Cipher = TDE_CIPHER_AES_128;
+int			TdeKeyLength = KEY_DATA_SIZE_128;
+
+/* Custom GUC variable */
+static const struct config_enum_entry cipher_options[] = {
+	{"aes_128", TDE_CIPHER_AES_128, false},
+	{"aes_256", TDE_CIPHER_AES_256, false},
+};
+
+static void
+assign_keys_size(int newval, void *extra)
+{
+	switch (newval)
+	{
+		case TDE_CIPHER_AES_128:
+			TdeKeyLength = KEY_DATA_SIZE_128;
+			break;
+		case TDE_CIPHER_AES_256:
+			TdeKeyLength = KEY_DATA_SIZE_256;
+			break;
+		default:
+			Assert(false);
+	}
+}
 
 void
 TdeGucInit(void)
@@ -51,4 +76,17 @@ TdeGucInit(void)
 							 NULL	/* show_hook */
 		);
 
+	DefineCustomEnumVariable("pg_tde.cipher",	/* name */
+							 "TDE encryption algorithm.",	/* short_desc */
+							 NULL,	/* long_desc */
+							 &Cipher,	/* value address */
+							 TDE_CIPHER_AES_128, /* boot value */
+							 cipher_options, /* options */
+							 PGC_SUSET, /* context */
+							 0, /* flags */
+							 NULL,	/* check_hook */
+							 assign_keys_size,	/* assign_hook */
+							 NULL	/* show_hook */
+		);
+
 }
diff --git a/contrib/pg_tde/src/smgr/pg_tde_smgr.c b/contrib/pg_tde/src/smgr/pg_tde_smgr.c
index 1f55a57b6e937..46aa9131b2341 100644
--- a/contrib/pg_tde/src/smgr/pg_tde_smgr.c
+++ b/contrib/pg_tde/src/smgr/pg_tde_smgr.c
@@ -11,6 +11,7 @@
 #include "encryption/enc_aes.h"
 #include "encryption/enc_tde.h"
 #include "pg_tde_event_capture.h"
+#include "pg_tde_guc.h"
 #include "smgr/pg_tde_smgr.h"
 
 typedef enum TDEMgrRelationEncryptionStatus
@@ -95,7 +96,7 @@ tde_smgr_create_key(const RelFileLocatorBackend *smgr_rlocator)
 {
 	InternalKey *key = palloc_object(InternalKey);
 
-	pg_tde_generate_internal_key(key);
+	pg_tde_generate_internal_key(key, TdeKeyLength);
 
 	if (RelFileLocatorBackendIsTemp(*smgr_rlocator))
 		tde_smgr_save_temp_key(&smgr_rlocator->locator, key);
@@ -113,7 +114,7 @@ tde_smgr_create_key_redo(const RelFileLocator *rlocator)
 {
 	InternalKey key;
 
-	pg_tde_generate_internal_key(&key);
+	pg_tde_generate_internal_key(&key, TdeKeyLength);
 
 	pg_tde_save_smgr_key(*rlocator, &key);
 }
@@ -234,7 +235,7 @@ tde_mdwritev(SMgrRelation reln, ForkNumber forknum, BlockNumber blocknum,
 
 			CalcBlockIv(forknum, bn, tdereln->relKey.base_iv, iv);
 
-			AesEncrypt(tdereln->relKey.key, iv, ((unsigned char **) buffers)[i], BLCKSZ, local_buffers[i]);
+			AesEncrypt(tdereln->relKey.key, tdereln->relKey.key_len, iv, ((unsigned char **) buffers)[i], BLCKSZ, local_buffers[i]);
 		}
 
 		mdwritev(reln, forknum, blocknum,
@@ -298,7 +299,7 @@ tde_mdextend(SMgrRelation reln, ForkNumber forknum, BlockNumber blocknum,
 
 		CalcBlockIv(forknum, blocknum, tdereln->relKey.base_iv, iv);
 
-		AesEncrypt(tdereln->relKey.key, iv, ((unsigned char *) buffer), BLCKSZ, local_blocks);
+		AesEncrypt(tdereln->relKey.key, tdereln->relKey.key_len, iv, ((unsigned char *) buffer), BLCKSZ, local_blocks);
 
 		mdextend(reln, forknum, blocknum, local_blocks, skipFsync);
 
@@ -353,7 +354,7 @@ tde_mdreadv(SMgrRelation reln, ForkNumber forknum, BlockNumber blocknum,
 
 		CalcBlockIv(forknum, bn, tdereln->relKey.base_iv, iv);
 
-		AesDecrypt(tdereln->relKey.key, iv, ((unsigned char **) buffers)[i], BLCKSZ, ((unsigned char **) buffers)[i]);
+		AesDecrypt(tdereln->relKey.key, tdereln->relKey.key_len, iv, ((unsigned char **) buffers)[i], BLCKSZ, ((unsigned char **) buffers)[i]);
 	}
 }
 
diff --git a/src/bin/pg_basebackup/pg_basebackup.c b/src/bin/pg_basebackup/pg_basebackup.c
index e45d09f7151ea..8504f1d8b8e30 100644
--- a/src/bin/pg_basebackup/pg_basebackup.c
+++ b/src/bin/pg_basebackup/pg_basebackup.c
@@ -687,7 +687,7 @@ StartLogStreamer(char *startpos, uint32 timeline, char *sysidentifier,
 			exit(1);
 		}
 		pg_tde_save_server_key(principalKey, false);
-		TDEXLogSmgrInitWrite(true);
+		TDEXLogSmgrInitWrite(true, 16);
 	}
 #endif
 
diff --git a/src/bin/pg_resetwal/pg_resetwal.c b/src/bin/pg_resetwal/pg_resetwal.c
index 3cebcb6adabe2..3e5fa01c93b9c 100644
--- a/src/bin/pg_resetwal/pg_resetwal.c
+++ b/src/bin/pg_resetwal/pg_resetwal.c
@@ -521,7 +521,7 @@ main(int argc, char *argv[])
 	 * We are doing a write initialization only here and not at the startup because we 
 	 * want to be sure that everything is checked and ready for writing at this point.
 	 */
-	TDEXLogSmgrInitWrite(false);
+	TDEXLogSmgrInitWrite(false, 16);
 #endif
 
 	/*