v2-0007-WIP-Introduce-third-party-OAuth-flow-plugins.patch

application/octet-stream

Filename: v2-0007-WIP-Introduce-third-party-OAuth-flow-plugins.patch
Type: application/octet-stream
Part: 7
Message: Re: [oauth] Stabilize the libpq-oauth ABI (and allow alternative implementations?)
From e20b555aba57ac7e810a0b3e7821690f7e41072f Mon Sep 17 00:00:00 2001
From: Jacob Champion <jacob.champion@enterprisedb.com>
Date: Wed, 3 Dec 2025 15:47:23 -0800
Subject: [PATCH v2 7/7] WIP: Introduce third-party OAuth flow plugins?

This experimental commit promotes the pg_start_oauthbearer API to a
public header (libpq-oauth.h) and adds a PGOAUTHMODULE environment
variable that overrides the load path for the plugin, allowing users to
provide their own. The libpq_oauth_init function is now optional, and
will remain undocumented. (Modules that don't provide it are marked as
user-defined.)

This is a relatively small amount of implementation change, but
unfortunately the tests have a large amount of code motion to be able to
share logic between the test executable and plugin. I might need to
split that into multiple squash! commits to make it more easily
reviewable.

TODO: figure out PGDLLEXPORT, which we do not currently provide publicly
TODO: add a public init() API so that no one tries to implement
      libpq_oauth_init()?
TODO: lock down PGOAUTHMODULE as necessary to avoid introducing exciting
      new vulnerabilities
TODO: how hard would it be to support Windows here?

Discussion: https://postgr.es/m/CAOYmi%2BmrGg%2Bn_X2MOLgeWcj3v_M00gR8uz_D7mM8z%3DdX1JYVbg%40mail.gmail.com
---
 src/interfaces/libpq/meson.build              |   1 +
 src/interfaces/libpq/Makefile                 |   2 +
 src/interfaces/libpq-oauth/oauth-curl.h       |  24 --
 src/interfaces/libpq/fe-auth-oauth.h          |   2 +-
 src/interfaces/libpq/libpq-oauth.h            |  52 +++
 src/interfaces/libpq-oauth/oauth-curl.c       |   2 +-
 src/interfaces/libpq/fe-auth-oauth.c          |  81 ++--
 src/test/modules/oauth_validator/meson.build  |  15 +
 src/test/modules/oauth_validator/Makefile     |  10 +-
 .../oauth_validator/oauth_test_common.h       |  26 ++
 src/test/modules/oauth_validator/oauth_flow.c |  69 ++++
 .../oauth_validator/oauth_hook_client.c       | 319 +--------------
 .../oauth_validator/oauth_test_common.c       | 374 ++++++++++++++++++
 .../modules/oauth_validator/t/002_client.pl   |  41 +-
 14 files changed, 649 insertions(+), 369 deletions(-)
 delete mode 100644 src/interfaces/libpq-oauth/oauth-curl.h
 create mode 100644 src/interfaces/libpq/libpq-oauth.h
 create mode 100644 src/test/modules/oauth_validator/oauth_test_common.h
 create mode 100644 src/test/modules/oauth_validator/oauth_flow.c
 create mode 100644 src/test/modules/oauth_validator/oauth_test_common.c

