PG-1213 Add SQL function for checking if WAL record is encrypted

contrib/pg_tde/Makefile

@@ -1,7 +1,7 @@
 PGFILEDESC = "pg_tde access method"
 MODULE_big = pg_tde
 EXTENSION = pg_tde
-DATA = pg_tde--1.0.sql
+DATA = pg_tde--1.0.sql pg_tde--1.0--2.0.sql
 
 # Since meson supports skipping test suites this is a make only feature
 ifndef TDE_MODE

contrib/pg_tde/pg_tde--1.0--2.0.sql

@@ -0,0 +1,5 @@
+CREATE FUNCTION pg_tde_is_wal_record_encrypted(lsn pg_lsn)
+RETURNS BOOLEAN
+LANGUAGE C
+AS 'MODULE_PATHNAME';
+REVOKE ALL ON FUNCTION pg_tde_is_wal_record_encrypted(pg_lsn) FROM PUBLIC;

contrib/pg_tde/pg_tde.control

@@ -1,4 +1,4 @@
 comment = 'pg_tde access method'
-default_version = '1.0'
+default_version = '2.0'
 module_pathname = '$libdir/pg_tde'
 relocatable = false

contrib/pg_tde/src/include/pg_tde.h

@@ -2,7 +2,7 @@
 #define PG_TDE_H
 
 #define PG_TDE_NAME "pg_tde"
-#define PG_TDE_VERSION "1.0.0"
+#define PG_TDE_VERSION "2.0.0"
 #define PG_TDE_VERSION_STRING PG_TDE_NAME " " PG_TDE_VERSION
 
 #define PG_TDE_DATA_DIR	"pg_tde"

contrib/pg_tde/src/pg_tde.c

@@ -15,10 +15,12 @@
 #include "storage/shmem.h"
 #include "utils/builtins.h"
 #include "utils/percona.h"
+#include "utils/pg_lsn.h"
 
 #include "access/pg_tde_tdemap.h"
 #include "access/pg_tde_xlog.h"
 #include "access/pg_tde_xlog_smgr.h"
+#include "access/pg_tde_xlog_keys.h"
 #include "catalog/tde_global_space.h"
 #include "catalog/tde_principal_key.h"
 #include "encryption/enc_aes.h"
@@ -41,6 +43,7 @@ static shmem_request_hook_type prev_shmem_request_hook = NULL;
 PG_FUNCTION_INFO_V1(pg_tde_extension_initialize);
 PG_FUNCTION_INFO_V1(pg_tde_version);
 PG_FUNCTION_INFO_V1(pg_tdeam_handler);
+PG_FUNCTION_INFO_V1(pg_tde_is_wal_record_encrypted);
 
 static void
 tde_shmem_request(void)
@@ -162,3 +165,23 @@ pg_tdeam_handler(PG_FUNCTION_ARGS)
 {
 	PG_RETURN_POINTER(GetHeapamTableAmRoutine());
 }
+
+/*
+ * Returns true if the WAL record at the given LSN is encrypted.
+ */
+Datum
+pg_tde_is_wal_record_encrypted(PG_FUNCTION_ARGS)
+{
+	XLogRecPtr lsn = PG_GETARG_LSN(0);
+	WALKeyCacheRec *keys = pg_tde_fetch_wal_keys(InvalidXLogRecPtr);
+
+	if (!keys)
+		PG_RETURN_BOOL(false);
+
+	for (WALKeyCacheRec *curr_key = keys; curr_key != NULL; curr_key = curr_key->next){
+		if (curr_key->start_lsn <= lsn && curr_key->end_lsn > lsn)
+			PG_RETURN_BOOL(curr_key->key.type == TDE_KEY_TYPE_WAL_ENCRYPTED);
+	}
+
+	PG_RETURN_BOOL(false);
+}

contrib/pg_tde/t/expected/wal_encrypt.out

@@ -13,6 +13,12 @@ SELECT key_name, provider_name, provider_id FROM pg_tde_server_key_info();
           |               |            
 (1 row)
 
+SELECT pg_tde_is_wal_record_encrypted(pg_current_wal_lsn());
+ pg_tde_is_wal_record_encrypted 
+--------------------------------
+ f
+(1 row)
+
 SELECT pg_tde_create_key_using_global_key_provider('server-key', 'file-keyring-010');
  pg_tde_create_key_using_global_key_provider 
 ---------------------------------------------
