Skip to content

Ensure decrypted internal keys never swapped #185

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
wants to merge 2 commits into
base: TDE_REL_17_STABLE
Choose a base branch
from
Draft

Ensure decrypted internal keys never swapped #185

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
6a8bd08
Ensure decrypted internal keys never swapped
dAdAbird Mar 31, 2025
23cbf8b
Fetch smgr key from TDE cache instead of TDESMgrRelation
dAdAbird Apr 10, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
281 changes: 136 additions & 145 deletions contrib/pg_tde/src/access/pg_tde_tdemap.c
View file
Open in desktop

Large diffs are not rendered by default.

61 changes: 37 additions & 24 deletions contrib/pg_tde/src/access/pg_tde_xlog_encrypt.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -70,12 +70,7 @@ typedef struct EncryptionStateData

static EncryptionStateData *EncryptionState = NULL;

/* TODO: can be swapped out to the disk */
static InternalKey EncryptionKey =
{
.rel_type = MAP_ENTRY_EMPTY,
.start_lsn = InvalidXLogRecPtr,
};
static InternalKey *EncryptionKey = NULL;
static void *EncryptionCryptCtx = NULL;

static int XLOGChooseNumBuffers(void);
Expand Down Expand Up @@ -165,48 +160,67 @@ TDEXLogWriteEncryptedPages(int fd, const void *buf, size_t count, off_t offset,
TimeLineID tli, XLogSegNo segno)
{
char iv_prefix[16];
InternalKey *key = &EncryptionKey;
char *enc_buff = EncryptionState->segBuf;

Assert(count <= TDEXLogEncryptBuffSize());

#ifdef TDE_XLOG_DEBUG
elog(DEBUG1, "write encrypted WAL, size: %lu, offset: %ld [%lX], seg: %X/%X, key_start_lsn: %X/%X",
count, offset, offset, LSN_FORMAT_ARGS(segno), LSN_FORMAT_ARGS(key->start_lsn));
count, offset, offset, LSN_FORMAT_ARGS(segno), LSN_FORMAT_ARGS(EncryptionKey->start_lsn));
#endif

CalcXLogPageIVPrefix(tli, segno, key->base_iv, iv_prefix);
CalcXLogPageIVPrefix(tli, segno, EncryptionKey->base_iv, iv_prefix);
PG_TDE_ENCRYPT_DATA(iv_prefix, offset,
(char *) buf, count,
enc_buff, key, &EncryptionCryptCtx);
enc_buff, EncryptionKey, &EncryptionCryptCtx);

return pg_pwrite(fd, enc_buff, count, offset);
}

#endif /* !FRONTEND */

void
TDEXLogSmgrInit(void)
static void
tde_init_wal_encryption_key(void)
{
#ifndef FRONTEND
/* TODO: move to the separate func, it's not an SMGR init */
InternalKey *key = pg_tde_read_last_wal_key();

/* TDOO: clean-up this mess */
if ((!key && EncryptXLog) || (key &&
((key->rel_type & TDE_KEY_TYPE_WAL_ENCRYPTED && !EncryptXLog) ||
(key->rel_type & TDE_KEY_TYPE_WAL_UNENCRYPTED && EncryptXLog))))
{
pg_tde_create_wal_key(&EncryptionKey, &GLOBAL_SPACE_RLOCATOR(XLOG_TDE_OID),
(EncryptXLog ? TDE_KEY_TYPE_WAL_ENCRYPTED : TDE_KEY_TYPE_WAL_UNENCRYPTED));
key = pg_tde_create_wal_key(&GLOBAL_SPACE_RLOCATOR(XLOG_TDE_OID),
(EncryptXLog ? TDE_KEY_TYPE_WAL_ENCRYPTED : TDE_KEY_TYPE_WAL_UNENCRYPTED));
}
else if (key)
{
EncryptionKey = *key;
pfree(key);
pg_atomic_write_u64(&EncryptionState->enc_key_lsn, EncryptionKey.start_lsn);
pg_atomic_write_u64(&EncryptionState->enc_key_lsn, key->start_lsn);
}


/*
* The relation cache might be moved to the new address space when growing
* later on. It happens when there are a lot of tde_heap relations during
* the recovery. So copy the key to the own locked mem. We could rewrite
* the EncryptionKey pointer after the grow, but it'd be more complicated
* and error-prone. We can sacrifice a couple of Kb.
*
* TODO: genralize and use `tde_alloc_safe_mem` instead?
*/
if (key)
{
long pageSize = sysconf(_SC_PAGESIZE);

EncryptionKey = (InternalKey *) MemoryContextAllocAligned(TopMemoryContext, pageSize, pageSize, MCXT_ALLOC_ZERO);
memcpy(EncryptionKey, key, sizeof(InternalKey));
}
}

