v40-0002-squash-Add-OAUTHBEARER-SASL-mechanism.patch

From de155343c816fa9ee8301522060576ee796857a4 Mon Sep 17 00:00:00 2001
From: Jacob Champion <jacob.champion@enterprisedb.com>
Date: Mon, 16 Dec 2024 13:57:14 -0800
Subject: [PATCH v40 2/3] squash! Add OAUTHBEARER SASL mechanism

Add require_auth=oauth support.
---
 doc/src/sgml/libpq.sgml                       |   9 +
 src/interfaces/libpq/fe-auth-oauth.c          |   7 +
 src/interfaces/libpq/fe-auth.c                |  22 +++
 src/interfaces/libpq/fe-connect.c             | 183 ++++++++++++++++--
 src/interfaces/libpq/libpq-int.h              |   2 +
 src/test/authentication/t/001_password.pl     |  18 +-
 .../modules/oauth_validator/t/001_server.pl   |  62 +++++-
 7 files changed, 278 insertions(+), 25 deletions(-)

diff --git a/doc/src/sgml/libpq.sgml b/doc/src/sgml/libpq.sgml
index 93266c4a0d0..fe0cbb0c800 100644
--- a/doc/src/sgml/libpq.sgml
+++ b/doc/src/sgml/libpq.sgml
@@ -1385,6 +1385,15 @@ postgresql://%2Fvar%2Flib%2Fpostgresql/dbname
           </listitem>
          </varlistentry>
 
+         <varlistentry>
+          <term><literal>oauth</literal></term>
+          <listitem>
+           <para>
+            The server must request an OAuth bearer token from the client.
+           </para>
+          </listitem>
+         </varlistentry>
+
          <varlistentry>
           <term><literal>none</literal></term>
           <listitem>
diff --git a/src/interfaces/libpq/fe-auth-oauth.c b/src/interfaces/libpq/fe-auth-oauth.c
index cb290c5c113..71940fb9f54 100644
--- a/src/interfaces/libpq/fe-auth-oauth.c
+++ b/src/interfaces/libpq/fe-auth-oauth.c
@@ -934,6 +934,13 @@ oauth_exchange(void *opaq, bool final,
 			*outputlen = strlen(*output);
 			state->step = FE_OAUTH_BEARER_SENT;
 
+			/*
+			 * For the purposes of require_auth, our side of authentication is
+			 * done at this point; the server will either accept the
+			 * connection or send an error. Unlike SCRAM, there is no
+			 * additional server data to check upon success.
+			 */
+			conn->client_finished_auth = true;
 			return SASL_CONTINUE;
 
 		case FE_OAUTH_BEARER_SENT:
diff --git a/src/interfaces/libpq/fe-auth.c b/src/interfaces/libpq/fe-auth.c
index 6289b8b60a7..421278dacf8 100644
--- a/src/interfaces/libpq/fe-auth.c
+++ b/src/interfaces/libpq/fe-auth.c
@@ -552,6 +552,28 @@ pg_SASL_init(PGconn *conn, int payloadlen, bool *async)
 		goto error;
 	}
 
+	/* Make sure require_auth is satisfied. */
+	if (conn->require_auth)
+	{
+		bool		allowed = false;
+
+		for (int i = 0; i < lengthof(conn->allowed_sasl_mechs); i++)
+		{
+			if (conn->sasl == conn->allowed_sasl_mechs[i])
+			{
+				allowed = true;
+				break;
+			}
+		}
+
+		if (!allowed)
+		{
+			libpq_append_conn_error(conn, "authentication method requirement \"%s\" failed: server requested %s authentication",
+									conn->require_auth, selected_mechanism);
+			goto error;
+		}
+	}
+
 	if (conn->channel_binding[0] == 'r' &&	/* require */
 		strcmp(selected_mechanism, SCRAM_SHA_256_PLUS_NAME) != 0)
 	{
diff --git a/src/interfaces/libpq/fe-connect.c b/src/interfaces/libpq/fe-connect.c
index 28a26a1d362..ff1322cb07d 100644
--- a/src/interfaces/libpq/fe-connect.c
+++ b/src/interfaces/libpq/fe-connect.c
@@ -1129,6 +1129,57 @@ libpq_prng_init(PGconn *conn)
 	pg_prng_seed(&conn->prng_state, rseed);
 }
 
