v5-0003-libpq-Introduce-PQAUTHDATA_OAUTH_BEARER_TOKEN_V2.patch

From 751df22f99de2075f276750f2bccf6548eb63a76 Mon Sep 17 00:00:00 2001
From: Jacob Champion <jacob.champion@enterprisedb.com>
Date: Mon, 1 Dec 2025 15:07:26 -0800
Subject: [PATCH v5 3/7] libpq: Introduce PQAUTHDATA_OAUTH_BEARER_TOKEN_V2

For the libpq-oauth module to eventually make use of the
PGoauthBearerRequest API, it needs some additional functionality: the
derived Issuer ID for the authorization server needs to be provided by
libpq, and error messages need to be built without relying on PGconn
internals. These features seem useful for application hooks, too, so
that they don't each have to reinvent the wheel.

The original plan was for additions to PGoauthBearerRequest to be made
without a version bump to the PGauthData type, and then applications
would simply check a LIBPQ_HAS_* macro at compile time to decide whether
they could use the new features. That theoretically works for
applications linked against libpq, since it's not safe to downgrade
libpq from the version you've compiled against.

That strategy won't work for plugins, though, due to a complication
first noticed during the libpq-oauth module split: it's normal for a
plugin on disk to be *newer* than the libpq that's loading it, because
you might have upgraded your installation while an application was
running. (In other words, a plugin architecture causes the compile-time
and run-time dependency arrows to point in opposite directions, so
plugins won't be able to rely on the LIBPQ_HAS_* macros to determine
what APIs are available to them.)

Instead, add a new struct which extends the "v1" PGoauthBearerRequest.
When an application hook receives PQAUTHDATA_OAUTH_BEARER_TOKEN_V2, it
may safely downcast the request pointer it receives in its callbacks to
make use of the new functionality. libpq will first try the new version,
then fall back to the old before giving up.

Reviewed-by: Chao Li <li.evan.chao@gmail.com>
Reviewed-by: Zsolt Parragi <zsolt.parragi@percona.com>
Discussion: https://postgr.es/m/CAOYmi%2BmrGg%2Bn_X2MOLgeWcj3v_M00gR8uz_D7mM8z%3DdX1JYVbg%40mail.gmail.com
---
 doc/src/sgml/libpq.sgml                       | 83 ++++++++++++++++
 src/tools/pgindent/typedefs.list              |  1 +
 src/interfaces/libpq/libpq-fe.h               | 31 +++++-
 src/interfaces/libpq/fe-auth-oauth.c          | 87 +++++++++++------
 .../oauth_validator/oauth_hook_client.c       | 97 ++++++++++++++++++-
 .../modules/oauth_validator/t/002_client.pl   | 43 +++++++-
 6 files changed, 306 insertions(+), 36 deletions(-)

diff --git a/doc/src/sgml/libpq.sgml b/doc/src/sgml/libpq.sgml
index 57349ef8f84..d64f63433c0 100644
--- a/doc/src/sgml/libpq.sgml
+++ b/doc/src/sgml/libpq.sgml
@@ -10362,6 +10362,7 @@ PQauthDataHook_type PQgetAuthDataHook(void);
         <indexterm><primary>PQAUTHDATA_PROMPT_OAUTH_DEVICE</primary></indexterm>
        </term>
        <listitem>
+        <para><emphasis>Available in PostgreSQL 18 and later.</emphasis></para>
         <para>
          Replaces the default user prompt during the builtin device
          authorization client flow. <replaceable>data</replaceable> points to
@@ -10414,6 +10415,7 @@ typedef struct _PGpromptOAuthDevice
         <indexterm><primary>PQAUTHDATA_OAUTH_BEARER_TOKEN</primary></indexterm>
        </term>
        <listitem>
+        <para><emphasis>Available in PostgreSQL 18 and later.</emphasis></para>
         <para>
          Adds a custom implementation of a flow, replacing the builtin flow if
          it is <link linkend="configure-option-with-libcurl">installed</link>.
@@ -10421,6 +10423,13 @@ typedef struct _PGpromptOAuthDevice
          user/issuer/scope combination, if one is available without blocking, or
          else set up an asynchronous callback to retrieve one.
         </para>
