From 296acdb1551db523009ce7201daa03ef3e33f182 Mon Sep 17 00:00:00 2001
From: Zsolt Parragi <zsolt.parragi@percona.com>
Date: Thu, 18 Dec 2025 08:25:46 +0000
Subject: [PATCH] Adding hooks to HBA parsing, option c

This commit showcases how we can add a generic, non oauth dependent hook
to parsing hba entries, and also adds a simple test to the existing
oauth_validator test suite.
---
 src/backend/libpq/hba.c                       | 34 ++++++++----
 src/include/libpq/hba.h                       | 31 +++++++++++
 .../modules/oauth_validator/t/002_client.pl   | 51 ++++++++++++++++++
 src/test/modules/oauth_validator/validator.c  | 53 +++++++++++++++++++
 4 files changed, 160 insertions(+), 9 deletions(-)

diff --git a/src/backend/libpq/hba.c b/src/backend/libpq/hba.c
index 4c259f58d77..409801ec6d8 100644
--- a/src/backend/libpq/hba.c
+++ b/src/backend/libpq/hba.c
@@ -125,6 +125,11 @@ static const char *const UserAuthName[] =
 StaticAssertDecl(lengthof(UserAuthName) == USER_AUTH_LAST + 1,
 				 "UserAuthName[] must match the UserAuth enum");
 
+/*
+ * Hook for plugins to extend pg_hba.conf option parsing.
+ */
+hba_parse_option_hook_type hba_parse_option_hook = NULL;
+
 
 static List *tokenize_expand_file(List *tokens, const char *outer_filename,
 								  const char *inc_filename, int elevel,
@@ -2507,15 +2512,26 @@ parse_hba_auth_opt(char *name, char *val, HbaLine *hbaline,
 	}
 	else
 	{
-		ereport(elevel,
-				(errcode(ERRCODE_CONFIG_FILE_ERROR),
-				 errmsg("unrecognized authentication option name: \"%s\"",
-						name),
-				 errcontext("line %d of configuration file \"%s\"",
-							line_num, file_name)));
-		*err_msg = psprintf("unrecognized authentication option name: \"%s\"",
-							name);
-		return false;
+		bool		handled = false;
+
+		if (hba_parse_option_hook)
+		{
+			handled = (*hba_parse_option_hook) (name, val, hbaline,
+												elevel, err_msg);
+		}
+
+		if (!handled)
+		{
+			ereport(elevel,
+					(errcode(ERRCODE_CONFIG_FILE_ERROR),
+					 errmsg("unrecognized authentication option name: \"%s\"",
+							name),
+					 errcontext("line %d of configuration file \"%s\"",
+								line_num, file_name)));
+			*err_msg = psprintf("unrecognized authentication option name: \"%s\"",
+								name);
+			return false;
+		}
 	}
 	return true;
 }
diff --git a/src/include/libpq/hba.h b/src/include/libpq/hba.h
index 7b93ba4a709..bd4658c72c8 100644
--- a/src/include/libpq/hba.h
+++ b/src/include/libpq/hba.h
@@ -172,6 +172,37 @@ typedef struct TokenizedAuthLine
 /* avoid including libpq/libpq-be.h here */
 typedef struct Port Port;
 
+/*
+ * Hook for plugins to extend pg_hba.conf option parsing.
+ *
+ * This hook is called by parse_hba_auth_opt() when it encounters an option
+ * name that it doesn't recognize. Plugins can use this to parse custom
+ * authentication.
+ *
+ * Parameters:
+ *   name     - The option name being parsed (e.g., "custom_option")
+ *   val      - The option value (may be NULL for boolean-style options)
+ *   hbaline  - The HbaLine structure being populated. Plugins should not
+ *              modify standard fields, but can use this to check auth_method,
+ *              conntype, etc. to validate option applicability.
+ *   elevel   - Error level for reporting (LOG, ERROR, etc.)
+ *   err_msg  - Output parameter for error messages. Set this to a palloc'd
+ *              string if returning false due to a validation error.
+ *
+ * Return value:
+ *   true  - The hook recognized and successfully handled this option.
+ *   false - The hook doesn't recognize this option, or encountered an error.
+ *           If an error occurred, the hook should set *err_msg and/or call
+ *           ereport().
+ */
+typedef bool (*hba_parse_option_hook_type) (const char *name,
+											const char *val,
+											HbaLine *hbaline,
+											int elevel,
+											char **err_msg);
+
+extern PGDLLIMPORT hba_parse_option_hook_type hba_parse_option_hook;
+
 extern bool load_hba(void);
 extern bool load_ident(void);
 extern const char *hba_authname(UserAuth auth_method);
diff --git a/src/test/modules/oauth_validator/t/002_client.pl b/src/test/modules/oauth_validator/t/002_client.pl
index e6c91fc911c..6576d6e41e0 100644
--- a/src/test/modules/oauth_validator/t/002_client.pl
+++ b/src/test/modules/oauth_validator/t/002_client.pl
@@ -29,6 +29,8 @@ $node->init;
 $node->append_conf('postgresql.conf', "log_connections = all\n");
 $node->append_conf('postgresql.conf',
 	"oauth_validator_libraries = 'validator'\n");
+$node->append_conf('postgresql.conf',
+	"shared_preload_libraries = 'validator'\n");
 # Needed to inspect postmaster log after connection failure:
 $node->append_conf('postgresql.conf', "log_min_messages = debug2");
 $node->start;
