v42-0003-libpq-handle-asynchronous-actions-during-SASL.patch

From bda684d19cc7a9a24c84f72176b6a84cc08ebecd Mon Sep 17 00:00:00 2001
From: Jacob Champion <jacob.champion@enterprisedb.com>
Date: Wed, 8 Jan 2025 09:30:05 -0800
Subject: [PATCH v42 3/6] libpq: handle asynchronous actions during SASL

This adds the ability for a SASL mechanism to signal to PQconnectPoll()
that some arbitrary work must be done, external to the Postgres
connection, before authentication can continue. The intent is for the
upcoming OAUTHBEARER mechanism to make use of this functionality.

To ensure that threads are not blocked waiting for the SASL mechanism to
make long-running calls, the mechanism communicates with the top-level
client via the "altsock": a file or socket descriptor, opaque to this
layer of libpq, which is signaled when work is ready to be done again.
This socket temporarily takes the place of the standard connection
descriptor, so PQsocket() clients should continue to operate correctly
using their existing polling implementations.

A mechanism should set an authentication callback (conn->async_auth())
and a cleanup callback (conn->cleanup_async_auth()), return SASL_ASYNC
during the exchange, and assign conn->altsock during the first call to
async_auth(). When the cleanup callback is called, either because
authentication has succeeded or because the connection is being
dropped, the altsock must be released and disconnected from the PGconn.
---
 src/interfaces/libpq/fe-auth-sasl.h  |  11 ++-
 src/interfaces/libpq/fe-auth-scram.c |   6 +-
 src/interfaces/libpq/fe-auth.c       | 114 +++++++++++++++++++--------
 src/interfaces/libpq/fe-auth.h       |   3 +-
 src/interfaces/libpq/fe-connect.c    |  73 ++++++++++++++++-
 src/interfaces/libpq/fe-misc.c       |  35 +++++---
 src/interfaces/libpq/libpq-fe.h      |   2 +
 src/interfaces/libpq/libpq-int.h     |   6 ++
 8 files changed, 201 insertions(+), 49 deletions(-)

diff --git a/src/interfaces/libpq/fe-auth-sasl.h b/src/interfaces/libpq/fe-auth-sasl.h
index f0c62139092..f06f547c07d 100644
--- a/src/interfaces/libpq/fe-auth-sasl.h
+++ b/src/interfaces/libpq/fe-auth-sasl.h
@@ -30,6 +30,7 @@ typedef enum
 	SASL_COMPLETE = 0,
 	SASL_FAILED,
 	SASL_CONTINUE,
+	SASL_ASYNC,
 } SASLStatus;
 
 /*
@@ -77,6 +78,8 @@ typedef struct pg_fe_sasl_mech
 	 *
 	 *	state:	   The opaque mechanism state returned by init()
 	 *
+	 *	final:	   true if the server has sent a final exchange outcome
+	 *
 	 *	input:	   The challenge data sent by the server, or NULL when
 	 *			   generating a client-first initial response (that is, when
 	 *			   the server expects the client to send a message to start
@@ -101,12 +104,18 @@ typedef struct pg_fe_sasl_mech
 	 *
 	 *	SASL_CONTINUE:	The output buffer is filled with a client response.
 	 *					Additional server challenge is expected
+	 *	SASL_ASYNC:		Some asynchronous processing external to the
+	 *					connection needs to be done before a response can be
+	 *					generated. The mechanism is responsible for setting up
+	 *					conn->async_auth/cleanup_async_auth appropriately
+	 *					before returning.
 	 *	SASL_COMPLETE:	The SASL exchange has completed successfully.
 	 *	SASL_FAILED:	The exchange has failed and the connection should be
 	 *					dropped.
 	 *--------
 	 */