+        <note>
+         <para>
+          For <productname>PostgreSQL</productname> releases 19 and later,
+          applications should prefer
+          <link linkend="libpq-oauth-authdata-oauth-bearer-token-v2"><literal>PQAUTHDATA_OAUTH_BEARER_TOKEN_V2</literal></link>.
+         </para>
+        </note>
         <para>
          <replaceable>data</replaceable> points to an instance
          of <symbol>PGoauthBearerRequest</symbol>, which should be filled in
@@ -10516,6 +10525,80 @@ typedef struct PGoauthBearerRequest
         </para>
        </listitem>
       </varlistentry>
+
+      <varlistentry id="libpq-oauth-authdata-oauth-bearer-token-v2">
+       <term>
+        <symbol>PQAUTHDATA_OAUTH_BEARER_TOKEN_V2</symbol>
+        <indexterm>
+         <primary>PQAUTHDATA_OAUTH_BEARER_TOKEN_V2</primary>
+         <secondary>PQAUTHDATA_OAUTH_BEARER_TOKEN</secondary>
+        </indexterm>
+       </term>
+       <listitem>
+        <para><emphasis>Available in PostgreSQL 19 and later.</emphasis></para>
+        <para>
+         Provides all the functionality of
+         <link linkend="libpq-oauth-authdata-oauth-bearer-token"><literal>PQAUTHDATA_OAUTH_BEARER_TOKEN</literal></link>,
+         and adds the ability to set custom error messages and inspect the
+         OAuth issuer ID that the client expects to use.
+        </para>
+        <para>
+         <replaceable>data</replaceable> points to an instance
+         of <symbol>PGoauthBearerRequestV2</symbol>, which should be filled in
+         by the implementation:
+<synopsis>
+typedef struct
+{
+    PGoauthBearerRequest v1;    /* see the PGoauthBearerRequest struct, above */
+
+    /* Hook inputs (constant across all calls) */
+    const char *issuer;            /* the issuer identifier (RFC 9207) in use */
+
+    /* Hook outputs */
+    const char *error;             /* hook-defined error message */
+} PGoauthBearerRequestV2;
+</synopsis>
+        </para>
+        <para>
+         When a <literal>PQAUTHDATA_OAUTH_BEARER_TOKEN_V2</literal> hook is in
+         use, <application>libpq</application> additionally guarantees that the
+         <literal>request</literal> pointer that is provided to the
+         <replaceable>v1.async</replaceable> and <replaceable>v1.cleanup</replaceable>
+         callbacks may be safely cast to <literal>(PGoauthBearerRequestV2&nbsp;*)</literal>.
+         Implementations must otherwise follow the v1 API, as described above,
+         using the <replaceable>v1</replaceable> struct member.
+        </para>
+        <warning>
+         <para>
+          Casting to <literal>(PGoauthBearerRequestV2&nbsp;*)</literal> is
+          <emphasis>only</emphasis> safe when the hook type is
+          <literal>PQAUTHDATA_OAUTH_BEARER_TOKEN_V2</literal>. Applications may
+          crash or misbehave if a hook attempts to use v2 features with the v1
+          hook type.
+         </para>
+        </warning>
+        <para>
+         In addition to the functionality of the version 1 API, the v2 struct
+         provides an additional input and output for the hook:
+        </para>
+        <para>
+         <replaceable>issuer</replaceable> contains the issuer identifier, as
+         defined in <ulink url="https://datatracker.ietf.org/doc/html/rfc9207">RFC 9207</ulink>,
+         that is in use for the current connection. This identifier is
+         determined by <xref linkend="libpq-connect-oauth-issuer"/>.
+         To avoid mix-up attacks, custom flows should ensure that any discovery
+         metadata provided by the authorization server matches this issuer ID.
+        </para>
+        <para>
+         <replaceable>error</replaceable> may be set to point to a custom
+         error message when a flow fails. The message will be included as part
+         of <link linkend="libpq-PQerrorMessage"><literal>PQerrorMessage()</literal></link>.
+         Hooks must free any error message allocations during the
+         <replaceable>v1.cleanup</replaceable> callback.
+        </para>
+       </listitem>
+      </varlistentry>
+
      </variablelist>
     </para>
    </sect3>
diff --git a/src/tools/pgindent/typedefs.list b/src/tools/pgindent/typedefs.list
index 77e3c04144e..3250564d4ff 100644
--- a/src/tools/pgindent/typedefs.list
+++ b/src/tools/pgindent/typedefs.list
@@ -1926,6 +1926,7 @@ PGdataValue
 PGlobjfuncs
 PGnotify
 PGoauthBearerRequest