+/*
+ * Fills the connection's allowed_sasl_mechs list with all supported SASL
+ * mechanisms.
+ */
+static inline void
+fill_allowed_sasl_mechs(PGconn *conn)
+{
+	/*---
+	 * We only support two mechanisms at the moment, so rather than deal with a
+	 * linked list, conn->allowed_sasl_mechs is an array of static length. We
+	 * rely on the compile-time assertion here to keep us honest.
+	 *
+	 * To add a new mechanism to require_auth,
+	 * - update the length of conn->allowed_sasl_mechs,
+	 * - add the new pg_fe_sasl_mech pointer to this function, and
+	 * - handle the new mechanism name in the require_auth portion of
+	 *   pqConnectOptions2(), below.
+	 */
+	StaticAssertDecl(lengthof(conn->allowed_sasl_mechs) == 2,
+					 "fill_allowed_sasl_mechs() must be updated when resizing conn->allowed_sasl_mechs[]");
+
+	conn->allowed_sasl_mechs[0] = &pg_scram_mech;
+	conn->allowed_sasl_mechs[1] = &pg_oauth_mech;
+}
+
+/*
+ * Clears the connection's allowed_sasl_mechs list.
+ */
+static inline void
+clear_allowed_sasl_mechs(PGconn *conn)
+{
+	for (int i = 0; i < lengthof(conn->allowed_sasl_mechs); i++)
+		conn->allowed_sasl_mechs[i] = NULL;
+}
+
+/*
+ * Helper routine that searches the static allowed_sasl_mechs list for a
+ * specific mechanism.
+ */
+static inline int
+index_of_allowed_sasl_mech(PGconn *conn, const pg_fe_sasl_mech *mech)
+{
+	for (int i = 0; i < lengthof(conn->allowed_sasl_mechs); i++)
+	{
+		if (conn->allowed_sasl_mechs[i] == mech)
+			return i;
+	}
+
+	return -1;
+}
+
 /*
  *		pqConnectOptions2
  *
@@ -1370,17 +1421,19 @@ pqConnectOptions2(PGconn *conn)
 		bool		negated = false;
 
 		/*
-		 * By default, start from an empty set of allowed options and add to
-		 * it.
+		 * By default, start from an empty set of allowed methods and
+		 * mechanisms, and add to it.
 		 */
 		conn->auth_required = true;
 		conn->allowed_auth_methods = 0;
