#412 Add OAuth bearer token injection via PQsetAuthDataHook

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(

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} $<TARGET_OBJECTS:postgres_ext_library>
     PARENT_SCOPE)

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

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 &parameter) {
+	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",

src/postgres_oauth.cpp

@@ -0,0 +1,115 @@
+#include "postgres_oauth.hpp"
+
+#include <cstdlib>
+#include <cstring>
+#include <mutex>
+#include <string>
+
+#ifdef _WIN32
+#include <windows.h>
+#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<OAuthTokenState *>(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<std::mutex> 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<PGoauthBearerRequest *>(data);
+
+			auto *ts = static_cast<OAuthTokenState *>(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<std::mutex> 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<std::mutex> lock(oauth_token_mutex);
+	if (!oauth_token_value.empty()) {
+		SecureZero(&oauth_token_value[0], oauth_token_value.size());
+	}
+	oauth_token_value.clear();
+}
+
+} // namespace duckdb

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;