diff --git a/src/interfaces/libpq/meson.build b/src/interfaces/libpq/meson.build
index b259c998fa2..ac71548b059 100644
--- a/src/interfaces/libpq/meson.build
+++ b/src/interfaces/libpq/meson.build
@@ -128,6 +128,7 @@ pkgconfig.generate(
 install_headers(
   'libpq-fe.h',
   'libpq-events.h',
+  'libpq-oauth.h',
 )
 
 install_headers(
diff --git a/src/interfaces/libpq/Makefile b/src/interfaces/libpq/Makefile
index 9fe321147fc..1d0cd22f6f0 100644
--- a/src/interfaces/libpq/Makefile
+++ b/src/interfaces/libpq/Makefile
@@ -155,6 +155,7 @@ $(top_builddir)/src/port/pg_config_paths.h:
 install: all installdirs install-lib
 	$(INSTALL_DATA) $(srcdir)/libpq-fe.h '$(DESTDIR)$(includedir)'
 	$(INSTALL_DATA) $(srcdir)/libpq-events.h '$(DESTDIR)$(includedir)'
+	$(INSTALL_DATA) $(srcdir)/libpq-oauth.h '$(DESTDIR)$(includedir)'
 	$(INSTALL_DATA) $(srcdir)/libpq-int.h '$(DESTDIR)$(includedir_internal)'
 	$(INSTALL_DATA) $(srcdir)/fe-auth-sasl.h '$(DESTDIR)$(includedir_internal)'
 	$(INSTALL_DATA) $(srcdir)/pqexpbuffer.h '$(DESTDIR)$(includedir_internal)'
@@ -177,6 +178,7 @@ installdirs: installdirs-lib
 uninstall: uninstall-lib
 	rm -f '$(DESTDIR)$(includedir)/libpq-fe.h'
 	rm -f '$(DESTDIR)$(includedir)/libpq-events.h'
+	rm -f '$(DESTDIR)$(includedir)/libpq-oauth.h'
 	rm -f '$(DESTDIR)$(includedir_internal)/libpq-int.h'
 	rm -f '$(DESTDIR)$(includedir_internal)/fe-auth-sasl.h'
 	rm -f '$(DESTDIR)$(includedir_internal)/pqexpbuffer.h'
diff --git a/src/interfaces/libpq-oauth/oauth-curl.h b/src/interfaces/libpq-oauth/oauth-curl.h
deleted file mode 100644
index 47704689586..00000000000
--- a/src/interfaces/libpq-oauth/oauth-curl.h
+++ /dev/null
@@ -1,24 +0,0 @@
-/*-------------------------------------------------------------------------
- *
- * oauth-curl.h
- *
- *	  Definitions for OAuth Device Authorization module
- *
- * Portions Copyright (c) 1996-2025, PostgreSQL Global Development Group
- * Portions Copyright (c) 1994, Regents of the University of California
- *
- * src/interfaces/libpq-oauth/oauth-curl.h
- *
- *-------------------------------------------------------------------------
- */
-
-#ifndef OAUTH_CURL_H
-#define OAUTH_CURL_H
-
-#include "libpq-fe.h"
-
-/* Exported flow callback. */
-extern PGDLLEXPORT int pg_start_oauthbearer(PGconn *conn,
-											PGoauthBearerRequestV2 *request);
-
-#endif							/* OAUTH_CURL_H */
diff --git a/src/interfaces/libpq/fe-auth-oauth.h b/src/interfaces/libpq/fe-auth-oauth.h
index b9aed879e64..30d3ff6741e 100644
--- a/src/interfaces/libpq/fe-auth-oauth.h
+++ b/src/interfaces/libpq/fe-auth-oauth.h
@@ -35,7 +35,7 @@ typedef struct
 	void	   *async_ctx;
 
 	bool		builtin;
-	void	   *builtin_flow;
+	void	   *flow_module;
 } fe_oauth_state;
 
 extern void pqClearOAuthToken(PGconn *conn);
diff --git a/src/interfaces/libpq/libpq-oauth.h b/src/interfaces/libpq/libpq-oauth.h
new file mode 100644
index 00000000000..2a62b330b1c
--- /dev/null
+++ b/src/interfaces/libpq/libpq-oauth.h
@@ -0,0 +1,52 @@
+/*-------------------------------------------------------------------------
+ *
+ * libpq-oauth.h
+ *	  This file contains structs and functions used by custom OAuth plugins.
+ *
+ * Copyright (c) 2025, PostgreSQL Global Development Group
+ *
+ * src/interfaces/libpq/libpq-oauth.h
+ *
+ *-------------------------------------------------------------------------
+ */
+
+#ifndef LIBPQ_OAUTH_H
+#define LIBPQ_OAUTH_H
+
+#include "libpq-fe.h"
+
+#ifdef __cplusplus
+extern "C"
+{
+#endif
+
+/* XXX can't rely on c.h, but duplicating this is asking for trouble */
+#ifndef PGDLLEXPORT
+#ifdef _WIN32
+#define PGDLLEXPORT __declspec (dllexport)
+#elif defined(__has_attribute)
+#if __has_attribute(visibility)
+#define PGDLLEXPORT __attribute__((visibility("default")))
+#else
+#define PGDLLEXPORT
+#endif
+#else
+#define PGDLLEXPORT
+#endif
+#endif
+
+/*
+ * V1 API
+ *
+ * Flow plugins must provide an implementation of this callback.
+ *
+ * TODO: provide a magic struct that allows backwards but not forwards compat?
+ */
+extern PGDLLEXPORT int pg_start_oauthbearer(PGconn *conn,
+											PGoauthBearerRequestV2 *request);
+
+#ifdef __cplusplus
+}
+#endif
+
+#endif							/* LIBPQ_OAUTH_H */
diff --git a/src/interfaces/libpq-oauth/oauth-curl.c b/src/interfaces/libpq-oauth/oauth-curl.c
index b94458c860f..587c9c882e6 100644
--- a/src/interfaces/libpq-oauth/oauth-curl.c
+++ b/src/interfaces/libpq-oauth/oauth-curl.c
@@ -29,8 +29,8 @@
 #endif
 
 #include "common/jsonapi.h"
+#include "libpq-oauth.h"
 #include "mb/pg_wchar.h"
-#include "oauth-curl.h"
 
 #ifdef USE_DYNAMIC_OAUTH
 
diff --git a/src/interfaces/libpq/fe-auth-oauth.c b/src/interfaces/libpq/fe-auth-oauth.c
index da6a2c4a2f7..88cf1b45938 100644
--- a/src/interfaces/libpq/fe-auth-oauth.c
+++ b/src/interfaces/libpq/fe-auth-oauth.c
@@ -17,6 +17,8 @@
 
 #ifdef USE_DYNAMIC_OAUTH
 #include <dlfcn.h>
+#else
+#include "libpq-oauth.h"
 #endif
 
 #include "common/base64.h"
@@ -887,15 +889,33 @@ use_builtin_flow(PGconn *conn, fe_oauth_state *state, PGoauthBearerRequestV2 *re
 	 * On the other platforms, load the module using only the basename, to
 	 * rely on the runtime linker's standard search behavior.
 	 */
-	const char *const module_name =
+	const char *module_name =
 #if defined(__darwin__)
 		LIBDIR "/libpq-oauth" DLSUFFIX;
 #else
 		"libpq-oauth" DLSUFFIX;
 #endif
 
-	state->builtin_flow = dlopen(module_name, RTLD_NOW | RTLD_LOCAL);
-	if (!state->builtin_flow)
+	/*-
+	 * Additionally, the user may override the module path explicitly to be
+	 * able to provide their own module, via PGOAUTHMODULE.
+	 *
+	 * TODO: error messages below need to be rethought when this is set
+	 * TODO: have to think about _all_ the security ramifications of this. What
+	 * existing protections in LD_LIBRARY_PATH (and/or SIP) are we potentially
+	 * bypassing? Should we check the permissions of the file somehow...?
+	 * TODO: maybe disallow anything not underneath LIBDIR? or PKGLIBDIR?
+	 * Should it have a naming convention?
+	 */
+	{
+		const char *env = getenv("PGOAUTHMODULE");
+
+		if (env && env[0])
+			module_name = env;
+	}
+
+	state->flow_module = dlopen(module_name, RTLD_NOW | RTLD_LOCAL);
+	if (!state->flow_module)
 	{
 		/*
 		 * For end users, this probably isn't an error condition, it just
@@ -910,8 +930,7 @@ use_builtin_flow(PGconn *conn, fe_oauth_state *state, PGoauthBearerRequestV2 *re
 		return 0;
 	}
 
-	if ((init = dlsym(state->builtin_flow, "libpq_oauth_init")) == NULL
-		|| (start_flow = dlsym(state->builtin_flow, "pg_start_oauthbearer")) == NULL)
+	if ((start_flow = dlsym(state->flow_module, "pg_start_oauthbearer")) == NULL)
 	{
 		/*
 		 * This is more of an error condition than the one above, but due to
@@ -920,7 +939,7 @@ use_builtin_flow(PGconn *conn, fe_oauth_state *state, PGoauthBearerRequestV2 *re
 		if (oauth_unsafe_debugging_enabled())
 			fprintf(stderr, "failed dlsym for libpq-oauth: %s\n", dlerror());
 
-		dlclose(state->builtin_flow);
+		dlclose(state->flow_module);
 		return 0;
 	}
 
@@ -930,34 +949,46 @@ use_builtin_flow(PGconn *conn, fe_oauth_state *state, PGoauthBearerRequestV2 *re
 	 */
 
 	/*
-	 * We need to inject necessary function pointers into the module. This
-	 * only needs to be done once -- even if the pointers are constant,
-	 * assigning them while another thread is executing the flows feels like
-	 * tempting fate.
+	 * Our libpq-oauth.so provides a special initialization function for libpq
+	 * integration. It's not a problem if we don't find this; it just means
+	 * that a user-defined PGOAUTHMODULE is being used.
 	 */
-	if ((lockerr = pthread_mutex_lock(&init_mutex)) != 0)
+	init = dlsym(state->flow_module, "libpq_oauth_init");
+
+	if (!init)
+		state->builtin = false; /* adjust our error messages */
+	else
 	{
-		/* Should not happen... but don't continue if it does. */
-		Assert(false);
+		/*
+		 * We need to inject necessary function pointers into the module. This
+		 * only needs to be done once -- even if the pointers are constant,
+		 * assigning them while another thread is executing the flows feels
+		 * like tempting fate.
+		 */
+		if ((lockerr = pthread_mutex_lock(&init_mutex)) != 0)
+		{
+			/* Should not happen... but don't continue if it does. */
+			Assert(false);
 
-		libpq_append_conn_error(conn, "failed to lock mutex (%d)", lockerr);
-		return 0;
-	}
+			libpq_append_conn_error(conn, "failed to lock mutex (%d)", lockerr);
+			return 0;
+		}
 
-	if (!initialized)
-	{
-		init(pg_g_threadlock,
+		if (!initialized)
+		{
+			init(pg_g_threadlock,
 #ifdef ENABLE_NLS
-			 libpq_gettext
+				 libpq_gettext
 #else
-			 NULL
+				 NULL
 #endif
-			);
+				);
 
-		initialized = true;
-	}
+			initialized = true;
+		}
 
-	pthread_mutex_unlock(&init_mutex);
+		pthread_mutex_unlock(&init_mutex);
+	}
 
 	return (start_flow(conn, request) == 0) ? 1 : -1;
 }
diff --git a/src/test/modules/oauth_validator/meson.build b/src/test/modules/oauth_validator/meson.build
index a6f937fd7d7..1d898270220 100644
--- a/src/test/modules/oauth_validator/meson.build
+++ b/src/test/modules/oauth_validator/meson.build
@@ -50,6 +50,7 @@ test_install_libs += magic_validator
 
 oauth_hook_client_sources = files(
   'oauth_hook_client.c',
+  'oauth_test_common.c',
 )
 
 if host_system == 'windows'
@@ -67,6 +68,19 @@ oauth_hook_client = executable('oauth_hook_client',
 )
 testprep_targets += oauth_hook_client
 
+oauth_flow = shared_module('oauth_flow',
+  files(
+    'oauth_flow.c',
+    'oauth_test_common.c',
+  ),
+  include_directories: [postgres_inc],
+  dependencies: [frontend_shlib_code, libpq],
+  kwargs: default_lib_args + {
+    'install': false,
+  },
+)
+testprep_targets += oauth_flow
+
 tests += {
   'name': 'oauth_validator',
   'sd': meson.current_source_dir(),
@@ -80,6 +94,7 @@ tests += {
       'PYTHON': python.full_path(),
       'with_libcurl': oauth_flow_supported ? 'yes' : 'no',
       'with_python': 'yes',
+      'flow_module_path': oauth_flow.full_path(),
     },
   },
 }