+		clear_allowed_sasl_mechs(conn);
 
 		for (first = true, more = true; more; first = false)
 		{
 			char	   *method,
 					   *part;
-			uint32		bits;
+			uint32		bits = 0;
+			const pg_fe_sasl_mech *mech = NULL;
 
 			part = parse_comma_separated_list(&s, &more);
 			if (part == NULL)
@@ -1396,11 +1449,12 @@ pqConnectOptions2(PGconn *conn)
 				if (first)
 				{
 					/*
-					 * Switch to a permissive set of allowed options, and
-					 * subtract from it.
+					 * Switch to a permissive set of allowed methods and
+					 * mechanisms, and subtract from it.
 					 */
 					conn->auth_required = false;
 					conn->allowed_auth_methods = -1;
+					fill_allowed_sasl_mechs(conn);
 				}
 				else if (!negated)
 				{
@@ -1425,6 +1479,10 @@ pqConnectOptions2(PGconn *conn)
 				return false;
 			}
 
+			/*
+			 * First group: methods that can be handled solely with the
+			 * authentication request codes.
+			 */
 			if (strcmp(method, "password") == 0)
 			{
 				bits = (1 << AUTH_REQ_PASSWORD);
@@ -1443,13 +1501,26 @@ pqConnectOptions2(PGconn *conn)
 				bits = (1 << AUTH_REQ_SSPI);
 				bits |= (1 << AUTH_REQ_GSS_CONT);
 			}
+
+			/*
+			 * Next group: SASL mechanisms. All of these use the same request
+			 * codes, so the list of allowed mechanisms is tracked separately.
+			 *
+			 * fill_allowed_sasl_mechs() must be updated when adding a new
+			 * mechanism here!
+			 */
 			else if (strcmp(method, "scram-sha-256") == 0)
 			{
-				/* This currently assumes that SCRAM is the only SASL method. */
-				bits = (1 << AUTH_REQ_SASL);
-				bits |= (1 << AUTH_REQ_SASL_CONT);
-				bits |= (1 << AUTH_REQ_SASL_FIN);
+				mech = &pg_scram_mech;
 			}
+			else if (strcmp(method, "oauth") == 0)
+			{
+				mech = &pg_oauth_mech;
+			}
+
+			/*
+			 * Final group: meta-options.
+			 */
 			else if (strcmp(method, "none") == 0)
 			{
 				/*
@@ -1485,20 +1556,68 @@ pqConnectOptions2(PGconn *conn)
 				return false;
 			}
 
-			/* Update the bitmask. */
-			if (negated)
+			if (mech)
 			{
-				if ((conn->allowed_auth_methods & bits) == 0)
-					goto duplicate;
+				/*
+				 * Update the mechanism set only. The method bitmask will be
+				 * updated for SASL further down.
+				 */
+				Assert(!bits);
 
-				conn->allowed_auth_methods &= ~bits;
+				if (negated)
+				{
+					/* Remove the existing mechanism from the list. */
+					i = index_of_allowed_sasl_mech(conn, mech);
+					if (i < 0)
+						goto duplicate;
+
+					conn->allowed_sasl_mechs[i] = NULL;
+				}
+				else
+				{
+					/*
+					 * Find a space to put the new mechanism (after making
+					 * sure it's not already there).
+					 */
+					i = index_of_allowed_sasl_mech(conn, mech);
+					if (i >= 0)
+						goto duplicate;
+
+					i = index_of_allowed_sasl_mech(conn, NULL);
+					if (i < 0)
+					{
+						/* Should not happen; the pointer list is corrupted. */
+						Assert(false);
+
+						conn->status = CONNECTION_BAD;
+						libpq_append_conn_error(conn,
+												"internal error: no space in allowed_sasl_mechs");
+						free(part);
+						return false;
+					}
+
+					conn->allowed_sasl_mechs[i] = mech;
+				}
 			}
 			else
 			{
-				if ((conn->allowed_auth_methods & bits) == bits)
-					goto duplicate;
+				/* Update the method bitmask. */
+				Assert(bits);
 
-				conn->allowed_auth_methods |= bits;
+				if (negated)
+				{
+					if ((conn->allowed_auth_methods & bits) == 0)
+						goto duplicate;
+
+					conn->allowed_auth_methods &= ~bits;
+				}
+				else
+				{
+					if ((conn->allowed_auth_methods & bits) == bits)
+						goto duplicate;
+
+					conn->allowed_auth_methods |= bits;
+				}
 			}
 
 			free(part);
@@ -1517,6 +1636,36 @@ pqConnectOptions2(PGconn *conn)
 			free(part);
 			return false;
 		}
+
+		/*
+		 * Finally, allow SASL authentication requests if (and only if) we've
+		 * allowed any mechanisms.
+		 */
+		{
+			bool		allowed = false;
+			const uint32 sasl_bits =
+				(1 << AUTH_REQ_SASL)
+				| (1 << AUTH_REQ_SASL_CONT)
+				| (1 << AUTH_REQ_SASL_FIN);
+
+			for (i = 0; i < lengthof(conn->allowed_sasl_mechs); i++)
+			{
+				if (conn->allowed_sasl_mechs[i])
+				{
+					allowed = true;
+					break;
+				}
+			}
+
+			/*
+			 * For the standard case, add the SASL bits to the (default-empty)
+			 * set if needed. For the negated case, remove them.
+			 */
+			if (!negated && allowed)
+				conn->allowed_auth_methods |= sasl_bits;
+			else if (negated && !allowed)
+				conn->allowed_auth_methods &= ~sasl_bits;
+		}
 	}
 
 	/*
diff --git a/src/interfaces/libpq/libpq-int.h b/src/interfaces/libpq/libpq-int.h
index 6e6ab2c1f51..fd571119033 100644
--- a/src/interfaces/libpq/libpq-int.h
+++ b/src/interfaces/libpq/libpq-int.h
@@ -511,6 +511,8 @@ struct pg_conn
 								 * the server? */
 	uint32		allowed_auth_methods;	/* bitmask of acceptable AuthRequest
 										 * codes */
