From 9726ac39442dd3af479903bfab62894765a6ab2f Mon Sep 17 00:00:00 2001
From: Jacob Champion <jacob.champion@enterprisedb.com>
Date: Tue, 17 Mar 2026 12:01:28 -0700
Subject: [PATCH v3 2/2] WIP: oauth: Allow validators to register custom HBA
 options

(lacks user documentation)

Two new API entry points for validator callbacks:
- RegisterOAuthHBAOptions
- GetOAuthHBAOption

Registering options "foo" and "bar" allows a user to set validator.foo
and validator.bar on an `oauth` HBA line.

The bulk of the patch is not the conceptually simple API implementation,
but guardrails on the simple API to make sure it doesn't bind our hands
in the future, either for callback architecture or HBA syntax.

Suggested-by: Zsolt Parragi <zsolt.parragi@percona.com>
Suggested-by: VASUKI M <vasukianand0119@gmail.com>
Investigated-by: Zsolt Parragi <zsolt.parragi@percona.com>
---
 src/include/libpq/hba.h                       |   2 +
 src/include/libpq/oauth.h                     |  15 +-
 src/backend/libpq/auth-oauth.c                | 225 ++++++++++++++++++
 src/backend/libpq/hba.c                       |  25 ++
 .../modules/oauth_validator/t/001_server.pl   |  97 ++++++++
 src/test/modules/oauth_validator/validator.c  |  48 +++-
 6 files changed, 407 insertions(+), 5 deletions(-)

diff --git a/src/include/libpq/hba.h b/src/include/libpq/hba.h
index c4570ce9b3f..e8898561c8c 100644
--- a/src/include/libpq/hba.h
+++ b/src/include/libpq/hba.h
@@ -140,6 +140,8 @@ typedef struct HbaLine
 	char	   *oauth_scope;
 	char	   *oauth_validator;
 	bool		oauth_skip_usermap;
+	List	   *oauth_opt_keys;
+	List	   *oauth_opt_vals;
 } HbaLine;
 
 typedef struct IdentLine
diff --git a/src/include/libpq/oauth.h b/src/include/libpq/oauth.h
index 60f493acddd..86f463a284e 100644
--- a/src/include/libpq/oauth.h
+++ b/src/include/libpq/oauth.h
@@ -96,6 +96,17 @@ typedef struct OAuthValidatorCallbacks
 	ValidatorValidateCB validate_cb;
 } OAuthValidatorCallbacks;
 
+/*
+ * A validator can register a list of custom option names during its startup_cb,
+ * then later retrieve the user settings for each during validation. This
+ * enables per-HBA-line configuration. For more information, refer to the OAuth
+ * validator modules documentation.
+ */
+extern void RegisterOAuthHBAOptions(ValidatorModuleState *state, int num,
+									const char *opts[]);
+extern const char *GetOAuthHBAOption(const ValidatorModuleState *state,
+									 const char *optname);
+
 /*
  * Type of the shared library symbol _PG_oauth_validator_module_init which is
  * required for all validator modules.  This function will be invoked during
@@ -107,9 +118,7 @@ extern PGDLLEXPORT const OAuthValidatorCallbacks *_PG_oauth_validator_module_ini
 /* Implementation */
 extern PGDLLIMPORT const pg_be_sasl_mech pg_be_oauth_mech;
 
-/*
- * Ensure a validator named in the HBA is permitted by the configuration.
- */
 extern bool check_oauth_validator(HbaLine *hbaline, int elevel, char **err_msg);
+extern bool valid_oauth_hba_option_name(const char *name);
 
 #endif							/* PG_OAUTH_H */
diff --git a/src/backend/libpq/auth-oauth.c b/src/backend/libpq/auth-oauth.c
index 6d4be70aada..032b49e1743 100644
--- a/src/backend/libpq/auth-oauth.c
+++ b/src/backend/libpq/auth-oauth.c
@@ -25,6 +25,7 @@
 #include "libpq/hba.h"
 #include "libpq/oauth.h"
 #include "libpq/sasl.h"
