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

Add user function for checking if WAL record with given LSN is
encrypted. Function assumes that provided LSN belongs to current
timeline, however it allows to specify timeline ID as second argument.

Author: artemgavrilov

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/expected/version.out

@@ -2,7 +2,7 @@ CREATE EXTENSION pg_tde;
 SELECT pg_tde_version();
  pg_tde_version 
 ----------------
- pg_tde 1.0.0
+ pg_tde 2.0.0
 (1 row)
 
 DROP EXTENSION pg_tde;

contrib/pg_tde/meson.build

@@ -78,6 +78,7 @@ endif
 install_data(
   'pg_tde.control',
   'pg_tde--1.0.sql',
+  'pg-tde--1.0--2.0.sql',
   kwargs: contrib_data_args,
 )
 

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, tli integer DEFAULT 0)
+RETURNS BOOLEAN
+LANGUAGE C
+AS 'MODULE_PATHNAME';
+REVOKE ALL ON FUNCTION pg_tde_is_wal_record_encrypted(pg_lsn, integer) 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,39 @@ 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);
+	int			tli = PG_GETARG_INT32(1);
+	WalLocation loc;
+	WALKeyCacheRec *keys;
+
+	if (tli == 0)
+		tli = GetWALInsertionTimeLine();
+
+	/* Load all keys for the given timeline */
+	loc = (WalLocation)
+	{
+		.tli = tli,.lsn = 0
+	};
+
+	keys = pg_tde_fetch_wal_keys(loc);
+	if (!keys)
+		PG_RETURN_BOOL(false);
+
+	loc.lsn = lsn;
+
+	for (WALKeyCacheRec *curr_key = keys; curr_key != NULL; curr_key = curr_key->next)
+	{
+		if (wal_location_cmp(loc, curr_key->start) >= 0 &&
+			wal_location_cmp(loc, curr_key->end) < 0)
+			PG_RETURN_BOOL(curr_key->key.type == WAL_KEY_TYPE_ENCRYPTED);
+	}
+
+	PG_RETURN_BOOL(false);
+}

contrib/pg_tde/t/expected/basic.out

@@ -20,7 +20,7 @@ CREATE EXTENSION pg_tde;
 SELECT extname, extversion FROM pg_extension WHERE extname = 'pg_tde';
  extname | extversion 
 ---------+------------
- pg_tde  | 1.0
+ pg_tde  | 2.0
 (1 row)
 
 CREATE TABLE test_enc (id SERIAL, k INTEGER, PRIMARY KEY (id)) USING tde_heap;

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 
 ---------------------------------------------
@@ -53,6 +59,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;
@@ -62,6 +74,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;
@@ -71,6 +89,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 
@@ -79,6 +103,12 @@ SHOW pg_tde.wal_encrypt;
 (1 row)
 
 INSERT INTO test_wal (k) VALUES (7), (8);
+SELECT pg_tde_is_wal_record_encrypted('0/15883E8'::pg_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,9 @@
 	'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 +63,9 @@
 
 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 +75,9 @@
 
 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 +87,19 @@
 
 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);");