Skip to content

Move SMGR code out of pg_tde_tdemap.c #361

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.

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

Merged
merged 3 commits into from
May 28, 2025
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
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
309 changes: 69 additions & 240 deletions contrib/pg_tde/src/access/pg_tde_tdemap.c

Large diffs are not rendered by default.

3 changes: 2 additions & 1 deletion contrib/pg_tde/src/access/pg_tde_xlog.c
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,7 @@

#include "access/pg_tde_xlog.h"
#include "encryption/enc_tde.h"
#include "smgr/pg_tde_smgr.h"

static void tdeheap_rmgr_redo(XLogReaderState *record);
static void tdeheap_rmgr_desc(StringInfo buf, XLogReaderState *record);
Expand Down Expand Up @@ -52,7 +53,7 @@ tdeheap_rmgr_redo(XLogReaderState *record)
{
XLogRelKey *xlrec = (XLogRelKey *) XLogRecGetData(record);

pg_tde_create_smgr_key_perm_redo(&xlrec->rlocator);
tde_smgr_create_key_redo(&xlrec->rlocator);
}
else if (info == XLOG_TDE_ADD_PRINCIPAL_KEY)
{
Expand Down
6 changes: 2 additions & 4 deletions contrib/pg_tde/src/common/pg_tde_utils.c
Original file line number Diff line number Diff line change
Expand Up @@ -16,8 +16,7 @@

#ifndef FRONTEND
#include "fmgr.h"
#include "catalog/pg_class.h"
#include "access/pg_tde_tdemap.h"
#include "smgr/pg_tde_smgr.h"
#include "access/relation.h"
#include "utils/rel.h"

Expand All @@ -28,7 +27,6 @@ pg_tde_is_encrypted(PG_FUNCTION_ARGS)
Oid relationOid = PG_GETARG_OID(0);
LOCKMODE lockmode = AccessShareLock;
Relation rel = relation_open(relationOid, lockmode);
RelFileLocatorBackend rlocator = {.locator = rel->rd_locator,.backend = rel->rd_backend};
bool result;

if (!RELKIND_HAS_STORAGE(rel->rd_rel->relkind))
Expand All @@ -42,7 +40,7 @@ pg_tde_is_encrypted(PG_FUNCTION_ARGS)
errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
errmsg("we cannot check if temporary relations from other backends are encrypted"));

result = IsSMGRRelationEncrypted(rlocator);
result = tde_smgr_rel_is_encrypted(RelationGetSmgr(rel));

relation_close(rel, lockmode);

Expand Down
25 changes: 25 additions & 0 deletions contrib/pg_tde/src/encryption/enc_tde.c
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,13 @@
#include "encryption/enc_aes.h"
#include "storage/bufmgr.h"

#ifdef FRONTEND
#include "pg_tde_fe.h"
#endif

#include <openssl/rand.h>
#include <openssl/err.h>

#define AES_BLOCK_SIZE 16
#define NUM_AES_BLOCKS_IN_BATCH 200
#define DATA_BYTES_PER_AES_BATCH (NUM_AES_BLOCKS_IN_BATCH * AES_BLOCK_SIZE)
Expand All @@ -23,6 +30,24 @@ iv_prefix_debug(const char *iv_prefix, char *out_hex)
}
#endif

void
pg_tde_generate_internal_key(InternalKey *int_key, TDEMapEntryType entry_type)
{
int_key->type = entry_type;
int_key->start_lsn = InvalidXLogRecPtr;

if (!RAND_bytes(int_key->key, INTERNAL_KEY_LEN))
ereport(ERROR,
errcode(ERRCODE_INTERNAL_ERROR),
errmsg("could not generate internal key: %s",
ERR_error_string(ERR_get_error(), NULL)));
if (!RAND_bytes(int_key->base_iv, INTERNAL_KEY_IV_LEN))
ereport(ERROR,
errcode(ERRCODE_INTERNAL_ERROR),
errmsg("could not generate IV: %s",
ERR_error_string(ERR_get_error(), NULL)));
}

