pg_rewind: Handle partial writes with remote source

Author: dAdAbird

fetools/pg18/pg_rewind/libpq_source.c

@@ -17,6 +17,9 @@
 #include "pg_rewind.h"
 #include "port/pg_bswap.h"
 #include "rewind_source.h"
+#include "tde_ops.h"
+
+#include "pg_tde.h"
 
 /*
  * Files are fetched MAX_CHUNK_SIZE bytes at a time, and with a
@@ -31,6 +34,7 @@ typedef struct
 	const char *path;			/* path relative to data directory root */
 	off_t		offset;
 	size_t		length;
+	bool		encrypt;
 } fetch_range_request;
 
 typedef struct
@@ -71,6 +75,10 @@ static char *libpq_fetch_file(rewind_source *source, const char *path,
 static XLogRecPtr libpq_get_current_wal_insert_lsn(rewind_source *source);
 static void libpq_destroy(rewind_source *source);
 
+static void libpq_queue_process_fetch_range(rewind_source *source, const char *path,
+											bool needs_encrypt, off_t off, size_t len);
+static void libpq_fetch_tde_keys(rewind_source *source);
+
 /*
  * Create a new libpq source.
  *
@@ -100,6 +108,8 @@ init_libpq_source(PGconn *conn)
 	initStringInfo(&src->offsets);
 	initStringInfo(&src->lengths);
 
+	libpq_fetch_tde_keys(&src->common);
+
 	return &src->common;
 }
 
@@ -345,7 +355,7 @@ libpq_queue_fetch_file(rewind_source *source, const char *path, size_t len)
 	 * fetch-requests are for a whole file.
 	 */
 	open_target_file(path, true);
-	libpq_queue_fetch_range(source, path, 0, Max(len, MAX_CHUNK_SIZE));
+	libpq_queue_process_fetch_range(source, path, false, 0, Max(len, MAX_CHUNK_SIZE));
 }
 
 /*
@@ -354,6 +364,17 @@ libpq_queue_fetch_file(rewind_source *source, const char *path, size_t len)
 static void
 libpq_queue_fetch_range(rewind_source *source, const char *path, off_t off,
 						size_t len)
+{
+	libpq_queue_process_fetch_range(source, path, true, off, len);
+}
+
+/*
+ * A workhorse for libpq_queue_fetch_range.
+ * `needs_encrypt` indicates if file's blocks may need re-encryption.
+ */
+static void
+libpq_queue_process_fetch_range(rewind_source *source, const char *path,
+								bool needs_encrypt, off_t off, size_t len)
 {
 	libpq_source *src = (libpq_source *) source;
 
@@ -406,6 +427,7 @@ libpq_queue_fetch_range(rewind_source *source, const char *path, off_t off,
 		src->request_queue[src->num_requests].path = path;
 		src->request_queue[src->num_requests].offset = off;
 		src->request_queue[src->num_requests].length = thislen;
+		src->request_queue[src->num_requests].encrypt = needs_encrypt;
 		src->num_requests++;
 
 		off += thislen;
@@ -420,6 +442,7 @@ static void
 libpq_finish_fetch(rewind_source *source)
 {
 	process_queued_fetch_requests((libpq_source *) source);
+	flush_current_key();
 }
 
 static void
@@ -592,6 +615,19 @@ process_queued_fetch_requests(libpq_source *src)
 
 			open_target_file(filename, false);
 
+			if (rq->encrypt)
+			{
+				Assert(chunksize % BLCKSZ == 0);
+
+				ensure_tde_keys(filename);
+
+				for (int i = 0; i < chunksize / BLCKSZ; i++)
+				{
+					unsigned char *data = (unsigned char *) chunk + BLCKSZ * i;
+
+					encrypt_block(data, chunkoff + BLCKSZ * i);
+				}
+			}
 			write_target_range(chunk, chunkoff, chunksize);
 		}
 
@@ -682,3 +718,52 @@ libpq_destroy(rewind_source *source)
 
 	/* NOTE: we don't close the connection here, as it was not opened by us. */
 }