diff --git a/src/test/modules/oauth_validator/Makefile b/src/test/modules/oauth_validator/Makefile
index 05b9f06ed73..e22098dacb8 100644
--- a/src/test/modules/oauth_validator/Makefile
+++ b/src/test/modules/oauth_validator/Makefile
@@ -14,11 +14,13 @@ PGFILEDESC = "validator - test OAuth validator module"
 
 PROGRAM = oauth_hook_client
 PGAPPICON = win32
-OBJS = $(WIN32RES) oauth_hook_client.o
+OBJS = $(WIN32RES) oauth_hook_client.o oauth_test_common.o
 
 PG_CPPFLAGS = -I$(libpq_srcdir)
 PG_LIBS_INTERNAL += $(libpq_pgport)
 
+EXTRA_CLEAN = oauth_flow$(DLSUFFIX) oauth_flow.o
+
 NO_INSTALLCHECK = 1
 
 TAP_TESTS = 1
@@ -33,8 +35,14 @@ top_builddir = ../../../..
 include $(top_builddir)/src/Makefile.global
 include $(top_srcdir)/contrib/contrib-global.mk
 
+all: oauth_flow$(DLSUFFIX)
+
+oauth_flow$(DLSUFFIX): oauth_flow.o oauth_test_common.o
+	$(CC) $(CFLAGS) $^ $(LDFLAGS) $(libpq_pgport_shlib) $(LDFLAGS_SL) -shared -o $@
+
 export PYTHON
 export with_libcurl
 export with_python
+export flow_module_path := $(abs_top_builddir)/$(subdir)/oauth_flow$(DLSUFFIX)
 
 endif