-	SASLStatus	(*exchange) (void *state, char *input, int inputlen,
+	SASLStatus	(*exchange) (void *state, bool final,
+							 char *input, int inputlen,
 							 char **output, int *outputlen);
 
 	/*--------
diff --git a/src/interfaces/libpq/fe-auth-scram.c b/src/interfaces/libpq/fe-auth-scram.c
index 557e9c568b6..fe18615197f 100644
--- a/src/interfaces/libpq/fe-auth-scram.c
+++ b/src/interfaces/libpq/fe-auth-scram.c
@@ -24,7 +24,8 @@
 /* The exported SCRAM callback mechanism. */
 static void *scram_init(PGconn *conn, const char *password,
 						const char *sasl_mechanism);
-static SASLStatus scram_exchange(void *opaq, char *input, int inputlen,
+static SASLStatus scram_exchange(void *opaq, bool final,
+								 char *input, int inputlen,
 								 char **output, int *outputlen);
 static bool scram_channel_bound(void *opaq);
 static void scram_free(void *opaq);
@@ -205,7 +206,8 @@ scram_free(void *opaq)
  * Exchange a SCRAM message with backend.
  */
 static SASLStatus
-scram_exchange(void *opaq, char *input, int inputlen,
+scram_exchange(void *opaq, bool final,
+			   char *input, int inputlen,
 			   char **output, int *outputlen)
 {
 	fe_scram_state *state = (fe_scram_state *) opaq;
diff --git a/src/interfaces/libpq/fe-auth.c b/src/interfaces/libpq/fe-auth.c
index 70753d8ec29..a9a23eb97ee 100644
--- a/src/interfaces/libpq/fe-auth.c
+++ b/src/interfaces/libpq/fe-auth.c
@@ -430,7 +430,7 @@ pg_SSPI_startup(PGconn *conn, int use_negotiate, int payloadlen)
  * Initialize SASL authentication exchange.
  */
 static int
-pg_SASL_init(PGconn *conn, int payloadlen)
+pg_SASL_init(PGconn *conn, int payloadlen, bool *async)
 {
 	char	   *initialresponse = NULL;
 	int			initialresponselen;
@@ -448,7 +448,7 @@ pg_SASL_init(PGconn *conn, int payloadlen)
 		goto error;
 	}
 
-	if (conn->sasl_state)
+	if (conn->sasl_state && !conn->async_auth)
 	{
 		libpq_append_conn_error(conn, "duplicate SASL authentication request");
 		goto error;
@@ -607,26 +607,48 @@ pg_SASL_init(PGconn *conn, int payloadlen)
 
 	Assert(conn->sasl);
 
-	/*
-	 * Initialize the SASL state information with all the information gathered
-	 * during the initial exchange.
-	 *
-	 * Note: Only tls-unique is supported for the moment.
-	 */
-	conn->sasl_state = conn->sasl->init(conn,
-										password,
-										selected_mechanism);
 	if (!conn->sasl_state)
-		goto oom_error;
+	{
+		/*
+		 * Initialize the SASL state information with all the information
+		 * gathered during the initial exchange.
+		 *
+		 * Note: Only tls-unique is supported for the moment.
+		 */
+		conn->sasl_state = conn->sasl->init(conn,
+											password,
+											selected_mechanism);
+		if (!conn->sasl_state)
+			goto oom_error;
+	}
+	else
+	{
+		/*
+		 * This is only possible if we're returning from an async loop.
+		 * Disconnect it now.
+		 */
+		Assert(conn->async_auth);
+		conn->async_auth = NULL;
+	}
 
 	/* Get the mechanism-specific Initial Client Response, if any */
-	status = conn->sasl->exchange(conn->sasl_state,
+	status = conn->sasl->exchange(conn->sasl_state, false,
 								  NULL, -1,
 								  &initialresponse, &initialresponselen);
 
 	if (status == SASL_FAILED)
 		goto error;
 
+	if (status == SASL_ASYNC)
+	{
+		/*
+		 * The mechanism should have set up the necessary callbacks; all we
+		 * need to do is signal the caller.
+		 */
+		*async = true;
+		return STATUS_OK;
+	}
+
 	/*
 	 * Build a SASLInitialResponse message, and send it.
 	 */
@@ -671,7 +693,7 @@ oom_error:
  * the protocol.
  */
 static int
-pg_SASL_continue(PGconn *conn, int payloadlen, bool final)
+pg_SASL_continue(PGconn *conn, int payloadlen, bool final, bool *async)
 {
 	char	   *output;
 	int			outputlen;
@@ -701,11 +723,25 @@ pg_SASL_continue(PGconn *conn, int payloadlen, bool final)
 	/* For safety and convenience, ensure the buffer is NULL-terminated. */
 	challenge[payloadlen] = '\0';
 
-	status = conn->sasl->exchange(conn->sasl_state,
+	status = conn->sasl->exchange(conn->sasl_state, final,
 								  challenge, payloadlen,
 								  &output, &outputlen);
 	free(challenge);			/* don't need the input anymore */
 
+	if (status == SASL_ASYNC)
+	{
+		/*
+		 * The mechanism should have set up the necessary callbacks; all we
+		 * need to do is signal the caller.
+		 */
+		*async = true;
+
+		/*
+		 * The mechanism may optionally generate some output to send before
+		 * switching over to async auth, so continue onwards.
+		 */
+	}
+
 	if (final && status == SASL_CONTINUE)
 	{
 		if (outputlen != 0)
@@ -1013,12 +1049,18 @@ check_expected_areq(AuthRequest areq, PGconn *conn)
  * it. We are responsible for reading any remaining extra data, specific
  * to the authentication method. 'payloadlen' is the remaining length in
  * the message.
+ *
+ * If *async is set to true on return, the client doesn't yet have enough
+ * information to respond, and the caller must temporarily switch to
+ * conn->async_auth() to continue driving the exchange.
  */
 int
-pg_fe_sendauth(AuthRequest areq, int payloadlen, PGconn *conn)
+pg_fe_sendauth(AuthRequest areq, int payloadlen, PGconn *conn, bool *async)
 {
 	int			oldmsglen;
 
+	*async = false;
+
 	if (!check_expected_areq(areq, conn))
 		return STATUS_ERROR;
 
@@ -1176,7 +1218,7 @@ pg_fe_sendauth(AuthRequest areq, int payloadlen, PGconn *conn)
 			 * The request contains the name (as assigned by IANA) of the
 			 * authentication mechanism.
 			 */
-			if (pg_SASL_init(conn, payloadlen) != STATUS_OK)
+			if (pg_SASL_init(conn, payloadlen, async) != STATUS_OK)
 			{
 				/* pg_SASL_init already set the error message */
 				return STATUS_ERROR;
@@ -1185,23 +1227,33 @@ pg_fe_sendauth(AuthRequest areq, int payloadlen, PGconn *conn)
 
 		case AUTH_REQ_SASL_CONT:
 		case AUTH_REQ_SASL_FIN:
-			if (conn->sasl_state == NULL)
 			{
-				appendPQExpBufferStr(&conn->errorMessage,
-									 "fe_sendauth: invalid authentication request from server: AUTH_REQ_SASL_CONT without AUTH_REQ_SASL\n");
-				return STATUS_ERROR;
-			}
-			oldmsglen = conn->errorMessage.len;
-			if (pg_SASL_continue(conn, payloadlen,
-								 (areq == AUTH_REQ_SASL_FIN)) != STATUS_OK)
-			{
-				/* Use this message if pg_SASL_continue didn't supply one */
-				if (conn->errorMessage.len == oldmsglen)
+				bool		final = false;
+
+				if (conn->sasl_state == NULL)
+				{
 					appendPQExpBufferStr(&conn->errorMessage,
-										 "fe_sendauth: error in SASL authentication\n");
-				return STATUS_ERROR;
+										 "fe_sendauth: invalid authentication request from server: AUTH_REQ_SASL_CONT without AUTH_REQ_SASL\n");
+					return STATUS_ERROR;
+				}
+				oldmsglen = conn->errorMessage.len;
+
+				if (areq == AUTH_REQ_SASL_FIN)
+					final = true;
+
+				if (pg_SASL_continue(conn, payloadlen, final, async) != STATUS_OK)
+				{
+					/*
+					 * Append a generic error message unless pg_SASL_continue
+					 * did set a more specific one already.
+					 */
+					if (conn->errorMessage.len == oldmsglen)
+						appendPQExpBufferStr(&conn->errorMessage,
+											 "fe_sendauth: error in SASL authentication\n");
+					return STATUS_ERROR;
+				}
+				break;
 			}
-			break;
 
 		default:
 			libpq_append_conn_error(conn, "authentication method %u not supported", areq);
diff --git a/src/interfaces/libpq/fe-auth.h b/src/interfaces/libpq/fe-auth.h
index df0a68b0b21..1d4991f8996 100644
--- a/src/interfaces/libpq/fe-auth.h
+++ b/src/interfaces/libpq/fe-auth.h
@@ -19,7 +19,8 @@
 
 
 /* Prototypes for functions in fe-auth.c */
-extern int	pg_fe_sendauth(AuthRequest areq, int payloadlen, PGconn *conn);
+extern int	pg_fe_sendauth(AuthRequest areq, int payloadlen, PGconn *conn,
+						   bool *async);
 extern char *pg_fe_getusername(uid_t user_id, PQExpBuffer errorMessage);
 extern char *pg_fe_getauthname(PQExpBuffer errorMessage);
 
diff --git a/src/interfaces/libpq/fe-connect.c b/src/interfaces/libpq/fe-connect.c
index ccbcbb7acda..a90f261cdb7 100644
--- a/src/interfaces/libpq/fe-connect.c
+++ b/src/interfaces/libpq/fe-connect.c
@@ -501,6 +501,19 @@ pqDropConnection(PGconn *conn, bool flushInput)
 	conn->cmd_queue_recycle = NULL;
 
 	/* Free authentication/encryption state */
+	if (conn->cleanup_async_auth)
+	{
+		/*
+		 * Any in-progress async authentication should be torn down first so
+		 * that cleanup_async_auth() can depend on the other authentication
+		 * state if necessary.
+		 */
+		conn->cleanup_async_auth(conn);
+		conn->cleanup_async_auth = NULL;
+	}
+	conn->async_auth = NULL;
+	conn->altsock = PGINVALID_SOCKET;	/* cleanup_async_auth() should have
+										 * done this, but make sure. */
 #ifdef ENABLE_GSS
 	{
 		OM_uint32	min_s;
@@ -2847,6 +2860,7 @@ PQconnectPoll(PGconn *conn)
 		case CONNECTION_NEEDED:
 		case CONNECTION_GSS_STARTUP:
 		case CONNECTION_CHECK_TARGET:
+		case CONNECTION_AUTHENTICATING:
 			break;
 
 		default:
@@ -3882,6 +3896,7 @@ keep_going:						/* We will come back to here until there is
 				int			avail;
 				AuthRequest areq;
 				int			res;
+				bool		async;
 
 				/*
 				 * Scan the message from current point (note that if we find
@@ -4070,7 +4085,17 @@ keep_going:						/* We will come back to here until there is
 				 * Note that conn->pghost must be non-NULL if we are going to
 				 * avoid the Kerberos code doing a hostname look-up.
 				 */
-				res = pg_fe_sendauth(areq, msgLength, conn);
+				res = pg_fe_sendauth(areq, msgLength, conn, &async);
+
+				if (async && (res == STATUS_OK))
+				{
+					/*
+					 * We'll come back later once we're ready to respond.
+					 * Don't consume the request yet.
+					 */
+					conn->status = CONNECTION_AUTHENTICATING;
+					goto keep_going;
+				}
 
 				/*
 				 * OK, we have processed the message; mark data consumed.  We
@@ -4107,6 +4132,49 @@ keep_going:						/* We will come back to here until there is
 				goto keep_going;
 			}
 
+		case CONNECTION_AUTHENTICATING:
+			{
+				PostgresPollingStatusType status;
+
+				if (!conn->async_auth || !conn->cleanup_async_auth)
+				{
+					/* programmer error; should not happen */
+					libpq_append_conn_error(conn, "async authentication has no handler");
+					goto error_return;
+				}
+
+				/* Drive some external authentication work. */
+				status = conn->async_auth(conn);
+
+				if (status == PGRES_POLLING_FAILED)
+					goto error_return;
+
+				if (status == PGRES_POLLING_OK)
+				{
+					/* Done. Tear down the async implementation. */
+					conn->cleanup_async_auth(conn);
+					conn->cleanup_async_auth = NULL;
+					Assert(conn->altsock == PGINVALID_SOCKET);
+
+					/*
+					 * Reenter the authentication exchange with the server. We
+					 * didn't consume the message that started external
+					 * authentication, so it'll be reprocessed as if we just
+					 * received it.
+					 */
+					conn->status = CONNECTION_AWAITING_RESPONSE;
+
+					goto keep_going;
+				}
+
+				/*
+				 * Caller needs to poll some more. conn->async_auth() should
+				 * have assigned an altsock to poll on.
+				 */
+				Assert(conn->altsock != PGINVALID_SOCKET);
+				return status;
+			}
+
 		case CONNECTION_AUTH_OK:
 			{
 				/*
@@ -4788,6 +4856,7 @@ pqMakeEmptyPGconn(void)
 	conn->verbosity = PQERRORS_DEFAULT;
 	conn->show_context = PQSHOW_CONTEXT_ERRORS;
 	conn->sock = PGINVALID_SOCKET;
+	conn->altsock = PGINVALID_SOCKET;
 	conn->Pfdebug = NULL;
 
 	/*
@@ -7439,6 +7508,8 @@ PQsocket(const PGconn *conn)
 {
 	if (!conn)
 		return -1;
+	if (conn->altsock != PGINVALID_SOCKET)
+		return conn->altsock;
 	return (conn->sock != PGINVALID_SOCKET) ? conn->sock : -1;
 }
 
diff --git a/src/interfaces/libpq/fe-misc.c b/src/interfaces/libpq/fe-misc.c
index 2c60eb5b569..d78445c70af 100644
--- a/src/interfaces/libpq/fe-misc.c
+++ b/src/interfaces/libpq/fe-misc.c
@@ -1049,34 +1049,43 @@ pqWriteReady(PGconn *conn)
  * or both.  Returns >0 if one or more conditions are met, 0 if it timed
  * out, -1 if an error occurred.
  *
- * If SSL is in use, the SSL buffer is checked prior to checking the socket
- * for read data directly.
+ * If an altsock is set for asynchronous authentication, that will be used in
+ * preference to the "server" socket. Otherwise, if SSL is in use, the SSL
+ * buffer is checked prior to checking the socket for read data directly.
  */
 static int
 pqSocketCheck(PGconn *conn, int forRead, int forWrite, pg_usec_time_t end_time)
 {
 	int			result;
+	pgsocket	sock;
 
 	if (!conn)
 		return -1;
-	if (conn->sock == PGINVALID_SOCKET)
+
+	if (conn->altsock != PGINVALID_SOCKET)
+		sock = conn->altsock;
+	else
 	{
-		libpq_append_conn_error(conn, "invalid socket");
-		return -1;
-	}
+		sock = conn->sock;
+		if (sock == PGINVALID_SOCKET)
+		{
+			libpq_append_conn_error(conn, "invalid socket");
+			return -1;
+		}
 
 #ifdef USE_SSL
-	/* Check for SSL library buffering read bytes */
-	if (forRead && conn->ssl_in_use && pgtls_read_pending(conn))
-	{
-		/* short-circuit the select */
-		return 1;
-	}
+		/* Check for SSL library buffering read bytes */
+		if (forRead && conn->ssl_in_use && pgtls_read_pending(conn))
+		{
+			/* short-circuit the select */
+			return 1;
+		}
 #endif
+	}
 
 	/* We will retry as long as we get EINTR */
 	do
-		result = PQsocketPoll(conn->sock, forRead, forWrite, end_time);
+		result = PQsocketPoll(sock, forRead, forWrite, end_time);
 	while (result < 0 && SOCK_ERRNO == EINTR);
 
 	if (result < 0)
diff --git a/src/interfaces/libpq/libpq-fe.h b/src/interfaces/libpq/libpq-fe.h
index cce9ce60c55..a3491faf0c3 100644
--- a/src/interfaces/libpq/libpq-fe.h
+++ b/src/interfaces/libpq/libpq-fe.h
@@ -103,6 +103,8 @@ typedef enum
 	CONNECTION_CHECK_STANDBY,	/* Checking if server is in standby mode. */
 	CONNECTION_ALLOCATED,		/* Waiting for connection attempt to be
 								 * started.  */
+	CONNECTION_AUTHENTICATING,	/* Authentication is in progress with some
+								 * external system. */
 } ConnStatusType;
 
 typedef enum
diff --git a/src/interfaces/libpq/libpq-int.h b/src/interfaces/libpq/libpq-int.h
index e0d5b5fe0be..2546f9f8a50 100644
--- a/src/interfaces/libpq/libpq-int.h
+++ b/src/interfaces/libpq/libpq-int.h
@@ -513,6 +513,12 @@ struct pg_conn
 										 * know which auth response we're
 										 * sending */
 
+	/* Callbacks for external async authentication */
+	PostgresPollingStatusType (*async_auth) (PGconn *conn);
+	void		(*cleanup_async_auth) (PGconn *conn);
+	pgsocket	altsock;		/* alternative socket for client to poll */
+
+
 	/* Transient state needed while establishing connection */
 	PGTargetServerType target_server_type;	/* desired session properties */
 	PGLoadBalanceType load_balance_type;	/* desired load balancing
-- 
2.34.1