+#include "miscadmin.h"
 #include "storage/fd.h"
 #include "storage/ipc.h"
 #include "utils/json.h"
@@ -40,10 +41,15 @@ static int	oauth_exchange(void *opaq, const char *input, int inputlen,
 
 static void load_validator_library(const char *libname);
 static void shutdown_validator_library(void *arg);
+static bool check_validator_hba_options(Port *port, const char **logdetail);
 
 static ValidatorModuleState *validator_module_state;
 static const OAuthValidatorCallbacks *ValidatorCallbacks;
 
+static MemoryContext ValidatorMemoryContext;
+static List *ValidatorOptions;
+static bool ValidatorOptionsChecked;
+
 /* Mechanism declaration */
 const pg_be_sasl_mech pg_be_oauth_mech = {
 	.get_mechanisms = oauth_get_mechanisms,
@@ -108,6 +114,9 @@ oauth_init(Port *port, const char *selected_mech, const char *shadow_pass)
 				errcode(ERRCODE_PROTOCOL_VIOLATION),
 				errmsg("client selected an invalid SASL authentication mechanism"));
 
+	/* Save our memory context for later use by client API calls. */
+	ValidatorMemoryContext = CurrentMemoryContext;
+
 	ctx = palloc0_object(struct oauth_ctx);
 
 	ctx->state = OAUTH_STATE_INIT;
@@ -279,6 +288,16 @@ oauth_exchange(void *opaq, const char *input, int inputlen,
 				errmsg("malformed OAUTHBEARER message"),
 				errdetail("Message contains additional data after the final terminator."));
 
+	/*
+	 * Make sure all custom HBA options are understood by the validator before
+	 * continuing, since we couldn't check them during server start/reload.
+	 */
+	if (!check_validator_hba_options(ctx->port, logdetail))
+	{
+		ctx->state = OAUTH_STATE_FINISHED;
+		return PG_SASL_EXCHANGE_FAILURE;
+	}
+
 	if (!validate(ctx->port, auth, logdetail))
 	{
 		generate_error_response(ctx, output, outputlen);
@@ -807,6 +826,9 @@ shutdown_validator_library(void *arg)
 {
 	if (ValidatorCallbacks->shutdown_cb != NULL)
 		ValidatorCallbacks->shutdown_cb(validator_module_state);
+
+	/* The backing memory for this is about to disappear. */
+	ValidatorOptions = NIL;
 }
 
 /*
@@ -892,3 +914,206 @@ done:
 
 	return (*err_msg == NULL);
 }
+
+/*
+ * Client APIs for validator implementations
+ *
+ * Since we're currently not threaded, we only allow one validator in the
+ * process at a time. So we can make use of globals for now instead of looking
+ * up information using the state pointer. We probably shouldn't assume that the
+ * module hasn't temporarily changed memory contexts on us, though; functions
+ * here should defensively use an appropriate context when making global
+ * allocations.
+ */
+
+/*
+ * Adds to the list of allowed validator.* HBA options. Used during the
+ * startup_cb.
+ */
+void
+RegisterOAuthHBAOptions(ValidatorModuleState *state, int num,
+						const char *opts[])
+{
+	MemoryContext oldcontext;
+
+	if (!state)
+	{
+		Assert(false);
+		return;
+	}
+
+	oldcontext = MemoryContextSwitchTo(ValidatorMemoryContext);
+
+	for (int i = 0; i < num; i++)
+	{
+		if (!valid_oauth_hba_option_name(opts[i]))
+		{
+			/*
+			 * The user can't set this option in the HBA, so GetOAuthHBAOption
+			 * would always return NULL.
+			 */
+			ereport(WARNING,
+					errmsg("HBA option name \"%s\" is invalid and will be ignored",
+						   opts[i]),
+			/* translator: the second %s is a function name */
+					errcontext("validator module \"%s\", in call to %s",
+							   MyProcPort->hba->oauth_validator,
+							   "RegisterOAuthHBAOptions"));
+			continue;
+		}
+
+		ValidatorOptions = lappend(ValidatorOptions, pstrdup(opts[i]));
+	}
+
+	MemoryContextSwitchTo(oldcontext);
+
+	/*
+	 * Wait to validate the HBA against the registered options until later
+	 * (see check_validator_hba_options()).
+	 *
+	 * Delaying allows the validator to make multiple registration calls, to
+	 * append to the list; it lets us make the check in a place where we can
+	 * report the error without leaking details to the client; and it avoids
+	 * exporting the order of operations between HBA matching and the
+	 * startup_cb call as an API guarantee. (The last issue may become
+	 * relevant with a threaded model.)
+	 */
+}
+
+/*
+ * Restrict the names available to custom HBA options, so that we don't
+ * accidentally prevent future syntax extensions to HBA files.
+ */
+bool
+valid_oauth_hba_option_name(const char *name)
+{
+	/*
+	 * This list is not incredibly principled, since the goal is just to bound
+	 * compatibility guarantees for our HBA parser. Alphanumerics seem
+	 * obviously fine, and it's difficult to argue against the punctuation
+	 * that's already included in some HBA option names and identifiers.
+	 */
+	static const char *name_allowed_set =
+		"abcdefghijklmnopqrstuvwxyz"
+		"ABCDEFGHIJKLMNOPQRSTUVWXYZ"
+		"0123456789_-";
+
+	size_t		span;
+
+	if (!name[0])
+		return false;
+
+	span = strspn(name, name_allowed_set);
+	return name[span] == '\0';
+}
+
+/*
+ * Verifies that all validator.* HBA options specified by the user were actually
+ * registered by the validator library in use.
+ */
+static bool
+check_validator_hba_options(Port *port, const char **logdetail)
+{
+	HbaLine    *hba = port->hba;
+
+	foreach_ptr(char, key, hba->oauth_opt_keys)
+	{
+		bool		found = false;
+
+		/* O(n^2) shouldn't be a problem here in practice. */
+		foreach_ptr(char, optname, ValidatorOptions)
+		{
+			if (strcmp(key, optname) == 0)
+			{
+				found = true;
+				break;
+			}
+		}
+
+		if (!found)
+		{
+			/*
+			 * Bad option name. Mirror the error messages in hba.c here,
+			 * keeping in mind that the original "validator." prefix was
+			 * stripped from the key during parsing.
+			 *
+			 * Since this is affecting live connections, which is unusual for
+			 * HBA, be noisy with a WARNING. (Warnings aren't sent to clients
+			 * prior to successful authentication, so this won't disclose the
+			 * server config.) It'll duplicate some of the information in the
+			 * logdetail, but that should make it hard to miss the connection
+			 * between the two.
+			 */
+			char	   *name = psprintf("validator.%s", key);
+
+			*logdetail = psprintf(_("unrecognized authentication option name: \"%s\""),
+								  name);
+			ereport(WARNING,
+					errcode(ERRCODE_CONFIG_FILE_ERROR),
+					errmsg("unrecognized authentication option name: \"%s\"",
+						   name),
+			/* translator: the first %s is the name of the module */
+					errdetail("The installed validator module (\"%s\") did not define an option named \"%s\".",
+							  hba->oauth_validator, key),
+					errhint("All OAuth connections matching this line will fail. Correct the option and reload the server configuration."),
+					errcontext("line %d of configuration file \"%s\"",
+							   hba->linenumber, hba->sourcefile));
+
+			return false;
+		}
+	}
+
+	ValidatorOptionsChecked = true; /* unfetter GetOAuthHBAOption() */
+	return true;
+}
+
+/*
+ * Retrieves the setting for a validator.* HBA option, or NULL if not found.
+ * This may only be used during the validate_cb and shutdown_cb.
+ */
+const char *
+GetOAuthHBAOption(const ValidatorModuleState *state, const char *optname)
+{
+	HbaLine    *hba = MyProcPort->hba;
+	ListCell   *lc_k;
+	ListCell   *lc_v;
+	const char *ret = NULL;
+
+	if (!ValidatorOptionsChecked)
+	{
+		/*
+		 * Prevent the startup_cb from retrieving HBA options that it has just
+		 * registered. This probably seems strange -- why refuse to hand out
+		 * information we already know? -- but this lets us reserve the
+		 * ability to perform the startup_cb call earlier, before we know
+		 * which HBA line is matched by a connection, without breaking this
+		 * API.
+		 */
+		return NULL;
+	}
+
+	if (!state || !hba)
+	{
+		Assert(false);
+		return NULL;
+	}
+
+	Assert(list_length(hba->oauth_opt_keys) == list_length(hba->oauth_opt_vals));
+
+	forboth(lc_k, hba->oauth_opt_keys, lc_v, hba->oauth_opt_vals)
+	{
+		const char *key = lfirst(lc_k);
+		const char *val = lfirst(lc_v);
+
+		if (strcmp(key, optname) == 0)
+		{
+			/*
+			 * Don't return yet -- when regular HBA options are specified more
+			 * than once, the last one wins. Do the same for these options.
+			 */
+			ret = val;
+		}
+	}
+
+	return ret;
+}
diff --git a/src/backend/libpq/hba.c b/src/backend/libpq/hba.c
index 87ee541e880..0569e8fced8 100644
--- a/src/backend/libpq/hba.c
+++ b/src/backend/libpq/hba.c
@@ -2497,6 +2497,31 @@ parse_hba_auth_opt(char *name, char *val, HbaLine *hbaline,
 		REQUIRE_AUTH_OPTION(uaOAuth, "validator", "oauth");
 		hbaline->oauth_validator = pstrdup(val);
 	}
+	else if (strncmp(name, "validator.", strlen("validator.")) == 0)
+	{
+		const char *key = name + strlen("validator.");
+
+		/*
+		 * Validator modules may register their own per-HBA-line options.
+		 * Unfortunately, since we don't want to require these modules to be
+		 * loaded into the postmaster, we don't know if the options are valid
+		 * yet and must store them for later. Perform only a basic syntax
+		 * check here.
+		 */
+		if (!valid_oauth_hba_option_name(key))
+		{
+			ereport(elevel,
+					(errcode(ERRCODE_CONFIG_FILE_ERROR),
+					 errmsg("invalid OAuth validator option name: \"%s\"", name),
+					 errcontext("line %d of configuration file \"%s\"",
+								line_num, file_name)));
+			return false;
+		}
+
+		REQUIRE_AUTH_OPTION(uaOAuth, name, "oauth");
+		hbaline->oauth_opt_keys = lappend(hbaline->oauth_opt_keys, pstrdup(key));
+		hbaline->oauth_opt_vals = lappend(hbaline->oauth_opt_vals, pstrdup(val));
+	}
 	else if (strcmp(name, "delegate_ident_mapping") == 0)
 	{
 		REQUIRE_AUTH_OPTION(uaOAuth, "delegate_ident_mapping", "oauth");
diff --git a/src/test/modules/oauth_validator/t/001_server.pl b/src/test/modules/oauth_validator/t/001_server.pl
index 28d123b833e..0756e6038c4 100644
--- a/src/test/modules/oauth_validator/t/001_server.pl
+++ b/src/test/modules/oauth_validator/t/001_server.pl
@@ -579,10 +579,29 @@ $node->connect_fails(
 
 $bgconn->query_safe("ALTER SYSTEM RESET oauth_validator.error_detail");
 $bgconn->query_safe("ALTER SYSTEM RESET oauth_validator.internal_error");
+
+# We complain when bad option names are registered, but connections may proceed
+# (since users can't set those options in the HBA anyway).
+$bgconn->query_safe("ALTER SYSTEM RESET oauth_validator.authn_id");
+$bgconn->query_safe("ALTER SYSTEM RESET oauth_validator.authorize_tokens");
+$bgconn->query_safe("ALTER SYSTEM SET oauth_validator.invalid_hba TO true");
+
 $node->reload;
 $log_start =
   $node->wait_for_log(qr/reloading configuration files/, $log_start);
 
+$node->connect_ok(
+	"$common_connstr user=test",
+	"bad registered HBA option",
+	expected_stderr =>
+	  qr@Visit https://example\.com/ and enter the code: postgresuser@,
+	log_like => [
+		qr/WARNING:\s+HBA option name "bad option name" is invalid and will be ignored/,
+		qr/CONTEXT:\s+validator module "validator", in call to RegisterOAuthHBAOptions/,
+	]);
+
+$bgconn->query_safe("ALTER SYSTEM RESET oauth_validator.invalid_hba");
+
 #
 # Test user mapping.
 #
@@ -651,6 +670,84 @@ $node->reload;
 $log_start =
   $node->wait_for_log(qr/reloading configuration files/, $log_start);
 
+$bgconn->quit;    # the tests below restart the server
+
+#
+# Test validator-specific HBA options.
+#
+
+unlink($node->data_dir . '/pg_hba.conf');
+$node->append_conf(
+	'pg_hba.conf', qq{
+local all test    oauth issuer="$issuer" scope="openid postgres" delegate_ident_mapping=1 \\
+                        validator.authn_id="ignored" validator.authn_id="other-identity"
+local all testalt oauth issuer="$issuer" scope="openid postgres" validator.log="testalt message"
+});
+
+$node->reload;
+$log_start =
+  $node->wait_for_log(qr/reloading configuration files/, $log_start);
+
+$node->connect_ok(
+	"$common_connstr user=test",
+	"custom HBA setting (test)",
+	expected_stderr =>
+	  qr@Visit https://example\.com/ and enter the code: postgresuser@,
+	log_like => [qr/connection authenticated: identity="other-identity"/]);
+$node->connect_ok(
+	"$common_connstr user=testalt",
+	"custom HBA setting (testalt)",
+	expected_stderr =>
+	  qr@Visit https://example\.com/ and enter the code: postgresuser@,
+	log_like => [
+		qr/LOG:\s+testalt message/,
+		qr/connection authenticated: identity="testalt"/,
+	]);
+
+# bad syntax
+unlink($node->data_dir . '/pg_hba.conf');
+$node->append_conf(
+	'pg_hba.conf', qq{
+local all testalt oauth issuer="$issuer" scope="openid postgres" validator.=1
+});
+
+$log_start = -s $node->logfile;
+$node->restart(fail_ok => 1);
+$node->log_check("empty HBA option name",
+	$log_start,
+	log_like => [qr/invalid OAuth validator option name: "validator\."/]);
+
+unlink($node->data_dir . '/pg_hba.conf');
+$node->append_conf(
+	'pg_hba.conf', qq{
+local all testalt oauth issuer="$issuer" scope="openid postgres" validator.@@=1
+});
+
+$log_start = -s $node->logfile;
+$node->restart(fail_ok => 1);
+$node->log_check("invalid HBA option name",
+	$log_start,
+	log_like => [qr/invalid OAuth validator option name: "validator\.@@"/]);
+
+# unknown settings (validation is deferred to connect time)
+unlink($node->data_dir . '/pg_hba.conf');
+$node->append_conf(
+	'pg_hba.conf', qq{
+local all testalt oauth issuer="$issuer" scope="openid postgres" \\
+                        validator.log=ignored validator.bad=1
+});
+$node->restart;
+
+$node->connect_fails(
+	"$common_connstr user=testalt",
+	"bad HBA setting",
+	expected_stderr => qr/OAuth bearer authentication failed/,
+	log_like => [
+		qr/WARNING:\s+unrecognized authentication option name: "validator\.bad"/,
+		qr/FATAL:\s+OAuth bearer authentication failed/,
+		qr/DETAIL:\s+unrecognized authentication option name: "validator\.bad"/,
+	]);
+
 #
 # Test multiple validators.
 #
diff --git a/src/test/modules/oauth_validator/validator.c b/src/test/modules/oauth_validator/validator.c
index 353e0e0d32a..85fb4c08bf2 100644
--- a/src/test/modules/oauth_validator/validator.c
+++ b/src/test/modules/oauth_validator/validator.c
@@ -42,13 +42,21 @@ static char *authn_id = NULL;
 static bool authorize_tokens = true;
 static char *error_detail = NULL;
 static bool internal_error = false;
+static bool invalid_hba = false;
+
+/* HBA options */
+static const char *hba_opts[] = {
+	"authn_id",					/* overrides the default authn_id */
+	"log",						/* logs an arbitrary string */
+};
 
 /*---
  * Extension entry point. Sets up GUCs for use by tests:
  *
  * - oauth_validator.authn_id	Sets the user identifier to return during token
  *								validation. Defaults to the username in the
- *								startup packet.
+ *								startup packet, or the validator.authn_id HBA
+ *								option if it is set.
  *
  * - oauth_validator.authorize_tokens
  *								Sets whether to successfully validate incoming
@@ -96,6 +104,14 @@ _PG_init(void)
 							 PGC_SIGHUP,
 							 0,
 							 NULL, NULL, NULL);
+	DefineCustomBoolVariable("oauth_validator.invalid_hba",
+							 "Should the validator register an invalid option?",
+							 NULL,
+							 &invalid_hba,
+							 false,
+							 PGC_SIGHUP,
+							 0,
+							 NULL, NULL, NULL);
 
 	MarkGUCPrefixReserved("oauth_validator");
 }
@@ -124,6 +140,29 @@ validator_startup(ValidatorModuleState *state)
 	if (state->sversion != PG_VERSION_NUM)
 		elog(ERROR, "oauth_validator: sversion set to %d", state->sversion);
 
+	/*
+	 * Test the behavior of custom HBA options. Registered options should not
+	 * be retrievable during startup (we want to discourage modules from
+	 * relying on the relative order of client connections and the
+	 * startup_cb).
+	 */
+	RegisterOAuthHBAOptions(state, lengthof(hba_opts), hba_opts);
+	for (int i = 0; i < lengthof(hba_opts); i++)
+	{
+		if (GetOAuthHBAOption(state, hba_opts[i]))
+			elog(ERROR,
+				 "oauth_validator: GetOAuthValidatorOption(\"%s\") was non-NULL during startup_cb",
+				 hba_opts[i]);
+	}
+
+	if (invalid_hba)
+	{
+		/* Register a bad option, which should print a WARNING to the logs. */
+		const char *invalid = "bad option name";
+
+		RegisterOAuthHBAOptions(state, 1, &invalid);
+	}
+
 	state->private_data = PRIVATE_COOKIE;
 }
 
@@ -141,7 +180,7 @@ validator_shutdown(ValidatorModuleState *state)
 
 /*
  * Validator implementation. Logs the incoming data and authorizes the token by
- * default; the behavior can be modified via the module's GUC settings.
+ * default; the behavior can be modified via the module's GUC and HBA settings.
  */
 static bool
 validate_token(const ValidatorModuleState *state,
@@ -153,6 +192,9 @@ validate_token(const ValidatorModuleState *state,
 		elog(ERROR, "oauth_validator: private state cookie changed to %p in validate",
 			 state->private_data);
 
+	if (GetOAuthHBAOption(state, "log"))
+		elog(LOG, "%s", GetOAuthHBAOption(state, "log"));
+
 	elog(LOG, "oauth_validator: token=\"%s\", role=\"%s\"", token, role);
 	elog(LOG, "oauth_validator: issuer=\"%s\", scope=\"%s\"",
 		 MyProcPort->hba->oauth_issuer,
@@ -165,6 +207,8 @@ validate_token(const ValidatorModuleState *state,
 	res->authorized = authorize_tokens;
 	if (authn_id)
 		res->authn_id = pstrdup(authn_id);
+	else if (GetOAuthHBAOption(state, "authn_id"))
+		res->authn_id = pstrdup(GetOAuthHBAOption(state, "authn_id"));
 	else
 		res->authn_id = pstrdup(role);
 
-- 
2.34.1

