pg_rewind: Re-encrypt fsm/vm forks for partial writes

We rewrite the internal key for relations when partially re-encrypting
blocks. That makes its FSM and VM fork unreadable as they are still
encrypted with the old key. To fix that, we re-encrypt such forks with
the proper key after we finish processing the main fork file.

As pg_rewind processes files in the order of operation types (see
file_action_t) and whole-file copies occur before any partial writes,
we assume that for files already in the target datadir, we rewrite
them in-place.

Author: dAdAbird

fetools/pg18/pg_rewind/libpq_source.c

@@ -625,7 +625,7 @@ process_queued_fetch_requests(libpq_source *src)
 				{
 					unsigned char *data = (unsigned char *) chunk + BLCKSZ * i;
 
-					encrypt_block(data, chunkoff + BLCKSZ * i);
+					encrypt_block(data, chunkoff + BLCKSZ * i, MAIN_FORKNUM);
 				}
 			}
 			write_target_range(chunk, chunkoff, chunksize);

fetools/pg18/pg_rewind/local_source.c

@@ -177,7 +177,7 @@ local_queue_fetch_range(rewind_source *source, const char *path, off_t off,
 			pg_fatal("unexpected EOF while reading file \"%s\"", srcpath);
 
 		/* Re-encrypt blocks with a proper key if neeed. */
-		encrypt_block((unsigned char *) buf.data, begin);
+		encrypt_block((unsigned char *) buf.data, begin, MAIN_FORKNUM);
 
 		write_target_range(buf.data, begin, readlen);
 		begin += readlen;

fetools/pg18/pg_rewind/tde_ops.c

@@ -35,12 +35,79 @@ static current_file_data current_tde_file =
 static char tde_tmp_scource[MAXPGPATH] = "/tmp/pg_tde_rewindXXXXXX";
 static bool source_has_tde = false;
 
+static void
+recrypt_fork(ForkNumber fork)
+{
+	int			srcfd;
+	int			trgfd;
+	char		srcpath[MAXPGPATH];
+	PGIOAlignedBlock buf;
+	size_t		written_len;
+	RelPathStr	rp = relpathperm(current_tde_file.rlocator, fork);
+
+	snprintf(srcpath, sizeof(srcpath), "%s/%s", datadir_target, rp.str);
+
+	/* check if fork exists, nothing to do if it does not */
+	if (access(srcpath, F_OK) != 0)
+		return;
+
+	srcfd = open(srcpath, O_RDONLY | PG_BINARY, 0);
+	if (srcfd < 0)
+	{
+		/*
+		 * Server can recover from wrecked VM/FSM, hence only warnings here
+		 * and in the rest of the function
+		 */
+		pg_log_warning("could not open file for reading \"%s\": %m", srcpath);
+		return;
+	}
+
+	trgfd = open(srcpath, O_WRONLY | PG_BINARY, 0);
+	if (trgfd < 0)
+	{
+		pg_log_warning("could not open file for writing \"%s\": %m", srcpath);
+		close(srcfd);
+		return;
+	}
+
+	written_len = 0;
+	for (;;)
+	{
+		ssize_t		read_len;
+
+		read_len = read(srcfd, buf.data, sizeof(buf));
+
+		if (read_len < 0)
+			pg_fatal("could not read file \"%s\": %m", srcpath);
+		else if (read_len == 0)
+			break;				/* EOF reached */
+
+		encrypt_block((unsigned char *) buf.data, written_len, fork);
+
+		if (write(trgfd, buf.data, read_len) != read_len)
+		{
+			pg_log_warning("could not write block to fork file \"%s\": %m", srcpath);
+			break;
+		}
+		written_len += read_len;
+	}
+
+	close(srcfd);
+	close(trgfd);
+}
+
+
 void
 flush_current_key(void)
 {
 	if (current_tde_file.source_key == NULL)
 		return;
 
+	pg_log_debug("ensure forks encryption for \"%s\"", current_tde_file.path);
+
+	recrypt_fork(FSM_FORKNUM);
+	recrypt_fork(VISIBILITYMAP_FORKNUM);
+
 	pg_log_debug("update internal key for \"%s\"", current_tde_file.path);
 	pg_tde_set_data_dir(tde_tmp_scource);
 	pg_tde_save_smgr_key(current_tde_file.rlocator, current_tde_file.target_key, true);
@@ -96,7 +163,7 @@ ensure_tde_keys(const char *relpath)
 }
 
 void
-encrypt_block(unsigned char *buf, off_t file_offset)
+encrypt_block(unsigned char *buf, off_t file_offset, ForkNumber fork)
 {
 	BlockNumber blkno;
 
@@ -109,8 +176,8 @@ encrypt_block(unsigned char *buf, off_t file_offset)
 	blkno = file_offset / BLCKSZ + current_tde_file.segNo * RELSEG_SIZE;
 
 	pg_log_debug("re-encrypt block in %s, offset: %ld, blockNum: %u", current_tde_file.path, (long) file_offset, blkno);
-	tde_decrypt_smgr_block(current_tde_file.source_key, MAIN_FORKNUM, blkno, buf, buf);
-	tde_encrypt_smgr_block(current_tde_file.target_key, MAIN_FORKNUM, blkno, buf, buf);
+	tde_decrypt_smgr_block(current_tde_file.source_key, fork, blkno, buf, buf);
+	tde_encrypt_smgr_block(current_tde_file.target_key, fork, blkno, buf, buf);
 }
 
 

fetools/pg18/pg_rewind/tde_ops.h

@@ -1,9 +1,11 @@
 #ifndef PG_REWIND_TDE_FILE_H
 #define PG_REWIND_TDE_FILE_H
 
+#include "common/relpath.h"
+
 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 encrypt_block(unsigned char *buf, off_t file_offset, ForkNumber fork);
 
 extern void destroy_tde_tmp_dir(void);
 extern void write_tmp_source_file(const char *fname, char *buf, size_t size);

meson.build

@@ -320,6 +320,7 @@ tap_tests = [
   't/pg_rewind_basic.pl',
   't/pg_rewind_databases.pl',
   't/pg_rewind_enc_copy_blocks.pl',
+  't/pg_rewind_enc_fsm.pl',
   't/pg_rewind_enc_unchanged_rel.pl',
   't/pg_rewind_extrafiles.pl',
   't/pg_rewind_growing_files.pl',

t/pg_rewind_enc_fsm.pl

@@ -0,0 +1,72 @@
+
+# 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 tbl1 (id INTEGER GENERATED ALWAYS AS IDENTITY PRIMARY KEY, f1 TEXT) USING tde_heap"
+	);
+	primary_psql(
+		"INSERT INTO tbl1 (f1) SELECT repeat('abcdeF', 1000) FROM generate_series(1, 1000)"
+	);
+	primary_psql("CHECKPOINT");
+
+	RewindTest::promote_standby();
+
+	# Trigger updated blocks in FSM
+	standby_psql(
+		"DELETE FROM tbl1 WHERE id % 15 = 0;"
+	);
+	standby_psql(
+		"INSERT INTO tbl1 (f1) SELECT repeat('ghijk', 100) FROM generate_series(1, 1000)"
+	);
+
+
+	RewindTest::run_pg_rewind($test_mode);
+
+	ok( !$RewindTest::node_primary->log_contains(
+			'; zeroing out page'
+	),
+	'verify there are no corrupted _fsm relations');
+
+	check_query(
+		'SELECT count(*) FROM tbl1',
+		qq(1934
+),
+		'check table');
+
+	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();