PG-1411 Make pg_resetwal work with TDE

contrib/pg_tde/src/access/pg_tde_xlog_smgr.c

@@ -240,6 +240,26 @@ tdeheap_xlog_seg_write(int fd, const void *buf, size_t count, off_t offset,
 	if (EncryptXLog)
 		return TDEXLogWriteEncryptedPages(fd, buf, count, offset, tli, segno);
 	else
+#else
+	InternalKey *key = pg_tde_read_last_wal_key();
+
+	if (key && key->type == TDE_KEY_TYPE_WAL_ENCRYPTED)
+	{
+		TDEPrincipalKey *principal_key;
+		InternalKey encryptionKey;
+
+		char		iv_prefix[16];
+		char		enc_buff[XLOG_BLCKSZ];
+		void	   *crypt_ctx = NULL;
+
+		CalcXLogPageIVPrefix(tli, segno, key->base_iv, iv_prefix);
+		pg_tde_stream_crypt(iv_prefix, offset,
+							(char *) buf, count,
+							enc_buff, key, &crypt_ctx);
+
+		return pg_pwrite(fd, enc_buff, count, offset);
+	}
+	else
 #endif
 		return pg_pwrite(fd, buf, count, offset);
 }

src/bin/pg_resetwal/.gitignore

@@ -1,2 +1,5 @@
 /pg_resetwal
 /tmp_check/
+
+# Source files copied from src/backend/access/transam/
+/xlogreader.c

src/bin/pg_resetwal/Makefile

@@ -15,14 +15,29 @@ subdir = src/bin/pg_resetwal
 top_builddir = ../../..
 include $(top_builddir)/src/Makefile.global
 
+override CPPFLAGS := -DFRONTEND $(CPPFLAGS)
 LDFLAGS_INTERNAL += -L$(top_builddir)/src/fe_utils -lpgfeutils
 
 OBJS = \
 	$(WIN32RES) \
+	xlogreader.o \
 	pg_resetwal.o
 
+ifeq ($(enable_percona_ext),yes)
+
+OBJS += \
+	$(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 $(CPPFLAGS)
+endif
+
 all: pg_resetwal
 
+xlogreader.c: % : $(top_srcdir)/src/backend/access/transam/%
+	rm -f $@ && $(LN_S) $< .
+
 pg_resetwal: $(OBJS) | submake-libpgport
 	$(CC) $(CFLAGS) $^ $(LDFLAGS) $(LDFLAGS_EX) $(LIBS) -o $@$(X)
 

src/bin/pg_resetwal/meson.build

@@ -10,10 +10,22 @@ if host_system == 'windows'
     '--FILEDESC', 'pg_resetwal - reset PostgreSQL WAL log'])
 endif
 
+link_w  = []
+include_dirs = [postgres_inc]
+
+if percona_ext == true
+  link_w = [pg_tde_frontend]
+  include_dirs = [postgres_inc, pg_tde_inc]
+  pg_resetwal_sources += xlogreader_sources
+endif
+
 pg_resetwal = executable('pg_resetwal',
   pg_resetwal_sources,
   dependencies: [frontend_code],
+  c_args: ['-DFRONTEND'], # needed for xlogreader et al
   kwargs: default_bin_args,
+  include_directories: include_dirs,
+  link_with: link_w
 )
 bin_targets += pg_resetwal
 

src/bin/pg_resetwal/pg_resetwal.c

@@ -59,6 +59,14 @@
 #include "pg_getopt.h"
 #include "storage/large_object.h"
 
+#ifdef PERCONA_EXT
+#include "pg_tde.h"
+#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
+
 static ControlFileData ControlFile; /* pg_control values */
 static XLogSegNo newXlogSegNo;	/* new XLOG segment # */
 static bool guessed = false;	/* T if we had to guess at any values */
@@ -344,6 +352,15 @@ main(int argc, char *argv[])
 	}
 #endif
 
+#ifdef PERCONA_EXT
+	{
+		char tde_path[MAXPGPATH];
+		snprintf(tde_path, sizeof(tde_path), "%s/%s", DataDir, PG_TDE_DATA_DIR);
+		pg_tde_fe_init(tde_path);
+		TDEXLogSmgrInit();
+	}
+#endif
+
 	get_restricted_token();
 
 	/* Set mask based on PGDATA permissions */
@@ -1134,14 +1151,25 @@ WriteEmptyXLOG(void)
 		pg_fatal("could not open file \"%s\": %m", path);
 
 	errno = 0;
+#ifdef PERCONA_EXT
+	if (xlog_smgr->seg_write(fd, buffer.data, XLOG_BLCKSZ, 0,
+						 ControlFile.checkPointCopy.ThisTimeLineID,
+						 newXlogSegNo) != XLOG_BLCKSZ)
+#else
 	if (write(fd, buffer.data, XLOG_BLCKSZ) != XLOG_BLCKSZ)
+#endif
 	{
 		/* if write didn't set errno, assume problem is no disk space */
 		if (errno == 0)
 			errno = ENOSPC;
 		pg_fatal("could not write file \"%s\": %m", path);
 	}
 
+#ifdef PERCONA_EXT
+	/* If we used xlog smgr, we need to update the file offset */
+	lseek(fd, XLOG_BLCKSZ, SEEK_CUR);
+#endif
+
 	/* Fill the rest of the file with zeroes */
 	memset(buffer.data, 0, XLOG_BLCKSZ);
 	for (nbytes = XLOG_BLCKSZ; nbytes < WalSegSz; nbytes += XLOG_BLCKSZ)