libpq: Handle asynchronous actions during SASL

This adds the ability for a SASL mechanism to signal PQconnectPoll()
that some arbitrary work, external to the Postgres connection, is
required for authentication to continue.  There is no consumer for
this capability as part of this commit, it is infrastructure which
is required for future work on supporting the OAUTHBEARER mechanism.

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.  The altsock temporarily replaces the regular connection
descriptor, so existing PQsocket() clients should continue to operate
correctly using their existing polling implementations.

For a mechanism to use this it should set an authentication callback,
conn->async_auth(), and a cleanup callback, conn->cleanup_async_auth(),
and return SASL_ASYNC during the exchange.  It should then 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 object.

This was extracted from the larger OAUTHBEARER patchset which has
been developed, and reviewed by many, over several years and it is
thus likely that some reviewer credit of much earlier versions has
been accidentally omitted.

Author: Jacob Champion <jacob.champion@enterprisedb.com>
Reviewed-by: Daniel Gustafsson <daniel@yesql.se>
Reviewed-by: Peter Eisentraut <peter@eisentraut.org>
Reviewed-by: Antonin Houska <ah@cybertec.at>
Discussion: https://postgr.es/m/CAOYmi+kJqzo6XsR9TEhvVfeVNQ-TyFM5LATypm9yoQVYk=4Wrw@mail.gmail.com

1 files changed, 89 insertions, 31 deletions

diff --git a/src/interfaces/libpq/fe-auth.c b/src/interfaces/libpq/fe-auth.c
index 70753d8ec29..761ee8f88f7 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,54 @@ 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.
+		 *
+		 * In non-assertion builds, this postcondition is enforced at time of
+		 * use in PQconnectPoll().
+		 */
+		Assert(conn->async_auth);
+		Assert(conn->cleanup_async_auth);
+
+		*async = true;
+		return STATUS_OK;
+	}
+
 	/*
 	 * Build a SASLInitialResponse message, and send it.
 	 */
@@ -671,7 +699,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 +729,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 +1055,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 +1224,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 +1233,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);