v47-0009-fixup-Add-OAUTHBEARER-SASL-mechanism.patch

From 27135876559a1c5b5cac7e68c43403f257f391fd Mon Sep 17 00:00:00 2001
From: Jacob Champion <jacob.champion@enterprisedb.com>
Date: Tue, 4 Feb 2025 17:00:47 -0800
Subject: [PATCH v47 09/11] fixup! Add OAUTHBEARER SASL mechanism

---
 src/interfaces/libpq/fe-auth-oauth-curl.c     | 131 +++++++++++++++---
 .../oauth_validator/oauth_hook_client.c       |  35 ++++-
 .../modules/oauth_validator/t/001_server.pl   |  15 ++
 3 files changed, 159 insertions(+), 22 deletions(-)

diff --git a/src/interfaces/libpq/fe-auth-oauth-curl.c b/src/interfaces/libpq/fe-auth-oauth-curl.c
index 02c5b50afcd..96c5096e4ca 100644
--- a/src/interfaces/libpq/fe-auth-oauth-curl.c
+++ b/src/interfaces/libpq/fe-auth-oauth-curl.c
@@ -167,9 +167,7 @@ struct async_ctx
 {
 	enum OAuthStep step;		/* where are we in the flow? */
 
-#ifdef HAVE_SYS_EPOLL_H
-	int			timerfd;		/* a timerfd for signaling async timeouts */
-#endif
+	int			timerfd;		/* descriptor for signaling async timeouts */
 	pgsocket	mux;			/* the multiplexer socket containing all
 								 * descriptors tracked by libcurl, plus the
 								 * timerfd */
@@ -275,10 +273,8 @@ free_async_ctx(PGconn *conn, struct async_ctx *actx)
 
 	if (actx->mux != PGINVALID_SOCKET)
 		close(actx->mux);
-#ifdef HAVE_SYS_EPOLL_H
 	if (actx->timerfd >= 0)
 		close(actx->timerfd);
-#endif
 
 	free(actx);
 }
@@ -1089,8 +1085,9 @@ parse_access_token(struct async_ctx *actx, struct token *tok)
  * select() on instead of the Postgres socket during OAuth negotiation.
  *
  * This is just an epoll set or kqueue abstracting multiple other descriptors.
- * A timerfd is always part of the set when using epoll; it's just disabled
- * when we're not using it.
+ * For epoll, the timerfd is always part of the set; it's just disabled when
+ * we're not using it. For kqueue, the "timerfd" is actually a second kqueue
+ * instance which is only added to the set when needed.
  */
 static bool
 setup_multiplexer(struct async_ctx *actx)
@@ -1128,6 +1125,19 @@ setup_multiplexer(struct async_ctx *actx)
 		return false;
 	}
 
+	/*
+	 * Originally, we set EVFILT_TIMER directly on the top-level multiplexer.
+	 * This makes it difficult to implement timer_expired(), though, so now we
+	 * set EVFILT_TIMER on a separate actx->timerfd, which is chained to
+	 * actx->mux while the timer is active.
+	 */
+	actx->timerfd = kqueue();
+	if (actx->timerfd < 0)
+	{
+		actx_error(actx, "failed to create timer kqueue: %m");
+		return false;
+	}
+
 	return true;
 #endif
 
@@ -1286,9 +1296,12 @@ register_socket(CURL *curl, curl_socket_t socket, int what, void *ctx,
 
 /*
  * Enables or disables the timer in the multiplexer set. The timeout value is
- * in milliseconds (negative values disable the timer). Rather than continually
- * adding and removing the timer, we keep it in the set at all times and just
- * disarm it when it's not needed.
+ * in milliseconds (negative values disable the timer).
+ *
+ * For epoll, rather than continually adding and removing the timer, we keep it
+ * in the set at all times and just disarm it when it's not needed. For kqueue,
+ * the timer is removed completely when disabled to prevent stale timeouts from
+ * remaining in the queue.
  */
 static bool
 set_timer(struct async_ctx *actx, long timeout)
@@ -1320,20 +1333,87 @@ set_timer(struct async_ctx *actx, long timeout)
 		actx_error(actx, "setting timerfd to %ld: %m", timeout);
 		return false;
 	}
+
+	return true;
 #endif
 #ifdef HAVE_SYS_EVENT_H
 	struct kevent ev;
 
+	/* Enable/disable the timer itself. */
 	EV_SET(&ev, 1, EVFILT_TIMER, timeout < 0 ? EV_DELETE : (EV_ADD | EV_ONESHOT),
 		   0, timeout, 0);