diff --git a/src/test/modules/oauth_validator/oauth_test_common.h b/src/test/modules/oauth_validator/oauth_test_common.h
new file mode 100644
index 00000000000..33e72e30440
--- /dev/null
+++ b/src/test/modules/oauth_validator/oauth_test_common.h
@@ -0,0 +1,26 @@
+/*-------------------------------------------------------------------------
+ *
+ * oauth_test_common.h
+ *	  Shared functionality for oauth_hook_client and oauth_flow
+ *
+ * Copyright (c) 2025, PostgreSQL Global Development Group
+ *
+ *-------------------------------------------------------------------------
+ */
+
+#ifndef OAUTH_TEST_COMMON_H
+#define OAUTH_TEST_COMMON_H
+
+/*
+ * Only public headers can be here, since oauth_flow.c is trying to test only
+ * the public API.
+ */
+#include "libpq-fe.h"
+
+extern int	stress_async;		/* for oauth_hook_client */
+
+extern char *oauth_test_parse_argv(int argc, char *argv[], int for_plugin);
+extern int	oauth_test_authdata_hook(PGauthData type, PGconn *conn, void *data);
+extern int	oauth_test_start_flow(PGconn *conn, PGoauthBearerRequestV2 *request);
+
+#endif							/* OAUTH_TEST_COMMON_H */
diff --git a/src/test/modules/oauth_validator/oauth_flow.c b/src/test/modules/oauth_validator/oauth_flow.c
new file mode 100644
index 00000000000..8068a45ae29
--- /dev/null
+++ b/src/test/modules/oauth_validator/oauth_flow.c
@@ -0,0 +1,69 @@
+/*-------------------------------------------------------------------------
+ *
+ * oauth_flow.c
+ *	  Test plugin for clientside OAuth flows
+ *
+ * Copyright (c) 2025, PostgreSQL Global Development Group
+ *
+ *-------------------------------------------------------------------------
+ */
+
+#include <stdlib.h>
+#include <string.h>
+
+/* Since we want to test the public API, only include public headers here. */
+#include "libpq-fe.h"
+#include "libpq-oauth.h"
+#include "oauth_test_common.h"
+
+static void
+load_test_flags(void)
+{
+	int			argc;
+	char	  **argv;
+	char	   *env = getenv("OAUTH_TEST_FLAGS");
+	int			flag_count;
+	int			i;
+
+	if (!env || !env[0])
+	{
+		fprintf(stderr, "OAUTH_TEST_FLAGS must be set\n");
+		exit(1);
+	}
+
+	flag_count = 1;
+	for (char *c = env; *c; c++)
+	{
+		if (*c == '\x01')
+			flag_count++;
+	}
+
+	argc = flag_count + 1;
+	argv = malloc(sizeof(*argv) * (argc + 1));
+	if (!argv)
+	{
+		fprintf(stderr, "out of memory");
+		exit(1);
+	}
+
+	argv[0] = "[plugin test]";
+	for (i = 1; i < flag_count; i++)
+	{
+		argv[i] = env;
+
+		env = strchr(env, '\x01');
+		*env++ = '\0';
+	}
+	argv[flag_count] = env;
+	argv[argc] = NULL;
+
+	oauth_test_parse_argv(argc, argv, 1 /* plugin */ );
+}
+
+int
+pg_start_oauthbearer(PGconn *conn, PGoauthBearerRequestV2 *request)
+{
+	load_test_flags();
+
+	return oauth_test_start_flow(conn, request);
+}
diff --git a/src/test/modules/oauth_validator/oauth_hook_client.c b/src/test/modules/oauth_validator/oauth_hook_client.c
index b1a3b014079..e0b022a268c 100644
--- a/src/test/modules/oauth_validator/oauth_hook_client.c
+++ b/src/test/modules/oauth_validator/oauth_hook_client.c
@@ -18,144 +18,18 @@
 
 #include <sys/socket.h>
 
-#include "getopt_long.h"
 #include "libpq-fe.h"
-#include "pqexpbuffer.h"
 
