From f8876fb63d52564ab6afcd7c4f678b2fded3610a Mon Sep 17 00:00:00 2001 From: 0814celsus Date: Mon, 20 Apr 2026 17:45:15 +0200 Subject: [PATCH] Add OAuth bearer token injection via PQsetAuthDataHook Implements the PG18 PQsetAuthDataHook API to inject OAuth bearer tokens during OAUTHBEARER authentication. This allows connecting to OAuth-enabled PostgreSQL servers with a pre-obtained token, bypassing the device authorization flow. Token sources (checked in priority order): 1. DuckDB setting: SET pg_oauth_token = '' 2. Environment variable: PGOAUTHTOKEN The hook is registered once at extension load time. --- CMakeLists.txt | 3 +- src/CMakeLists.txt | 3 +- src/include/postgres_oauth.hpp | 26 +++++ src/postgres_extension.cpp | 16 ++++ src/postgres_oauth.cpp | 115 +++++++++++++++++++++++ test/sql/storage/attach_oauth_token.test | 45 +++++++++ 6 files changed, 206 insertions(+), 2 deletions(-) create mode 100644 src/include/postgres_oauth.hpp create mode 100644 src/postgres_oauth.cpp create mode 100644 test/sql/storage/attach_oauth_token.test diff --git a/CMakeLists.txt b/CMakeLists.txt index 034505d50..5f99434aa 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -22,7 +22,8 @@ endif() include_directories( include database-connector/src/include - ${OPENSSL_INCLUDE_DIR}) + ${OPENSSL_INCLUDE_DIR} + ${PostgreSQL_INCLUDE_DIRS}) if(WIN32) include_directories( diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 5b99d6e43..0ea6847bf 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -22,7 +22,8 @@ add_library( postgres_storage.cpp postgres_text_reader.cpp postgres_utils.cpp - postgres_logging.cpp) + postgres_logging.cpp + postgres_oauth.cpp) set(ALL_OBJECT_FILES ${ALL_OBJECT_FILES} $ PARENT_SCOPE) diff --git a/src/include/postgres_oauth.hpp b/src/include/postgres_oauth.hpp new file mode 100644 index 000000000..262428341 --- /dev/null +++ b/src/include/postgres_oauth.hpp @@ -0,0 +1,26 @@ +//===----------------------------------------------------------------------===// +// DuckDB +// +// postgres_oauth.hpp +// +// +//===----------------------------------------------------------------------===// + +#pragma once + +#include "duckdb.hpp" + +namespace duckdb { + +//! Initializes the PQsetAuthDataHook for OAuth bearer token injection. +//! The hook provides tokens from either the PGOAUTHTOKEN environment variable +//! or the pg_oauth_token DuckDB setting. +void PostgresInitOAuthHook(); + +//! Sets the in-memory OAuth token (used by the pg_oauth_token setting) +void PostgresSetOAuthToken(const string &token); + +//! Clears the in-memory OAuth token +void PostgresClearOAuthToken(); + +} // namespace duckdb diff --git a/src/postgres_extension.cpp b/src/postgres_extension.cpp index 208092ff7..af54cd202 100644 --- a/src/postgres_extension.cpp +++ b/src/postgres_extension.cpp @@ -22,6 +22,7 @@ #include "duckdb/common/error_data.hpp" #include "postgres_logging.hpp" #include "postgres_hstore.hpp" +#include "postgres_oauth.hpp" using namespace duckdb; @@ -154,7 +155,18 @@ static std::string CreatePoolNote(const std::string &option) { "\"FROM postgres_configure_pool(catalog_name='my_attached_postgres_db', " + option + ")\""; } +static void SetOAuthToken(ClientContext &context, SetScope scope, Value ¶meter) { + if (parameter.IsNull() || StringValue::Get(parameter).empty()) { + PostgresClearOAuthToken(); + } else { + PostgresSetOAuthToken(StringValue::Get(parameter)); + } +} + static void LoadInternal(ExtensionLoader &loader) { + // Register the OAuth bearer token hook before any connections are made + PostgresInitOAuthHook(); + PostgresScanFunction postgres_fun; loader.RegisterFunction(postgres_fun); @@ -231,6 +243,10 @@ static void LoadInternal(ExtensionLoader &loader) { LogicalType::VARCHAR, Value(), SetPostgresNullByteReplacement); config.AddExtensionOption("pg_debug_show_queries", "DEBUG SETTING: print all queries sent to Postgres to stdout", LogicalType::BOOLEAN, Value::BOOLEAN(false), SetPostgresDebugQueryPrint); + config.AddExtensionOption("pg_oauth_token", + "OAuth bearer token for PostgreSQL OAUTHBEARER authentication. " + "Takes priority over the PGOAUTHTOKEN environment variable", + LogicalType::VARCHAR, Value(), SetOAuthToken); config.AddExtensionOption("pg_use_text_protocol", "Whether or not to use TEXT protocol to read data. This is slower, but provides better " "compatibility with non-Postgres systems", diff --git a/src/postgres_oauth.cpp b/src/postgres_oauth.cpp new file mode 100644 index 000000000..bc9a3b6c7 --- /dev/null +++ b/src/postgres_oauth.cpp @@ -0,0 +1,115 @@ +#include "postgres_oauth.hpp" + +#include +#include +#include +#include + +#ifdef _WIN32 +#include +#define SecureZero(ptr, len) SecureZeroMemory(ptr, len) +#else +#define SecureZero(ptr, len) explicit_bzero(ptr, len) +#endif + +extern "C" { +#include "libpq-fe.h" +} + +namespace duckdb { + +//! Global mutex protecting the token string +static std::mutex oauth_token_mutex; +//! The in-memory token set via the DuckDB setting +static std::string oauth_token_value; + +//! Previous hook in the chain (if any) +static PQauthDataHook_type prev_hook = nullptr; + +struct OAuthTokenState { + char *token_copy; +}; + +static void OAuthTokenCleanup(PGconn *conn, PGoauthBearerRequest *request) { + auto *ts = static_cast(request->user); + if (ts) { + if (ts->token_copy) { + SecureZero(ts->token_copy, strlen(ts->token_copy)); + free(ts->token_copy); + } + free(ts); + request->user = nullptr; + } + request->token = nullptr; +} + +static int OAuthBearerTokenHook(PGauthData type, PGconn *conn, void *data) { + if (type == PQAUTHDATA_OAUTH_BEARER_TOKEN) { + const char *token = nullptr; + + // Priority 1: DuckDB setting (more secure, in-process memory only) + std::string token_str; + { + std::lock_guard lock(oauth_token_mutex); + token_str = oauth_token_value; + } + + // Priority 2: Environment variable PGOAUTHTOKEN + if (token_str.empty()) { + const char *env_token = std::getenv("PGOAUTHTOKEN"); + if (env_token && env_token[0] != '\0') { + token_str = env_token; + } + } + + if (!token_str.empty()) { + auto *request = static_cast(data); + + auto *ts = static_cast(malloc(sizeof(OAuthTokenState))); + if (!ts) { + return -1; + } + + ts->token_copy = strdup(token_str.c_str()); + if (!ts->token_copy) { + free(ts); + return -1; + } + + request->token = ts->token_copy; + request->cleanup = OAuthTokenCleanup; + request->user = ts; + request->async = nullptr; + return 1; + } + } + + // Fall through to previous hook + if (prev_hook) { + return prev_hook(type, conn, data); + } + return 0; +} + +void PostgresInitOAuthHook() { + prev_hook = PQgetAuthDataHook(); + PQsetAuthDataHook(OAuthBearerTokenHook); +} + +void PostgresSetOAuthToken(const string &token) { + std::lock_guard lock(oauth_token_mutex); + if (!oauth_token_value.empty()) { + SecureZero(&oauth_token_value[0], oauth_token_value.size()); + } + oauth_token_value = token; +} + +void PostgresClearOAuthToken() { + std::lock_guard lock(oauth_token_mutex); + if (!oauth_token_value.empty()) { + SecureZero(&oauth_token_value[0], oauth_token_value.size()); + } + oauth_token_value.clear(); +} + +} // namespace duckdb diff --git a/test/sql/storage/attach_oauth_token.test b/test/sql/storage/attach_oauth_token.test new file mode 100644 index 000000000..c1c4ca6a0 --- /dev/null +++ b/test/sql/storage/attach_oauth_token.test @@ -0,0 +1,45 @@ +# name: test/sql/storage/attach_oauth_token.test +# description: Test OAuth bearer token injection via PQsetAuthDataHook +# group: [storage] + +require postgres_scanner + +require-env POSTGRES_TEST_DATABASE_AVAILABLE + +require-env PGOAUTHTOKEN + +require-env PG_OAUTH_DSN + +# Explicitly load the extension so settings are registered +statement ok +LOAD postgres_scanner; + +# Test 1: Connect using the pg_oauth_token DuckDB setting +statement ok +SET pg_oauth_token = '${PGOAUTHTOKEN}'; + +statement ok +ATTACH '${PG_OAUTH_DSN}' AS pg_oauth (TYPE POSTGRES); + +query I +SELECT current_user FROM postgres_query('pg_oauth', 'SELECT current_user'); +---- +admin + +statement ok +DETACH pg_oauth; + +# Test 2: Clear the setting and verify env var fallback (PGOAUTHTOKEN) works +statement ok +SET pg_oauth_token = ''; + +statement ok +ATTACH '${PG_OAUTH_DSN}' AS pg_oauth_env (TYPE POSTGRES); + +query I +SELECT current_user FROM postgres_query('pg_oauth_env', 'SELECT current_user'); +---- +admin + +statement ok +DETACH pg_oauth_env;