PG-1603 Make pg_basebackup work with encrypted WAL

When WAL is streamed during the backup (default mode), it comes in
unencrypted. But we need keys to encrypt it. For now, we expect that
the user would put `pg_tde` dir containing the `1664_key` and
`1664_providers` into the destination directory before starting the
backup. We encrypt the streamed WAL according to internal keys. No
`pg_tde` dir means no streamed WAL encryption.

Author: dAdAbird

contrib/pg_tde/src/access/pg_tde_xlog_smgr.c

@@ -329,10 +329,6 @@ tdeheap_xlog_seg_read(int fd, void *buf, size_t count, off_t offset,
 					  TimeLineID tli, XLogSegNo segno, int segSize)
 {
 	ssize_t		readsz;
-	WALKeyCacheRec *keys = pg_tde_get_wal_cache_keys();
-	XLogRecPtr	write_key_lsn;
-	WalLocation data_end = {.tli = tli};
-	WalLocation data_start = {.tli = tli};
 
 #ifdef TDE_XLOG_DEBUG
 	elog(DEBUG1, "read from a WAL segment, size: %lu offset: %ld [%lX], seg: %u_%X/%X",
@@ -344,6 +340,23 @@ tdeheap_xlog_seg_read(int fd, void *buf, size_t count, off_t offset,
 	if (readsz <= 0)
 		return readsz;
 
+	TDEXLogCryptBuffer(buf, count, offset, tli, segno, segSize);
+
+	return readsz;
+}
+
+/*
+ * [De]Crypt buffer if needed based on provided segment offset, number and TLI
+ */
+void
+TDEXLogCryptBuffer(void *buf, size_t count, off_t offset,
+				   TimeLineID tli, XLogSegNo segno, int segSize)
+{
+	WALKeyCacheRec *keys = pg_tde_get_wal_cache_keys();
+	XLogRecPtr	write_key_lsn;
+	WalLocation data_end = {.tli = tli};
+	WalLocation data_start = {.tli = tli};
+
 	if (!keys)
 	{
 		WalLocation start = {.tli = 1,.lsn = 0};
@@ -377,7 +390,7 @@ tdeheap_xlog_seg_read(int fd, void *buf, size_t count, off_t offset,
 	}
 
 	XLogSegNoOffsetToRecPtr(segno, offset, segSize, data_start.lsn);
-	XLogSegNoOffsetToRecPtr(segno, offset + readsz, segSize, data_end.lsn);
+	XLogSegNoOffsetToRecPtr(segno, offset + count, segSize, data_end.lsn);
 
 	/*
 	 * TODO: this is higly ineffective. We should get rid of linked list and
@@ -414,7 +427,7 @@ tdeheap_xlog_seg_read(int fd, void *buf, size_t count, off_t offset,
 				/* We have reached the end of the segment */
 				if (dec_end == 0)
 				{
-					dec_end = offset + readsz;
+					dec_end = offset + count;
 				}
 
 				dec_sz = dec_end - dec_off;
@@ -433,8 +446,6 @@ tdeheap_xlog_seg_read(int fd, void *buf, size_t count, off_t offset,
 			}
 		}
 	}
-
-	return readsz;
 }
 
 union u128cast

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

@@ -13,4 +13,7 @@ extern void TDEXLogSmgrInit(void);
 extern void TDEXLogSmgrInitWrite(bool encrypt_xlog);
 extern void TDEXLogSmgrInitWriteReuseKey(void);
 
+extern void TDEXLogCryptBuffer(void *buf, size_t count, off_t offset,
+							   TimeLineID tli, XLogSegNo segno, int segSize);
+
 #endif							/* PG_TDE_XLOGSMGR_H */

contrib/pg_tde/t/RewindTest.pm

@@ -44,6 +44,7 @@ use PostgreSQL::Test::Cluster;
 use PostgreSQL::Test::RecursiveCopy;
 use PostgreSQL::Test::Utils;
 use Test::More;
+use pgtde;
 
 our @EXPORT = qw(
   $node_primary
@@ -199,7 +200,7 @@ sub create_standby
 	$node_standby =
 	  PostgreSQL::Test::Cluster->new(
 		'standby' . ($extra_name ? "_${extra_name}" : ''));
-	$node_primary->backup('my_backup');
+	PGTDE::backup($node_primary, 'my_backup');
 	$node_standby->init_from_backup($node_primary, 'my_backup');
 	my $connstr_primary = $node_primary->connstr();
 

contrib/pg_tde/t/pgtde.pm

@@ -109,4 +109,17 @@ sub compare_results
 	return compare($expected_filename_with_path, $out_filename_with_path);
 }
 
+sub backup
+{
+	my ($node, $backup_name, %params) = @_;
+	my $backup_dir = $node->backup_dir . '/' . $backup_name;
+
+	mkdir $backup_dir or die "mkdir($backup_dir) failed: $!";
+
+	PostgreSQL::Test::RecursiveCopy::copypath($node->data_dir . '/pg_tde',
+		$backup_dir . '/pg_tde');
+
+	$node->backup($backup_name, %params);
+}
+
 1;

src/bin/pg_basebackup/Makefile

@@ -44,6 +44,17 @@ BBOBJS = \
 	bbstreamer_tar.o \
 	bbstreamer_zstd.o
 
+ifeq ($(enable_percona_ext),yes)
+
+OBJS += \
+	xlogreader.o \
+	$(top_srcdir)/src/fe_utils/simple_list.o \
+	$(top_builddir)/src/libtde/libtdexlog.a \
+	$(top_builddir)/src/libtde/libtde.a
+
+override CPPFLAGS := -I$(top_srcdir)/contrib/pg_tde/src/include -I$(top_srcdir)/contrib/pg_tde/src/libkmip/libkmip/include -DFRONTEND $(CPPFLAGS)
+endif
+
 all: pg_basebackup pg_createsubscriber pg_receivewal pg_recvlogical
 
 pg_basebackup: $(BBOBJS) $(OBJS) | submake-libpq submake-libpgport submake-libpgfeutils
@@ -58,6 +69,9 @@ pg_receivewal: pg_receivewal.o $(OBJS) | submake-libpq submake-libpgport submake
 pg_recvlogical: pg_recvlogical.o $(OBJS) | submake-libpq submake-libpgport submake-libpgfeutils
 	$(CC) $(CFLAGS) pg_recvlogical.o $(OBJS) $(LDFLAGS) $(LDFLAGS_EX) $(LIBS) -o $@$(X)
 
+xlogreader.c: % : $(top_srcdir)/src/backend/access/transam/%
+	rm -f $@ && $(LN_S) $< .
+
 install: all installdirs
 	$(INSTALL_PROGRAM) pg_basebackup$(X) '$(DESTDIR)$(bindir)/pg_basebackup$(X)'
 	$(INSTALL_PROGRAM) pg_createsubscriber$(X) '$(DESTDIR)$(bindir)/pg_createsubscriber$(X)'
@@ -76,7 +90,7 @@ uninstall:
 clean distclean:
 	rm -f pg_basebackup$(X) pg_createsubscriber$(X) pg_receivewal$(X) pg_recvlogical$(X) \
 		$(BBOBJS) pg_createsubscriber.o pg_receivewal.o pg_recvlogical.o \
-		$(OBJS)
+		$(OBJS) xlogreader.c
 	rm -rf tmp_check
 
 check:

src/bin/pg_basebackup/bbstreamer_file.c

@@ -18,6 +18,10 @@
 #include "common/logging.h"
 #include "common/string.h"
 
+#ifdef PERCONA_EXT
+#include "pg_tde.h"
+#endif
+
 typedef struct bbstreamer_plain_writer
 {
 	bbstreamer	base;
@@ -297,7 +301,8 @@ should_allow_existing_directory(const char *pathname)
 		strcmp(filename, "pg_xlog") == 0 ||
 		strcmp(filename, "archive_status") == 0 ||
 		strcmp(filename, "summaries") == 0 ||
-		strcmp(filename, "pg_tblspc") == 0)
+		strcmp(filename, "pg_tblspc") == 0 ||
+		strcmp(filename, PG_TDE_DATA_DIR) == 0) 
 		return true;
 
 	if (strspn(filename, "0123456789") == strlen(filename))