-	if (kevent(actx->mux, &ev, 1, NULL, 0, NULL) < 0 && errno != ENOENT)
+	if (kevent(actx->timerfd, &ev, 1, NULL, 0, NULL) < 0 && errno != ENOENT)
 	{
 		actx_error(actx, "setting kqueue timer to %ld: %m", timeout);
 		return false;
 	}
-#endif
+
+	/*
+	 * Add/remove the timer to/from the mux. (In contrast with epoll, if we
+	 * allowed the timer to remain registered here after being disabled, the
+	 * mux queue would retain any previous stale timeout notifications and
+	 * remain readable.)
+	 */
+	EV_SET(&ev, actx->timerfd, EVFILT_READ, timeout < 0 ? EV_DELETE : EV_ADD,
+		   0, 0, 0);
+	if (kevent(actx->mux, &ev, 1, NULL, 0, NULL) < 0 && errno != ENOENT)
+	{
+		actx_error(actx, "could not update timer on kqueue: %m");
+		return false;
+	}
 
 	return true;
+#endif
+
+	actx_error(actx, "libpq does not support timers on this platform");
+	return false;
+}
+
+/*
+ * Returns 1 if the timeout in the multiplexer set has expired since the last
+ * call to set_timer(), 0 if the timer is still running, or -1 (with an
+ * actx_error() report) if the timer cannot be queried.
+ */
+static int
+timer_expired(struct async_ctx *actx)
+{
+#if HAVE_SYS_EPOLL_H
+	struct itimerspec spec = {0};
+
+	if (timerfd_gettime(actx->timerfd, &spec) < 0)
+	{
+		actx_error(actx, "getting timerfd value: %m");
+		return -1;
+	}
+
+	/*
+	 * This implementation assumes we're using single-shot timers. If you
+	 * change to using intervals, you'll need to reimplement this function
+	 * too, possibly with the read() or select() interfaces for timerfd.
+	 */
+	Assert(spec.it_interval.tv_sec == 0
+		   && spec.it_interval.tv_nsec == 0);
+
+	/* If the remaining time to expiration is zero, we're done. */
+	return (spec.it_value.tv_sec == 0
+			&& spec.it_value.tv_nsec == 0);
+#endif
+#ifdef HAVE_SYS_EVENT_H
+	int			res;
+
+	/* Is the timer queue ready? */
+	res = PQsocketPoll(actx->timerfd, 1 /* forRead */ , 0, 0);
+	if (res < 0)
+	{
+		actx_error(actx, "checking kqueue for timeout: %m");
+		return -1;
+	}
+
+	return (res > 0);
+#endif
+
+	actx_error(actx, "libpq does not support timers on this platform");
+	return -1;
 }
 
 /*
@@ -2510,9 +2590,7 @@ pg_fe_run_oauth_flow_impl(PGconn *conn)
 		}
 
 		actx->mux = PGINVALID_SOCKET;
-#ifdef HAVE_SYS_EPOLL_H
 		actx->timerfd = -1;
-#endif
 
 		/* Should we enable unsafe features? */
 		actx->debugging = oauth_unsafe_debugging_enabled();
@@ -2556,10 +2634,28 @@ pg_fe_run_oauth_flow_impl(PGconn *conn)
 						/* not done yet */
 						return status;
 					}
+
+					break;
 				}
 
 			case OAUTH_STEP_WAIT_INTERVAL:
-				/* TODO check that the timer has expired */
+
+				/*
+				 * The client application is supposed to wait until our timer
+				 * expires before calling PQconnectPoll() again, but that
+				 * might not happen. To avoid sending a token request early,
+				 * check the timer before continuing.
+				 */
+				if (!timer_expired(actx))
+				{
+					conn->altsock = actx->timerfd;
+					return PGRES_POLLING_READING;
+				}
+
+				/* Disable the expired timer. */
+				if (!set_timer(actx, -1))
+					goto error_return;
+
 				break;
 		}
 
@@ -2633,15 +2729,12 @@ pg_fe_run_oauth_flow_impl(PGconn *conn)
 				if (!set_timer(actx, actx->authz.interval * 1000))
 					goto error_return;
 
-#ifdef HAVE_SYS_EPOLL_H
-
 				/*
 				 * No Curl requests are running, so we can simplify by having
 				 * the client wait directly on the timerfd rather than the
-				 * multiplexer. (This isn't possible for kqueue.)
+				 * multiplexer.
 				 */
 				conn->altsock = actx->timerfd;