@@ -115,6 +117,55 @@ test(
 	expected_stdout => qr/connection succeeded/,
 	log_like => [qr/oauth_validator: token="my-token", role="$user"/]);
 
+# Test custom HBA option parsing hook
+my $log_start_custom = -s $node->logfile;
+unlink($node->data_dir . '/pg_hba.conf');
+$node->append_conf(
+	'pg_hba.conf', qq{
+local all test oauth issuer="$issuer" scope="$scope" test_custom_claim="my_custom_value"
+});
+$node->reload;
+$node->wait_for_log(qr/reloading configuration files/, $log_start_custom);
+
+$node->wait_for_log(
+	qr/oauth_validator: parsed custom HBA option test_custom_claim="my_custom_value"/,
+	$log_start_custom);
+
+test(
+	"custom HBA option is parsed and used",
+	flags => [
+		"--token", "test-token",
+		"--expected-uri", "$issuer/.well-known/openid-configuration",
+		"--expected-scope", $scope,
+	],
+	expected_stdout => qr/connection succeeded/,
+	log_like => [qr/oauth_validator: custom_claim="my_custom_value"/]);
+
+# Test that unknown HBA options still fail
+my $log_start_unknown = -s $node->logfile;
+unlink($node->data_dir . '/pg_hba.conf');
+$node->append_conf(
+	'pg_hba.conf', qq{
+local all test oauth issuer="$issuer" scope="$scope" unknown_option="value"
+});
+$node->reload;
+$node->wait_for_log(qr/reloading configuration files/, $log_start_unknown);
+
+# Check that the server logged the error about the unknown option
+$node->wait_for_log(
+	qr/unrecognized authentication option name: "unknown_option"/,
+	$log_start_unknown);
+pass("unknown HBA option is rejected");
+
+# Restore working configuration
+unlink($node->data_dir . '/pg_hba.conf');
+$node->append_conf(
+	'pg_hba.conf', qq{
+local all test oauth issuer="$issuer" scope="$scope"
+});
+$node->reload;
+$node->wait_for_log(qr/reloading configuration files/);
+
 if ($ENV{with_libcurl} ne 'yes')
 {
 	# libpq should help users out if no OAuth support is built in.
diff --git a/src/test/modules/oauth_validator/validator.c b/src/test/modules/oauth_validator/validator.c
index 42b69646fbb..0115f228cea 100644
--- a/src/test/modules/oauth_validator/validator.c
+++ b/src/test/modules/oauth_validator/validator.c
@@ -14,6 +14,7 @@
 #include "postgres.h"
 
 #include "fmgr.h"
+#include "libpq/hba.h"
 #include "libpq/oauth.h"
 #include "miscadmin.h"
 #include "utils/guc.h"
@@ -41,6 +42,50 @@ static const OAuthValidatorCallbacks validator_callbacks = {
 static char *authn_id = NULL;
 static bool authorize_tokens = true;
 
+static char *custom_claim = NULL;
+
+static hba_parse_option_hook_type prev_hba_parse_option_hook = NULL;
+
+static bool
+validator_hba_parse_option(const char *name, const char *val,
+						   HbaLine *hbaline, int elevel, char **err_msg)
+{
+	int			line_num = hbaline->linenumber;
+	char	   *file_name = hbaline->sourcefile;
+
+	if (prev_hba_parse_option_hook)
+	{
+		if ((*prev_hba_parse_option_hook) (name, val, hbaline,
+										   elevel, err_msg))
+			return true;
+	}
+
+	if (strcmp(name, "test_custom_claim") == 0)
+	{
+		if (val == NULL || val[0] == '\0')
+		{
+			ereport(elevel,
+					(errcode(ERRCODE_CONFIG_FILE_ERROR),
+					 errmsg("test_custom_claim requires a value"),
+					 errcontext("line %d of configuration file \"%s\"",
+								line_num, file_name)));
+			*err_msg = pstrdup("test_custom_claim requires a value");
+			return false;
+		}
+
+		if (custom_claim)
+			pfree(custom_claim);
+		custom_claim = pstrdup(val);
+
+		elog(LOG, "oauth_validator: parsed custom HBA option test_custom_claim=\"%s\"",
+			 custom_claim);
+
+		return true;
+	}
+
+	return false;
+}
+
 /*---
  * Extension entry point. Sets up GUCs for use by tests:
  *
@@ -55,6 +100,8 @@ static bool authorize_tokens = true;
 void
 _PG_init(void)
 {
+	elog(LOG, "oauth_validator: _PG_init() called, installing HBA parse option hook");
+
 	DefineCustomStringVariable("oauth_validator.authn_id",
 							   "Authenticated identity to use for future connections",
 							   NULL,
@@ -73,6 +120,9 @@ _PG_init(void)
 							 NULL, NULL, NULL);
 
 	MarkGUCPrefixReserved("oauth_validator");
+
+	prev_hba_parse_option_hook = hba_parse_option_hook;
+	hba_parse_option_hook = validator_hba_parse_option;
 }
 
 /*
@@ -133,6 +183,9 @@ validate_token(const ValidatorModuleState *state,
 		 MyProcPort->hba->oauth_issuer,
 		 MyProcPort->hba->oauth_scope);
 
+	if (custom_claim)
+		elog(LOG, "oauth_validator: custom_claim=\"%s\"", custom_claim);
+
 	res->authorized = authorize_tokens;
 	if (authn_id)
 		res->authn_id = pstrdup(authn_id);
-- 
2.43.0