/*
* Encrypts/decrypts `data` with a given `key`. The result is written to `out`.
*
Expand Down
9 changes: 4 additions & 5 deletions contrib/pg_tde/src/include/access/pg_tde_tdemap.h
Original file line number Diff line number Diff line change
Expand Up @@ -83,8 +83,6 @@ extern WALKeyCacheRec *pg_tde_fetch_wal_keys(XLogRecPtr start_lsn);
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_smgr_key_perm_redo(const RelFileLocator *newrlocator);
extern void pg_tde_create_wal_key(InternalKey *rel_key_data, const RelFileLocator *newrlocator, TDEMapEntryType flags);

#define PG_TDE_MAP_FILENAME "%d_keys"
Expand All @@ -95,9 +93,10 @@ pg_tde_set_db_file_path(Oid dbOid, char *path)
join_path_components(path, pg_tde_get_data_dir(), psprintf(PG_TDE_MAP_FILENAME, dbOid));
}

extern bool IsSMGRRelationEncrypted(RelFileLocatorBackend rel);
extern InternalKey *GetSMGRRelationKey(RelFileLocatorBackend rel);
extern void DeleteSMGRRelationKey(RelFileLocatorBackend rel);
extern void pg_tde_save_smgr_key(RelFileLocator rel, const InternalKey *key, bool write_xlog);
extern bool pg_tde_has_smgr_key(RelFileLocator rel);
extern InternalKey *pg_tde_get_smgr_key(RelFileLocator rel);
extern void pg_tde_free_key_map_entry(RelFileLocator rel);

extern int pg_tde_count_relations(Oid dbOid);

Expand Down
1 change: 1 addition & 0 deletions contrib/pg_tde/src/include/encryption/enc_tde.h
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@

#include "access/pg_tde_tdemap.h"

extern void pg_tde_generate_internal_key(InternalKey *int_key, TDEMapEntryType entry_type);
extern void pg_tde_stream_crypt(const char *iv_prefix, uint32 start_offset, const char *data, uint32 data_len, char *out, InternalKey *key, void **ctxPtr);

#endif /* ENC_TDE_H */
5 changes: 5 additions & 0 deletions contrib/pg_tde/src/include/smgr/pg_tde_smgr.h
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,11 @@
#ifndef PG_TDE_SMGR_H
#define PG_TDE_SMGR_H

#include "storage/relfilelocator.h"
#include "storage/smgr.h"

extern void RegisterStorageMgr(void);
extern void tde_smgr_create_key_redo(const RelFileLocator *rlocator);
extern bool tde_smgr_rel_is_encrypted(SMgrRelation reln);

#endif /* PG_TDE_SMGR_H */
160 changes: 145 additions & 15 deletions contrib/pg_tde/src/smgr/pg_tde_smgr.c
Original file line number Diff line number Diff line change
@@ -1,10 +1,13 @@
#include "smgr/pg_tde_smgr.h"
#include "postgres.h"

#include "smgr/pg_tde_smgr.h"
#include "storage/smgr.h"
#include "storage/md.h"
#include "catalog/catalog.h"
#include "encryption/enc_aes.h"
#include "encryption/enc_tde.h"
#include "access/pg_tde_tdemap.h"
#include "utils/hsearch.h"
#include "pg_tde_event_capture.h"

typedef enum TDEMgrRelationEncryptionStatus
Expand Down Expand Up @@ -43,26 +46,80 @@ typedef struct TDESMgrRelation
InternalKey relKey;
} TDESMgrRelation;

typedef struct
{
RelFileLocator rel;
InternalKey key;
} TempRelKeyEntry;

#define INIT_TEMP_RELS 16

/*
* Each backend has a hashtable that stores the keys for all temproary tables.
*/
static HTAB *TempRelKeys = NULL;

static SMgrId OurSMgrId = MaxSMgrId;

