From 408c5dce04a99d2d89bff3dde18d4d66fd229fe6 Mon Sep 17 00:00:00 2001
From: Andreas Karlsson <andreas.karlsson@percona.com>
Date: Thu, 13 Mar 2025 14:45:20 +0100
Subject: [PATCH] PG-1416 Throw error when no principal key for encrypted table

By first looking up the principal key and then the relation key we hid
errors when we have lost the principal key but some relations are still
encrypted. Which could for example lead to trying to read an encrypted
table as it was unencrypted causing errors like the following:

  ERROR:  invalid page in block 0 of relation base/5/16448

This does not solve the much scarier issue when we get another principal
key back from the key server but we at least get better error messages
for this common case.
---
 contrib/pg_tde/meson.build                    |   1 +
 contrib/pg_tde/src/access/pg_tde_tdemap.c     |  35 ++---
 contrib/pg_tde/src/common/pg_tde_utils.c      |  36 ++---
 contrib/pg_tde/src/smgr/pg_tde_smgr.c         |   9 --
 contrib/pg_tde/t/010_change_key_provider.pl   | 144 ++++++++++++++++++
 .../t/expected/010_change_key_provider.out    |  65 ++++++++
 6 files changed, 235 insertions(+), 55 deletions(-)
 create mode 100644 contrib/pg_tde/t/010_change_key_provider.pl
 create mode 100644 contrib/pg_tde/t/expected/010_change_key_provider.out

