diff --git a/contrib/pg_tde/src/catalog/tde_principal_key.c b/contrib/pg_tde/src/catalog/tde_principal_key.c index b168bb578f4..c3579880734 100644 --- a/contrib/pg_tde/src/catalog/tde_principal_key.c +++ b/contrib/pg_tde/src/catalog/tde_principal_key.c @@ -60,7 +60,6 @@ typedef struct TdePrincipalKeySharedState LWLockPadded *Locks; dshash_table_handle hashHandle; void *rawDsaArea; /* DSA area pointer */ - } TdePrincipalKeySharedState; typedef struct TdePrincipalKeylocalState @@ -84,10 +83,6 @@ static dshash_parameters principal_key_dsh_params = { static TdePrincipalKeylocalState principalKeyLocalState; static void principal_key_info_attach_shmem(void); -static Size initialize_shared_state(void *start_address); -static void initialize_objects_in_dsa_area(dsa_area *dsa, void *raw_dsa_area); -static Size required_shared_mem_size(void); -static void shared_memory_shutdown(int code, Datum arg); static void clear_principal_key_cache(Oid databaseId); static inline dshash_table *get_principal_key_hash(void); static TDEPrincipalKey *get_principal_key_from_cache(Oid dbOid); @@ -111,32 +106,6 @@ PG_FUNCTION_INFO_V1(pg_tde_set_server_key_using_global_key_provider); static void pg_tde_set_principal_key_internal(Oid providerOid, Oid dbOid, const char *principal_key_name, const char *provider_name, bool ensure_new_key); -static const TDEShmemSetupRoutine principal_key_info_shmem_routine = { - .init_shared_state = initialize_shared_state, - .init_dsa_area_objects = initialize_objects_in_dsa_area, - .required_shared_mem_size = required_shared_mem_size, - .shmem_kill = shared_memory_shutdown -}; - -void -InitializePrincipalKeyInfo(void) -{ - ereport(LOG, errmsg("Initializing TDE principal key info")); - RegisterShmemRequest(&principal_key_info_shmem_routine); -} - -/* - * Lock to guard internal/principal key. Usually, this lock has to be held until - * the caller fetches an internal_key or rotates the principal. - */ -LWLock * -tde_lwlock_enc_keys(void) -{ - Assert(principalKeyLocalState.sharedPrincipalKeyState); - - return &principalKeyLocalState.sharedPrincipalKeyState->Locks[TDE_LWLOCK_ENC_KEY].lock; -} - /* * Request some pages so we can fit the DSA header, empty hash table plus some * extra. Additional memory to grow the hash map will be allocated as needed @@ -147,8 +116,8 @@ tde_lwlock_enc_keys(void) */ #define CACHE_DSA_INITIAL_SIZE (4096 * 64) -static Size -required_shared_mem_size(void) +Size +PrincipalKeyShmemSize(void) { Size sz = CACHE_DSA_INITIAL_SIZE; @@ -156,41 +125,75 @@ required_shared_mem_size(void) return MAXALIGN(sz); } -/* - * Initialize the shared area for Principal key info. - * This includes locks and cache area for principal key info - */ - -static Size -initialize_shared_state(void *start_address) +void +PrincipalKeyShmemInit(void) { - TdePrincipalKeySharedState *sharedState = (TdePrincipalKeySharedState *) start_address; + bool found; + char *free_start; + Size required_shmem_size = PrincipalKeyShmemSize(); - ereport(LOG, errmsg("initializing shared state for principal key")); + LWLockAcquire(AddinShmemInitLock, LW_EXCLUSIVE); - sharedState->Locks = GetNamedLWLockTranche(TDE_TRANCHE_NAME); + /* Create or attach to the shared memory state */ + ereport(NOTICE, errmsg("PrincipalKeyShmemInit: requested %ld bytes", required_shmem_size)); + free_start = ShmemInitStruct("pg_tde", required_shmem_size, &found); - principalKeyLocalState.sharedPrincipalKeyState = sharedState; - principalKeyLocalState.sharedHash = NULL; + if (!found) + { + TdePrincipalKeySharedState *sharedState; + Size sz; + Size dsa_area_size; + dsa_area *dsa; + dshash_table *dsh; - return sizeof(TdePrincipalKeySharedState); -} + /* Now place shared state structure */ + sharedState = (TdePrincipalKeySharedState *) free_start; + sz = MAXALIGN(sizeof(TdePrincipalKeySharedState)); + free_start += sz; + Assert(sz <= required_shmem_size); -static void -initialize_objects_in_dsa_area(dsa_area *dsa, void *raw_dsa_area) -{ - dshash_table *dsh; - TdePrincipalKeySharedState *sharedState = principalKeyLocalState.sharedPrincipalKeyState; + /* Create DSA area */ + dsa_area_size = required_shmem_size - sz; + Assert(dsa_area_size > 0); + + ereport(LOG, errmsg("creating DSA area of size %lu", dsa_area_size)); + + dsa = dsa_create_in_place(free_start, + dsa_area_size, + LWLockNewTrancheId(), 0); + dsa_pin(dsa); + + /* Limit area size during population to get a nice error */ + dsa_set_size_limit(dsa, dsa_area_size); + + principal_key_dsh_params.tranche_id = LWLockNewTrancheId(); + dsh = dshash_create(dsa, &principal_key_dsh_params, NULL); + + dsa_set_size_limit(dsa, -1); - ereport(LOG, errmsg("initializing dsa area objects for principal key")); + sharedState->Locks = GetNamedLWLockTranche(TDE_TRANCHE_NAME); + sharedState->hashHandle = dshash_get_hash_table_handle(dsh); + sharedState->rawDsaArea = free_start; - Assert(sharedState != NULL); + principalKeyLocalState.sharedPrincipalKeyState = sharedState; + principalKeyLocalState.sharedHash = NULL; - sharedState->rawDsaArea = raw_dsa_area; - principal_key_dsh_params.tranche_id = LWLockNewTrancheId(); - dsh = dshash_create(dsa, &principal_key_dsh_params, NULL); - sharedState->hashHandle = dshash_get_hash_table_handle(dsh); - dshash_detach(dsh); + dshash_detach(dsh); + } + + LWLockRelease(AddinShmemInitLock); +} + +/* + * Lock to guard internal/principal key. Usually, this lock has to be held until + * the caller fetches an internal_key or rotates the principal. + */ +LWLock * +tde_lwlock_enc_keys(void) +{ + Assert(principalKeyLocalState.sharedPrincipalKeyState); + + return &principalKeyLocalState.sharedPrincipalKeyState->Locks[TDE_LWLOCK_ENC_KEY].lock; } /* @@ -217,12 +220,6 @@ principal_key_info_attach_shmem(void) MemoryContextSwitchTo(oldcontext); } -static void -shared_memory_shutdown(int code, Datum arg) -{ - principalKeyLocalState.sharedPrincipalKeyState = NULL; -} - void set_principal_key_with_keyring(const char *key_name, const char *provider_name, Oid providerOid, Oid dbOid, bool ensure_new_key) diff --git a/contrib/pg_tde/src/common/pg_tde_shmem.c b/contrib/pg_tde/src/common/pg_tde_shmem.c index 697a3e23b68..2d32893ff3a 100644 --- a/contrib/pg_tde/src/common/pg_tde_shmem.c +++ b/contrib/pg_tde/src/common/pg_tde_shmem.c @@ -11,122 +11,10 @@ #include "postgres.h" -#include "lib/dshash.h" -#include "nodes/pg_list.h" -#include "storage/ipc.h" -#include "storage/lwlock.h" -#include "storage/shmem.h" - #include "common/pg_tde_shmem.h" -static void tde_shmem_shutdown(int code, Datum arg); - -List *registeredShmemRequests = NIL; -bool shmemInited = false; - -void -RegisterShmemRequest(const TDEShmemSetupRoutine *routine) -{ - Assert(shmemInited == false); - registeredShmemRequests = lappend(registeredShmemRequests, (void *) routine); -} - -Size -TdeRequiredSharedMemorySize(void) -{ - Size sz = 0; - ListCell *lc; - - foreach(lc, registeredShmemRequests) - { - TDEShmemSetupRoutine *routine = (TDEShmemSetupRoutine *) lfirst(lc); - - if (routine->required_shared_mem_size) - sz = add_size(sz, routine->required_shared_mem_size()); - } - return MAXALIGN(sz); -} - int TdeRequiredLocksCount(void) { return TDE_LWLOCK_COUNT; } - -void -TdeShmemInit(void) -{ - bool found; - char *free_start; - Size required_shmem_size = TdeRequiredSharedMemorySize(); - - LWLockAcquire(AddinShmemInitLock, LW_EXCLUSIVE); - /* Create or attach to the shared memory state */ - ereport(NOTICE, errmsg("TdeShmemInit: requested %ld bytes", required_shmem_size)); - free_start = ShmemInitStruct("pg_tde", required_shmem_size, &found); - - if (!found) - { - /* First time through ... */ - dsa_area *dsa; - ListCell *lc; - Size used_size = 0; - Size dsa_area_size; - - /* Now place all shared state structures */ - foreach(lc, registeredShmemRequests) - { - Size sz = 0; - TDEShmemSetupRoutine *routine = (TDEShmemSetupRoutine *) lfirst(lc); - - if (routine->init_shared_state) - { - sz = routine->init_shared_state(free_start); - used_size += MAXALIGN(sz); - free_start += MAXALIGN(sz); - Assert(used_size <= required_shmem_size); - } - } - /* Create DSA area */ - dsa_area_size = required_shmem_size - used_size; - Assert(dsa_area_size > 0); - - ereport(LOG, errmsg("creating DSA area of size %lu", dsa_area_size)); - dsa = dsa_create_in_place(free_start, - dsa_area_size, - LWLockNewTrancheId(), 0); - dsa_pin(dsa); - dsa_set_size_limit(dsa, dsa_area_size); - - /* Initialize all DSA area objects */ - foreach(lc, registeredShmemRequests) - { - TDEShmemSetupRoutine *routine = (TDEShmemSetupRoutine *) lfirst(lc); - - if (routine->init_dsa_area_objects) - routine->init_dsa_area_objects(dsa, free_start); - } - ereport(LOG, errmsg("setting no limit to DSA area of size %lu", dsa_area_size)); - - dsa_set_size_limit(dsa, -1); /* Let it grow outside the shared - * memory */ - - shmemInited = true; - } - LWLockRelease(AddinShmemInitLock); - on_shmem_exit(tde_shmem_shutdown, (Datum) 0); -} - -static void -tde_shmem_shutdown(int code, Datum arg) -{ - ListCell *lc; - - foreach(lc, registeredShmemRequests) - { - TDEShmemSetupRoutine *routine = (TDEShmemSetupRoutine *) lfirst(lc); - - if (routine->shmem_kill) - routine->shmem_kill(code, arg); - } -} diff --git a/contrib/pg_tde/src/include/catalog/tde_principal_key.h b/contrib/pg_tde/src/include/catalog/tde_principal_key.h index 26dd6f22bef..185a9229295 100644 --- a/contrib/pg_tde/src/include/catalog/tde_principal_key.h +++ b/contrib/pg_tde/src/include/catalog/tde_principal_key.h @@ -42,7 +42,8 @@ typedef struct XLogPrincipalKeyRotate #define SizeoOfXLogPrincipalKeyRotate offsetof(XLogPrincipalKeyRotate, buff) -extern void InitializePrincipalKeyInfo(void); +extern void PrincipalKeyShmemInit(void); +extern Size PrincipalKeyShmemSize(void); #ifndef FRONTEND extern void principal_key_startup_cleanup(Oid databaseId); diff --git a/contrib/pg_tde/src/include/common/pg_tde_shmem.h b/contrib/pg_tde/src/include/common/pg_tde_shmem.h index 31dfb3a6917..ba36d10d246 100644 --- a/contrib/pg_tde/src/include/common/pg_tde_shmem.h +++ b/contrib/pg_tde/src/include/common/pg_tde_shmem.h @@ -8,9 +8,6 @@ #ifndef PG_TDE_SHMEM_H #define PG_TDE_SHMEM_H -#include "postgres.h" -#include "utils/dsa.h" - #define TDE_TRANCHE_NAME "pg_tde_tranche" typedef enum @@ -22,39 +19,6 @@ typedef enum TDE_LWLOCK_COUNT } TDELockTypes; -typedef struct TDEShmemSetupRoutine -{ - /* - * init_shared_state gets called at the time of extension load you can - * initialize the data structures required to be placed in shared memory - * in this callback The callback must return the size of the shared memory - * area acquired. The argument to the function is the start of the shared - * memory address that can be used to store the shared data structures. - */ - Size (*init_shared_state) (void *raw_dsa_area); - - /* - * shmem_startup gets called at the time of postmaster shutdown - */ - void (*shmem_kill) (int code, Datum arg); - - /* - * The callback must return the size of the shared memory acquired. - */ - Size (*required_shared_mem_size) (void); - - /* - * Gets called after all shared memory structures are initialized and here - * you can create shared memory hash tables or any other shared objects - * that needs to live in DSA area. - */ - void (*init_dsa_area_objects) (dsa_area *dsa, void *raw_dsa_area); -} TDEShmemSetupRoutine; - -/* Interface to register the shared memory requests */ -extern void RegisterShmemRequest(const TDEShmemSetupRoutine *routine); -extern void TdeShmemInit(void); -extern Size TdeRequiredSharedMemorySize(void); extern int TdeRequiredLocksCount(void); #endif /* PG_TDE_SHMEM_H */ diff --git a/contrib/pg_tde/src/pg_tde.c b/contrib/pg_tde/src/pg_tde.c index 4157b72d6f7..974b3358a22 100644 --- a/contrib/pg_tde/src/pg_tde.c +++ b/contrib/pg_tde/src/pg_tde.c @@ -55,13 +55,15 @@ PG_FUNCTION_INFO_V1(pg_tdeam_handler); static void tde_shmem_request(void) { - Size sz = TdeRequiredSharedMemorySize(); + Size sz = 0; int required_locks = TdeRequiredLocksCount(); + sz = add_size(sz, PrincipalKeyShmemSize()); sz = add_size(sz, TDEXLogEncryptStateSize()); if (prev_shmem_request_hook) prev_shmem_request_hook(); + RequestAddinShmemSpace(sz); RequestNamedLWLockTranche(TDE_TRANCHE_NAME, required_locks); ereport(LOG, errmsg("tde_shmem_request: requested %ld bytes", sz)); @@ -73,8 +75,8 @@ tde_shmem_startup(void) if (prev_shmem_startup_hook) prev_shmem_startup_hook(); - TdeShmemInit(); - InitializeKeyProviderInfo(); + KeyProviderShmemInit(); + PrincipalKeyShmemInit(); TDEXLogShmemInit(); TDEXLogSmgrInit(); } @@ -100,7 +102,6 @@ _PG_init(void) AesInit(); TdeGucInit(); TdeEventCaptureInit(); - InitializePrincipalKeyInfo(); InstallFileKeyring(); InstallVaultV2Keyring(); InstallKmipKeyring();