-static int	handle_auth_data(PGauthData type, PGconn *conn, void *data);
-static PostgresPollingStatusType async_cb(PGconn *conn,
-										  PGoauthBearerRequest *req,
-										  pgsocket *altsock);
-static PostgresPollingStatusType misbehave_cb(PGconn *conn,
-											  PGoauthBearerRequest *req,
-											  pgsocket *altsock);
-
-static void
-usage(char *argv[])
-{
-	printf("usage: %s [flags] CONNINFO\n\n", argv[0]);
-
-	printf("recognized flags:\n");
-	printf("  -h, --help              show this message\n");
-	printf("  -v VERSION              select the hook API version (default 2)\n");
-	printf("  --expected-scope SCOPE  fail if received scopes do not match SCOPE\n");
-	printf("  --expected-uri URI      fail if received configuration link does not match URI\n");
-	printf("  --expected-issuer ISS   fail if received issuer does not match ISS (v2 only)\n");
-	printf("  --misbehave=MODE        have the hook fail required postconditions\n"
-		   "                          (MODEs: no-hook, fail-async, no-token, no-socket)\n");
-	printf("  --no-hook               don't install OAuth hooks\n");
-	printf("  --hang-forever          don't ever return a token (combine with connect_timeout)\n");
-	printf("  --token TOKEN           use the provided TOKEN value\n");
-	printf("  --error ERRMSG          fail instead, with the given ERRMSG (v2 only)\n");
-	printf("  --stress-async          busy-loop on PQconnectPoll rather than polling\n");
-}
-
-/* --options */
-static bool no_hook = false;
-static bool hang_forever = false;
-static bool stress_async = false;
-static const char *expected_uri = NULL;
-static const char *expected_issuer = NULL;
-static const char *expected_scope = NULL;
-static const char *misbehave_mode = NULL;
-static char *token = NULL;
-static char *errmsg = NULL;
-static int	hook_version = PQAUTHDATA_OAUTH_BEARER_TOKEN_V2;
+#include "oauth_test_common.h"
 
 int
 main(int argc, char *argv[])
 {
-	static const struct option long_options[] = {
-		{"help", no_argument, NULL, 'h'},
-
-		{"expected-scope", required_argument, NULL, 1000},
-		{"expected-uri", required_argument, NULL, 1001},
-		{"no-hook", no_argument, NULL, 1002},
-		{"token", required_argument, NULL, 1003},
-		{"hang-forever", no_argument, NULL, 1004},
-		{"misbehave", required_argument, NULL, 1005},
-		{"stress-async", no_argument, NULL, 1006},
-		{"expected-issuer", required_argument, NULL, 1007},
-		{"error", required_argument, NULL, 1008},
-		{0}
-	};
-
-	const char *conninfo;
+	const char *conninfo = oauth_test_parse_argv(argc, argv, 0 /* hook */ );
 	PGconn	   *conn;
-	int			c;
-
-	while ((c = getopt_long(argc, argv, "hv:", long_options, NULL)) != -1)
-	{
-		switch (c)
-		{
-			case 'h':
-				usage(argv);
-				return 0;
-
-			case 'v':
-				if (strcmp(optarg, "1") == 0)
-					hook_version = PQAUTHDATA_OAUTH_BEARER_TOKEN;
-				else if (strcmp(optarg, "2") == 0)
-					hook_version = PQAUTHDATA_OAUTH_BEARER_TOKEN_V2;
-				else
-				{
-					usage(argv);
-					return 1;
-				}
-				break;
-
-			case 1000:			/* --expected-scope */
-				expected_scope = optarg;
-				break;
-
-			case 1001:			/* --expected-uri */
-				expected_uri = optarg;
-				break;
-
-			case 1002:			/* --no-hook */
-				no_hook = true;
-				break;
-
-			case 1003:			/* --token */
-				token = optarg;
-				break;
-
-			case 1004:			/* --hang-forever */
-				hang_forever = true;
-				break;
-
-			case 1005:			/* --misbehave */
-				misbehave_mode = optarg;
-				break;
-
-			case 1006:			/* --stress-async */
-				stress_async = true;
-				break;
-
-			case 1007:			/* --expected-issuer */
-				expected_issuer = optarg;
-				break;
-
-			case 1008:			/* --error */
-				errmsg = optarg;
-				break;
-
-			default:
-				usage(argv);
-				return 1;
-		}
-	}
-
-	if (argc != optind + 1)
-	{
-		usage(argv);
-		return 1;
-	}
-
-	conninfo = argv[optind];
 
 	/* Set up our OAuth hooks. */
-	PQsetAuthDataHook(handle_auth_data);
+	PQsetAuthDataHook(oauth_test_authdata_hook);
 
 	/* Connect. (All the actual work is in the hook.) */
 	if (stress_async)
@@ -193,190 +67,3 @@ main(int argc, char *argv[])
 	PQfinish(conn);
 	return 0;
 }