#endif /* !FRONTEND */

void
TDEXLogInit(void)
{
#ifndef FRONTEND
tde_init_wal_encryption_key();
pg_tde_set_db_file_path(GLOBAL_SPACE_RLOCATOR(XLOG_TDE_OID).dbOid, EncryptionState->db_map_path);

#endif
Expand All @@ -224,15 +238,14 @@ tdeheap_xlog_seg_write(int fd, const void *buf, size_t count, off_t offset,
*
* This func called with WALWriteLock held, so no need in any extra sync.
*/
if (EncryptionKey.rel_type & TDE_KEY_TYPE_GLOBAL &&
pg_atomic_read_u64(&EncryptionState->enc_key_lsn) == 0)
if (EncryptionKey && pg_atomic_read_u64(&EncryptionState->enc_key_lsn) == 0)
{
XLogRecPtr lsn;

XLogSegNoOffsetToRecPtr(segno, offset, wal_segment_size, lsn);

pg_tde_wal_last_key_set_lsn(lsn, EncryptionState->db_map_path);
EncryptionKey.start_lsn = lsn;
EncryptionKey->start_lsn = lsn;
pg_atomic_write_u64(&EncryptionState->enc_key_lsn, lsn);
}

Expand Down
2 changes: 1 addition & 1 deletion contrib/pg_tde/src/include/access/pg_tde_tdemap.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -100,7 +100,7 @@ extern WALKeyCacheRec *pg_tde_get_wal_cache_keys(void);
extern void pg_tde_wal_last_key_set_lsn(XLogRecPtr lsn, const char *keyfile_path);

extern InternalKey *pg_tde_create_smgr_key(const RelFileLocatorBackend *newrlocator);
extern void pg_tde_create_wal_key(InternalKey *rel_key_data, const RelFileLocator *newrlocator, uint32 flags);
extern InternalKey *pg_tde_create_wal_key(const RelFileLocator *newrlocator, uint32 flags);
extern void pg_tde_free_key_map_entry(const RelFileLocator *rlocator, off_t offset);
extern void pg_tde_write_key_map_entry_redo(const TDEMapEntry *write_map_entry, TDESignedPrincipalKeyInfo *signed_key_info);

Expand Down
2 changes: 1 addition & 1 deletion contrib/pg_tde/src/include/access/pg_tde_xlog_encrypt.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,6 @@

extern Size TDEXLogEncryptStateSize(void);
extern void TDEXLogShmemInit(void);
extern void TDEXLogSmgrInit(void);
extern void TDEXLogInit(void);

#endif /* PG_TDE_XLOGENCRYPT_H */
2 changes: 1 addition & 1 deletion contrib/pg_tde/src/pg_tde.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -92,7 +92,7 @@ tde_shmem_startup(void)
AesInit();

TDEXLogShmemInit();
TDEXLogSmgrInit();
TDEXLogInit();
}

void
Expand Down
28 changes: 6 additions & 22 deletions contrib/pg_tde/src/smgr/pg_tde_smgr.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,6 @@ typedef struct TDESMgrRelationData
struct _MdfdVec *md_seg_fds[MAX_FORKNUM + 1];

bool encrypted_relation;
InternalKey relKey;
} TDESMgrRelationData;