diff --git a/contrib/pg_tde/meson.build b/contrib/pg_tde/meson.build
index 84c28af2667..053aa80f9b9 100644
--- a/contrib/pg_tde/meson.build
+++ b/contrib/pg_tde/meson.build
@@ -101,6 +101,7 @@ tap_tests = [
       't/007_tde_heap.pl',
       't/008_key_rotate_tablespace.pl',
       't/009_wal_encrypt.pl',
+      't/010_change_key_provider.pl',
 ]
 
 tests += {
diff --git a/contrib/pg_tde/src/access/pg_tde_tdemap.c b/contrib/pg_tde/src/access/pg_tde_tdemap.c
index bdb6a47d3ce..52ab5c3103f 100644
--- a/contrib/pg_tde/src/access/pg_tde_tdemap.c
+++ b/contrib/pg_tde/src/access/pg_tde_tdemap.c
@@ -981,6 +981,18 @@ pg_tde_get_key_from_file(const RelFileLocator *rlocator, uint32 key_type)
 
 	Assert(rlocator);
 
+	/* Get the file paths */
+	pg_tde_set_db_file_paths(rlocator->dbOid, db_map_path, db_keydata_path);
+
+	if (access(db_map_path, F_OK) == -1)
+		return NULL;
+
+	/* Read the map entry and get the index of the relation key */
+	key_index = pg_tde_process_map_entry(rlocator, key_type, db_map_path, &offset, false);
+
+	if (key_index == -1)
+		return NULL;
+
 	/*
 	 * Get/generate a principal key, create the key for relation and get the
 	 * encrypted key with bytes to write
@@ -996,27 +1008,8 @@ pg_tde_get_key_from_file(const RelFileLocator *rlocator, uint32 key_type)
 	LWLockAcquire(lock_pk, LW_SHARED);
 	principal_key = GetPrincipalKey(rlocator->dbOid, LW_SHARED);
 	if (principal_key == NULL)
-	{
-		LWLockRelease(lock_pk);
-		return NULL;
-	}
-
-	/* Get the file paths */
-	pg_tde_set_db_file_paths(rlocator->dbOid, db_map_path, db_keydata_path);
-
-	if (access(db_map_path, F_OK) == -1)
-	{
-		LWLockRelease(lock_pk);
-		return NULL;
-	}
-	/* Read the map entry and get the index of the relation key */
-	key_index = pg_tde_process_map_entry(rlocator, key_type, db_map_path, &offset, false);
-
-	if (key_index == -1)
-	{
-		LWLockRelease(lock_pk);
-		return NULL;
-	}
+		ereport(ERROR,
+				(errmsg("failed to retrieve principal key. Create one using pg_tde_set_principal_key before using encrypted tables.")));
 
 	enc_rel_key_data = pg_tde_read_keydata(db_keydata_path, key_index, principal_key);
 	LWLockRelease(lock_pk);
diff --git a/contrib/pg_tde/src/common/pg_tde_utils.c b/contrib/pg_tde/src/common/pg_tde_utils.c
index 1c1e122005f..5377a9faca6 100644
--- a/contrib/pg_tde/src/common/pg_tde_utils.c
+++ b/contrib/pg_tde/src/common/pg_tde_utils.c
@@ -34,35 +34,21 @@ Datum
 pg_tde_is_encrypted(PG_FUNCTION_ARGS)
 {
 	Oid			tableOid = PG_GETARG_OID(0);
-	Oid			dbOid = MyDatabaseId;
-	TDEPrincipalKey *principalKey = NULL;
+	LOCKMODE	lockmode = AccessShareLock;
+	Relation	rel = relation_open(tableOid, lockmode);
+	RelFileLocatorBackend rlocator = {.locator = rel->rd_locator,.backend = rel->rd_backend};
+	InternalKey *key;
 
-	LWLockAcquire(tde_lwlock_enc_keys(), LW_SHARED);
-	principalKey = GetPrincipalKey(dbOid, LW_SHARED);
-	LWLockRelease(tde_lwlock_enc_keys());
+	if (RelFileLocatorBackendIsTemp(rlocator) && !rel->rd_islocaltemp)
+		ereport(ERROR,
+				(errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
+				 errmsg("we cannot check if temporary relations from other backends are encrypted")));
 
-	if (principalKey == NULL)
-	{
-		PG_RETURN_BOOL(false);
-	}
-
-	{
-		LOCKMODE	lockmode = AccessShareLock;
-		Relation	rel = relation_open(tableOid, lockmode);
-		InternalKey *key;
-		RelFileLocatorBackend rlocator = {.locator = rel->rd_locator,.backend = rel->rd_backend};
+	key = GetSMGRRelationKey(rlocator);
 
-		if (RelFileLocatorBackendIsTemp(rlocator) && !rel->rd_islocaltemp)
-			ereport(ERROR,
-					(errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
-					 errmsg("we cannot check if temporary relations from other backends are encrypted")));
+	relation_close(rel, lockmode);
 
-		key = GetSMGRRelationKey(rlocator);
-
-		relation_close(rel, lockmode);
-
-		PG_RETURN_BOOL(key != NULL);
-	}
+	PG_RETURN_BOOL(key != NULL);
 }
 
 /*
diff --git a/contrib/pg_tde/src/smgr/pg_tde_smgr.c b/contrib/pg_tde/src/smgr/pg_tde_smgr.c
index ec9af0681f2..065de80a92a 100644
--- a/contrib/pg_tde/src/smgr/pg_tde_smgr.c
+++ b/contrib/pg_tde/src/smgr/pg_tde_smgr.c
@@ -40,7 +40,6 @@ tde_smgr_get_key(SMgrRelation reln, RelFileLocator *old_locator, bool can_create
 {
 	TdeCreateEvent *event;
 	InternalKey *key;
-	TDEPrincipalKey *pk;
 
 	if (IsCatalogRelationOid(reln->smgr_rlocator.locator.relNumber))
 	{
@@ -48,14 +47,6 @@ tde_smgr_get_key(SMgrRelation reln, RelFileLocator *old_locator, bool can_create
 		return NULL;
 	}
 
-	LWLockAcquire(tde_lwlock_enc_keys(), LW_SHARED);
-	pk = GetPrincipalKey(reln->smgr_rlocator.locator.dbOid, LW_SHARED);
-	LWLockRelease(tde_lwlock_enc_keys());
-	if (pk == NULL)
-	{
-		return NULL;
-	}
-
 	event = GetCurrentTdeCreateEvent();
 
 	/* see if we have a key for the relation, and return if yes */
diff --git a/contrib/pg_tde/t/010_change_key_provider.pl b/contrib/pg_tde/t/010_change_key_provider.pl
new file mode 100644
index 00000000000..94ac25da537
--- /dev/null
+++ b/contrib/pg_tde/t/010_change_key_provider.pl
@@ -0,0 +1,144 @@
+#!/usr/bin/perl
+
+use strict;
+use warnings;
+use File::Basename;
+use File::Compare;
+use File::Copy;
+use Test::More;
+use lib 't';
+use pgtde;
+
+# Get file name and CREATE out file name and dirs WHERE requried
+PGTDE::setup_files_dir(basename($0));
+
+# CREATE new PostgreSQL node and do initdb
+my $node = PGTDE->pgtde_init_pg();
+my $pgdata = $node->data_dir;
+
+# UPDATE postgresql.conf to include/load pg_tde library
+open my $conf, '>>', "$pgdata/postgresql.conf";
+print $conf "shared_preload_libraries = 'pg_tde'\n";
+close $conf;
+
+unlink('/tmp/change_key_provider_1.per');
+unlink('/tmp/change_key_provider_2.per');
+unlink('/tmp/change_key_provider_3.per');
+
+# Start server
+my $rt_value = $node->start;
+ok($rt_value == 1, "Start Server");
+
+# CREATE EXTENSION IF NOT EXISTS and change out file permissions
+my ($cmdret, $stdout, $stderr) = $node->psql('postgres', 'CREATE EXTENSION IF NOT EXISTS pg_tde;', extra_params => ['-a']);
+ok($cmdret == 0, "CREATE PGTDE EXTENSION");
+PGTDE::append_to_file($stdout);
+
+$stdout = $node->safe_psql('postgres', "SELECT pg_tde_add_key_provider_file('file-vault', '/tmp/change_key_provider_1.per');", extra_params => ['-a']);
+PGTDE::append_to_file($stdout);
+$stdout = $node->safe_psql('postgres', "SELECT pg_tde_list_all_key_providers();", extra_params => ['-a']);
+PGTDE::append_to_file($stdout);
+$stdout = $node->safe_psql('postgres', "SELECT pg_tde_set_principal_key('test-key', 'file-vault');", extra_params => ['-a']);
+PGTDE::append_to_file($stdout);
+
+$stdout = $node->safe_psql('postgres', 'CREATE TABLE test_enc (id serial, k integer, PRIMARY KEY (id)) USING tde_heap;', extra_params => ['-a']);
+PGTDE::append_to_file($stdout);
+$stdout = $node->safe_psql('postgres', 'INSERT INTO test_enc (k) VALUES (5), (6);', extra_params => ['-a']);
+PGTDE::append_to_file($stdout);
+
+$stdout = $node->safe_psql('postgres', "SELECT pg_tde_verify_principal_key();", extra_params => ['-a']);
+PGTDE::append_to_file($stdout);
+$stdout = $node->safe_psql('postgres', "SELECT pg_tde_is_encrypted('test_enc');", extra_params => ['-a']);
+PGTDE::append_to_file($stdout);
+$stdout = $node->safe_psql('postgres', 'SELECT * FROM test_enc ORDER BY id;', extra_params => ['-a']);
+PGTDE::append_to_file($stdout);
+
+# Change provider and move file
+PGTDE::append_to_file("-- mv /tmp/change_key_provider_1.per /tmp/change_key_provider_2.per");
+move('/tmp/change_key_provider_1.per', '/tmp/change_key_provider_2.per');
+$stdout = $node->safe_psql('postgres', "SELECT pg_tde_change_key_provider_file('file-vault', '/tmp/change_key_provider_2.per');", extra_params => ['-a']);
+PGTDE::append_to_file($stdout);
+$stdout = $node->safe_psql('postgres', "SELECT pg_tde_list_all_key_providers();", extra_params => ['-a']);
+PGTDE::append_to_file($stdout);
+
+$stdout = $node->safe_psql('postgres', "SELECT pg_tde_verify_principal_key();", extra_params => ['-a']);
+PGTDE::append_to_file($stdout);
+$stdout = $node->safe_psql('postgres', "SELECT pg_tde_is_encrypted('test_enc');", extra_params => ['-a']);
+PGTDE::append_to_file($stdout);
+$stdout = $node->safe_psql('postgres', 'SELECT * FROM test_enc ORDER BY id;', extra_params => ['-a']);
+PGTDE::append_to_file($stdout);
+
+# Restart the server
+PGTDE::append_to_file("-- server restart");
+$rt_value = $node->stop();
+$rt_value = $node->start();
+
+# Verify
+$stdout = $node->safe_psql('postgres', "SELECT pg_tde_verify_principal_key();", extra_params => ['-a']);
+PGTDE::append_to_file($stdout);
+$stdout = $node->safe_psql('postgres', "SELECT pg_tde_is_encrypted('test_enc');", extra_params => ['-a']);
+PGTDE::append_to_file($stdout);
+$stdout = $node->safe_psql('postgres', 'SELECT * FROM test_enc ORDER BY id;', extra_params => ['-a']);
+PGTDE::append_to_file($stdout);
+
+# Change provider and do not move file
+$stdout = $node->safe_psql('postgres', "SELECT pg_tde_change_key_provider_file('file-vault', '/tmp/change_key_provider_3.per');", extra_params => ['-a']);
+PGTDE::append_to_file($stdout);
+$stdout = $node->safe_psql('postgres', "SELECT pg_tde_list_all_key_providers();", extra_params => ['-a']);
+PGTDE::append_to_file($stdout);
+
+(undef, $stdout, $stderr) = $node->psql('postgres', "SELECT pg_tde_verify_principal_key();", extra_params => ['-a']);
+PGTDE::append_to_file($stdout);
+PGTDE::append_to_file($stderr);
+$stdout = $node->safe_psql('postgres', "SELECT pg_tde_is_encrypted('test_enc');", extra_params => ['-a']);
+PGTDE::append_to_file($stdout);
+$stdout = $node->safe_psql('postgres', 'SELECT * FROM test_enc ORDER BY id;', extra_params => ['-a']);
+PGTDE::append_to_file($stdout);
+
+# Restart the server
+PGTDE::append_to_file("-- server restart");
+$rt_value = $node->stop();
+$rt_value = $node->start();
+
+# Verify
+(undef, $stdout, $stderr) = $node->psql('postgres', "SELECT pg_tde_verify_principal_key();", extra_params => ['-a']);
+PGTDE::append_to_file($stdout);
+PGTDE::append_to_file($stderr);
+(undef, $stdout, $stderr) = $node->psql('postgres', "SELECT pg_tde_is_encrypted('test_enc');", extra_params => ['-a']);
+PGTDE::append_to_file($stdout);
+PGTDE::append_to_file($stderr);
+(undef, $stdout, $stderr) = $node->psql('postgres', 'SELECT * FROM test_enc ORDER BY id;', extra_params => ['-a']);
+PGTDE::append_to_file($stdout);
+PGTDE::append_to_file($stderr);
+
+PGTDE::append_to_file("-- mv /tmp/change_key_provider_2.per /tmp/change_key_provider_3.per");
+move('/tmp/change_key_provider_2.per', '/tmp/change_key_provider_3.per');
+
+# Restart the server
+PGTDE::append_to_file("-- server restart");
+$rt_value = $node->stop();
+$rt_value = $node->start();
+
+# Verify
+$stdout = $node->safe_psql('postgres', "SELECT pg_tde_verify_principal_key();", extra_params => ['-a']);
+PGTDE::append_to_file($stdout);
+$stdout = $node->safe_psql('postgres', "SELECT pg_tde_is_encrypted('test_enc');", extra_params => ['-a']);
+PGTDE::append_to_file($stdout);
+$stdout = $node->safe_psql('postgres', 'SELECT * FROM test_enc ORDER BY id;', extra_params => ['-a']);
+PGTDE::append_to_file($stdout);
+
+# DROP EXTENSION
+(undef, $stdout, $stderr) = $node->psql('postgres', 'DROP EXTENSION pg_tde CASCADE;', extra_params => ['-a']);
+PGTDE::append_to_file($stdout);
+PGTDE::append_to_file($stderr);
+# Stop the server
+$node->stop();
+
+# compare the expected and out file
+my $compare = PGTDE->compare_results();
+
+# Test/check if expected and result/out file match. If Yes, test passes.
+is($compare, 0, "Compare Files: $PGTDE::expected_filename_with_path and $PGTDE::out_filename_with_path files.");
+
+# Done testing for this testcase file.
+done_testing();
diff --git a/contrib/pg_tde/t/expected/010_change_key_provider.out b/contrib/pg_tde/t/expected/010_change_key_provider.out
new file mode 100644
index 00000000000..3f0f35f7243
--- /dev/null
+++ b/contrib/pg_tde/t/expected/010_change_key_provider.out
@@ -0,0 +1,65 @@
+CREATE EXTENSION IF NOT EXISTS pg_tde;
+SELECT pg_tde_add_key_provider_file('file-vault', '/tmp/change_key_provider_1.per');
+1
+SELECT pg_tde_list_all_key_providers();
+(1,file-vault,file,"{""type"" : ""file"", ""path"" : ""/tmp/change_key_provider_1.per""}")
+SELECT pg_tde_set_principal_key('test-key', 'file-vault');
+t
+CREATE TABLE test_enc (id serial, k integer, PRIMARY KEY (id)) USING tde_heap;
+INSERT INTO test_enc (k) VALUES (5), (6);
+SELECT pg_tde_verify_principal_key();
+
+SELECT pg_tde_is_encrypted('test_enc');
+t
+SELECT * FROM test_enc ORDER BY id;
+1|5
+2|6
+-- mv /tmp/change_key_provider_1.per /tmp/change_key_provider_2.per
+SELECT pg_tde_change_key_provider_file('file-vault', '/tmp/change_key_provider_2.per');
+1
+SELECT pg_tde_list_all_key_providers();
+(1,file-vault,file,"{""type"" : ""file"", ""path"" : ""/tmp/change_key_provider_2.per""}")
+SELECT pg_tde_verify_principal_key();
+
+SELECT pg_tde_is_encrypted('test_enc');
+t
+SELECT * FROM test_enc ORDER BY id;
+1|5
+2|6
+-- server restart
+SELECT pg_tde_verify_principal_key();
+
+SELECT pg_tde_is_encrypted('test_enc');
+t
+SELECT * FROM test_enc ORDER BY id;
+1|5
+2|6
+SELECT pg_tde_change_key_provider_file('file-vault', '/tmp/change_key_provider_3.per');
+1
+SELECT pg_tde_list_all_key_providers();
+(1,file-vault,file,"{""type"" : ""file"", ""path"" : ""/tmp/change_key_provider_3.per""}")
+SELECT pg_tde_verify_principal_key();
+psql:<stdin>:1: ERROR:  Failed to retrieve key from keyring
+SELECT pg_tde_is_encrypted('test_enc');
+t
+SELECT * FROM test_enc ORDER BY id;
+1|5
+2|6
+-- server restart
+SELECT pg_tde_verify_principal_key();
+psql:<stdin>:1: ERROR:  Failed to retrieve key from keyring
+SELECT pg_tde_is_encrypted('test_enc');
+psql:<stdin>:1: ERROR:  failed to retrieve principal key. Create one using pg_tde_set_principal_key before using encrypted tables.
+SELECT * FROM test_enc ORDER BY id;
+psql:<stdin>:1: ERROR:  failed to retrieve principal key. Create one using pg_tde_set_principal_key before using encrypted tables.
+-- mv /tmp/change_key_provider_2.per /tmp/change_key_provider_3.per
+-- server restart
+SELECT pg_tde_verify_principal_key();
+
+SELECT pg_tde_is_encrypted('test_enc');
+t
+SELECT * FROM test_enc ORDER BY id;
+1|5
+2|6
+DROP EXTENSION pg_tde CASCADE;
+psql:<stdin>:1: NOTICE:  drop cascades to table test_enc