+
+static void
+libpq_fetch_tde_keys(rewind_source *source)
+{
+	PGconn	   *conn = ((libpq_source *) source)->conn;
+	PGresult   *res;
+
+	res = PQexec(conn, "SELECT pg_ls_dir('" PG_TDE_DATA_DIR "', true, false)");
+
+	if (PQresultStatus(res) != PGRES_TUPLES_OK)
+		pg_fatal("could not fetch file list: %s",
+				 PQresultErrorMessage(res));
+
+	/* no tde dir, nothing to do */
+	if (PQnfields(res) == 0)
+	{
+		PQclear(res);
+		return;
+	}
+
+	init_tde();
+
+	for (int i = 0; i < PQntuples(res); i++)
+	{
+		char	   *path;
+		char	   *tde_file_buf;
+		size_t		size;
+		char		target_path[MAXPGPATH];
+
+		if (PQgetisnull(res, i, 0))
+		{
+			/*
+			 * The file was removed from the server while the query was
+			 * running. Ignore it.
+			 */
+			continue;
+		}
+
+		path = PQgetvalue(res, i, 0);
+
+		snprintf(target_path, MAXPGPATH, "%s/%s", PG_TDE_DATA_DIR, path);
+		tde_file_buf = libpq_fetch_file(source, target_path, &size);
+
+		write_tmp_source_file(path, tde_file_buf, size);
+		pg_free(tde_file_buf);
+	}
+
+	PQclear(res);
+}

fetools/pg18/pg_rewind/local_source.c

@@ -210,7 +210,7 @@ local_fetch_tde_keys(rewind_source *source)
 	if (!directory_exists(tde_source_dir))
 		return;
 
-	create_tde_tmp_dir();
+	init_tde();
 	copy_tmp_tde_files(tde_source_dir);
 }
 

fetools/pg18/pg_rewind/tde_ops.c

@@ -15,6 +15,7 @@
 #include "pg_tde.h"
 
 static void copy_dir(const char *src, const char *dst);
+static void create_tde_tmp_dir(void);
 
 typedef struct
 {
@@ -32,6 +33,7 @@ static current_file_data current_tde_file =
 
 /* Dir for an operational copy of source's tde files (_keys, etc)  */
 static char tde_tmp_scource[MAXPGPATH] = "/tmp/pg_tde_rewindXXXXXX";
+static bool source_has_tde = false;
 
 void
 flush_current_key(void)
@@ -55,13 +57,18 @@ ensure_tde_keys(const char *relpath)
 	RelFileLocator rlocator;
 	unsigned int segNo;
 
+	/* no TDE on source, nothing to do */
+	if (!source_has_tde)
+		return;
+
 	/* the same file, nothing to do */
 	if (strcmp(current_tde_file.path, relpath) == 0)
 		return;
 
 	flush_current_key();
 
-	path_rlocator(relpath, &rlocator, &segNo);
+	if (!path_rlocator(relpath, &rlocator, &segNo))
+		return;
 
 	pg_tde_set_data_dir(tde_tmp_scource);
 	current_tde_file.source_key = pg_tde_get_smgr_key(rlocator);
@@ -107,14 +114,12 @@ encrypt_block(unsigned char *buf, off_t file_offset)
 }
 
 
-void
+static void
 create_tde_tmp_dir(void)
 {
 	if (mkdtemp(tde_tmp_scource) == NULL)
 		pg_fatal("could not create temporary directory \"%s\": %m", tde_tmp_scource);
 
-	atexit(destroy_tde_tmp_dir);
-
 	pg_log_debug("created temporary pg_tde directory: %s", tde_tmp_scource);
 }
 
@@ -195,6 +200,14 @@ copy_dir(const char *src, const char *dst)
 		pg_fatal("could not close directory \"%s\": %m", src);
 }
 