+PGoauthBearerRequestV2
 PGpipelineStatus
 PGpromptOAuthDevice
 PGresAttDesc
diff --git a/src/interfaces/libpq/libpq-fe.h b/src/interfaces/libpq/libpq-fe.h
index 1b46032b6f9..f06e7a972c3 100644
--- a/src/interfaces/libpq/libpq-fe.h
+++ b/src/interfaces/libpq/libpq-fe.h
@@ -66,6 +66,8 @@ extern "C"
 /* Features added in PostgreSQL v19: */
 /* Indicates presence of PQgetThreadLock */
 #define LIBPQ_HAS_GET_THREAD_LOCK 1
+/* Indicates presence of the PQAUTHDATA_OAUTH_BEARER_TOKEN_V2 authdata hook */
+#define LIBPQ_HAS_OAUTH_BEARER_TOKEN_V2 1
 
 /*
  * Option flags for PQcopyResult
@@ -197,7 +199,9 @@ typedef enum
 {
 	PQAUTHDATA_PROMPT_OAUTH_DEVICE, /* user must visit a device-authorization
 									 * URL */
-	PQAUTHDATA_OAUTH_BEARER_TOKEN,	/* server requests an OAuth Bearer token */
+	PQAUTHDATA_OAUTH_BEARER_TOKEN,	/* server requests an OAuth Bearer token
+									 * (v2 is preferred; see below) */
+	PQAUTHDATA_OAUTH_BEARER_TOKEN_V2,	/* newest API for OAuth Bearer tokens */
 } PGauthData;
 
 /* PGconn encapsulates a connection to the backend.
@@ -735,6 +739,7 @@ extern int	PQenv2encoding(void);
 
 /* === in fe-auth.c === */
 