static void tde_smgr_save_temp_key(const RelFileLocator *newrlocator, const InternalKey *key);
static InternalKey *tde_smgr_get_temp_key(const RelFileLocator *rel);
static bool tde_smgr_has_temp_key(const RelFileLocator *rel);
static void tde_smgr_remove_temp_key(const RelFileLocator *rel);
static void CalcBlockIv(ForkNumber forknum, BlockNumber bn, const unsigned char *base_iv, unsigned char *iv);

static InternalKey *
tde_smgr_create_key(const RelFileLocatorBackend *smgr_rlocator)
{
InternalKey *key = palloc_object(InternalKey);

pg_tde_generate_internal_key(key, TDE_KEY_TYPE_SMGR);

if (RelFileLocatorBackendIsTemp(*smgr_rlocator))
tde_smgr_save_temp_key(&smgr_rlocator->locator, key);
else
pg_tde_save_smgr_key(smgr_rlocator->locator, key, true);

return key;
}

void
tde_smgr_create_key_redo(const RelFileLocator *rlocator)
{
InternalKey key;

if (pg_tde_has_smgr_key(*rlocator))
return;

pg_tde_generate_internal_key(&key, TDE_KEY_TYPE_SMGR);

pg_tde_save_smgr_key(*rlocator, &key, false);
}

static bool
tde_smgr_is_encrypted(const RelFileLocatorBackend *smgr_rlocator)
{
/* Do not try to encrypt/decrypt catalog tables */
if (IsCatalogRelationOid(smgr_rlocator->locator.relNumber))
return false;

return IsSMGRRelationEncrypted(*smgr_rlocator);
if (RelFileLocatorBackendIsTemp(*smgr_rlocator))
return tde_smgr_has_temp_key(&smgr_rlocator->locator);
else
return pg_tde_has_smgr_key(smgr_rlocator->locator);
}

static InternalKey *
tde_smgr_get_key(const RelFileLocatorBackend *smgr_rlocator)
{
/* Do not try to encrypt/decrypt catalog tables */
if (IsCatalogRelationOid(smgr_rlocator->locator.relNumber))
return NULL;
if (RelFileLocatorBackendIsTemp(*smgr_rlocator))
return tde_smgr_get_temp_key(&smgr_rlocator->locator);
else
return pg_tde_get_smgr_key(smgr_rlocator->locator);
}

return GetSMGRRelationKey(*smgr_rlocator);
static void
tde_smgr_remove_key(const RelFileLocatorBackend *smgr_rlocator)
{
if (RelFileLocatorBackendIsTemp(*smgr_rlocator))
tde_smgr_remove_temp_key(&smgr_rlocator->locator);
else
pg_tde_free_key_map_entry(smgr_rlocator->locator);
}

static bool
Expand All @@ -86,13 +143,25 @@ tde_smgr_should_encrypt(const RelFileLocatorBackend *smgr_rlocator, RelFileLocat
.backend = smgr_rlocator->backend,
};

return IsSMGRRelationEncrypted(old_smgr_locator);
return tde_smgr_is_encrypted(&old_smgr_locator);
}
}

return false;
}

bool
tde_smgr_rel_is_encrypted(SMgrRelation reln)
{
TDESMgrRelation *tdereln = (TDESMgrRelation *) reln;

if (reln->smgr_which != OurSMgrId)
Copy link
Collaborator

Choose a reason for hiding this comment

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

This is the first time we actually use our SMgrId afaics. Should we use it more? Why haven't we needed it before?

Copy link
Collaborator Author

Choose a reason for hiding this comment

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

Because the other calls go through the smgr*() functions while this one does not.

return false;

return tdereln->encryption_status == RELATION_KEY_AVAILABLE ||
tdereln->encryption_status == RELATION_KEY_NOT_AVAILABLE;
}

