From 79a6fffbb4ef9e1afd38a9687742941ab4d99417 Mon Sep 17 00:00:00 2001 From: Andreas Karlsson Date: Fri, 27 Feb 2026 00:19:09 +0100 Subject: [PATCH v1] PoC: Add support for per-connection OAuth hooks to libpq The current global hook makes life hard for authors of language drivers based on libpq. TODO - Clean up code - Write documentation - Write tests --- src/interfaces/libpq-oauth/oauth-curl.c | 3 +-- src/interfaces/libpq/exports.txt | 2 ++ src/interfaces/libpq/fe-auth-oauth.c | 6 ++++- src/interfaces/libpq/fe-auth.c | 26 +++++++++++++++++++ src/interfaces/libpq/libpq-fe.h | 3 +++ src/interfaces/libpq/libpq-int.h | 10 +++++++ .../oauth_validator/oauth_hook_client.c | 16 +++++++++--- 7 files changed, 60 insertions(+), 6 deletions(-) diff --git a/src/interfaces/libpq-oauth/oauth-curl.c b/src/interfaces/libpq-oauth/oauth-curl.c index 691e7ec1d9f..1967de4a31f 100644 --- a/src/interfaces/libpq-oauth/oauth-curl.c +++ b/src/interfaces/libpq-oauth/oauth-curl.c @@ -2621,9 +2621,8 @@ prompt_user(struct async_ctx *actx, PGconn *conn) .verification_uri_complete = actx->authz.verification_uri_complete, .expires_in = actx->authz.expires_in, }; - PQauthDataHook_type hook = PQgetAuthDataHook(); - res = hook(PQAUTHDATA_PROMPT_OAUTH_DEVICE, conn, &prompt); + res = PQrunConnAuthDataHook(PQAUTHDATA_PROMPT_OAUTH_DEVICE, conn, &prompt); if (!res) { diff --git a/src/interfaces/libpq/exports.txt b/src/interfaces/libpq/exports.txt index dbbae642d76..98c630b30a9 100644 --- a/src/interfaces/libpq/exports.txt +++ b/src/interfaces/libpq/exports.txt @@ -210,3 +210,5 @@ PQgetAuthDataHook 207 PQdefaultAuthDataHook 208 PQfullProtocolVersion 209 appendPQExpBufferVA 210 +PQsetConnAuthDataHook 211 +PQrunConnAuthDataHook 212 diff --git a/src/interfaces/libpq/fe-auth-oauth.c b/src/interfaces/libpq/fe-auth-oauth.c index 67879d64b39..3ca7e594111 100644 --- a/src/interfaces/libpq/fe-auth-oauth.c +++ b/src/interfaces/libpq/fe-auth-oauth.c @@ -998,7 +998,11 @@ setup_token_request(PGconn *conn, fe_oauth_state *state) Assert(request.openid_configuration); /* The client may have overridden the OAuth flow. */ - res = PQauthDataHook(PQAUTHDATA_OAUTH_BEARER_TOKEN, conn, &request); + if (conn->authDataHooks.proc) + res = conn->authDataHooks.proc(conn->authDataHooks.arg, PQAUTHDATA_OAUTH_BEARER_TOKEN, conn, &request); + else + res = PQauthDataHook(PQAUTHDATA_OAUTH_BEARER_TOKEN, conn, &request); + if (res > 0) { PGoauthBearerRequest *request_copy; diff --git a/src/interfaces/libpq/fe-auth.c b/src/interfaces/libpq/fe-auth.c index f05aaea9651..0533c572aeb 100644 --- a/src/interfaces/libpq/fe-auth.c +++ b/src/interfaces/libpq/fe-auth.c @@ -1602,3 +1602,29 @@ PQdefaultAuthDataHook(PGauthData type, PGconn *conn, void *data) { return 0; /* handle nothing */ } + +PQConnAuthDataHook +PQsetConnAuthDataHook(PGconn *conn, PQConnAuthDataHook proc, void *arg) +{ + PQConnAuthDataHook old; + + if (conn == NULL) + return NULL; + + old = conn->authDataHooks.proc; + if (proc) + { + conn->authDataHooks.proc = proc; + conn->authDataHooks.arg = arg; + } + return old; +} + +int +PQrunConnAuthDataHook(PGauthData type, PGconn *conn, void *data) +{ + if (conn->authDataHooks.proc) + return conn->authDataHooks.proc(conn->authDataHooks.arg, type, conn, data); + + return PQauthDataHook(type, conn, data); +} \ No newline at end of file diff --git a/src/interfaces/libpq/libpq-fe.h b/src/interfaces/libpq/libpq-fe.h index 905f2f33ab8..ea48706b31e 100644 --- a/src/interfaces/libpq/libpq-fe.h +++ b/src/interfaces/libpq/libpq-fe.h @@ -812,9 +812,12 @@ extern char *PQencryptPasswordConn(PGconn *conn, const char *passwd, const char extern PGresult *PQchangePassword(PGconn *conn, const char *user, const char *passwd); typedef int (*PQauthDataHook_type) (PGauthData type, PGconn *conn, void *data); +typedef int (*PQConnAuthDataHook) (void *arg, PGauthData type, PGconn *conn, void *data); extern void PQsetAuthDataHook(PQauthDataHook_type hook); extern PQauthDataHook_type PQgetAuthDataHook(void); extern int PQdefaultAuthDataHook(PGauthData type, PGconn *conn, void *data); +extern PQConnAuthDataHook PQsetConnAuthDataHook(PGconn *conn, PQConnAuthDataHook proc, void *arg); +extern int PQrunConnAuthDataHook(PGauthData type, PGconn *conn, void *data); /* === in encnames.c === */ diff --git a/src/interfaces/libpq/libpq-int.h b/src/interfaces/libpq/libpq-int.h index bd7eb59f5f8..6fce6f3add3 100644 --- a/src/interfaces/libpq/libpq-int.h +++ b/src/interfaces/libpq/libpq-int.h @@ -152,6 +152,13 @@ typedef struct void *noticeProcArg; } PGNoticeHooks; +/* Fields needed for OAuth handling */ +typedef struct +{ + PQConnAuthDataHook proc; + void *arg; +} PGAuthDataHooks; + typedef struct PGEvent { PGEventProc proc; /* the function to call on events */ @@ -453,6 +460,9 @@ struct pg_conn /* Callback procedures for notice message processing */ PGNoticeHooks noticeHooks; + /* Callback procedures for OAuth authentication */ + PGAuthDataHooks authDataHooks; + /* Event procs registered via PQregisterEventProc */ PGEvent *events; /* expandable array of event data */ int nEvents; /* number of active events */ diff --git a/src/test/modules/oauth_validator/oauth_hook_client.c b/src/test/modules/oauth_validator/oauth_hook_client.c index 60dd1dcdaa0..612dd2503a5 100644 --- a/src/test/modules/oauth_validator/oauth_hook_client.c +++ b/src/test/modules/oauth_validator/oauth_hook_client.c @@ -22,6 +22,7 @@ #include "libpq-fe.h" static int handle_auth_data(PGauthData type, PGconn *conn, void *data); +static int handle_auth_data_async(void *arg, PGauthData type, PGconn *conn, void *data); static PostgresPollingStatusType async_cb(PGconn *conn, PGoauthBearerRequest *req, pgsocket *altsock); @@ -125,9 +126,6 @@ main(int argc, char *argv[]) conninfo = argv[optind]; - /* Set up our OAuth hooks. */ - PQsetAuthDataHook(handle_auth_data); - /* Connect. (All the actual work is in the hook.) */ if (stress_async) { @@ -141,6 +139,9 @@ main(int argc, char *argv[]) conn = PQconnectStart(conninfo); + /* Set up our OAuth hooks. */ + PQsetConnAuthDataHook(conn, handle_auth_data_async, NULL); + do { res = PQconnectPoll(conn); @@ -148,6 +149,9 @@ main(int argc, char *argv[]) } else { + /* Set up our OAuth hooks. */ + PQsetAuthDataHook(handle_auth_data); + /* Perform a standard synchronous connection. */ conn = PQconnectdb(conninfo); } @@ -225,6 +229,12 @@ handle_auth_data(PGauthData type, PGconn *conn, void *data) return 1; } +static int +handle_auth_data_async(void *arg, PGauthData type, PGconn *conn, void *data) +{ + return handle_auth_data(type, conn, data); +} + static PostgresPollingStatusType async_cb(PGconn *conn, PGoauthBearerRequest *req, pgsocket *altsock) { -- 2.47.3