+/* Authdata for PQAUTHDATA_PROMPT_OAUTH_DEVICE */
 typedef struct _PGpromptOAuthDevice
 {
 	const char *verification_uri;	/* verification URI to visit */
@@ -755,6 +760,7 @@ typedef struct _PGpromptOAuthDevice
 #define PQ_SOCKTYPE int
 #endif
 
+/* Authdata for PQAUTHDATA_OAUTH_BEARER_TOKEN */
 typedef struct PGoauthBearerRequest
 {
 	/* Hook inputs (constant across all calls) */
@@ -788,7 +794,8 @@ typedef struct PGoauthBearerRequest
 
 	/*
 	 * Callback to clean up custom allocations. A hook implementation may use
-	 * this to free request->token and any resources in request->user.
+	 * this to free request->token and any resources in request->user. V2
+	 * implementations should additionally free request->error, if set.
 	 *
 	 * This is technically optional, but highly recommended, because there is
 	 * no other indication as to when it is safe to free the token.
@@ -813,6 +820,26 @@ typedef struct PGoauthBearerRequest
 
 #undef PQ_SOCKTYPE
 
+/* Authdata for PQAUTHDATA_OAUTH_BEARER_TOKEN_V2 */
+typedef struct
+{
+	PGoauthBearerRequest v1;	/* see the PGoauthBearerRequest struct, above */
+
+	/* Hook inputs (constant across all calls) */
+	const char *issuer;			/* the issuer identifier (RFC 9207) in use, as
+								 * derived from the connection's oauth_issuer */
+
+	/* Hook outputs */
+
+	/*
+	 * Hook-defined error message which will be included in the connection's
+	 * PQerrorMessage() output when the flow fails. libpq does not take
+	 * ownership of this pointer; any allocations should be freed during the
+	 * cleanup callback.
+	 */
+	const char *error;
+} PGoauthBearerRequestV2;
+
 extern char *PQencryptPassword(const char *passwd, const char *user);
 extern char *PQencryptPasswordConn(PGconn *conn, const char *passwd, const char *user, const char *algorithm);
 extern PGresult *PQchangePassword(PGconn *conn, const char *user, const char *passwd);
diff --git a/src/interfaces/libpq/fe-auth-oauth.c b/src/interfaces/libpq/fe-auth-oauth.c
index 67879d64b39..b64d05a13e3 100644
--- a/src/interfaces/libpq/fe-auth-oauth.c
+++ b/src/interfaces/libpq/fe-auth-oauth.c
@@ -675,6 +675,25 @@ cleanup:
 	return success;
 }
 
+/*
+ * Helper for handling flow failures. If anything was put into request->error,
+ * it's added to conn->errorMessage here.
+ */
+static void
+report_user_flow_error(PGconn *conn, const PGoauthBearerRequestV2 *request)
+{
+	appendPQExpBufferStr(&conn->errorMessage,
+						 libpq_gettext("user-defined OAuth flow failed"));
+
+	if (request->error)
+	{
+		appendPQExpBufferStr(&conn->errorMessage, ": ");
+		appendPQExpBufferStr(&conn->errorMessage, request->error);
+	}
+
+	appendPQExpBufferChar(&conn->errorMessage, '\n');
+}
+
 /*
  * Callback implementation of conn->async_auth() for a user-defined OAuth flow.
  * Delegates the retrieval of the token to the application's async callback.
@@ -687,20 +706,23 @@ static PostgresPollingStatusType
 run_user_oauth_flow(PGconn *conn)
 {
 	fe_oauth_state *state = conn->sasl_state;
-	PGoauthBearerRequest *request = state->async_ctx;
+	PGoauthBearerRequestV2 *request = state->async_ctx;
 	PostgresPollingStatusType status;
 
-	if (!request->async)
+	if (!request->v1.async)
 	{
 		libpq_append_conn_error(conn,
 								"user-defined OAuth flow provided neither a token nor an async callback");
 		return PGRES_POLLING_FAILED;
 	}
 
-	status = request->async(conn, request, &conn->altsock);
+	status = request->v1.async(conn,
+							   (PGoauthBearerRequest *) request,
+							   &conn->altsock);
+
 	if (status == PGRES_POLLING_FAILED)
 	{
-		libpq_append_conn_error(conn, "user-defined OAuth flow failed");
+		report_user_flow_error(conn, request);
 		return status;
 	}
 	else if (status == PGRES_POLLING_OK)
@@ -710,14 +732,14 @@ run_user_oauth_flow(PGconn *conn)
 		 * onto the original string, since it may not be safe for us to free()
 		 * it.)
 		 */
-		if (!request->token)
+		if (!request->v1.token)
 		{
 			libpq_append_conn_error(conn,
 									"user-defined OAuth flow did not provide a token");
 			return PGRES_POLLING_FAILED;
 		}
 
-		conn->oauth_token = strdup(request->token);
+		conn->oauth_token = strdup(request->v1.token);
 		if (!conn->oauth_token)
 		{
 			libpq_append_conn_error(conn, "out of memory");
@@ -739,19 +761,20 @@ run_user_oauth_flow(PGconn *conn)
 }
 
 /*
- * Cleanup callback for the async user flow. Delegates most of its job to the
- * user-provided cleanup implementation, then disconnects the altsock.
+ * Cleanup callback for the async user flow. Delegates most of its job to
+ * PGoauthBearerRequestV2.cleanup(), then disconnects the altsock and frees the
+ * request itself.
  */
 static void
 cleanup_user_oauth_flow(PGconn *conn)
 {
 	fe_oauth_state *state = conn->sasl_state;
-	PGoauthBearerRequest *request = state->async_ctx;
+	PGoauthBearerRequestV2 *request = state->async_ctx;
 
 	Assert(request);
 
-	if (request->cleanup)
-		request->cleanup(conn, request);
+	if (request->v1.cleanup)
+		request->v1.cleanup(conn, (PGoauthBearerRequest *) request);
 	conn->altsock = PGINVALID_SOCKET;
 
 	free(request);
@@ -975,8 +998,8 @@ use_builtin_flow(PGconn *conn, fe_oauth_state *state)
  * token for presentation to the server.
  *
  * If the application has registered a custom flow handler using
- * PQAUTHDATA_OAUTH_BEARER_TOKEN, it may either return a token immediately (e.g.
- * if it has one cached for immediate use), or set up for a series of
+ * PQAUTHDATA_OAUTH_BEARER_TOKEN[_V2], it may either return a token immediately
+ * (e.g. if it has one cached for immediate use), or set up for a series of
  * asynchronous callbacks which will be managed by run_user_oauth_flow().
  *
  * If the default handler is used instead, a Device Authorization flow is used
@@ -990,27 +1013,37 @@ static bool
 setup_token_request(PGconn *conn, fe_oauth_state *state)
 {
 	int			res;
-	PGoauthBearerRequest request = {
-		.openid_configuration = conn->oauth_discovery_uri,
-		.scope = conn->oauth_scope,
+	PGoauthBearerRequestV2 request = {
+		.v1 = {
+			.openid_configuration = conn->oauth_discovery_uri,
+			.scope = conn->oauth_scope,
+		},
+		.issuer = conn->oauth_issuer_id,
 	};
 
-	Assert(request.openid_configuration);
+	Assert(request.v1.openid_configuration);
+	Assert(request.issuer);
+
+	/*
+	 * The client may have overridden the OAuth flow. Try the v2 hook first,
+	 * then fall back to the v1 implementation.
+	 */
+	res = PQauthDataHook(PQAUTHDATA_OAUTH_BEARER_TOKEN_V2, conn, &request);
+	if (res == 0)
+		res = PQauthDataHook(PQAUTHDATA_OAUTH_BEARER_TOKEN, conn, &request);
 
-	/* The client may have overridden the OAuth flow. */
-	res = PQauthDataHook(PQAUTHDATA_OAUTH_BEARER_TOKEN, conn, &request);
 	if (res > 0)
 	{
-		PGoauthBearerRequest *request_copy;
+		PGoauthBearerRequestV2 *request_copy;
 
-		if (request.token)
+		if (request.v1.token)
 		{
 			/*
 			 * We already have a token, so copy it into the conn. (We can't
 			 * hold onto the original string, since it may not be safe for us
 			 * to free() it.)
 			 */
-			conn->oauth_token = strdup(request.token);
+			conn->oauth_token = strdup(request.v1.token);
 			if (!conn->oauth_token)
 			{
 				libpq_append_conn_error(conn, "out of memory");
@@ -1018,8 +1051,8 @@ setup_token_request(PGconn *conn, fe_oauth_state *state)
 			}
 
 			/* short-circuit */
-			if (request.cleanup)
-				request.cleanup(conn, &request);
+			if (request.v1.cleanup)
+				request.v1.cleanup(conn, (PGoauthBearerRequest *) &request);
 			return true;
 		}
 
@@ -1038,7 +1071,7 @@ setup_token_request(PGconn *conn, fe_oauth_state *state)
 	}
 	else if (res < 0)
 	{
-		libpq_append_conn_error(conn, "user-defined OAuth flow failed");
+		report_user_flow_error(conn, &request);
 		goto fail;
 	}
 	else if (!use_builtin_flow(conn, state))
@@ -1050,8 +1083,8 @@ setup_token_request(PGconn *conn, fe_oauth_state *state)
 	return true;
 
 fail:
-	if (request.cleanup)
-		request.cleanup(conn, &request);
+	if (request.v1.cleanup)
+		request.v1.cleanup(conn, (PGoauthBearerRequest *) &request);
 	return false;
 }
 
diff --git a/src/test/modules/oauth_validator/oauth_hook_client.c b/src/test/modules/oauth_validator/oauth_hook_client.c
index 60dd1dcdaa0..5bddc0f807a 100644
--- a/src/test/modules/oauth_validator/oauth_hook_client.c
+++ b/src/test/modules/oauth_validator/oauth_hook_client.c
@@ -20,6 +20,7 @@
 
 #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,
@@ -36,13 +37,16 @@ usage(char *argv[])
 
 	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");
 }
 
@@ -51,9 +55,12 @@ 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;
 
 int
 main(int argc, char *argv[])
@@ -68,6 +75,8 @@ main(int argc, char *argv[])
 		{"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}
 	};
 
@@ -75,7 +84,7 @@ main(int argc, char *argv[])
 	PGconn	   *conn;
 	int			c;
 
-	while ((c = getopt_long(argc, argv, "h", long_options, NULL)) != -1)
+	while ((c = getopt_long(argc, argv, "hv:", long_options, NULL)) != -1)
 	{
 		switch (c)
 		{
@@ -83,6 +92,18 @@ main(int argc, char *argv[])
 				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;
@@ -111,6 +132,14 @@ main(int argc, char *argv[])
 				stress_async = true;
 				break;
 
+			case 1007:			/* --expected-issuer */
+				expected_issuer = optarg;
+				break;
+
+			case 1008:			/* --error */
+				errmsg = optarg;
+				break;
+
 			default:
 				usage(argv);
 				return 1;
@@ -167,16 +196,24 @@ main(int argc, char *argv[])
 
 /*
  * PQauthDataHook implementation. Replaces the default client flow by handling
- * PQAUTHDATA_OAUTH_BEARER_TOKEN.
+ * PQAUTHDATA_OAUTH_BEARER_TOKEN[_V2].
  */
 static int
 handle_auth_data(PGauthData type, PGconn *conn, void *data)
 {
-	PGoauthBearerRequest *req = data;
+	PGoauthBearerRequest *req;
+	PGoauthBearerRequestV2 *req2 = NULL;
+
+	Assert(hook_version == PQAUTHDATA_OAUTH_BEARER_TOKEN ||
+		   hook_version == PQAUTHDATA_OAUTH_BEARER_TOKEN_V2);
 
-	if (no_hook || (type != PQAUTHDATA_OAUTH_BEARER_TOKEN))
+	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. */
@@ -221,6 +258,44 @@ handle_auth_data(PGauthData type, PGconn *conn, void *data)
 		}
 	}
 
+	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;
+		}
+
+		req2->error = errmsg;
+		return -1;
+	}
+
 	req->token = token;
 	return 1;
 }