+void
+init_tde(void)
+{
+	source_has_tde = true;
+	create_tde_tmp_dir();
+	atexit(destroy_tde_tmp_dir);
+}
+
 void
 copy_tmp_tde_files(const char *from)
 {
@@ -209,6 +222,9 @@ fetch_tde_dir(void)
 	if (dry_run)
 		return;
 
+	if (!source_has_tde)
+		return;
+
 	snprintf(target_tde_dir, MAXPGPATH, "%s/%s", datadir_target, PG_TDE_DATA_DIR);
 
 	rmtree(target_tde_dir, false);

fetools/pg18/pg_rewind/tde_ops.h

@@ -5,10 +5,10 @@ extern void flush_current_key(void);
 extern void ensure_tde_keys(const char *relpath);
 extern void encrypt_block(unsigned char *buf, off_t file_offset);
 
-extern void create_tde_tmp_dir(void);
 extern void destroy_tde_tmp_dir(void);
 extern void write_tmp_source_file(const char *fname, char *buf, size_t size);
 extern void fetch_tde_dir(void);
 extern void copy_tmp_tde_files(const char *from);
+extern void init_tde(void);
 
 #endif							/* PG_REWIND_TDE_FILE_H */

meson.build

@@ -319,6 +319,7 @@ tap_tests = [
   't/pg_resetwal_corrupted.pl',
   't/pg_rewind_basic.pl',
   't/pg_rewind_databases.pl',
+  't/pg_rewind_enc_copy_blocks.pl',
   't/pg_rewind_extrafiles.pl',
   't/pg_rewind_growing_files.pl',
   't/pg_rewind_keep_recycled_wals.pl',

t/RewindTest.pm

@@ -165,6 +165,16 @@ shared_preload_libraries = 'pg_tde'
 		"SELECT pg_tde_set_server_key_using_global_key_provider('global-db-principal-key', 'file-keyring-wal');"
 	);
 
+	$node_primary->safe_psql('postgres',
+		"SELECT pg_tde_add_database_key_provider_file('file-keyring','${tde_keyring_file}');"
+	);
+	$node_primary->safe_psql('postgres',
+		"SELECT pg_tde_create_key_using_database_key_provider('test-db-key', 'file-keyring');"
+	);
+	$node_primary->safe_psql('postgres',
+		"SELECT pg_tde_set_key_using_database_key_provider('test-db-key', 'file-keyring');"
+	);
+
 	$node_primary->append_conf(
 		'postgresql.conf', q{
 pg_tde.wal_encrypt = on

t/pg_rewind_enc_copy_blocks.pl

@@ -0,0 +1,82 @@
+
+# Copyright (c) 2021-2024, PostgreSQL Global Development Group
+
+use strict;
+use warnings FATAL => 'all';
+use PostgreSQL::Test::Utils;
+use Test::More;
+
+use FindBin;
+use lib $FindBin::RealBin;
+
+use RewindTest;
+
+sub run_test
+{
+	my $test_mode = shift;
+	my $extra_name = shift;
+	my $extra_conf = shift;
+
+	my $cluster_name = $test_mode;
+
+	$cluster_name = $cluster_name . $extra_name if defined $extra_name;
+
+	RewindTest::setup_cluster($cluster_name, [], $extra_conf);
+	RewindTest::start_primary();
+	RewindTest::create_standby($cluster_name);
+
+	primary_psql(
+		"CREATE TABLE tail_t (id INTEGER GENERATED ALWAYS AS IDENTITY PRIMARY KEY, f1 TEXT) USING tde_heap"
+	);
+	primary_psql(
+		"INSERT INTO tail_t (f1) SELECT repeat('abcdeF', 1000) FROM generate_series(1, 1000)"
+	);
+	primary_psql(
+		"CREATE TABLE block_t (id INTEGER GENERATED ALWAYS AS IDENTITY PRIMARY KEY, f1 TEXT) USING tde_heap"
+	);
+	primary_psql(
+		"INSERT INTO block_t (f1) SELECT repeat('abcdeF', 1000) FROM generate_series(1, 1000)"
+	);
+	primary_psql("CHECKPOINT");
+
+	RewindTest::promote_standby();
+
+	# Makes pg_rewind to copy some blocks of the relation
+	# (mixing data encrypted with different keys on the target).
+	primary_psql("UPDATE block_t SET f1='YYYYYYY' WHERE id % 10 = 0;");
+
+	# Insert some data making rewind to copy the tail of this relation
+	# (mixing data encrypted with different keys on the target).
+	standby_psql(
+		"INSERT INTO tail_t (f1) SELECT repeat('ghijk', 100) FROM generate_series(1, 1000)"
+	);
+	standby_psql("CHECKPOINT");
+
+
+	RewindTest::run_pg_rewind($test_mode);
+
+	check_query(
+		'SELECT count(*) FROM tail_t',
+		qq(2000
+),
+		'tail-copy');
+
+	check_query(
+		'SELECT count(*) FROM block_t',
+		qq(1000
+),
+		'blocks-copy');
+
+	RewindTest::clean_rewind_test();
+	return;
+}
+
+# Run the test in both modes
+run_test('local');
+run_test('remote');
+run_test('archive');
+
+my @conf_params = ("pg_tde.cipher = 'aes_256'");
+run_test('local', "_aes_256", \@conf_params);
+
+done_testing();