v48-0006-fixup-Add-OAUTHBEARER-SASL-mechanism.patch

From 954341052b44ab0f89e679926c9b66a085d2f064 Mon Sep 17 00:00:00 2001
From: Jacob Champion <jacob.champion@enterprisedb.com>
Date: Thu, 6 Feb 2025 20:53:33 -0800
Subject: [PATCH v48 6/8] fixup! Add OAUTHBEARER SASL mechanism

---
 doc/src/sgml/oauth-validators.sgml            | 22 +++++----
 src/backend/libpq/auth-oauth.c                | 20 ++++++--
 src/include/libpq/oauth.h                     | 48 ++++++++++++++++++-
 .../modules/oauth_validator/fail_validator.c  | 15 ++++--
 src/test/modules/oauth_validator/validator.c  | 28 +++++++----
 5 files changed, 105 insertions(+), 28 deletions(-)

diff --git a/doc/src/sgml/oauth-validators.sgml b/doc/src/sgml/oauth-validators.sgml
index c8bbac7b462..eb8c4431c2d 100644
--- a/doc/src/sgml/oauth-validators.sgml
+++ b/doc/src/sgml/oauth-validators.sgml
@@ -297,13 +297,15 @@
    validator module a function named
    <function>_PG_oauth_validator_module_init</function> must be provided. The
    return value of the function must be a pointer to a struct of type
-   <structname>OAuthValidatorCallbacks</structname>, which contains pointers to
-   the module's token validation functions. The returned
+   <structname>OAuthValidatorCallbacks</structname>, which contains a magic
+   number and pointers to the module's token validation functions. The returned
    pointer must be of server lifetime, which is typically achieved by defining
    it as a <literal>static const</literal> variable in global scope.
 <programlisting>
 typedef struct OAuthValidatorCallbacks
 {
+    uint32        magic;            /* must be set to PG_OAUTH_VALIDATOR_MAGIC */
+
     ValidatorStartupCB startup_cb;
     ValidatorShutdownCB shutdown_cb;
     ValidatorValidateCB validate_cb;
@@ -348,14 +350,16 @@ typedef void (*ValidatorStartupCB) (ValidatorModuleState *state);
     previous calls will be available in <structfield>state->private_data</structfield>.
 
 <programlisting>
-typedef ValidatorModuleResult *(*ValidatorValidateCB) (ValidatorModuleState *state, const char *token, const char *role);
+typedef bool (*ValidatorValidateCB) (const ValidatorModuleState *state,
+                                     const char *token, const char *role,
+                                     ValidatorModuleResult *result);
 </programlisting>
 
     <replaceable>token</replaceable> will contain the bearer token to validate.
     The server has ensured that the token is well-formed syntactically, but no
     other validation has been performed.  <replaceable>role</replaceable> will
     contain the role the user has requested to log in as.  The callback must
-    return a palloc'd <literal>ValidatorModuleResult</literal> struct, which is
+    set output parameters in the <literal>result</literal> struct, which is
     defined as below:
 
 <programlisting>
@@ -375,17 +379,17 @@ typedef struct ValidatorModuleResult
     determined.
    </para>
    <para>
-    The caller assumes ownership of the returned memory allocation, the
-    validator module should not in any way access the memory after it has been
-    returned.  A validator may instead return NULL to signal an internal
-    error.
+    A validator may return <literal>false</literal> to signal an internal error,
+    in which case any result parameters are ignored and the connection fails.
+    Otherwise the validator should return <literal>true</literal> to indicate
+    that it has processed the token and made an authorization decision.
    </para>
    <para>
     The behavior after <function>validate_cb</function> returns depends on the
     specific HBA setup.  Normally, the <structfield>authn_id</structfield> user
     name must exactly match the role that the user is logging in as.  (This
     behavior may be modified with a usermap.)  But when authenticating against
-    an HBA rule with <literal>trust_validator_authz</literal> turned on, the
+    an HBA rule with <literal>delegate_ident_mapping</literal> turned on, the
     server will not perform any checks on the value of
     <structfield>authn_id</structfield> at all; in this case it is up to the
     validator to ensure that the token carries enough privileges for the user to
diff --git a/src/backend/libpq/auth-oauth.c b/src/backend/libpq/auth-oauth.c
index aa16977c643..e2b5d1ed913 100644
--- a/src/backend/libpq/auth-oauth.c
+++ b/src/backend/libpq/auth-oauth.c
@@ -656,9 +656,9 @@ validate(Port *port, const char *auth)
 				errmsg("validation of OAuth token requested without a validator loaded"));
 
 	/* Call the validation function from the validator module */