@@ -273,6 +348,20 @@ 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;
+			req2->error = errmsg;
+		}
+
 		return PGRES_POLLING_FAILED;
 	}
 	else if (strcmp(misbehave_mode, "no-token") == 0)
diff --git a/src/test/modules/oauth_validator/t/002_client.pl b/src/test/modules/oauth_validator/t/002_client.pl
index d66cd0b4a5f..0bd5cbb9874 100644
--- a/src/test/modules/oauth_validator/t/002_client.pl
+++ b/src/test/modules/oauth_validator/t/002_client.pl
@@ -78,9 +78,9 @@ sub test
 	my $log_start = -s $node->logfile;
 	my ($stdout, $stderr) = run_command(\@cmd);
 
-	if (defined($params{expected_stdout}))
+	if ($params{expect_success})
 	{
-		like($stdout, $params{expected_stdout}, "$test_name: stdout matches");
+		like($stdout, qr/connection succeeded/, "$test_name: stdout matches");
 	}
 
 	if (defined($params{expected_stderr}))
@@ -110,11 +110,24 @@ test(
 	flags => [
 		"--token", "my-token",
 		"--expected-uri", "$issuer/.well-known/openid-configuration",
+		"--expected-issuer", "$issuer",
 		"--expected-scope", $scope,
 	],
-	expected_stdout => qr/connection succeeded/,
+	expect_success => 1,
 	log_like => [qr/oauth_validator: token="my-token", role="$user"/]);
 
