pg_rewind: Handle unchanged/truncated encrypted files

Encrypted files might remain unchanged or truncated by rewind on the
target. That leaves them encrypted with the target's key wich vanishes
after the pg_tde replace. So we need to ensure keys for such files.
Meaning if this is an encrypted relation, we save its target key into
the source's keys.

Author: dAdAbird

fetools/pg18/pg_rewind/filemap.c

@@ -487,6 +487,8 @@ action_to_str(file_action_t action)
 			return "CREATE";
 		case FILE_ACTION_REMOVE:
 			return "REMOVE";
+		case FILE_ACTION_ENSURE_TDE_KEY:
+			return "ENSURE_KEY";
 
 		default:
 			return "unknown";
@@ -591,7 +593,7 @@ isRelDataFile(const char *path)
 /*
  * Sets rlocator and segNo based on given path. Returns false if didn't find
  * a match.
- * 
+ *
  * Only concerned with files belonging to the main fork.
  */
 bool
@@ -719,13 +721,11 @@ decide_file_action(file_entry_t *entry)
 	if (strstr(path, ".DS_Store") != NULL)
 		return FILE_ACTION_NONE;
 
-	/* 
+	/*
 	 * Skip pg_tde key data but WAL-related stuff as WAL being replaced by
 	 * source's. We will handle the rest while re-encrypting data.
 	 */
-	if (strstr(path, "pg_tde/") != NULL &&
-		strstr(path, "pg_tde/wal_keys") == NULL &&
-		strstr(path, "pg_tde/1664_providers") == NULL)
+	if (strstr(path, "pg_tde/") != NULL)
 		return FILE_ACTION_NONE;
 
 	/*
@@ -847,14 +847,15 @@ decide_file_action(file_entry_t *entry)
 				 * in the target will be copied based on parsing the target
 				 * system's WAL, and any blocks modified in the source will be
 				 * updated after rewinding, when the source system's WAL is
-				 * replayed.
+				 * replayed. But we still have to sync source/target keys in
+				 * case it is encrypted.
 				 */
 				if (entry->target_size < entry->source_size)
 					return FILE_ACTION_COPY_TAIL;
 				else if (entry->target_size > entry->source_size)
 					return FILE_ACTION_TRUNCATE;
 				else
-					return FILE_ACTION_NONE;
+					return FILE_ACTION_ENSURE_TDE_KEY;
 			}
 			break;
 

fetools/pg18/pg_rewind/filemap.h

@@ -25,6 +25,9 @@ typedef enum
 								 * blocks based on the parsed WAL) */
 	FILE_ACTION_TRUNCATE,		/* truncate local file to 'newsize' bytes */
 	FILE_ACTION_REMOVE,			/* remove local file / directory / symlink */
+	FILE_ACTION_ENSURE_TDE_KEY, /* data file with no action, but we to check
+								 * if it is encrypted and sync source/target
+								 * keys */
 } file_action_t;
 
 typedef enum

fetools/pg18/pg_rewind/pg_rewind.c

@@ -31,6 +31,7 @@
 #include "pg_rewind.h"
 #include "rewind_source.h"
 #include "storage/bufpage.h"
+#include "tde_ops.h"
 
 #include "pg_tde.h"
 #include "access/pg_tde_fe_init.h"
@@ -612,12 +613,28 @@ perform_rewind(filemap_t *filemap, rewind_source *source,
 				/* nothing else to do */
 				break;
 
+			case FILE_ACTION_ENSURE_TDE_KEY:
+
+				/*
+				 * Partial rewrites will ensure the keys on their own.
+				 * Moreover, some partial updates, when the source is libpq,
+				 * may happen in the last turn, when source->finish_fetch() is
+				 * called. So running ensure_tde_keys for such files
+				 * prematurely will make them unreadable since the source key
+				 * would be updated before we use it to decrypt source data.
+				 */
+				if (entry->target_pages_to_overwrite.bitmapsize == 0)
+					ensure_tde_keys(entry->path);
+				break;
+
 			case FILE_ACTION_COPY:
 				source->queue_fetch_file(source, entry->path, entry->source_size);
 				break;
 
 			case FILE_ACTION_TRUNCATE:
 				truncate_target_file(entry->path, entry->source_size);
+				if (entry->target_pages_to_overwrite.bitmapsize == 0)
+					ensure_tde_keys(entry->path);
 				break;
 
 			case FILE_ACTION_COPY_TAIL:

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_unchanged_rel.pl',
   't/pg_rewind_extrafiles.pl',
   't/pg_rewind_growing_files.pl',
   't/pg_rewind_keep_recycled_wals.pl',

t/pg_rewind_enc_unchanged_rel.pl

@@ -0,0 +1,65 @@
+
+# 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 tbl (id INTEGER GENERATED ALWAYS AS IDENTITY PRIMARY KEY, f1 TEXT) USING tde_heap"
+	);
+	primary_psql(
+		"INSERT INTO tbl (f1) SELECT repeat('abcdeF', 1000) FROM generate_series(1, 1000)"
+	);
+	primary_psql("CHECKPOINT");
+
+	RewindTest::promote_standby();
+
+
+	# Makes an index relation to remain unchanged on target. So test that we
+	# preserve a target's internal key for this rel
+	standby_psql("CHECKPOINT");
+
+
+	RewindTest::run_pg_rewind($test_mode);
+
+	check_query(
+		'SELECT count(*) FROM tbl',
+		qq(1000
+),
+		'read-unchanged');
+
+
+	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();