-	ret = ValidatorCallbacks->validate_cb(validator_module_state,
-										  token, port->user_name);
-	if (ret == NULL)
+	ret = palloc0(sizeof(ValidatorModuleResult));
+	if (!ValidatorCallbacks->validate_cb(validator_module_state, token,
+										 port->user_name, ret))
 	{
 		ereport(LOG, errmsg("internal error in OAuth validator module"));
 		return false;
@@ -756,8 +756,22 @@ load_validator_library(const char *libname)
 	ValidatorCallbacks = (*validator_init) ();
 	Assert(ValidatorCallbacks);
 
+	/*
+	 * Check the magic number, to protect against break-glass scenarios where
+	 * the ABI must change within a major version. load_external_function()
+	 * already checks for compatibility across major versions.
+	 */
+	if (ValidatorCallbacks->magic != PG_OAUTH_VALIDATOR_MAGIC)
+		ereport(ERROR,
+				errmsg("%s module \"%s\": magic number mismatch",
+					   "OAuth validator", libname),
+				errdetail("Server has magic number 0x%08X, module has 0x%08X.",
+						  PG_OAUTH_VALIDATOR_MAGIC, ValidatorCallbacks->magic));
+
 	/* Allocate memory for validator library private state data */
 	validator_module_state = (ValidatorModuleState *) palloc0(sizeof(ValidatorModuleState));
+	validator_module_state->sversion = PG_VERSION_NUM;
+
 	if (ValidatorCallbacks->startup_cb != NULL)
 		ValidatorCallbacks->startup_cb(validator_module_state);
 
diff --git a/src/include/libpq/oauth.h b/src/include/libpq/oauth.h
index 4fcdda74305..7e249613e10 100644
--- a/src/include/libpq/oauth.h
+++ b/src/include/libpq/oauth.h
@@ -20,26 +20,72 @@ extern PGDLLIMPORT char *oauth_validator_libraries_string;
 
 typedef struct ValidatorModuleState
 {
+	/* Holds the server's PG_VERSION_NUM. Reserved for future extensibility. */
+	int			sversion;
+
+	/*
+	 * Private data pointer for use by a validator module. This can be used to
+	 * store state for the module that will be passed to each of its
+	 * callbacks.
+	 */
 	void	   *private_data;
 } ValidatorModuleState;
 
 typedef struct ValidatorModuleResult
 {
+	/*
+	 * Should be set to true if the token carries sufficient permissions for
+	 * the bearer to connect.
+	 */
 	bool		authorized;
+
+	/*
+	 * If the token authenticates the user, this should be set to a palloc'd
+	 * string containing the SYSTEM_USER to use for HBA mapping. Consider
+	 * setting this even if result->authorized is false so that DBAs may use
+	 * the logs to match end users to token failures.
+	 *
+	 * This is required if the module is not configured for ident mapping
+	 * delegation. See the validator module documentation for details.
+	 */
 	char	   *authn_id;
 } ValidatorModuleResult;
 
+/*
+ * Validator module callbacks
+ *
+ * These callback functions should be defined by validator modules and returned
+ * via _PG_oauth_validator_module_init().  ValidatorValidateCB is the only
+ * required callback. For more information about the purpose of each callback,
+ * refer to the OAuth validator modules documentation.
+ */
 typedef void (*ValidatorStartupCB) (ValidatorModuleState *state);
 typedef void (*ValidatorShutdownCB) (ValidatorModuleState *state);
-typedef ValidatorModuleResult *(*ValidatorValidateCB) (ValidatorModuleState *state, const char *token, const char *role);
+typedef bool (*ValidatorValidateCB) (const ValidatorModuleState *state,
+									 const char *token, const char *role,
+									 ValidatorModuleResult *result);
+
+/*
+ * Identifies the compiled ABI version of the validator module. Since the server
+ * already enforces the PG_MODULE_MAGIC number for modules across major
+ * versions, this is reserved for emergency use within a stable release line.
+ * May it never need to change.
+ */
+#define PG_OAUTH_VALIDATOR_MAGIC 0x20250207
 
 typedef struct OAuthValidatorCallbacks
 {
+	uint32		magic;			/* must be set to PG_OAUTH_VALIDATOR_MAGIC */
+
 	ValidatorStartupCB startup_cb;
 	ValidatorShutdownCB shutdown_cb;
 	ValidatorValidateCB validate_cb;
 } OAuthValidatorCallbacks;
 
+/*
+ * Type of the shared library symbol _PG_oauth_validator_module_init that is
+ * looked up when loading a validator module.
+ */
 typedef const OAuthValidatorCallbacks *(*OAuthValidatorModuleInit) (void);
 extern PGDLLEXPORT const OAuthValidatorCallbacks *_PG_oauth_validator_module_init(void);
 
diff --git a/src/test/modules/oauth_validator/fail_validator.c b/src/test/modules/oauth_validator/fail_validator.c
index f77a3e115c6..7b1e69518d9 100644
--- a/src/test/modules/oauth_validator/fail_validator.c
+++ b/src/test/modules/oauth_validator/fail_validator.c
@@ -19,12 +19,15 @@
 
 PG_MODULE_MAGIC;
 
-static ValidatorModuleResult *fail_token(ValidatorModuleState *state,
-										 const char *token,
-										 const char *role);
+static bool fail_token(const ValidatorModuleState *state,
+					   const char *token,
+					   const char *role,
+					   ValidatorModuleResult *result);
 
 /* Callback implementations (we only need the main one) */
 static const OAuthValidatorCallbacks validator_callbacks = {
+	PG_OAUTH_VALIDATOR_MAGIC,
+
 	.validate_cb = fail_token,
 };
 
@@ -34,8 +37,10 @@ _PG_oauth_validator_module_init(void)
 	return &validator_callbacks;
 }
 