static void
tde_mdwritev(SMgrRelation reln, ForkNumber forknum, BlockNumber blocknum,
const void **buffers, BlockNumber nblocks, bool skipFsync)
Expand Down Expand Up @@ -159,8 +228,8 @@ tde_mdunlink(RelFileLocatorBackend rlocator, ForkNumber forknum, bool isRedo)
*/
if (forknum == MAIN_FORKNUM || forknum == InvalidForkNumber)
{
if (IsSMGRRelationEncrypted(rlocator))
DeleteSMGRRelationKey(rlocator);
if (tde_smgr_is_encrypted(&rlocator))
tde_smgr_remove_key(&rlocator);
}
}

Expand Down Expand Up @@ -282,7 +351,7 @@ tde_mdcreate(RelFileLocator relold, SMgrRelation reln, ForkNumber forknum, bool
InternalKey *key = tde_smgr_get_key(&reln->smgr_rlocator);

if (!isRedo && !key && tde_smgr_should_encrypt(&reln->smgr_rlocator, &relold))
key = pg_tde_create_smgr_key(&reln->smgr_rlocator);
key = tde_smgr_create_key(&reln->smgr_rlocator);

if (key)
{
Expand Down Expand Up @@ -347,7 +416,68 @@ RegisterStorageMgr(void)
{
if (storage_manager_id != MdSMgrId)
elog(FATAL, "Another storage manager was loaded before pg_tde. Multiple storage managers is unsupported.");
storage_manager_id = smgr_register(&tde_smgr, sizeof(TDESMgrRelation));
OurSMgrId = smgr_register(&tde_smgr, sizeof(TDESMgrRelation));
storage_manager_id = OurSMgrId;
}

static void
tde_smgr_save_temp_key(const RelFileLocator *newrlocator, const InternalKey *key)
{
TempRelKeyEntry *entry;
bool found;

if (TempRelKeys == NULL)
{
HASHCTL ctl;

ctl.keysize = sizeof(RelFileLocator);
ctl.entrysize = sizeof(TempRelKeyEntry);
TempRelKeys = hash_create("pg_tde temporary relation keys",
INIT_TEMP_RELS,
&ctl,
HASH_ELEM | HASH_BLOBS);
}

entry = (TempRelKeyEntry *) hash_search(TempRelKeys,
newrlocator,
HASH_ENTER, &found);
Assert(!found);

entry->key = *key;
}

static InternalKey *
tde_smgr_get_temp_key(const RelFileLocator *rel)
{
TempRelKeyEntry *entry;

if (TempRelKeys == NULL)
return NULL;

entry = hash_search(TempRelKeys, rel, HASH_FIND, NULL);

if (entry)
{
InternalKey *key = palloc_object(InternalKey);

*key = entry->key;
return key;
}

return NULL;
}

static bool
tde_smgr_has_temp_key(const RelFileLocator *rel)
{
return TempRelKeys && hash_search(TempRelKeys, rel, HASH_FIND, NULL);
}

static void
tde_smgr_remove_temp_key(const RelFileLocator *rel)
{
Assert(TempRelKeys);
hash_search(TempRelKeys, rel, HASH_REMOVE, NULL);
}

/*
Expand Down
5 changes: 2 additions & 3 deletions src/bin/pg_checksums/pg_checksums.c
Original file line number Diff line number Diff line change
Expand Up @@ -138,10 +138,9 @@ pg_tde_init(const char *datadir)
static bool
is_pg_tde_encypted(Oid spcOid, Oid dbOid, RelFileNumber relNumber)
{
RelFileLocator locator = {.spcOid = spcOid, .dbOid = dbOid,.relNumber = relNumber};
RelFileLocatorBackend rlocator = {.locator = locator,.backend = INVALID_PROC_NUMBER};
RelFileLocator locator = {.spcOid = spcOid,.dbOid = dbOid,.relNumber = relNumber};

return IsSMGRRelationEncrypted(rlocator);
return pg_tde_has_smgr_key(locator);
}
#endif

Expand Down