-
-/*
- * PQauthDataHook implementation. Replaces the default client flow by handling
- * PQAUTHDATA_OAUTH_BEARER_TOKEN[_V2].
- */
-static int
-handle_auth_data(PGauthData type, PGconn *conn, void *data)
-{
-	PGoauthBearerRequest *req;
-	PGoauthBearerRequestV2 *req2 = NULL;
-
-	Assert(hook_version == PQAUTHDATA_OAUTH_BEARER_TOKEN ||
-		   hook_version == PQAUTHDATA_OAUTH_BEARER_TOKEN_V2);
-
-	if (no_hook || type != hook_version)
-		return 0;
-
-	req = data;
-	if (type == PQAUTHDATA_OAUTH_BEARER_TOKEN_V2)
-		req2 = data;
-
-	if (hang_forever)
-	{
-		/* Start asynchronous processing. */
-		req->async = async_cb;
-		return 1;
-	}
-
-	if (misbehave_mode)
-	{
-		if (strcmp(misbehave_mode, "no-hook") != 0)
-			req->async = misbehave_cb;
-		return 1;
-	}
-
-	if (expected_uri)
-	{
-		if (!req->openid_configuration)
-		{
-			fprintf(stderr, "expected URI \"%s\", got NULL\n", expected_uri);
-			return -1;
-		}
-
-		if (strcmp(expected_uri, req->openid_configuration) != 0)
-		{
-			fprintf(stderr, "expected URI \"%s\", got \"%s\"\n", expected_uri, req->openid_configuration);
-			return -1;
-		}
-	}
-
-	if (expected_scope)
-	{
-		if (!req->scope)
-		{
-			fprintf(stderr, "expected scope \"%s\", got NULL\n", expected_scope);
-			return -1;
-		}
-
-		if (strcmp(expected_scope, req->scope) != 0)
-		{
-			fprintf(stderr, "expected scope \"%s\", got \"%s\"\n", expected_scope, req->scope);
-			return -1;
-		}
-	}
-
-	if (expected_issuer)
-	{
-		if (!req2)
-		{
-			fprintf(stderr, "--expected-issuer cannot be combined with -v1\n");
-			return -1;
-		}
-
-		if (!req2->issuer)
-		{
-			fprintf(stderr, "expected issuer \"%s\", got NULL\n", expected_issuer);
-			return -1;
-		}
-
-		if (strcmp(expected_issuer, req2->issuer) != 0)
-		{
-			fprintf(stderr, "expected issuer \"%s\", got \"%s\"\n", expected_issuer, req2->issuer);
-			return -1;
-		}
-	}
-
-	if (errmsg)
-	{
-		if (token)
-		{
-			fprintf(stderr, "--error cannot be combined with --token\n");
-			return -1;
-		}
-		else if (!req2)
-		{
-			fprintf(stderr, "--error cannot be combined with -v1\n");
-			return -1;
-		}
-
-		appendPQExpBufferStr(req2->error, errmsg);
-		return -1;
-	}
-
-	req->token = token;
-	return 1;
-}
-
-static PostgresPollingStatusType
-async_cb(PGconn *conn, PGoauthBearerRequest *req, pgsocket *altsock)
-{
-	if (hang_forever)
-	{
-		/*
-		 * This code tests that nothing is interfering with libpq's handling
-		 * of connect_timeout.
-		 */
-		static pgsocket sock = PGINVALID_SOCKET;
-
-		if (sock == PGINVALID_SOCKET)
-		{
-			/* First call. Create an unbound socket to wait on. */
-#ifdef WIN32
-			WSADATA		wsaData;
-			int			err;
-
-			err = WSAStartup(MAKEWORD(2, 2), &wsaData);
-			if (err)
-			{
-				perror("WSAStartup failed");
-				return PGRES_POLLING_FAILED;
-			}
-#endif
-			sock = socket(AF_INET, SOCK_DGRAM, 0);
-			if (sock == PGINVALID_SOCKET)
-			{
-				perror("failed to create datagram socket");
-				return PGRES_POLLING_FAILED;
-			}
-		}
-
-		/* Make libpq wait on the (unreadable) socket. */
-		*altsock = sock;
-		return PGRES_POLLING_READING;
-	}
-
-	req->token = token;
-	return PGRES_POLLING_OK;
-}
-
-static PostgresPollingStatusType
-misbehave_cb(PGconn *conn, PGoauthBearerRequest *req, pgsocket *altsock)
-{
-	if (strcmp(misbehave_mode, "fail-async") == 0)
-	{
-		/* Just fail "normally". */
-		if (errmsg)
-		{
-			PGoauthBearerRequestV2 *req2;
-
-			if (hook_version == PQAUTHDATA_OAUTH_BEARER_TOKEN)
-			{
-				fprintf(stderr, "--error cannot be combined with -v1\n");
-				exit(1);
-			}
-
-			req2 = (PGoauthBearerRequestV2 *) req;
-			appendPQExpBufferStr(req2->error, errmsg);
-		}
-
-		return PGRES_POLLING_FAILED;
-	}
-	else if (strcmp(misbehave_mode, "no-token") == 0)
-	{
-		/* Callbacks must assign req->token before returning OK. */
-		return PGRES_POLLING_OK;
-	}
-	else if (strcmp(misbehave_mode, "no-socket") == 0)
-	{
-		/* Callbacks must assign *altsock before asking for polling. */
-		return PGRES_POLLING_READING;
-	}
-	else
-	{
-		fprintf(stderr, "unrecognized --misbehave mode: %s\n", misbehave_mode);
-		exit(1);
-	}
-}
diff --git a/src/test/modules/oauth_validator/oauth_test_common.c b/src/test/modules/oauth_validator/oauth_test_common.c
new file mode 100644
index 00000000000..f2a5b180a65
--- /dev/null
+++ b/src/test/modules/oauth_validator/oauth_test_common.c
@@ -0,0 +1,374 @@
+/*-------------------------------------------------------------------------
+ *
+ * oauth_test_common.c
+ *	  Shared functionality for oauth_hook_client and oauth_flow
+ *
+ * Copyright (c) 2025, PostgreSQL Global Development Group
+ *
+ *-------------------------------------------------------------------------
+ */
+
+#include "postgres_fe.h"
+
+#include <sys/socket.h>
+
+#include "getopt_long.h"
+#include "libpq-fe.h"
+#include "pqexpbuffer.h"
+
+#include "oauth_test_common.h"
+
+static PostgresPollingStatusType async_cb(PGconn *conn,
+										  PGoauthBearerRequest *req,
+										  pgsocket *altsock);
+static PostgresPollingStatusType misbehave_cb(PGconn *conn,
+											  PGoauthBearerRequest *req,
+											  pgsocket *altsock);
+
+/* --options */
+static bool no_hook = false;
+static bool hang_forever = false;
+static const char *expected_uri = NULL;
+static const char *expected_issuer = NULL;
+static const char *expected_scope = NULL;
+static const char *misbehave_mode = NULL;
+static char *token = NULL;
+static char *errmsg = NULL;
+static int	hook_version = PQAUTHDATA_OAUTH_BEARER_TOKEN_V2;
+
+/*
+ * XXX: stress_async is exported for the benefit of oauth_hook_client. Since
+ * we only use public headers (libpq-fe.h) for oauth_flow, it needs to be an int
+ * rather than a bool.
+ */
+int			stress_async = false;
+
+static void
+usage(char *argv[])
+{
+	printf("usage: %s [flags] CONNINFO\n\n", argv[0]);
+
+	printf("recognized flags:\n");
+	printf("  -h, --help              show this message\n");
+	printf("  -v VERSION              select the hook API version (default 2)\n");
+	printf("  --expected-scope SCOPE  fail if received scopes do not match SCOPE\n");
+	printf("  --expected-uri URI      fail if received configuration link does not match URI\n");
+	printf("  --expected-issuer ISS   fail if received issuer does not match ISS (v2 only)\n");
+	printf("  --misbehave=MODE        have the hook fail required postconditions\n"
+		   "                          (MODEs: no-hook, fail-async, no-token, no-socket)\n");
+	printf("  --no-hook               don't install OAuth hooks\n");
+	printf("  --hang-forever          don't ever return a token (combine with connect_timeout)\n");
+	printf("  --token TOKEN           use the provided TOKEN value\n");
+	printf("  --error ERRMSG          fail instead, with the given ERRMSG (v2 only)\n");
+	printf("  --stress-async          busy-loop on PQconnectPoll rather than polling\n");
+}
+
+char *
+oauth_test_parse_argv(int argc, char *argv[], int for_plugin)
+{
+	static const struct option long_options[] = {
+		{"help", no_argument, NULL, 'h'},
+
+		{"expected-scope", required_argument, NULL, 1000},
+		{"expected-uri", required_argument, NULL, 1001},
+		{"no-hook", no_argument, NULL, 1002},
+		{"token", required_argument, NULL, 1003},
+		{"hang-forever", no_argument, NULL, 1004},
+		{"misbehave", required_argument, NULL, 1005},
+		{"stress-async", no_argument, NULL, 1006},
+		{"expected-issuer", required_argument, NULL, 1007},
+		{"error", required_argument, NULL, 1008},
+		{0}
+	};
+
+	int			c;
+
+	if (for_plugin)
+	{
+		/* The "real" argv has already been parsed. Reset optind. */
+		optind = 1;
+	}
+
+	while ((c = getopt_long(argc, argv, "hv:", long_options, NULL)) != -1)
+	{
+		switch (c)
+		{
+			case 'h':
+				usage(argv);
+				exit(0);
+
+			case 'v':
+				if (strcmp(optarg, "1") == 0)
+					hook_version = PQAUTHDATA_OAUTH_BEARER_TOKEN;
+				else if (strcmp(optarg, "2") == 0)
+					hook_version = PQAUTHDATA_OAUTH_BEARER_TOKEN_V2;
+				else
+				{
+					usage(argv);
+					exit(1);
+				}
+				break;
+
+			case 1000:			/* --expected-scope */
+				expected_scope = optarg;
+				break;
+
+			case 1001:			/* --expected-uri */
+				expected_uri = optarg;
+				break;
+
+			case 1002:			/* --no-hook */
+				no_hook = true;
+				break;
+
+			case 1003:			/* --token */
+				token = optarg;
+				break;
+
+			case 1004:			/* --hang-forever */
+				hang_forever = true;
+				break;
+
+			case 1005:			/* --misbehave */
+				misbehave_mode = optarg;
+				break;
+
+			case 1006:			/* --stress-async */
+				stress_async = true;
+				break;
+
+			case 1007:			/* --expected-issuer */
+				expected_issuer = optarg;
+				break;
+
+			case 1008:			/* --error */
+				errmsg = optarg;
+				break;
+
+			default:
+				usage(argv);
+				exit(1);
+		}
+	}
+
+	if (argc != (for_plugin ? optind : optind + 1))
+	{
+		usage(argv);
+		exit(1);
+	}
+
+	return argv[optind];
+}
+
+/*
+ * PQauthDataHook implementation. Replaces the default client flow by handling
+ * PQAUTHDATA_OAUTH_BEARER_TOKEN[_V2].
+ */
+int
+oauth_test_authdata_hook(PGauthData type, PGconn *conn, void *data)
+{
+	PGoauthBearerRequest *req;
+	PGoauthBearerRequestV2 *req2 = NULL;
+
+	Assert(hook_version == PQAUTHDATA_OAUTH_BEARER_TOKEN ||
+		   hook_version == PQAUTHDATA_OAUTH_BEARER_TOKEN_V2);
+
+	if (no_hook || type != hook_version)
+		return 0;
+
+	req = data;
+	if (type == PQAUTHDATA_OAUTH_BEARER_TOKEN_V2)
+		req2 = data;
+
+	if (hang_forever)
+	{
+		/* Start asynchronous processing. */
+		req->async = async_cb;
+		return 1;
+	}
+
+	if (misbehave_mode)
+	{
+		if (strcmp(misbehave_mode, "no-hook") != 0)
+			req->async = misbehave_cb;
+		return 1;
+	}
+
+	if (expected_uri)
+	{
+		if (!req->openid_configuration)
+		{
+			fprintf(stderr, "expected URI \"%s\", got NULL\n", expected_uri);
+			return -1;
+		}
+
+		if (strcmp(expected_uri, req->openid_configuration) != 0)
+		{
+			fprintf(stderr, "expected URI \"%s\", got \"%s\"\n", expected_uri, req->openid_configuration);
+			return -1;
+		}
+	}
+
+	if (expected_scope)
+	{
+		if (!req->scope)
+		{
+			fprintf(stderr, "expected scope \"%s\", got NULL\n", expected_scope);
+			return -1;
+		}
+
+		if (strcmp(expected_scope, req->scope) != 0)
+		{
+			fprintf(stderr, "expected scope \"%s\", got \"%s\"\n", expected_scope, req->scope);
+			return -1;
+		}
+	}
+
+	if (expected_issuer)
+	{
+		if (!req2)
+		{
+			fprintf(stderr, "--expected-issuer cannot be combined with -v1\n");
+			return -1;
+		}
+
+		if (!req2->issuer)
+		{
+			fprintf(stderr, "expected issuer \"%s\", got NULL\n", expected_issuer);
+			return -1;
+		}
+
+		if (strcmp(expected_issuer, req2->issuer) != 0)
+		{
+			fprintf(stderr, "expected issuer \"%s\", got \"%s\"\n", expected_issuer, req2->issuer);
+			return -1;
+		}
+	}
+
+	if (errmsg)
+	{
+		if (token)
+		{
+			fprintf(stderr, "--error cannot be combined with --token\n");
+			return -1;
+		}
+		else if (!req2)
+		{
+			fprintf(stderr, "--error cannot be combined with -v1\n");
+			return -1;
+		}
+
+		appendPQExpBufferStr(req2->error, errmsg);
+		return -1;
+	}
+
+	req->token = token;
+	return 1;
+}
+
+/*
+ * Sets up a request for a plugin module (pg_start_oauthbearer()) rather than
+ * using the hook.
+ */
+int
+oauth_test_start_flow(PGconn *conn, PGoauthBearerRequestV2 *request)
+{
+	int			ret;
+
+	/*
+	 * We can still defer to the hook above to avoid copying code; we just
+	 * have to translate the return value.
+	 */
+	ret = oauth_test_authdata_hook(PQAUTHDATA_OAUTH_BEARER_TOKEN_V2, conn,
+								   request);
+
+	if (ret == 0)
+	{
+		/* This is a bug in the test. */
+		fprintf(stderr, "plugin tests cannot make use of -v1 or --no-hook\n");
+		exit(1);
+	}
+
+	return (ret == 1) ? 0 : -1;
+}
+
+static PostgresPollingStatusType
+async_cb(PGconn *conn, PGoauthBearerRequest *req, pgsocket *altsock)
+{
+	if (hang_forever)
+	{
+		/*
+		 * This code tests that nothing is interfering with libpq's handling
+		 * of connect_timeout.
+		 */
+		static pgsocket sock = PGINVALID_SOCKET;
+
+		if (sock == PGINVALID_SOCKET)
+		{
+			/* First call. Create an unbound socket to wait on. */
+#ifdef WIN32
+			WSADATA		wsaData;
+			int			err;
+
+			err = WSAStartup(MAKEWORD(2, 2), &wsaData);
+			if (err)
+			{
+				perror("WSAStartup failed");
+				return PGRES_POLLING_FAILED;
+			}
+#endif
+			sock = socket(AF_INET, SOCK_DGRAM, 0);
+			if (sock == PGINVALID_SOCKET)
+			{
+				perror("failed to create datagram socket");
+				return PGRES_POLLING_FAILED;
+			}
+		}
+
+		/* Make libpq wait on the (unreadable) socket. */
+		*altsock = sock;
+		return PGRES_POLLING_READING;
+	}
+
+	req->token = token;
+	return PGRES_POLLING_OK;
+}
+
+static PostgresPollingStatusType
+misbehave_cb(PGconn *conn, PGoauthBearerRequest *req, pgsocket *altsock)
+{
+	if (strcmp(misbehave_mode, "fail-async") == 0)
+	{
+		/* Just fail "normally". */
+		if (errmsg)
+		{
+			PGoauthBearerRequestV2 *req2;
+
+			if (hook_version == PQAUTHDATA_OAUTH_BEARER_TOKEN)
+			{
+				fprintf(stderr, "--error cannot be combined with -v1\n");
+				exit(1);
+			}
+
+			req2 = (PGoauthBearerRequestV2 *) req;
+			appendPQExpBufferStr(req2->error, errmsg);
+		}
+
+		return PGRES_POLLING_FAILED;
+	}
+	else if (strcmp(misbehave_mode, "no-token") == 0)
+	{
+		/* Callbacks must assign req->token before returning OK. */
+		return PGRES_POLLING_OK;
+	}
+	else if (strcmp(misbehave_mode, "no-socket") == 0)
+	{
+		/* Callbacks must assign *altsock before asking for polling. */
+		return PGRES_POLLING_READING;
+	}
+	else
+	{
+		fprintf(stderr, "unrecognized --misbehave mode: %s\n", misbehave_mode);
+		exit(1);
+	}
+}
diff --git a/src/test/modules/oauth_validator/t/002_client.pl b/src/test/modules/oauth_validator/t/002_client.pl
index f878a23d201..b23e579f1fd 100644
--- a/src/test/modules/oauth_validator/t/002_client.pl
+++ b/src/test/modules/oauth_validator/t/002_client.pl
@@ -1,6 +1,6 @@
 #
 # Exercises the API for custom OAuth client flows, using the oauth_hook_client
