PG-1603 Make pg_basebackup work with encrypted WAL

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)