Skip to content

Commit c8e5310

Browse files
artemgavrilovartemgavrilov
committed
PG-1213 Add function that returns WAL encryption ragnes
Add user function that shows LSN ranges for encrypted records in WAL.
1 parent 097235c commit c8e5310

File tree

6 files changed

+181
-1
lines changed

6 files changed

+181
-1
lines changed

‎contrib/pg_tde/documentation/docs/functions.md

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -432,3 +432,11 @@ SELECT pg_tde_is_wal_record_encrypted(
432432
'optional_timeline_id'
433433
);
434434
```
435+
436+
### pg_tde_get_wal_encryption_ranges
437+
438+
This function returns the ranges of WAL records that are encrypted. It returns a set of records with the start and end LSNs and TLI for each range.
439+
440+
```sql
441+
SELECT * FROM pg_tde_get_wal_encryption_ranges();
442+
```

‎contrib/pg_tde/meson.build

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -129,6 +129,7 @@ tap_tests = [
129129
't/unlogged_tables.pl',
130130
't/wal_archiving.pl',
131131
't/wal_encrypt.pl',
132+
't/wal_encryption_ranges.pl',
132133
't/wal_key_tli.pl',
133134
't/2pc_replication.pl',
134135
't/stream_rep.pl',

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

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,17 @@
1+
-- Function to check if a WAL record is encrypted
12
CREATE FUNCTION pg_tde_is_wal_record_encrypted(lsn pg_lsn, tli integer DEFAULT 0)
23
RETURNS BOOLEAN
34
LANGUAGE C
45
AS 'MODULE_PATHNAME';
56
REVOKE ALL ON FUNCTION pg_tde_is_wal_record_encrypted(pg_lsn, integer) FROM PUBLIC;
7+
8+
-- Function to get WAL encryption ranges
9+
CREATE FUNCTION pg_tde_get_wal_encryption_ranges
10+
(OUT start_tli integer,
11+
OUT start_lsn pg_lsn,
12+
OUT end_tli integer,
13+
OUT end_lsn pg_lsn)
14+
RETURNS SETOF RECORD
15+
LANGUAGE C
16+
AS 'MODULE_PATHNAME';
17+
REVOKE ALL ON FUNCTION pg_tde_get_wal_encryption_ranges() FROM PUBLIC;

‎contrib/pg_tde/src/pg_tde.c

Lines changed: 64 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,8 @@
3535

3636
PG_MODULE_MAGIC;
3737

38+
#define PG_TDE_LIST_WAL_KEYS_RANGES_COLS 4
39+
3840
static void pg_tde_init_data_dir(void);
3941

4042
static shmem_startup_hook_type prev_shmem_startup_hook = NULL;
@@ -44,6 +46,7 @@ PG_FUNCTION_INFO_V1(pg_tde_extension_initialize);
4446
PG_FUNCTION_INFO_V1(pg_tde_version);
4547
PG_FUNCTION_INFO_V1(pg_tdeam_handler);
4648
PG_FUNCTION_INFO_V1(pg_tde_is_wal_record_encrypted);
49+
PG_FUNCTION_INFO_V1(pg_tde_get_wal_encryption_ranges);
4750

4851
static void
4952
tde_shmem_request(void)
@@ -205,3 +208,64 @@ pg_tde_is_wal_record_encrypted(PG_FUNCTION_ARGS)
205208

206209
PG_RETURN_BOOL(false);
207210
}
211+
212+
/*
213+
* Returns WAL encryption ranges. WAL records within the LSN range are encrypted.
214+
*/
215+
Datum
216+
pg_tde_get_wal_encryption_ranges(PG_FUNCTION_ARGS)
217+
{
218+
Tuplestorestate *tupstore;
219+
TupleDesc tupdesc;
220+
ReturnSetInfo *rsinfo = (ReturnSetInfo *) fcinfo->resultinfo;
221+
MemoryContext per_query_ctx;
222+
MemoryContext oldcontext;
223+
WALKeyCacheRec *keys;
224+
WalLocation loc = {.tli = 0,.lsn = 0};
225+
226+
/* check to see if caller supports us returning a tuplestore */
227+
if (rsinfo == NULL || !IsA(rsinfo, ReturnSetInfo))
228+
ereport(ERROR,
229+
errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
230+
errmsg("set-valued function called in context that cannot accept a set"));
231+
if (!(rsinfo->allowedModes & SFRM_Materialize))
232+
ereport(ERROR,
233+
errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
234+
errmsg("materialize mode required, but it is not allowed in this context"));
235+
236+
/* Switch into long-lived context to construct returned data structures */
237+
per_query_ctx = rsinfo->econtext->ecxt_per_query_memory;
238+
oldcontext = MemoryContextSwitchTo(per_query_ctx);
239+
240+
/* Build a tuple descriptor for our result type */
241+
if (get_call_result_type(fcinfo, NULL, &tupdesc) != TYPEFUNC_COMPOSITE)
242+
elog(ERROR, "return type must be a row type");
243+
244+
tupstore = tuplestore_begin_heap(true, false, work_mem);
245+
rsinfo->returnMode = SFRM_Materialize;
246+
rsinfo->setResult = tupstore;
247+
rsinfo->setDesc = tupdesc;
248+
249+
MemoryContextSwitchTo(oldcontext);
250+
251+
keys = pg_tde_fetch_wal_keys(loc);
252+
253+
for (WALKeyCacheRec *curr_key = keys; curr_key != NULL; curr_key = curr_key->next)
254+
{
255+
Datum values[PG_TDE_LIST_WAL_KEYS_RANGES_COLS] = {0};
256+
bool nulls[PG_TDE_LIST_WAL_KEYS_RANGES_COLS] = {0};
257+
int i = 0;
258+
259+
if (curr_key->key.type != WAL_KEY_TYPE_ENCRYPTED)
260+
continue;
261+
262+
values[i++] = Int64GetDatum(curr_key->start.tli);
263+
values[i++] = Int64GetDatum(curr_key->start.lsn);
264+
values[i++] = Int64GetDatum(curr_key->end.tli);
265+
values[i++] = Int64GetDatum(curr_key->end.lsn);
266+
267+
tuplestore_putvalues(tupstore, tupdesc, values, nulls);
268+
}
269+
270+
return (Datum) 0;
271+
}

‎contrib/pg_tde/t/expected/wal_encrypt.out

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -103,7 +103,7 @@ SHOW pg_tde.wal_encrypt;
103103
(1 row)
104104

105105
INSERT INTO test_wal (k) VALUES (7), (8);
106-
SELECT pg_tde_is_wal_record_encrypted('0/15883E8'::pg_lsn);
106+
SELECT pg_tde_is_wal_record_encrypted(pg_current_wal_lsn());
107107
pg_tde_is_wal_record_encrypted
108108
--------------------------------
109109
t

‎contrib/pg_tde/t/wal_encryption_ranges.pl

Lines changed: 95 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,95 @@
1+
#!/usr/bin/perl
2+
3+
use strict;
4+
use warnings;
5+
use File::Basename;
6+
use Test::More;
7+
use lib 't';
8+
use pgtde;
9+
10+
PGTDE::setup_files_dir(basename($0));
11+
12+
unlink('/tmp/wal_encryption_ranges.per');
13+
14+
my $psql_out = '';
15+
16+
my $node = PostgreSQL::Test::Cluster->new('main');
17+
$node->init;
18+
$node->append_conf('postgresql.conf', "shared_preload_libraries = 'pg_tde'");
19+
$node->append_conf('postgresql.conf', "wal_level = 'logical'");
20+
$node->start;
21+
22+
# Create and configure pg_tde extension
23+
$node->psql('postgres', "CREATE EXTENSION pg_tde;");
24+
25+
$node->psql('postgres',
26+
"SELECT pg_tde_add_global_key_provider_file('file-keyring', '/tmp/wal_encryption_ranges.per');"
27+
);
28+
29+
$node->psql('postgres',
30+
"SELECT pg_tde_create_key_using_global_key_provider('server-key', 'file-keyring');"
31+
);
32+
$node->psql('postgres',
33+
"SELECT pg_tde_set_server_key_using_global_key_provider('server-key', 'file-keyring');"
34+
);
35+
36+
$node->psql('postgres', 'SELECT pg_tde_verify_server_key();');
37+
38+
# Create test table and enable WAL encryption
39+
$node->psql('postgres',
40+
'CREATE TABLE test_wal (id SERIAL, k INTEGER, PRIMARY KEY (id));');
41+
42+
$node->psql('postgres', 'ALTER SYSTEM SET pg_tde.wal_encrypt = on;');
43+
44+
$node->restart;
45+
46+
# Insert some data to generate WAL records and check that WAL records are encrypted
47+
$node->psql('postgres', 'INSERT INTO test_wal (k) VALUES (1), (2);');
48+
49+
my $enc_lsn = $node->safe_psql('postgres', "SELECT pg_current_wal_lsn();");
50+
$node->psql(
51+
'postgres',
52+
"SELECT pg_tde_is_wal_record_encrypted('$enc_lsn'::pg_lsn);",
53+
stdout => \$psql_out);
54+
is($psql_out, 't', "Check that WAL record is encrypted");
55+
56+
# Force PG to switch to a new WAL segment to avoid situation when we decrypt
57+
# non-full WAL page on WAL encryption off.
58+
$node->psql('postgres', 'SELECT pg_switch_wal();');
59+
$node->psql('postgres', 'ALTER SYSTEM SET pg_tde.wal_encrypt = off;');
60+
61+
$node->restart;
62+
63+
# Insert more data to generate WAL records and check that they are not encrypted
64+
$node->psql('postgres', 'INSERT INTO test_wal (k) VALUES (3), (4);');
65+
66+
my $dec_lsn = $node->safe_psql('postgres', "SELECT pg_current_wal_lsn();");
67+
$node->psql(
68+
'postgres',
69+
"SELECT pg_tde_is_wal_record_encrypted('$dec_lsn'::pg_lsn);",
70+
stdout => \$psql_out);
71+
is($psql_out, 'f', "Check that WAL record is not encrypted");
72+
73+
# Check that previously encrypted record is still encrypted
74+
$node->psql(
75+
'postgres',
76+
"SELECT pg_tde_is_wal_record_encrypted('$enc_lsn'::pg_lsn);",
77+
stdout => \$psql_out);
78+
is($psql_out, 't', "Check that WAL record is still encrypted");
79+
80+
# Check that WAL records that we recorded before match the encryption ranges. We use
81+
# relative comparisons to avoid issues with LSN stability across different runs of the test.
82+
my $ranges_count = $node->safe_psql(
83+
'postgres', "SELECT count(*) FROM pg_tde_get_wal_encryption_ranges()
84+
WHERE start_lsn <= '$enc_lsn'::pg_lsn
85+
AND end_lsn > '$enc_lsn'::pg_lsn
86+
AND end_lsn <= '$dec_lsn'::pg_lsn
87+
AND start_tli = 1
88+
AND end_tli = 1;");
89+
is($ranges_count, 1,
90+
"Check that WAL records correspond to expected encryption range");
91+
92+
$node->psql('postgres', 'DROP EXTENSION pg_tde;');
93+
$node->stop;
94+
95+
done_testing();

0 commit comments

Comments
 (0)