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);