pg_basebackup: encrypt streamed WAL with new key

Before, pg_basebackup would encrypt streamed WAL according to the keys
in pg_tde/wal_keys in the destination dir.

This commit introduces the number of changes:
pg_basebackup encrypts WAL only if the "-E --encrypt-wal" flag is
provided. In such a case, it would extract the principal key, truncate
pg_tde/wal_keys and encrypt WAL with a newly generated WAL key. We
still expect pg_tde/wal_keys and pg_tde/1664_providers in the
destination dir. In case these files are not provided, but "-E" is
specified, it fails with an error.

Author: dAdAbird

contrib/pg_tde/meson.build

@@ -109,6 +109,7 @@ tap_tests = [
   't/key_rotate_tablespace.pl',
   't/key_validation.pl',
   't/multiple_extensions.pl',
+  't/pg_basebackup.pl',
   't/pg_tde_change_key_provider.pl',
   't/pg_rewind_basic.pl',
   't/pg_rewind_databases.pl',

contrib/pg_tde/src/access/pg_tde_xlog_keys.c

@@ -715,7 +715,6 @@ pg_tde_save_server_key_redo(const TDESignedPrincipalKeyInfo *signed_key_info)
 }
 #endif
 
-#ifndef FRONTEND
 /*
  * Creates the key file and saves the principal key information.
  *
@@ -737,17 +736,18 @@ pg_tde_save_server_key(const TDEPrincipalKey *principal_key, bool write_xlog)
 
 	pg_tde_sign_principal_key_info(&signed_key_Info, principal_key);
 
+#ifndef FRONTEND
 	if (write_xlog)
 	{
 		XLogBeginInsert();
 		XLogRegisterData((char *) &signed_key_Info, sizeof(TDESignedPrincipalKeyInfo));
 		XLogInsert(RM_TDERMGR_ID, XLOG_TDE_ADD_PRINCIPAL_KEY);
 	}
+#endif
 
 	fd = pg_tde_open_wal_key_file_write(get_wal_key_file_path(), &signed_key_Info, true, &curr_pos);
 	CloseTransientFile(fd);
 }
-#endif
 
 /*
  * Get the principal key from the key file. The caller must hold

contrib/pg_tde/src/access/pg_tde_xlog_smgr.c

@@ -326,35 +326,21 @@ TDEXLogWriteEncryptedPages(int fd, const void *buf, size_t count, off_t offset,
 	return pg_pwrite(fd, enc_buff, count, offset);
 }
 
-static ssize_t
-tdeheap_xlog_seg_write(int fd, const void *buf, size_t count, off_t offset,
-					   TimeLineID tli, XLogSegNo segno, int segSize)
+/*
+ * Set the last (most recent) key's start location if not set.
+ */
+bool
+tde_ensure_xlog_key_location(WalLocation loc)
 {
 	bool		lastKeyUsable;
-	bool		afterWriteKey;
 #ifdef FRONTEND
 	bool		crashRecovery = false;
 #else
 	bool		crashRecovery = GetRecoveryState() == RECOVERY_STATE_CRASH;
 #endif
 
-	WalLocation loc = {.tli = tli};
-	WalLocation writeKeyLoc;
-
-	XLogSegNoOffsetToRecPtr(segno, offset, segSize, loc.lsn);
-
-	/*
-	 * Set the last (most recent) key's start LSN if not set.
-	 *
-	 * This func called with WALWriteLock held, so no need in any extra sync.
-	 */
-
-	writeKeyLoc.lsn = TDEXLogGetEncKeyLsn();
+	lastKeyUsable = (TDEXLogGetEncKeyLsn() != 0);
 	pg_read_barrier();
-	writeKeyLoc.tli = TDEXLogGetEncKeyTli();
-
-	lastKeyUsable = (writeKeyLoc.lsn != 0);
-	afterWriteKey = wal_location_cmp(writeKeyLoc, loc) <= 0;
 
 	if (EncryptionKey.type != WAL_KEY_TYPE_INVALID && !lastKeyUsable)
 	{
@@ -372,6 +358,30 @@ tdeheap_xlog_seg_write(int fd, const void *buf, size_t count, off_t offset,
 		}
 	}
 
+	return lastKeyUsable;
+}
+
+static ssize_t
+tdeheap_xlog_seg_write(int fd, const void *buf, size_t count, off_t offset,
+					   TimeLineID tli, XLogSegNo segno, int segSize)
+{
+	bool		lastKeyUsable;
+	bool		afterWriteKey;
+	WalLocation writeKeyLoc;
+	WalLocation loc = {.tli = tli};
+
+	XLogSegNoOffsetToRecPtr(segno, offset, segSize, loc.lsn);
+	lastKeyUsable = tde_ensure_xlog_key_location(loc);
+
+	/*
+	 * On backend this called with WALWriteLock held, so no need in any extra
+	 * sync.
+	 */
+	writeKeyLoc.lsn = TDEXLogGetEncKeyLsn();
+	pg_read_barrier();
+	writeKeyLoc.tli = TDEXLogGetEncKeyTli();
+	afterWriteKey = wal_location_cmp(writeKeyLoc, loc) <= 0;
+
 	if ((!afterWriteKey || !lastKeyUsable) && EncryptionKey.type != WAL_KEY_TYPE_INVALID)
 	{
 		return TDEXLogWriteEncryptedPagesOldKeys(fd, buf, count, offset, tli, segno, segSize);

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

@@ -7,6 +7,8 @@
 
 #include "postgres.h"
 
+#include "access/pg_tde_xlog_keys.h"
+
 extern Size TDEXLogEncryptStateSize(void);
 extern void TDEXLogShmemInit(void);
 extern void TDEXLogSmgrInit(void);
@@ -16,4 +18,6 @@ extern void TDEXLogSmgrInitWriteReuseKey(void);
 extern void TDEXLogCryptBuffer(const void *buf, void *out_buf, size_t count, off_t offset,
 							   TimeLineID tli, XLogSegNo segno, int segSize);
 
+extern bool tde_ensure_xlog_key_location(WalLocation loc);
+
 #endif							/* PG_TDE_XLOGSMGR_H */

contrib/pg_tde/t/pg_basebackup.pl

@@ -0,0 +1,49 @@
+
+# Copyright (c) 2021-2024, PostgreSQL Global Development Group
+
+use strict;
+use warnings FATAL => 'all';
+use Config;
+use File::Basename qw(basename dirname);
+use File::Path     qw(rmtree);
+use PostgreSQL::Test::Cluster;
+use PostgreSQL::Test::Utils;
+use Test::More;
+
+program_help_ok('pg_basebackup');
+program_version_ok('pg_basebackup');
+program_options_handling_ok('pg_basebackup');
+
+my $tempdir = PostgreSQL::Test::Utils::tempdir;
+
+my $node = PostgreSQL::Test::Cluster->new('main');
+
+# Initialize node without replication settings
+$node->init(
+	allows_streaming => 1,
+	extra => ['--data-checksums'],
+	auth_extra => [ '--create-role', 'backupuser' ]);
+$node->start;
+
+# Sanity checks for options with WAL encryption
+$node->command_fails_like(
+	[ 'pg_basebackup', '-D', "$tempdir/backup", '-E', '-Ft' ],
+	qr/can not encrypt WAL in tar mode/,
+	'encryption in tar mode');
+
+$node->command_fails_like(
+	[ 'pg_basebackup', '-D', "$tempdir/backup", '-E', '-X', 'fetch'],
+	qr/WAL encryption can only be used with WAL streaming/,
+	'encryption with WAL fetch');
+
+$node->command_fails_like(
+	[ 'pg_basebackup', '-D', "$tempdir/backup", '-E', '-X', 'none'],
+	qr/WAL encryption can only be used with WAL streaming/,
+	'encryption with WAL none');
+
+$node->command_fails_like(
+	[ 'pg_basebackup', '-D', "$tempdir/backup", '-E'],
+	qr/could not find global principal key/,
+	'encryption with no pg_tde dir');
+
+done_testing();

contrib/pg_tde/t/pgtde.pm

@@ -119,6 +119,8 @@ sub backup
 	PostgreSQL::Test::RecursiveCopy::copypath($node->data_dir . '/pg_tde',
 		$backup_dir . '/pg_tde');
 
+	push @{$params{backup_options}}, "-E";
+
 	$node->backup($backup_name, %params);
 }
 

src/bin/pg_basebackup/bbstreamer.h

@@ -204,7 +204,8 @@ extern bbstreamer *bbstreamer_gzip_writer_new(char *pathname, FILE *file,
 											  pg_compress_specification *compress);
 extern bbstreamer *bbstreamer_extractor_new(const char *basepath,
 											const char *(*link_map) (const char *),
-											void (*report_output_file) (const char *));
+											void (*report_output_file) (const char *),
+											bool encrypted_wal);
 
 extern bbstreamer *bbstreamer_gzip_decompressor_new(bbstreamer *next);
 extern bbstreamer *bbstreamer_lz4_compressor_new(bbstreamer *next,

src/bin/pg_basebackup/bbstreamer_file.c

@@ -38,6 +38,7 @@ typedef struct bbstreamer_extractor
 	void		(*report_output_file) (const char *);
 	char		filename[MAXPGPATH];
 	FILE	   *file;
+	bool		encryped_wal;
 } bbstreamer_extractor;
 
 static void bbstreamer_plain_writer_content(bbstreamer *streamer,
@@ -186,7 +187,8 @@ bbstreamer_plain_writer_free(bbstreamer *streamer)
 bbstreamer *
 bbstreamer_extractor_new(const char *basepath,
 						 const char *(*link_map) (const char *),
-						 void (*report_output_file) (const char *))
+						 void (*report_output_file) (const char *),
+						 bool encrypted_wal)
 {
 	bbstreamer_extractor *streamer;
 
@@ -196,6 +198,7 @@ bbstreamer_extractor_new(const char *basepath,
 	streamer->basepath = pstrdup(basepath);
 	streamer->link_map = link_map;
 	streamer->report_output_file = report_output_file;
+	streamer->encryped_wal = encrypted_wal;
 
 	return &streamer->base;
 }
@@ -240,9 +243,21 @@ bbstreamer_extractor_content(bbstreamer *streamer, bbstreamer_member *member,
 				extract_link(mystreamer->filename, linktarget);
 			}
 			else
+			{
+#ifdef PERCONA_EXT
+				/* 
+				 * A streamed WAL is encrypted with the newly generated WAL key,
+				 * hence we have to prevent these files from rewriting.
+				 */
+				if (mystreamer->encryped_wal &&
+						(strcmp(member->pathname, "pg_tde/wal_keys") == 0 ||
+						 strcmp(member->pathname, "pg_tde/1664_providers") == 0))
+					break;
+#endif
 				mystreamer->file =
 					create_file_for_extract(mystreamer->filename,
 											member->mode);
+			}
 
 			/* Report output file change. */
 			if (mystreamer->report_output_file)

src/bin/pg_basebackup/pg_basebackup.c

@@ -41,8 +41,12 @@
 #ifdef PERCONA_EXT
 #include "access/pg_tde_fe_init.h"
 #include "access/pg_tde_xlog_smgr.h"
+#include "access/pg_tde_xlog_keys.h"
 #include "access/xlog_smgr.h"
+#include "catalog/tde_principal_key.h"
 #include "pg_tde.h"
+
+#define GLOBAL_DATA_TDE_OID 1664
 #endif
 
 #define ERRCODE_DATA_CORRUPTED_BCP	"XX001"
@@ -145,6 +149,7 @@ static bool showprogress = false;
 static bool estimatesize = true;
 static int	verbose = 0;
 static IncludeWal includewal = STREAM_WAL;
+static bool encrypt_wal = false;
 static bool fastcheckpoint = false;
 static bool writerecoveryconf = false;
 static bool do_sync = true;
@@ -416,6 +421,9 @@ usage(void)
 	printf(_("      --waldir=WALDIR    location for the write-ahead log directory\n"));
 	printf(_("  -X, --wal-method=none|fetch|stream\n"
 			 "                         include required WAL files with specified method\n"));
+#ifdef PERCONA_EXT
+	printf(_("  -E, --encrypt-wal      encrypt streamed WAL\n"));
+#endif
 	printf(_("  -z, --gzip             compress tar output\n"));
 	printf(_("  -Z, --compress=[{client|server}-]METHOD[:DETAIL]\n"
 			 "                         compress on client or server as specified\n"));
@@ -568,6 +576,7 @@ LogStreamerMain(logstreamer_param *param)
 	stream.synchronous = false;
 	/* fsync happens at the end of pg_basebackup for all data */
 	stream.do_sync = false;
+	stream.encrypt = encrypt_wal;
 	stream.mark_done = true;
 	stream.partial_suffix = NULL;
 	stream.replication_slot = replication_slot;
@@ -662,12 +671,23 @@ StartLogStreamer(char *startpos, uint32 timeline, char *sysidentifier,
 			 "pg_xlog" : "pg_wal");
 
 #ifdef PERCONA_EXT
-	{
+	if (encrypt_wal) {
 		char tdedir[MAXPGPATH];
+		TDEPrincipalKey *principalKey;
 
 		snprintf(tdedir, sizeof(tdedir), "%s/%s", basedir, PG_TDE_DATA_DIR);
 		pg_tde_fe_init(tdedir);
 		TDEXLogSmgrInit();
+
+		principalKey = GetPrincipalKey(GLOBAL_DATA_TDE_OID, NULL);
+		if (!principalKey)
+		{
+			pg_log_error("could not find global principal key");
+			pg_log_error_hint("Copy PGDATA/pg_tde from the source to the backup destination dir.");
+			exit(1);
+		}
+		pg_tde_save_server_key(principalKey, false);
+		TDEXLogSmgrInitWrite(true);
 	}
 #endif
 
@@ -1187,7 +1207,8 @@ CreateBackupStreamer(char *archive_name, char *spclocation,
 			directory = get_tablespace_mapping(spclocation);
 		streamer = bbstreamer_extractor_new(directory,
 											get_tablespace_mapping,
-											progress_update_filename);
+											progress_update_filename,
+											encrypt_wal);
 	}
 	else
 	{
@@ -2393,6 +2414,9 @@ main(int argc, char **argv)
 		{"target", required_argument, NULL, 't'},
 		{"tablespace-mapping", required_argument, NULL, 'T'},
 		{"wal-method", required_argument, NULL, 'X'},
+#ifdef PERCONA_EXT
+		{"encrypt-wal", no_argument, NULL, 'e'},
+#endif
 		{"gzip", no_argument, NULL, 'z'},
 		{"compress", required_argument, NULL, 'Z'},
 		{"label", required_argument, NULL, 'l'},
@@ -2447,7 +2471,7 @@ main(int argc, char **argv)
 
 	atexit(cleanup_directories_atexit);
 
-	while ((c = getopt_long(argc, argv, "c:Cd:D:F:h:i:l:nNp:Pr:Rs:S:t:T:U:vwWX:zZ:",
+	while ((c = getopt_long(argc, argv, "c:Cd:D:EF:h:i:l:nNp:Pr:Rs:S:t:T:U:vwWX:zZ:",
 							long_options, &option_index)) != -1)
 	{
 		switch (c)
@@ -2560,6 +2584,11 @@ main(int argc, char **argv)
 					pg_fatal("invalid wal-method option \"%s\", must be \"fetch\", \"stream\", or \"none\"",
 							 optarg);
 				break;
+#ifdef PERCONA_EXT
+			case 'E':
+				encrypt_wal = true;
+				break;
+#endif
 			case 'z':
 				compression_algorithm = "gzip";
 				compression_detail = NULL;
@@ -2738,6 +2767,26 @@ main(int argc, char **argv)
 		exit(1);
 	}
 
+	/*
+	 * Sanity checks for WAL encryption.
+	 */
+	if (encrypt_wal)
+	{
+		if (includewal != STREAM_WAL)
+		{
+			pg_log_error("WAL encryption can only be used with WAL streaming");
+			pg_log_error_hint("Use -X stream with -E.");
+			exit(1);
+		}
+
+		if (format != 'p')
+		{
+			pg_log_error("can not encrypt WAL in tar mode");
+			pg_log_error_hint("Use -Fp with -E.");
+			exit(1);
+		}
+	}
+
 	/*
 	 * Sanity checks for replication slot options.
 	 */

src/bin/pg_basebackup/receivelog.c

@@ -28,6 +28,7 @@
 #ifdef PERCONA_EXT
 #include "access/pg_tde_fe_init.h"
 #include "access/pg_tde_xlog_smgr.h"
+#include "access/xlog_smgr.h"
 #include "catalog/tde_global_space.h"
 #endif
 
@@ -1131,8 +1132,13 @@ ProcessXLogDataMsg(PGconn *conn, StreamCtl *stream, char *copybuf, int len,
 		}
 
 #ifdef PERCONA_EXT
+		if (stream->encrypt)
 		{
 			void* enc_buf = copybuf + hdr_len + bytes_written;
+			WalLocation loc = {.tli = stream->timeline};
+
+			XLogSegNoOffsetToRecPtr(segno, xlogoff, WalSegSz, loc.lsn);
+			tde_ensure_xlog_key_location(loc);
 			TDEXLogCryptBuffer(enc_buf, enc_buf, bytes_to_write,
 							   xlogoff, stream->timeline, segno, WalSegSz);
 		}