-static ValidatorModuleResult *
-fail_token(ValidatorModuleState *state, const char *token, const char *role)
+static bool
+fail_token(const ValidatorModuleState *state,
+		   const char *token, const char *role,
+		   ValidatorModuleResult *res)
 {
 	elog(FATAL, "fail_validator: sentinel error");
 	pg_unreachable();
diff --git a/src/test/modules/oauth_validator/validator.c b/src/test/modules/oauth_validator/validator.c
index ef9bbb2866f..e218f5c8902 100644
--- a/src/test/modules/oauth_validator/validator.c
+++ b/src/test/modules/oauth_validator/validator.c
@@ -23,12 +23,15 @@ PG_MODULE_MAGIC;
 
 static void validator_startup(ValidatorModuleState *state);
 static void validator_shutdown(ValidatorModuleState *state);
-static ValidatorModuleResult *validate_token(ValidatorModuleState *state,
-											 const char *token,
-											 const char *role);
+static bool validate_token(const ValidatorModuleState *state,
+						   const char *token,
+						   const char *role,
+						   ValidatorModuleResult *result);
 
 /* Callback implementations (exercise all three) */
 static const OAuthValidatorCallbacks validator_callbacks = {
+	PG_OAUTH_VALIDATOR_MAGIC,
+
 	.startup_cb = validator_startup,
 	.shutdown_cb = validator_shutdown,
 	.validate_cb = validate_token
@@ -89,6 +92,13 @@ _PG_oauth_validator_module_init(void)
 static void
 validator_startup(ValidatorModuleState *state)
 {
+	/*
+	 * Make sure the server is correctly setting sversion. (Real modules
+	 * should not do this; it would defeat upgrade compatibility.)
+	 */
+	if (state->sversion != PG_VERSION_NUM)
+		elog(ERROR, "oauth_validator: sversion set to %d", state->sversion);
+
 	state->private_data = PRIVATE_COOKIE;
 }
 
@@ -108,18 +118,16 @@ 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.
  */
-static ValidatorModuleResult *
-validate_token(ValidatorModuleState *state, const char *token, const char *role)
+static bool
+validate_token(const ValidatorModuleState *state,
+			   const char *token, const char *role,
+			   ValidatorModuleResult *res)
 {
-	ValidatorModuleResult *res;
-
 	/* Check to make sure our private state still exists. */
 	if (state->private_data != PRIVATE_COOKIE)
 		elog(ERROR, "oauth_validator: private state cookie changed to %p in validate",
 			 state->private_data);
 
-	res = palloc(sizeof(ValidatorModuleResult));
-
 	elog(LOG, "oauth_validator: token=\"%s\", role=\"%s\"", token, role);
 	elog(LOG, "oauth_validator: issuer=\"%s\", scope=\"%s\"",
 		 MyProcPort->hba->oauth_issuer,
@@ -131,5 +139,5 @@ validate_token(ValidatorModuleState *state, const char *token, const char *role)
 	else
 		res->authn_id = pstrdup(role);
 
-	return res;
+	return true;
 }
-- 
2.34.1