-# test driver.
+# test driver and the oauth_flow custom plugin.
 #
 # Copyright (c) 2021-2025, PostgreSQL Global Development Group
 #
@@ -20,6 +20,10 @@ if (!$ENV{PG_TEST_EXTRA} || $ENV{PG_TEST_EXTRA} !~ /\boauth\b/)
 	  'Potentially unsafe test oauth not enabled in PG_TEST_EXTRA';
 }
 
+my $plugin_supported = (
+		 check_pg_config("#define HAVE_SYS_EVENT_H 1")
+	  or check_pg_config("#define HAVE_SYS_EPOLL_H 1"));
+
 #
 # Cluster Setup
 #
@@ -72,6 +76,8 @@ sub test
 		$flags = $params{flags};
 	}
 
+	# First run the oauth_hook_client, which uses PQauthDataHook to insert a new
+	# OAuth flow.
 	my @cmd = ("oauth_hook_client", @{$flags}, $common_connstr);
 	note "running '" . join("' '", @cmd) . "'";
 
@@ -103,6 +109,37 @@ sub test
 		$node->log_check("$test_name: log matches",
 			$log_start, log_like => $params{log_like});
 	}
+
+  SKIP:
+	{
+		last SKIP if $params{hook_only};
+		skip "OAuth modules are not supported on this platform"
+		  unless $plugin_supported;
+
+		# Run the same test with psql itself, loading the oauth_flow.so module.
+		local $ENV{PGOAUTHMODULE} = $ENV{flow_module_path};
+
+		# Flags are passed to the module via OAUTH_TEST_FLAGS, with 0x01 as a
+		# separator.
+		local $ENV{OAUTH_TEST_FLAGS} = join("\x01", @{$flags});
+
+		if ($params{expect_success})
+		{
+			$node->connect_ok(
+				$common_connstr,
+				"[plugin flow] $test_name",
+				expected_stderr => $params{expected_stderr},
+				log_like => $params{log_like});
+		}
+		else
+		{
+			$node->connect_fails(
+				$common_connstr,
+				"[plugin flow] $test_name",
+				expected_stderr => $params{expected_stderr},
+				log_like => $params{log_like});
+		}
+	}
 }
 
 test(
@@ -119,6 +156,7 @@ test(
 # Make sure the v1 hook continues to work.
 test(
 	"v1 synchronous hook can provide a token",
+	hook_only => 1,    # plugins don't support API v1
 	flags => [
 		"-v1",
 		"--token" => "my-token-v1",
@@ -133,6 +171,7 @@ if ($ENV{with_libcurl} ne 'yes')
 	# libpq should help users out if no OAuth support is built in.
 	test(
 		"fails without custom hook installed",
+		hook_only => 1,    # plugins can't use --no-hook
 		flags => ["--no-hook"],
 		expected_stderr =>
 		  qr/no OAuth flows are available \(try installing the libpq-oauth package\)/
-- 
2.34.1