From 7921343dcdaafbb0b08a8d8ff82be66a2e1d46d2 Mon Sep 17 00:00:00 2001 From: Muhammad Usama Date: Mon, 26 Aug 2024 20:58:06 +0500 Subject: [PATCH] PG-976: Fix for memory leaks (#264) This commit addresses several memory leaks primarily in the code related to obtaining the relation key and the storage manager (smgr). Additionally, fixes were made to the internal key cache maintained by each backend, as the keys were not properly added to the cache list. --- src/access/pg_tde_tdemap.c | 69 +++++++++++++---------- src/catalog/tde_global_space.c | 1 + src/catalog/tde_principal_key.c | 10 ++++ src/include/access/pg_tde_tdemap.h | 15 +---- src/include/smgr/pg_tde_smgr.h | 15 ++++- src/smgr/pg_tde_smgr.c | 88 ++++++++++++++++-------------- 6 files changed, 113 insertions(+), 85 deletions(-) diff --git a/src/access/pg_tde_tdemap.c b/src/access/pg_tde_tdemap.c index 5185731d..2254c0d4 100644 --- a/src/access/pg_tde_tdemap.c +++ b/src/access/pg_tde_tdemap.c @@ -79,6 +79,23 @@ typedef struct TDEMapFilePath char keydata_path[MAXPGPATH]; } TDEMapFilePath; +/* Relation key cache. + * + * TODO: For now it is just a linked list. Data can only be added w/o any + * ability to remove or change it. Also consider usage of more efficient data + * struct (hash map) in the shared memory(?) - currently allocated in the + * TopMemoryContext of the process. + */ +typedef struct RelKey +{ + Oid rel_id; + RelKeyData key; + struct RelKey *next; +} RelKey; + +/* Head of the key cache (linked list) */ +RelKey *tde_rel_key_map = NULL; + static int pg_tde_open_file_basic(char *tde_filename, int fileFlags, bool ignore_missing); static int pg_tde_file_header_write(char *tde_filename, int fd, TDEPrincipalKeyInfo *principal_key_info, off_t *bytes_written); static int pg_tde_file_header_read(char *tde_filename, int fd, TDEFileHeader *fheader, bool *is_new_file, off_t *bytes_read); @@ -149,12 +166,10 @@ pg_tde_create_key_map_entry(const RelFileLocator *newrlocator) * Add the encyrpted key to the key map data file structure. */ pg_tde_write_key_map_entry(newrlocator, enc_rel_key_data, &principal_key->keyInfo); - + pfree(enc_rel_key_data); return rel_key_data; } -/* Head of the key cache (linked list) */ -RelKey *tde_rel_key_map = NULL; /* * Returns TDE key for a given relation. @@ -166,13 +181,13 @@ GetRelationKey(RelFileLocator rel) { RelKey *curr; RelKeyData *key; - Oid rel_id = rel.relNumber; + for (curr = tde_rel_key_map; curr != NULL; curr = curr->next) { if (curr->rel_id == rel_id) { - return curr->key; + return &curr->key; } } @@ -180,26 +195,30 @@ GetRelationKey(RelFileLocator rel) if (key != NULL) { - pg_tde_put_key_into_map(rel.relNumber, key); + RelKeyData* cached_key = pg_tde_put_key_into_map(rel.relNumber, key); + pfree(key); + return cached_key; } - return key; + return key; /* returning NULL key */ } -void -pg_tde_put_key_into_map(Oid rel_id, RelKeyData *key) { - RelKey *new; - RelKey *prev = NULL; - - new = (RelKey *) MemoryContextAlloc(TopMemoryContext, sizeof(RelKey)); +RelKeyData * +pg_tde_put_key_into_map(Oid rel_id, RelKeyData *key) +{ + RelKey *new = (RelKey *) MemoryContextAlloc(TopMemoryContext, sizeof(RelKey)); new->rel_id = rel_id; - new->key = key; + memcpy(&new->key, key, sizeof(RelKeyData)); new->next = NULL; - if (prev == NULL) + if (tde_rel_key_map == NULL) tde_rel_key_map = new; else - prev->next = new; + { + new->next = tde_rel_key_map; + tde_rel_key_map = new; + } + return &new->key; } const char * @@ -221,20 +240,14 @@ tde_sprint_key(InternalKey *k) RelKeyData * tde_create_rel_key(Oid rel_id, InternalKey *key, TDEPrincipalKeyInfo *principal_key_info) { - RelKeyData *rel_key_data; - - rel_key_data = (RelKeyData *) MemoryContextAlloc(TopMemoryContext, sizeof(RelKeyData)); - - memcpy(&rel_key_data->principal_key_id, &principal_key_info->keyId, sizeof(TDEPrincipalKeyId)); - memcpy(&rel_key_data->internal_key, key, sizeof(InternalKey)); - rel_key_data->internal_key.ctx = NULL; + RelKeyData rel_key_data; + memcpy(&rel_key_data.principal_key_id, &principal_key_info->keyId, sizeof(TDEPrincipalKeyId)); + memcpy(&rel_key_data.internal_key, key, sizeof(InternalKey)); + rel_key_data.internal_key.ctx = NULL; /* Add to the decrypted key to cache */ - pg_tde_put_key_into_map(rel_id, rel_key_data); - - return rel_key_data; + return pg_tde_put_key_into_map(rel_id, &rel_key_data); } - /* * Encrypts a given key and returns the encrypted one. */ @@ -1030,7 +1043,7 @@ pg_tde_get_key_from_file(const RelFileLocator *rlocator) LWLockRelease(lock_files); rel_key_data = tde_decrypt_rel_key(principal_key, enc_rel_key_data, rlocator); - + pfree(enc_rel_key_data); return rel_key_data; } diff --git a/src/catalog/tde_global_space.c b/src/catalog/tde_global_space.c index d45d24a8..62ebeb84 100644 --- a/src/catalog/tde_global_space.c +++ b/src/catalog/tde_global_space.c @@ -181,6 +181,7 @@ init_keys(void) cache_internal_key(rel_key_data, TDE_INTERNAL_XLOG_KEY); pfree(rel_key_data); + pfree(enc_rel_key_data); pfree(mkey); } diff --git a/src/catalog/tde_principal_key.c b/src/catalog/tde_principal_key.c index ccad606e..0ad98c95 100644 --- a/src/catalog/tde_principal_key.c +++ b/src/catalog/tde_principal_key.c @@ -433,9 +433,16 @@ RotatePrincipalKey(TDEPrincipalKey *current_key, const char *new_key_name, const const keyInfo *keyInfo = NULL; GenericKeyring *keyring; bool is_rotated; + MemoryContext keyRotateCtx; + MemoryContext oldCtx; Assert(current_key != NULL); + keyRotateCtx = AllocSetContextCreate(CurrentMemoryContext, + "TDE key rotation temporary context", + ALLOCSET_DEFAULT_SIZES); + oldCtx = MemoryContextSwitchTo(keyRotateCtx); + /* * Let's set everything the same as the older principal key and * update only the required attributes. @@ -483,6 +490,9 @@ RotatePrincipalKey(TDEPrincipalKey *current_key, const char *new_key_name, const push_principal_key_to_cache(&new_principal_key); } + MemoryContextSwitchTo(oldCtx); + MemoryContextDelete(keyRotateCtx); + return is_rotated; } diff --git a/src/include/access/pg_tde_tdemap.h b/src/include/access/pg_tde_tdemap.h index db629b18..0a6b632d 100644 --- a/src/include/access/pg_tde_tdemap.h +++ b/src/include/access/pg_tde_tdemap.h @@ -27,19 +27,6 @@ typedef struct RelKeyData InternalKey internal_key; } RelKeyData; -/* Relation key cache. - * - * TODO: For now it is just a linked list. Data can only be added w/o any - * ability to remove or change it. Also consider usage of more efficient data - * struct (hash map) in the shared memory(?) - currently allocated in the - * TopMemoryContext of the process. - */ -typedef struct RelKey -{ - Oid rel_id; - RelKeyData *key; - struct RelKey *next; -} RelKey; typedef struct XLogRelKey { @@ -69,6 +56,6 @@ extern void pg_tde_set_db_file_paths(const RelFileLocator *rlocator, char *map_p const char * tde_sprint_key(InternalKey *k); -extern void pg_tde_put_key_into_map(Oid rel_id, RelKeyData *key); +extern RelKeyData *pg_tde_put_key_into_map(Oid rel_id, RelKeyData *key); #endif /*PG_TDE_MAP_H*/ diff --git a/src/include/smgr/pg_tde_smgr.h b/src/include/smgr/pg_tde_smgr.h index 359b34db..8c186881 100644 --- a/src/include/smgr/pg_tde_smgr.h +++ b/src/include/smgr/pg_tde_smgr.h @@ -1,4 +1,15 @@ -#pragma once +/*------------------------------------------------------------------------- + * + * pg_tde_smgr.h + * src/include/smgr/pg_tde_smgr.h + * + *------------------------------------------------------------------------- + */ -extern void RegisterStorageMgr(); +#ifndef PG_TDE_SMGR_H +#define PG_TDE_SMGR_H + +extern void RegisterStorageMgr(void); + +#endif /* PG_TDE_SMGR_H */ \ No newline at end of file diff --git a/src/smgr/pg_tde_smgr.c b/src/smgr/pg_tde_smgr.c index ee39a21d..964825d8 100644 --- a/src/smgr/pg_tde_smgr.c +++ b/src/smgr/pg_tde_smgr.c @@ -12,7 +12,7 @@ // TODO: implement proper IV // iv should be based on blocknum + relfile, available in the API -static char iv[16] = {0,}; +static unsigned char iv[16] = {0,}; static RelKeyData* tde_smgr_get_key(SMgrRelation reln) @@ -20,6 +20,8 @@ tde_smgr_get_key(SMgrRelation reln) // TODO: This recursion counter is a dirty hack until the metadata is in the catalog // As otherwise we would call GetPrincipalKey recursively and deadlock static int recursion = 0; + TdeCreateEvent *event; + RelKeyData *rkd; if(IsCatalogRelationOid(reln->smgr_rlocator.locator.relNumber)) { @@ -40,10 +42,10 @@ tde_smgr_get_key(SMgrRelation reln) return NULL; } - TdeCreateEvent* event = GetCurrentTdeCreateEvent(); + event = GetCurrentTdeCreateEvent(); // see if we have a key for the relation, and return if yes - RelKeyData* rkd = GetRelationKey(reln->smgr_rlocator.locator); + rkd = GetRelationKey(reln->smgr_rlocator.locator); if(rkd != NULL) { @@ -72,15 +74,11 @@ tde_smgr_get_key(SMgrRelation reln) return NULL; } -void +static void tde_mdwritev(SMgrRelation reln, ForkNumber forknum, BlockNumber blocknum, const void **buffers, BlockNumber nblocks, bool skipFsync) { AesInit(); - - char* local_blocks = malloc( BLCKSZ * (nblocks+1) ); - char* local_blocks_aligned = (char*)TYPEALIGN(PG_IO_ALIGN_SIZE, local_blocks); - const void** local_buffers = malloc ( sizeof(void*) * nblocks ); RelKeyData* rkd = tde_smgr_get_key(reln); @@ -90,31 +88,36 @@ tde_mdwritev(SMgrRelation reln, ForkNumber forknum, BlockNumber blocknum, return; } - - for(int i = 0; i < nblocks; ++i ) + else { - local_buffers[i] = &local_blocks_aligned[i*BLCKSZ]; - int out_len = BLCKSZ; - AesEncrypt(rkd->internal_key.key, iv, ((char**)buffers)[i], BLCKSZ, local_buffers[i], &out_len); - } + char *local_blocks = palloc(BLCKSZ * (nblocks + 1)); + char *local_blocks_aligned = (char *)TYPEALIGN(PG_IO_ALIGN_SIZE, local_blocks); + const void **local_buffers = palloc(sizeof(void *) * nblocks); + + for(int i = 0; i < nblocks; ++i ) + { + int out_len = BLCKSZ; + local_buffers[i] = &local_blocks_aligned[i * BLCKSZ]; + AesEncrypt(rkd->internal_key.key, iv, ((char**)buffers)[i], BLCKSZ, local_buffers[i], &out_len); + } - mdwritev(reln, forknum, blocknum, - local_buffers, nblocks, skipFsync); + mdwritev(reln, forknum, blocknum, + local_buffers, nblocks, skipFsync); - free(local_blocks); - free(local_buffers); + pfree(local_blocks); + pfree(local_buffers); + } } -void +static void tde_mdextend(SMgrRelation reln, ForkNumber forknum, BlockNumber blocknum, const void *buffer, bool skipFsync) { + RelKeyData *rkd; + AesInit(); - - char* local_blocks = malloc( BLCKSZ * (1+1) ); - char* local_blocks_aligned = (char*)TYPEALIGN(PG_IO_ALIGN_SIZE, local_blocks); - RelKeyData* rkd = tde_smgr_get_key(reln); + rkd = tde_smgr_get_key(reln); if(rkd == NULL) { @@ -122,30 +125,35 @@ tde_mdextend(SMgrRelation reln, ForkNumber forknum, BlockNumber blocknum, return; } + else + { - int out_len = BLCKSZ; - AesEncrypt(rkd->internal_key.key, iv, ((char*)buffer), BLCKSZ, local_blocks_aligned, &out_len); + char *local_blocks = palloc(BLCKSZ * (1 + 1)); + char *local_blocks_aligned = (char *)TYPEALIGN(PG_IO_ALIGN_SIZE, local_blocks); + int out_len = BLCKSZ; + AesEncrypt(rkd->internal_key.key, iv, ((char*)buffer), BLCKSZ, local_blocks_aligned, &out_len); - mdextend(reln, forknum, blocknum, local_blocks_aligned, skipFsync); + mdextend(reln, forknum, blocknum, local_blocks_aligned, skipFsync); - - free(local_blocks); + pfree(local_blocks); + } } -void +static void tde_mdreadv(SMgrRelation reln, ForkNumber forknum, BlockNumber blocknum, void **buffers, BlockNumber nblocks) { + RelKeyData *rkd; + int out_len = BLCKSZ; + AesInit(); mdreadv(reln, forknum, blocknum, buffers, nblocks); - RelKeyData* rkd = tde_smgr_get_key(reln); + rkd = tde_smgr_get_key(reln); if(rkd == NULL) - { return; - } for(int i = 0; i < nblocks; ++i) { @@ -163,21 +171,19 @@ tde_mdreadv(SMgrRelation reln, ForkNumber forknum, BlockNumber blocknum, break; } } - if(allZero) continue; - - int out_len = BLCKSZ; - AesDecrypt(rkd->internal_key.key, iv, ((char**)buffers)[i], BLCKSZ, ((char**)buffers)[i], &out_len); + if(allZero) + continue; + AesDecrypt(rkd->internal_key.key, iv, ((char **)buffers)[i], BLCKSZ, ((char **)buffers)[i], &out_len); } } -void +static void tde_mdcreate(SMgrRelation reln, ForkNumber forknum, bool isRedo) { // This is the only function that gets called during actual CREATE TABLE/INDEX (EVENT TRIGGER) // so we create the key here by loading it // Later calls then decide to encrypt or not based on the existence of the key - tde_smgr_get_key(reln); - + tde_smgr_get_key(reln); return mdcreate(reln, forknum, isRedo); } @@ -204,7 +210,7 @@ static const struct f_smgr tde_smgr = { .smgr_registersync = mdregistersync, }; -void RegisterStorageMgr() +void RegisterStorageMgr(void) { tde_smgr_id = smgr_register(&tde_smgr, 0); @@ -213,7 +219,7 @@ void RegisterStorageMgr() } #else -void RegisterStorageMgr() +void RegisterStorageMgr(void) { } #endif /* PERCONA_FORK */