From fd60ceb4c849c2a13b0201b57753506b60718e6e Mon Sep 17 00:00:00 2001 From: Jacob Champion Date: Thu, 6 Feb 2025 20:22:03 -0800 Subject: [PATCH v47 04/11] fixup! Add OAUTHBEARER SASL mechanism --- src/backend/libpq/auth-oauth.c | 13 +++++++++---- src/test/modules/oauth_validator/validator.c | 2 +- 2 files changed, 10 insertions(+), 5 deletions(-) diff --git a/src/backend/libpq/auth-oauth.c b/src/backend/libpq/auth-oauth.c index 6155d63a116..d910cbcb161 100644 --- a/src/backend/libpq/auth-oauth.c +++ b/src/backend/libpq/auth-oauth.c @@ -39,7 +39,7 @@ static int oauth_exchange(void *opaq, const char *input, int inputlen, char **output, int *outputlen, const char **logdetail); static void load_validator_library(const char *libname); -static void shutdown_validator_library(int code, Datum arg); +static void shutdown_validator_library(void *arg); static ValidatorModuleState *validator_module_state; static const OAuthValidatorCallbacks *ValidatorCallbacks; @@ -737,6 +737,7 @@ static void load_validator_library(const char *libname) { OAuthValidatorModuleInit validator_init; + MemoryContextCallback *mcb; Assert(libname && *libname); @@ -761,15 +762,19 @@ load_validator_library(const char *libname) if (ValidatorCallbacks->startup_cb != NULL) ValidatorCallbacks->startup_cb(validator_module_state); - before_shmem_exit(shutdown_validator_library, 0); + /* Shut down the library before cleaning up its state. */ + mcb = palloc0(sizeof(*mcb)); + mcb->func = shutdown_validator_library; + + MemoryContextRegisterResetCallback(CurrentMemoryContext, mcb); } /* * Call the validator module's shutdown callback, if one is provided. This is - * invoked via before_shmem_exit(). + * invoked during memory context reset. */ static void -shutdown_validator_library(int code, Datum arg) +shutdown_validator_library(void *arg) { if (ValidatorCallbacks->shutdown_cb != NULL) ValidatorCallbacks->shutdown_cb(validator_module_state); diff --git a/src/test/modules/oauth_validator/validator.c b/src/test/modules/oauth_validator/validator.c index bf94f091def..ef9bbb2866f 100644 --- a/src/test/modules/oauth_validator/validator.c +++ b/src/test/modules/oauth_validator/validator.c @@ -100,7 +100,7 @@ validator_shutdown(ValidatorModuleState *state) { /* Check to make sure our private state still exists. */ if (state->private_data != PRIVATE_COOKIE) - elog(ERROR, "oauth_validator: private state cookie changed to %p in shutdown", + elog(PANIC, "oauth_validator: private state cookie changed to %p in shutdown", state->private_data); } -- 2.34.1