src/bin/pg_basebackup/meson.build

@@ -12,10 +12,15 @@ common_sources = files(
   'walmethods.c',
 )
 
+common_sources += xlogreader_sources
+
 pg_basebackup_deps = [frontend_code, libpq, lz4, zlib, zstd]
 pg_basebackup_common = static_library('libpg_basebackup_common',
   common_sources,
+  c_args: ['-DFRONTEND'], # needed for xlogreader et al
+  link_with: pg_tde_frontend,
   dependencies: pg_basebackup_deps,
+  include_directories: pg_tde_inc,
   kwargs: internal_lib_args,
 )
 
@@ -34,6 +39,7 @@ pg_basebackup = executable('pg_basebackup',
   link_with: [pg_basebackup_common],
   dependencies: pg_basebackup_deps,
   kwargs: default_bin_args,
+  include_directories: pg_tde_inc,
 )
 bin_targets += pg_basebackup
 
@@ -71,6 +77,7 @@ pg_receivewal = executable('pg_receivewal',
   link_with: [pg_basebackup_common],
   dependencies: pg_basebackup_deps,
   kwargs: default_bin_args,
+  include_directories: pg_tde_inc,
 )
 bin_targets += pg_receivewal
 

src/bin/pg_basebackup/pg_basebackup.c