+# Make sure the v1 hook continues to work.
+test(
+	"v1 synchronous hook can provide a token",
+	flags => [
+		"-v1",
+		"--token" => "my-token-v1",
+		"--expected-uri" => "$issuer/.well-known/openid-configuration",
+		"--expected-scope" => $scope,
+	],
+	expect_success => 1,
+	log_like => [qr/oauth_validator: token="my-token-v1", role="$user"/]);
+
 if ($ENV{with_libcurl} ne 'yes')
 {
 	# libpq should help users out if no OAuth support is built in.
@@ -126,6 +139,15 @@ if ($ENV{with_libcurl} ne 'yes')
 	);
 }
 
+# v2 synchronous flows should be able to set custom error messages.
+test(
+	"basic synchronous hook can set error messages",
+	flags => [
+		"--error" => "a custom error message",
+	],
+	expected_stderr =>
+	  qr/user-defined OAuth flow failed: a custom error message/);
+
 # connect_timeout should work if the flow doesn't respond.
 $common_connstr = "$common_connstr connect_timeout=1";
 test(
@@ -163,6 +185,21 @@ foreach my $c (@cases)
 		"hook misbehavior: $c->{'flag'}",
 		flags => [ $c->{'flag'} ],
 		expected_stderr => $c->{'expected_error'});
+
+	test(
+		"hook misbehavior: $c->{'flag'} (v1)",
+		flags => [ '-v1', $c->{'flag'} ],
+		expected_stderr => $c->{'expected_error'});
 }
 
+# v2 async flows should be able to set error messages, too.
+test(
+	"asynchronous hook can set error messages",
+	flags => [
+		"--misbehave" => "fail-async",
+		"--error" => "async error message",
+	],
+	expected_stderr =>
+	  qr/user-defined OAuth flow failed: async error message/);
+
 done_testing();
-- 
2.34.1