+	const pg_fe_sasl_mech *allowed_sasl_mechs[2];	/* and acceptable SASL
+													 * mechanisms */
 	bool		client_finished_auth;	/* have we finished our half of the
 										 * authentication exchange? */
 	char		current_auth_response;	/* used by pqTraceOutputMessage to
diff --git a/src/test/authentication/t/001_password.pl b/src/test/authentication/t/001_password.pl
index 87e180af3d3..0838c080350 100644
--- a/src/test/authentication/t/001_password.pl
+++ b/src/test/authentication/t/001_password.pl
@@ -277,6 +277,16 @@ $node->connect_fails(
 	"require_auth methods cannot be duplicated, !none case",
 	expected_stderr =>
 	  qr/require_auth method "!none" is specified more than once/);
+$node->connect_fails(
+	"user=scram_role require_auth=scram-sha-256,scram-sha-256",
+	"require_auth methods cannot be duplicated, scram-sha-256 case",
+	expected_stderr =>
+	  qr/require_auth method "scram-sha-256" is specified more than once/);
+$node->connect_fails(
+	"user=scram_role require_auth=!scram-sha-256,!scram-sha-256",
+	"require_auth methods cannot be duplicated, !scram-sha-256 case",
+	expected_stderr =>
+	  qr/require_auth method "!scram-sha-256" is specified more than once/);
 
 # Unknown value defined in require_auth.
 $node->connect_fails(
@@ -394,11 +404,11 @@ $node->connect_fails(
 $node->connect_fails(
 	"user=scram_role require_auth=!scram-sha-256",
 	"SCRAM authentication forbidden, fails with SCRAM auth",
-	expected_stderr => qr/server requested SASL authentication/);
+	expected_stderr => qr/server requested SCRAM-SHA-256 authentication/);
 $node->connect_fails(
 	"user=scram_role require_auth=!password,!md5,!scram-sha-256",
 	"multiple authentication types forbidden, fails with SCRAM auth",
-	expected_stderr => qr/server requested SASL authentication/);
+	expected_stderr => qr/server requested SCRAM-SHA-256 authentication/);
 
 # Test that bad passwords are rejected.
 $ENV{"PGPASSWORD"} = 'badpass';
@@ -455,13 +465,13 @@ $node->connect_fails(
 	"user=scram_role require_auth=!scram-sha-256",
 	"password authentication forbidden, fails with SCRAM auth",
 	expected_stderr =>
-	  qr/authentication method requirement "!scram-sha-256" failed: server requested SASL authentication/
+	  qr/authentication method requirement "!scram-sha-256" failed: server requested SCRAM-SHA-256 authentication/
 );
 $node->connect_fails(
 	"user=scram_role require_auth=!password,!md5,!scram-sha-256",
 	"multiple authentication types forbidden, fails with SCRAM auth",
 	expected_stderr =>
-	  qr/authentication method requirement "!password,!md5,!scram-sha-256" failed: server requested SASL authentication/
+	  qr/authentication method requirement "!password,!md5,!scram-sha-256" failed: server requested SCRAM-SHA-256 authentication/
 );
 
 # Test SYSTEM_USER <> NULL with parallel workers.
diff --git a/src/test/modules/oauth_validator/t/001_server.pl b/src/test/modules/oauth_validator/t/001_server.pl
index 10d2d3da929..77a106b2cc2 100644
--- a/src/test/modules/oauth_validator/t/001_server.pl
+++ b/src/test/modules/oauth_validator/t/001_server.pl
@@ -124,6 +124,60 @@ $node->connect_fails(
 	  qr@server's discovery document at \Q$issuer/.well-known/oauth-authorization-server/alternate\E \(issuer "\Q$issuer/alternate\E"\) is incompatible with oauth_issuer \(\Q$issuer\E\)@
 );
 
+# Test require_auth settings against OAUTHBEARER.
+my @cases = (
+	{ require_auth => "oauth" },
+	{ require_auth => "oauth,scram-sha-256" },
+	{ require_auth => "password,oauth" },
+	{ require_auth => "none,oauth" },
+	{ require_auth => "!scram-sha-256" },
+	{ require_auth => "!none" },
+
+	{
+		require_auth => "!oauth",
+		failure => qr/server requested OAUTHBEARER authentication/
+	},
+	{
+		require_auth => "scram-sha-256",
+		failure => qr/server requested OAUTHBEARER authentication/
+	},
+	{
+		require_auth => "!password,!oauth",
+		failure => qr/server requested OAUTHBEARER authentication/
+	},
+	{
+		require_auth => "none",
+		failure => qr/server requested SASL authentication/
+	},
+	{
+		require_auth => "!oauth,!scram-sha-256",
+		failure => qr/server requested SASL authentication/
+	});
+
+$user = "test";
+foreach my $c (@cases)
+{
+	my $connstr =
+	  "user=$user dbname=postgres oauth_issuer=$issuer oauth_client_id=f02c6361-0635 require_auth=$c->{'require_auth'}";
+
+	if (defined $c->{'failure'})
+	{
+		$node->connect_fails(
+			$connstr,
+			"require_auth=$c->{'require_auth'} fails",
+			expected_stderr => $c->{'failure'});
+	}
+	else
+	{
+		$node->connect_ok(
+			$connstr,
+			"require_auth=$c->{'require_auth'} succeeds",
+			expected_stderr =>
+			  qr@Visit https://example\.com/ and enter the code: postgresuser@
+		);
+	}
+}
+
 # Make sure the client_id and secret are correctly encoded. $vschars contains
 # every allowed character for a client_id/_secret (the "VSCHAR" class).
 # $vschars_esc is additionally backslash-escaped for inclusion in a
@@ -134,15 +188,15 @@ my $vschars_esc =
   " !\"#\$%&\\'()*+,-./0123456789:;<=>?\@ABCDEFGHIJKLMNOPQRSTUVWXYZ[\\\\]^_`abcdefghijklmnopqrstuvwxyz{|}~";
 
 $node->connect_ok(
-	"user=$user dbname=postgres oauth_issuer=$issuer/alternate oauth_client_id='$vschars_esc'",
+	"user=$user dbname=postgres oauth_issuer=$issuer oauth_client_id='$vschars_esc'",
 	"escapable characters: client_id",
 	expected_stderr =>
-	  qr@Visit https://example\.org/ and enter the code: postgresuser@);
+	  qr@Visit https://example\.com/ and enter the code: postgresuser@);
 $node->connect_ok(
-	"user=$user dbname=postgres oauth_issuer=$issuer/alternate oauth_client_id='$vschars_esc' oauth_client_secret='$vschars_esc'",
+	"user=$user dbname=postgres oauth_issuer=$issuer oauth_client_id='$vschars_esc' oauth_client_secret='$vschars_esc'",
 	"escapable characters: client_id and secret",
 	expected_stderr =>
-	  qr@Visit https://example\.org/ and enter the code: postgresuser@);
+	  qr@Visit https://example\.com/ and enter the code: postgresuser@);
 
 #
 # Further tests rely on support for specific behaviors in oauth_server.py. To
-- 
2.34.1