@@ -54,6 +60,12 @@ SELECT slot_name FROM pg_create_logical_replication_slot('tde_slot', 'test_decod
 
 CREATE TABLE test_wal (id SERIAL, k INTEGER, PRIMARY KEY (id));
 INSERT INTO test_wal (k) VALUES (1), (2);
+SELECT pg_tde_is_wal_record_encrypted(pg_current_wal_lsn());
+ pg_tde_is_wal_record_encrypted 
+--------------------------------
+ t
+(1 row)
+
 ALTER SYSTEM SET pg_tde.wal_encrypt = off;
 -- server restart without wal encryption
 SHOW pg_tde.wal_encrypt;
@@ -63,6 +75,12 @@ SHOW pg_tde.wal_encrypt;
 (1 row)
 
 INSERT INTO test_wal (k) VALUES (3), (4);
+SELECT pg_tde_is_wal_record_encrypted(pg_current_wal_lsn());
+ pg_tde_is_wal_record_encrypted 
+--------------------------------
+ f
+(1 row)
+
 ALTER SYSTEM SET pg_tde.wal_encrypt = on;
 -- server restart with wal encryption
 SHOW pg_tde.wal_encrypt;
@@ -72,6 +90,12 @@ SHOW pg_tde.wal_encrypt;
 (1 row)
 
 INSERT INTO test_wal (k) VALUES (5), (6);
+SELECT pg_tde_is_wal_record_encrypted(pg_current_wal_lsn());
+ pg_tde_is_wal_record_encrypted 
+--------------------------------
+ t
+(1 row)
+
 -- server restart with still wal encryption
 SHOW pg_tde.wal_encrypt;
  pg_tde.wal_encrypt 
@@ -80,6 +104,12 @@ SHOW pg_tde.wal_encrypt;
 (1 row)
 
 INSERT INTO test_wal (k) VALUES (7), (8);
+SELECT pg_tde_is_wal_record_encrypted(pg_current_wal_lsn());
+ pg_tde_is_wal_record_encrypted 
+--------------------------------
+ t
+(1 row)
+
 SELECT data FROM pg_logical_slot_get_changes('tde_slot', NULL, NULL);
                            data                            
 -----------------------------------------------------------

contrib/pg_tde/t/wal_encrypt.pl

@@ -31,6 +31,8 @@
 	'SELECT key_name, provider_name, provider_id FROM pg_tde_server_key_info();'
 );
 
+PGTDE::psql($node, 'postgres', "SELECT pg_tde_is_wal_record_encrypted(pg_current_wal_lsn());");
+
 PGTDE::psql($node, 'postgres',
 	"SELECT pg_tde_create_key_using_global_key_provider('server-key', 'file-keyring-010');"
 );
@@ -60,6 +62,8 @@
 
 PGTDE::psql($node, 'postgres', 'INSERT INTO test_wal (k) VALUES (1), (2);');
 
+PGTDE::psql($node, 'postgres', "SELECT pg_tde_is_wal_record_encrypted(pg_current_wal_lsn());");
+
 PGTDE::psql($node, 'postgres', 'ALTER SYSTEM SET pg_tde.wal_encrypt = off;');
 
 PGTDE::append_to_result_file("-- server restart without wal encryption");
@@ -69,6 +73,8 @@
 
 PGTDE::psql($node, 'postgres', 'INSERT INTO test_wal (k) VALUES (3), (4);');
 
+PGTDE::psql($node, 'postgres', "SELECT pg_tde_is_wal_record_encrypted(pg_current_wal_lsn());");
+
 PGTDE::psql($node, 'postgres', 'ALTER SYSTEM SET pg_tde.wal_encrypt = on;');
 
 PGTDE::append_to_result_file("-- server restart with wal encryption");
@@ -78,13 +84,17 @@
 
 PGTDE::psql($node, 'postgres', 'INSERT INTO test_wal (k) VALUES (5), (6);');
 
+PGTDE::psql($node, 'postgres', "SELECT pg_tde_is_wal_record_encrypted(pg_current_wal_lsn());");
+
 PGTDE::append_to_result_file("-- server restart with still wal encryption");
 $node->restart;
 
 PGTDE::psql($node, 'postgres', "SHOW pg_tde.wal_encrypt;");
 
 PGTDE::psql($node, 'postgres', 'INSERT INTO test_wal (k) VALUES (7), (8);');
 
+PGTDE::psql($node, 'postgres', "SELECT pg_tde_is_wal_record_encrypted(pg_current_wal_lsn());");
+
 PGTDE::psql($node, 'postgres',
 	"SELECT data FROM pg_logical_slot_get_changes('tde_slot', NULL, NULL);");