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. And we encrypt streamed WAL according to the internal keys. No
`pg_tde` dir means no streamed WAL encryption.

Author: dAdAbird

contrib/pg_tde/src/access/pg_tde_xlog_smgr.c

@@ -252,10 +252,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;
-	XLogRecPtr	data_start;
-	XLogRecPtr	data_end;
 
 #ifdef TDE_XLOG_DEBUG
 	elog(DEBUG1, "read from a WAL segment, size: %lu offset: %ld [%lX], seg: %X/%X",
@@ -267,6 +263,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;
+	XLogRecPtr	data_start;
+	XLogRecPtr	data_end;
+
 	if (!keys)
 	{
 		/* cache is empty, try to read keys from disk */
@@ -294,7 +307,7 @@ tdeheap_xlog_seg_read(int fd, void *buf, size_t count, off_t offset,
 #endif
 
 	XLogSegNoOffsetToRecPtr(segno, offset, segSize, data_start);
-	XLogSegNoOffsetToRecPtr(segno, offset + readsz, segSize, data_end);
+	XLogSegNoOffsetToRecPtr(segno, offset + count, segSize, data_end);
 
 	/*
 	 * TODO: this is higly ineffective. We should get rid of linked list and
@@ -331,7 +344,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;
@@ -350,8 +363,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

@@ -11,4 +11,7 @@ extern Size TDEXLogEncryptStateSize(void);
 extern void TDEXLogShmemInit(void);
 extern void TDEXLogSmgrInit(void);
 
+extern void TDEXLogCryptBuffer(void *buf, size_t count, off_t offset,
+								   TimeLineID tli, XLogSegNo segno, int segSize);
+
 #endif							/* PG_TDE_XLOGSMGR_H */

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;
@@ -226,7 +230,9 @@ bbstreamer_extractor_content(bbstreamer *streamer, bbstreamer_member *member,
 
 			/* Dispatch based on file type. */
 			if (member->is_directory)
+			{
 				extract_directory(mystreamer->filename, member->mode);
+			}
 			else if (member->is_link)
 			{
 				const char *linktarget = member->linktarget;
@@ -236,9 +242,19 @@ bbstreamer_extractor_content(bbstreamer *streamer, bbstreamer_member *member,
 				extract_link(mystreamer->filename, linktarget);
 			}
 			else
+			{
+#ifdef PERCONA_EXT
+				/* 
+				 * Don't rewrite WAL keys and providers. User may have different
+				 * one on source and target.
+				 */
+				if (strncmp(member->pathname, "pg_tde/1664_", 12) == 0)
+					break;
+#endif
 				mystreamer->file =
 					create_file_for_extract(mystreamer->filename,
 											member->mode);
+			}
 
 			/* Report output file change. */
 			if (mystreamer->report_output_file)
@@ -297,7 +313,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
 {
@@ -621,6 +628,9 @@ StartLogStreamer(char *startpos, uint32 timeline, char *sysidentifier,
 	uint32		hi,
 				lo;
 	char		statusdir[MAXPGPATH];
+#ifdef PERCONA_EXT
+	char		tdedir[MAXPGPATH];
+#endif
 
 	param = pg_malloc0(sizeof(logstreamer_param));
 	param->timeline = timeline;
@@ -654,6 +664,12 @@ StartLogStreamer(char *startpos, uint32 timeline, char *sysidentifier,
 			 PQserverVersion(conn) < MINIMUM_VERSION_FOR_PG_WAL ?
 			 "pg_xlog" : "pg_wal");
 
+#ifdef PERCONA_EXT
+	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 +786,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 +2225,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)