v41-0005-squash-Add-OAUTHBEARER-SASL-mechanism.patch

From 18507c6978b7c357cb7d3d28743f37eef65643fa 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 v41 5/7] 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                |  7 ---
 src/interfaces/libpq/fe-connect.c             |  9 ++-
 src/interfaces/libpq/libpq-int.h              |  2 +-
 src/test/authentication/t/001_password.pl     |  8 +--
 .../modules/oauth_validator/t/001_server.pl   | 62 +++++++++++++++++--
 7 files changed, 86 insertions(+), 18 deletions(-)

diff --git a/doc/src/sgml/libpq.sgml b/doc/src/sgml/libpq.sgml
index aa92baf1fb3..9d132c6c4bd 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 f9133ad57e8..7a94de9c034 100644
--- a/src/interfaces/libpq/fe-auth-oauth.c
+++ b/src/interfaces/libpq/fe-auth-oauth.c
@@ -955,6 +955,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 dc7b1a2e725..27868506350 100644
--- a/src/interfaces/libpq/fe-auth.c
+++ b/src/interfaces/libpq/fe-auth.c
@@ -568,13 +568,6 @@ pg_SASL_init(PGconn *conn, int payloadlen, bool *async)
 
 		if (!allowed)
 		{
-			/*
-			 * TODO: this is dead code until a second SASL mechanism is added;
-			 * the connection can't have proceeded past check_expected_areq()
-			 * if no SASL methods are allowed.
-			 */
-			Assert(false);
-
 			libpq_append_conn_error(conn, "authentication method requirement \"%s\" failed: server requested %s authentication",
 									conn->require_auth, selected_mechanism);
 			goto error;
diff --git a/src/interfaces/libpq/fe-connect.c b/src/interfaces/libpq/fe-connect.c
index d2d967b86d5..6ba6432d750 100644
--- a/src/interfaces/libpq/fe-connect.c
+++ b/src/interfaces/libpq/fe-connect.c
@@ -1150,7 +1150,7 @@ static inline void
 fill_allowed_sasl_mechs(PGconn *conn)
 {
 	/*---
-	 * We only support one mechanism at the moment, so rather than deal with a
+	 * 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.
 	 *
@@ -1160,10 +1160,11 @@ fill_allowed_sasl_mechs(PGconn *conn)
 	 * - handle the new mechanism name in the require_auth portion of
 	 *   pqConnectOptions2(), below.
 	 */
-	StaticAssertDecl(lengthof(conn->allowed_sasl_mechs) == 1,
+	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;
 }
 
 /*
@@ -1525,6 +1526,10 @@ pqConnectOptions2(PGconn *conn)
 			{
 				mech = &pg_scram_mech;
 			}
+			else if (strcmp(method, "oauth") == 0)
+			{
+				mech = &pg_oauth_mech;
+			}
 
 			/*
 			 * Final group: meta-options.
diff --git a/src/interfaces/libpq/libpq-int.h b/src/interfaces/libpq/libpq-int.h
index 68645076227..7f0fcb9ee5a 100644
--- a/src/interfaces/libpq/libpq-int.h
+++ b/src/interfaces/libpq/libpq-int.h
@@ -511,7 +511,7 @@ struct pg_conn
 								 * the server? */
 	uint32		allowed_auth_methods;	/* bitmask of acceptable AuthRequest
 										 * codes */
-	const pg_fe_sasl_mech *allowed_sasl_mechs[1];	/* and acceptable SASL
+	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? */
diff --git a/src/test/authentication/t/001_password.pl b/src/test/authentication/t/001_password.pl
index 1357f806b6f..4ce22ccbdf2 100644
--- a/src/test/authentication/t/001_password.pl
+++ b/src/test/authentication/t/001_password.pl
@@ -404,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';
@@ -465,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 98dd532e133..96040e5ba95 100644
--- a/src/test/modules/oauth_validator/t/001_server.pl
+++ b/src/test/modules/oauth_validator/t/001_server.pl
@@ -134,6 +134,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
@@ -144,15 +198,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