@@ -38,7 +38,14 @@
 #include "receivelog.h"
 #include "streamutil.h"
 
-#define ERRCODE_DATA_CORRUPTED	"XX001"
+#ifdef PERCONA_EXT
+#include "access/pg_tde_fe_init.h"
+#include "access/pg_tde_xlog_smgr.h"
+#include "access/xlog_smgr.h"
+#include "pg_tde.h"
+#endif
+
+#define ERRCODE_DATA_CORRUPTED_BCP	"XX001"
 
 typedef struct TablespaceListCell
 {
@@ -654,6 +661,16 @@ StartLogStreamer(char *startpos, uint32 timeline, char *sysidentifier,
 			 PQserverVersion(conn) < MINIMUM_VERSION_FOR_PG_WAL ?
 			 "pg_xlog" : "pg_wal");
 
+#ifdef PERCONA_EXT
+	{
+		char tdedir[MAXPGPATH];
+
+		snprintf(tdedir, sizeof(tdedir), "%s/%s", basedir, PG_TDE_DATA_DIR);
+		pg_tde_fe_init(tdedir);
+		TDEXLogSmgrInit();
+	}
+#endif
+
 	/* Temporary replication slots are only supported in 10 and newer */
 	if (PQserverVersion(conn) < MINIMUM_VERSION_FOR_TEMP_SLOTS)
 		temp_replication_slot = false;
@@ -770,6 +787,14 @@ verify_dir_is_empty_or_create(char *dirname, bool *created, bool *found)
 		case 3:
 		case 4:
 
+#ifdef PERCONA_EXT
+			/* 
+			 * `pg_tde` may exists and contain keys and providers for the WAL
+			 * encryption
+			 */
+			if (strcmp(dirname, PG_TDE_DATA_DIR))
+				return;
+#endif
 			/*
 			 * Exists, not empty
 			 */
@@ -2201,7 +2226,7 @@ BaseBackup(char *compression_algorithm, char *compression_detail,
 		const char *sqlstate = PQresultErrorField(res, PG_DIAG_SQLSTATE);
 
 		if (sqlstate &&
-			strcmp(sqlstate, ERRCODE_DATA_CORRUPTED) == 0)
+			strcmp(sqlstate, ERRCODE_DATA_CORRUPTED_BCP) == 0)
 		{
 			pg_log_error("checksum error occurred");
 			checksum_failure = true;

src/bin/pg_basebackup/receivelog.c

@@ -25,6 +25,12 @@
 #include "receivelog.h"
 #include "streamutil.h"
 
+#ifdef PERCONA_EXT
+#include "access/pg_tde_fe_init.h"
+#include "access/pg_tde_xlog_smgr.h"
+#include "catalog/tde_global_space.h"
+#endif
+
 /* currently open WAL file */
 static Walfile *walfile = NULL;
 static bool reportFlushPosition = false;
@@ -1044,6 +1050,7 @@ ProcessXLogDataMsg(PGconn *conn, StreamCtl *stream, char *copybuf, int len,
 	int			bytes_left;
 	int			bytes_written;
 	int			hdr_len;
+	XLogSegNo 	segno;
 
 	/*
 	 * Once we've decided we don't want to receive any more, just ignore any
@@ -1071,6 +1078,8 @@ ProcessXLogDataMsg(PGconn *conn, StreamCtl *stream, char *copybuf, int len,
 	/* Extract WAL location for this block */
 	xlogoff = XLogSegmentOffset(*blockpos, WalSegSz);
 
+	XLByteToSeg(*blockpos, segno, WalSegSz);
+
 	/*
 	 * Verify that the initial location in the stream matches where we think
 	 * we are.
@@ -1121,6 +1130,11 @@ ProcessXLogDataMsg(PGconn *conn, StreamCtl *stream, char *copybuf, int len,
 			}
 		}
 
+#ifdef PERCONA_EXT
+		TDEXLogCryptBuffer(copybuf + hdr_len + bytes_written, bytes_to_write,
+						   xlogoff, stream->timeline, segno, WalSegSz);
+#endif
+
 		if (stream->walmethod->ops->write(walfile,
 										  copybuf + hdr_len + bytes_written,
 										  bytes_to_write) != bytes_to_write)