-#endif
 
 				actx->step = OAUTH_STEP_WAIT_INTERVAL;
 				actx->running = 1;
diff --git a/src/test/modules/oauth_validator/oauth_hook_client.c b/src/test/modules/oauth_validator/oauth_hook_client.c
index 12fe70c990b..fc003030ff8 100644
--- a/src/test/modules/oauth_validator/oauth_hook_client.c
+++ b/src/test/modules/oauth_validator/oauth_hook_client.c
@@ -40,14 +40,16 @@ usage(char *argv[])
 	printf(" --expected-uri URI		fail if received configuration link does not match URI\n");
 	printf(" --misbehave=MODE		have the hook fail required postconditions\n"
 		   "						(MODEs: no-hook, fail-async, no-token, no-socket)\n");
-	printf(" --no-hook				don't install OAuth hooks (connection will fail)\n");
+	printf(" --no-hook				don't install OAuth hooks\n");
 	printf(" --hang-forever			don't ever return a token (combine with connect_timeout)\n");
 	printf(" --token TOKEN			use the provided TOKEN value\n");
+	printf(" --stress-async			busy-loop on PQconnectPoll rather than polling\n");
 }
 
 /* --options */
 static bool no_hook = false;
 static bool hang_forever = false;
+static bool stress_async = false;
 static const char *expected_uri = NULL;
 static const char *expected_scope = NULL;
 static const char *misbehave_mode = NULL;
@@ -65,6 +67,7 @@ main(int argc, char *argv[])
 		{"token", required_argument, NULL, 1003},
 		{"hang-forever", no_argument, NULL, 1004},
 		{"misbehave", required_argument, NULL, 1005},
+		{"stress-async", no_argument, NULL, 1006},
 		{0}
 	};
 
@@ -104,6 +107,10 @@ main(int argc, char *argv[])
 				misbehave_mode = optarg;
 				break;
 
+			case 1006:			/* --stress-async */
+				stress_async = true;
+				break;
+
 			default:
 				usage(argv);
 				return 1;
@@ -122,10 +129,32 @@ main(int argc, char *argv[])
 	PQsetAuthDataHook(handle_auth_data);
 
 	/* Connect. (All the actual work is in the hook.) */
-	conn = PQconnectdb(conninfo);
+	if (stress_async)
+	{
+		/*
+		 * Perform an asynchronous connection, busy-looping on PQconnectPoll()
+		 * without actually waiting on socket events. This stresses code paths
+		 * that rely on asynchronous work to be done before continuing with
+		 * the next step in the flow.
+		 */
+		PostgresPollingStatusType res;
+
+		conn = PQconnectStart(conninfo);
+
+		do
+		{
+			res = PQconnectPoll(conn);
+		} while (res != PGRES_POLLING_FAILED && res != PGRES_POLLING_OK);
+	}
+	else
+	{
+		/* Perform a standard synchronous connection. */
+		conn = PQconnectdb(conninfo);
+	}
+
 	if (PQstatus(conn) != CONNECTION_OK)
 	{
-		fprintf(stderr, "Connection to database failed: %s\n",
+		fprintf(stderr, "connection to database failed: %s\n",
 				PQerrorMessage(conn));
 		PQfinish(conn);
 		return 1;
diff --git a/src/test/modules/oauth_validator/t/001_server.pl b/src/test/modules/oauth_validator/t/001_server.pl
index 80f52585896..f0b918390fd 100644
--- a/src/test/modules/oauth_validator/t/001_server.pl
+++ b/src/test/modules/oauth_validator/t/001_server.pl
@@ -389,6 +389,21 @@ $node->connect_fails(
 	  qr/failed to obtain access token: mutual TLS required for client \(invalid_client\)/
 );
 
+# Stress test: make sure our builtin flow operates correctly even if the client
+# application isn't respecting PGRES_POLLING_READING/WRITING signals returned
+# from PQconnectPoll().
+$base_connstr =
+  "$common_connstr port=" . $node->port . " host=" . $node->host;
+my @cmd = (
+	"oauth_hook_client", "--no-hook", "--stress-async",
+	connstr(stage => 'all', retries => 1, interval => 1));
+
+note "running '" . join("' '", @cmd) . "'";
+my ($stdout, $stderr) = run_command(\@cmd);
+
+like($stdout, qr/connection succeeded/, "stress-async: stdout matches");
+unlike($stderr, qr/connection to database failed/, "stress-async: stderr matches");
+
 #
 # This section of tests reconfigures the validator module itself, rather than
 # the OAuth server.
-- 
2.34.1