typedef TDESMgrRelationData *TDESMgrRelation;
Expand Down Expand Up @@ -81,7 +80,7 @@ tde_mdwritev(SMgrRelation reln, ForkNumber forknum, BlockNumber blocknum,
const void **buffers, BlockNumber nblocks, bool skipFsync)
{
TDESMgrRelation tdereln = (TDESMgrRelation) reln;
InternalKey *int_key = &tdereln->relKey;
InternalKey *int_key = GetSMGRRelationKey(tdereln->reln.smgr_rlocator);

if (!tdereln->encrypted_relation)
{
Expand Down Expand Up @@ -120,7 +119,7 @@ tde_mdextend(SMgrRelation reln, ForkNumber forknum, BlockNumber blocknum,
const void *buffer, bool skipFsync)
{
TDESMgrRelation tdereln = (TDESMgrRelation) reln;
InternalKey *int_key = &tdereln->relKey;
InternalKey *int_key = GetSMGRRelationKey(tdereln->reln.smgr_rlocator);

if (!tdereln->encrypted_relation)
{
Expand Down Expand Up @@ -149,7 +148,7 @@ tde_mdreadv(SMgrRelation reln, ForkNumber forknum, BlockNumber blocknum,
void **buffers, BlockNumber nblocks)
{
TDESMgrRelation tdereln = (TDESMgrRelation) reln;
InternalKey *int_key = &tdereln->relKey;
InternalKey *int_key = GetSMGRRelationKey(tdereln->reln.smgr_rlocator);

mdreadv(reln, forknum, blocknum, buffers, nblocks);

Expand Down Expand Up @@ -233,15 +232,7 @@ tde_mdcreate(RelFileLocator relold, SMgrRelation reln, ForkNumber forknum, bool
*/
key = tde_smgr_get_key(reln, event->alterAccessMethodMode ? NULL : &relold, true);

if (key)
{
tdereln->encrypted_relation = true;
tdereln->relKey = *key;
}
else
{
tdereln->encrypted_relation = false;
}
tdereln->encrypted_relation = key ? true : false;
}
}

Expand All @@ -254,15 +245,8 @@ tde_mdopen(SMgrRelation reln)
TDESMgrRelation tdereln = (TDESMgrRelation) reln;
InternalKey *key = tde_smgr_get_key(reln, NULL, false);

if (key)
{
tdereln->encrypted_relation = true;
tdereln->relKey = *key;
}
else
{
tdereln->encrypted_relation = false;
}
tdereln->encrypted_relation = key ? true : false;

mdopen(reln);
}

Expand Down
21 changes: 20 additions & 1 deletion contrib/pg_tde/t/009_wal_encrypt.pl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -91,9 +91,28 @@

$stdout = $node->safe_psql('postgres', "SELECT pg_drop_replication_slot('tde_slot');", extra_params => ['-a']);
PGTDE::append_to_file($stdout);
PGTDE::append_to_file("-- server restart with wal encryption and recovery with a lot of tde_heap relations");
$node->safe_psql('postgres',
q{
SELECT pg_tde_add_key_provider_file('file-keyring-010-2','/tmp/pg_tde_test_keyring010_2.per');
SELECT pg_tde_set_principal_key('test-db-principal-key','file-keyring-010-2');

do $$
DECLARE idx integer;
begin
for idx in 0..700 loop
EXECUTE format('CREATE TABLE t%s (c1 int) USING tde_heap', idx);
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Not that it matters in tests but in real production code I would have done.

EXECUTE format('CREATE TABLE %I (c1 int) USING tde_heap', 't' || idx);

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

So should I change it? I took it from the cache_alloc regression test. If to change, that it make sense to fix the regression test as well

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Up to you. I just wanted to share how real production code looks like.

end loop;
end; $$;
});

$rt_value = $node->restart();
ok($rt_value == 1, "Restart Server");

# TODO: add WAL content testing after the wal rework

# DROP EXTENSION
$stdout = $node->safe_psql('postgres', 'DROP EXTENSION pg_tde;', extra_params => ['-a']);
$stdout = $node->safe_psql('postgres', 'DROP EXTENSION pg_tde CASCADE;', extra_params => ['-a']);
PGTDE::append_to_file($stdout);
# Stop the server
$node->stop();
Expand Down
3 changes: 2 additions & 1 deletion contrib/pg_tde/t/expected/009_wal_encrypt.out
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -38,4 +38,5 @@ table public.test_wal: INSERT: id[integer]:6 k[integer]:6
COMMIT 742
SELECT pg_drop_replication_slot('tde_slot');

DROP EXTENSION pg_tde;
-- server restart with wal encryption and recovery with a lot of tde_heap relations
DROP EXTENSION pg_tde CASCADE;
2 changes: 1 addition & 1 deletion src/bin/pg_waldump/pg_waldump.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1156,7 +1156,7 @@ main(int argc, char **argv)
if (kringdir != NULL)
{
pg_tde_fe_init(kringdir);
TDEXLogSmgrInit();
TDEXLogInit();
}
#endif

Expand Down
Loading