Refactor SASL exchange in preparation for OAuth Bearer

From ae32698c528fce4829d4deccb922699d0b516f88 Mon Sep 17 00:00:00 2001
From: Daniel Gustafsson <dgustafsson@postgresql.org>
Date: Fri, 23 Feb 2024 11:09:54 +0100
Subject: [PATCH v2 1/2] Refactor SASL exchange to return tri-state status

The SASL exchange callback returned state in to output variables:
done and success.  This refactors that logic by introducing a new
return variable of type SASLStatus which makes the code easier to
read and understand, and prepares for future SASL exchanges which
operate asynchronously.

This was extracted from a larger patchset to introduce OAuthBearer
authentication and authorization.

Author: Jacob Champion <jacob.champion@enterprisedb.com>
Discussion: https://postgr.es/m/d1b467a78e0e36ed85a09adf979d04cf124a9d4b.camel@vmware.com
---
 src/interfaces/libpq/fe-auth-sasl.h  | 31 +++++++----
 src/interfaces/libpq/fe-auth-scram.c | 78 +++++++++++++---------------
 src/interfaces/libpq/fe-auth.c       | 28 +++++-----
 src/tools/pgindent/typedefs.list     |  1 +
 4 files changed, 71 insertions(+), 67 deletions(-)

diff --git a/src/interfaces/libpq/fe-auth-sasl.h b/src/interfaces/libpq/fe-auth-sasl.h
index ee5d1525b5..4eecf53a15 100644
--- a/src/interfaces/libpq/fe-auth-sasl.h
+++ b/src/interfaces/libpq/fe-auth-sasl.h
@@ -21,6 +21,17 @@
 
 #include "libpq-fe.h"
 
+/*
+ * Possible states for the SASL exchange, see the comment on exchange for an
+ * explanation of these.
+ */
+typedef enum
+{
+	SASL_COMPLETE = 0,
+	SASL_FAILED,
+	SASL_CONTINUE,
+} SASLStatus;
+
 /*
  * Frontend SASL mechanism callbacks.
  *
@@ -59,7 +70,8 @@ typedef struct pg_fe_sasl_mech
 	 * Produces a client response to a server challenge.  As a special case
 	 * for client-first SASL mechanisms, exchange() is called with a NULL
 	 * server response once at the start of the authentication exchange to
-	 * generate an initial response.
+	 * generate an initial response. Returns a SASLStatus indicating the
+	 * state and status of the exchange.
 	 *
 	 * Input parameters:
 	 *
@@ -79,22 +91,23 @@ typedef struct pg_fe_sasl_mech
 	 *
 	 *	output:	   A malloc'd buffer containing the client's response to
 	 *			   the server (can be empty), or NULL if the exchange should
-	 *			   be aborted.  (*success should be set to false in the
+	 *			   be aborted.  (The callback should return SASL_FAILED in the
 	 *			   latter case.)
 	 *
 	 *	outputlen: The length (0 or higher) of the client response buffer,
 	 *			   ignored if output is NULL.
 	 *
-	 *	done:      Set to true if the SASL exchange should not continue,
-	 *			   because the exchange is either complete or failed
+	 * Return value:
 	 *
-	 *	success:   Set to true if the SASL exchange completed successfully.
-	 *			   Ignored if *done is false.
+	 *	SASL_CONTINUE:	The output buffer is filled with a client response.
+	 *					Additional server challenge is expected
+	 *	SASL_COMPLETE:	The SASL exchange has completed successfully.
+	 *	SASL_FAILED:	The exchange has failed and the connection should be
+	 *					dropped.
 	 *--------
 	 */
-	void		(*exchange) (void *state, char *input, int inputlen,
-							 char **output, int *outputlen,
-							 bool *done, bool *success);
+	SASLStatus	(*exchange) (void *state, char *input, int inputlen,
+							 char **output, int *outputlen);
 
 	/*--------
 	 * channel_bound()
diff --git a/src/interfaces/libpq/fe-auth-scram.c b/src/interfaces/libpq/fe-auth-scram.c
index 04f0e5713d..0bb820e0d9 100644
--- a/src/interfaces/libpq/fe-auth-scram.c
+++ b/src/interfaces/libpq/fe-auth-scram.c
@@ -24,9 +24,8 @@
 /* The exported SCRAM callback mechanism. */
 static void *scram_init(PGconn *conn, const char *password,
 						const char *sasl_mechanism);
-static void scram_exchange(void *opaq, char *input, int inputlen,
-						   char **output, int *outputlen,
-						   bool *done, bool *success);
+static SASLStatus scram_exchange(void *opaq, char *input, int inputlen,
+								 char **output, int *outputlen);
 static bool scram_channel_bound(void *opaq);
 static void scram_free(void *opaq);
 
@@ -202,17 +201,14 @@ scram_free(void *opaq)
 /*
  * Exchange a SCRAM message with backend.
  */
-static void
+static SASLStatus
 scram_exchange(void *opaq, char *input, int inputlen,
-			   char **output, int *outputlen,
-			   bool *done, bool *success)
+			   char **output, int *outputlen)
 {
 	fe_scram_state *state = (fe_scram_state *) opaq;
 	PGconn	   *conn = state->conn;
 	const char *errstr = NULL;
 
-	*done = false;
-	*success = false;
 	*output = NULL;
 	*outputlen = 0;
 
@@ -225,12 +221,12 @@ scram_exchange(void *opaq, char *input, int inputlen,
 		if (inputlen == 0)
 		{
 			libpq_append_conn_error(conn, "malformed SCRAM message (empty message)");
-			goto error;
+			return SASL_FAILED;
 		}
 		if (inputlen != strlen(input))
 		{
 			libpq_append_conn_error(conn, "malformed SCRAM message (length mismatch)");
-			goto error;
+			return SASL_FAILED;
 		}
 	}
 
@@ -240,61 +236,59 @@ scram_exchange(void *opaq, char *input, int inputlen,
 			/* Begin the SCRAM handshake, by sending client nonce */
 			*output = build_client_first_message(state);
 			if (*output == NULL)
-				goto error;
+				return SASL_FAILED;
 
 			*outputlen = strlen(*output);
-			*done = false;
 			state->state = FE_SCRAM_NONCE_SENT;
-			break;
+			return SASL_CONTINUE;
 
 		case FE_SCRAM_NONCE_SENT:
 			/* Receive salt and server nonce, send response. */
 			if (!read_server_first_message(state, input))
-				goto error;
+				return SASL_FAILED;
 
 			*output = build_client_final_message(state);
 			if (*output == NULL)
-				goto error;
+				return SASL_FAILED;
 
 			*outputlen = strlen(*output);
-			*done = false;
 			state->state = FE_SCRAM_PROOF_SENT;
-			break;
+			return SASL_CONTINUE;
 
 		case FE_SCRAM_PROOF_SENT:
-			/* Receive server signature */
-			if (!read_server_final_message(state, input))
-				goto error;
-
-			/*
-			 * Verify server signature, to make sure we're talking to the
-			 * genuine server.
-			 */
-			if (!verify_server_signature(state, success, &errstr))
-			{
-				libpq_append_conn_error(conn, "could not verify server signature: %s", errstr);
-				goto error;
-			}
-
-			if (!*success)
 			{
-				libpq_append_conn_error(conn, "incorrect server signature");
+				bool		match;
+
+				/* Receive server signature */
+				if (!read_server_final_message(state, input))
+					return SASL_FAILED;
+
+				/*
+				 * Verify server signature, to make sure we're talking to the
+				 * genuine server.
+				 */
+				if (!verify_server_signature(state, &match, &errstr))
+				{
+					libpq_append_conn_error(conn, "could not verify server signature: %s", errstr);
+					return SASL_FAILED;
+				}
+
+				if (!match)
+				{
+					libpq_append_conn_error(conn, "incorrect server signature");
+				}
+				state->state = FE_SCRAM_FINISHED;
+				state->conn->client_finished_auth = true;
+				return match ? SASL_COMPLETE : SASL_FAILED;
 			}
-			*done = true;
-			state->state = FE_SCRAM_FINISHED;
-			state->conn->client_finished_auth = true;
-			break;
 
 		default:
 			/* shouldn't happen */
 			libpq_append_conn_error(conn, "invalid SCRAM exchange state");
-			goto error;
+			break;
 	}
-	return;
 
-error:
-	*done = true;
-	*success = false;
+	return SASL_FAILED;
 }
 
 /*
diff --git a/src/interfaces/libpq/fe-auth.c b/src/interfaces/libpq/fe-auth.c
index 1a8e4f6fbf..cf8af4c62e 100644
--- a/src/interfaces/libpq/fe-auth.c
+++ b/src/interfaces/libpq/fe-auth.c
@@ -423,11 +423,10 @@ pg_SASL_init(PGconn *conn, int payloadlen)
 {
 	char	   *initialresponse = NULL;
 	int			initialresponselen;
-	bool		done;
-	bool		success;
 	const char *selected_mechanism;
 	PQExpBufferData mechanism_buf;
 	char	   *password;
+	SASLStatus	status;
 
 	initPQExpBuffer(&mechanism_buf);
 
@@ -575,12 +574,11 @@ pg_SASL_init(PGconn *conn, int payloadlen)
 		goto oom_error;
 
 	/* Get the mechanism-specific Initial Client Response, if any */
-	conn->sasl->exchange(conn->sasl_state,
-						 NULL, -1,
-						 &initialresponse, &initialresponselen,
-						 &done, &success);
+	status = conn->sasl->exchange(conn->sasl_state,
+								  NULL, -1,
+								  &initialresponse, &initialresponselen);
 
-	if (done && !success)
+	if (status == SASL_FAILED)
 		goto error;
 
 	/*
@@ -629,10 +627,9 @@ pg_SASL_continue(PGconn *conn, int payloadlen, bool final)
 {
 	char	   *output;
 	int			outputlen;
-	bool		done;
-	bool		success;
 	int			res;
 	char	   *challenge;
+	SASLStatus	status;
 
 	/* Read the SASL challenge from the AuthenticationSASLContinue message. */
 	challenge = malloc(payloadlen + 1);
@@ -651,13 +648,12 @@ pg_SASL_continue(PGconn *conn, int payloadlen, bool final)
 	/* For safety and convenience, ensure the buffer is NULL-terminated. */
 	challenge[payloadlen] = '\0';
 
-	conn->sasl->exchange(conn->sasl_state,
-						 challenge, payloadlen,
-						 &output, &outputlen,
-						 &done, &success);
+	status = conn->sasl->exchange(conn->sasl_state,
+								  challenge, payloadlen,
+								  &output, &outputlen);
 	free(challenge);			/* don't need the input anymore */
 
-	if (final && !done)
+	if (final && status == SASL_CONTINUE)
 	{
 		if (outputlen != 0)
 			free(output);
@@ -670,7 +666,7 @@ pg_SASL_continue(PGconn *conn, int payloadlen, bool final)
 	 * If the exchange is not completed yet, we need to make sure that the
 	 * SASL mechanism has generated a message to send back.
 	 */
-	if (output == NULL && !done)
+	if (output == NULL && status == SASL_CONTINUE)
 	{
 		libpq_append_conn_error(conn, "no client response found after SASL exchange success");
 		return STATUS_ERROR;
@@ -692,7 +688,7 @@ pg_SASL_continue(PGconn *conn, int payloadlen, bool final)
 			return STATUS_ERROR;
 	}
 
-	if (done && !success)
+	if (status == SASL_FAILED)
 		return STATUS_ERROR;
 
 	return STATUS_OK;
diff --git a/src/tools/pgindent/typedefs.list b/src/tools/pgindent/typedefs.list
index fc8b15d0cf..2461567026 100644
--- a/src/tools/pgindent/typedefs.list
+++ b/src/tools/pgindent/typedefs.list
@@ -2423,6 +2423,7 @@ RuleLock
 RuleStmt
 RunningTransactions
 RunningTransactionsData
+SASLStatus
 SC_HANDLE
 SECURITY_ATTRIBUTES
 SECURITY_STATUS
-- 
2.32.1 (Apple Git-133)

From ed4a2e0c169e7b3a762aa93ee349b90d898962b2 Mon Sep 17 00:00:00 2001
From: Daniel Gustafsson <dgustafsson@postgresql.org>
Date: Fri, 23 Feb 2024 11:19:55 +0100
Subject: [PATCH v2 2/2] Explicitly require password for SCRAM exchange

This refactors the SASL init flow to set password_needed on the two
SCRAM exchanges currently supported. The code already required this
but was set up in such a way that all SASL exchanges required using
a password, a restriction which may not hold for all exchanges (the
example at hand being the proposed OAuthbearer exchange).

This was extracted from a larger patchset to introduce OAuthBearer
authentication and authorization.

Author: Jacob Champion <jacob.champion@enterprisedb.com>
Discussion: https://postgr.es/m/d1b467a78e0e36ed85a09adf979d04cf124a9d4b.camel@vmware.com
---
 src/interfaces/libpq/fe-auth.c | 28 +++++++++++++++-------------
 1 file changed, 15 insertions(+), 13 deletions(-)

diff --git a/src/interfaces/libpq/fe-auth.c b/src/interfaces/libpq/fe-auth.c
index cf8af4c62e..81ec08485d 100644
--- a/src/interfaces/libpq/fe-auth.c
+++ b/src/interfaces/libpq/fe-auth.c
@@ -425,7 +425,7 @@ pg_SASL_init(PGconn *conn, int payloadlen)
 	int			initialresponselen;
 	const char *selected_mechanism;
 	PQExpBufferData mechanism_buf;
-	char	   *password;
+	char	   *password = NULL;
 	SASLStatus	status;
 
 	initPQExpBuffer(&mechanism_buf);
@@ -446,8 +446,7 @@ pg_SASL_init(PGconn *conn, int payloadlen)
 	/*
 	 * Parse the list of SASL authentication mechanisms in the
 	 * AuthenticationSASL message, and select the best mechanism that we
-	 * support.  SCRAM-SHA-256-PLUS and SCRAM-SHA-256 are the only ones
-	 * supported at the moment, listed by order of decreasing importance.
+	 * support. Mechanisms are listed by order of decreasing importance.
 	 */
 	selected_mechanism = NULL;
 	for (;;)
@@ -487,6 +486,7 @@ pg_SASL_init(PGconn *conn, int payloadlen)
 				{
 					selected_mechanism = SCRAM_SHA_256_PLUS_NAME;
 					conn->sasl = &pg_scram_mech;
+					conn->password_needed = true;
 				}
 #else
 				/*
@@ -522,6 +522,7 @@ pg_SASL_init(PGconn *conn, int payloadlen)
 		{
 			selected_mechanism = SCRAM_SHA_256_NAME;
 			conn->sasl = &pg_scram_mech;
+			conn->password_needed = true;
 		}
 	}
 
@@ -545,18 +546,19 @@ pg_SASL_init(PGconn *conn, int payloadlen)
 
 	/*
 	 * First, select the password to use for the exchange, complaining if
-	 * there isn't one.  Currently, all supported SASL mechanisms require a
-	 * password, so we can just go ahead here without further distinction.
+	 * there isn't one and the selected SASL mechanism needs it.
 	 */
-	conn->password_needed = true;
-	password = conn->connhost[conn->whichhost].password;
-	if (password == NULL)
-		password = conn->pgpass;
-	if (password == NULL || password[0] == '\0')
+	if (conn->password_needed)
 	{
-		appendPQExpBufferStr(&conn->errorMessage,
-							 PQnoPasswordSupplied);
-		goto error;
+		password = conn->connhost[conn->whichhost].password;
+		if (password == NULL)
+			password = conn->pgpass;
+		if (password == NULL || password[0] == '\0')
+		{
+			appendPQExpBufferStr(&conn->errorMessage,
+								 PQnoPasswordSupplied);
+			goto error;
+		}
 	}
 
 	Assert(conn->sasl);
-- 
2.32.1 (Apple Git-133)

From 2ec88f69b9ad6e98403237a36a72ec5b7282f304 Mon Sep 17 00:00:00 2001
From: Daniel Gustafsson <dgustafsson@postgresql.org>
Date: Fri, 23 Feb 2024 11:19:55 +0100
Subject: [PATCH v3 2/2] Explicitly require password for SCRAM exchange

This refactors the SASL init flow to set password_needed on the two
SCRAM exchanges currently supported. The code already required this
but was set up in such a way that all SASL exchanges required using
a password, a restriction which may not hold for all exchanges (the
example at hand being the proposed OAuthbearer exchange).

This was extracted from a larger patchset to introduce OAuthBearer
authentication and authorization.

Author: Jacob Champion <jacob.champion@enterprisedb.com>
Discussion: https://postgr.es/m/d1b467a78e0e36ed85a09adf979d04cf124a9d4b.camel@vmware.com
---
 src/interfaces/libpq/fe-auth.c | 28 +++++++++++++++-------------
 1 file changed, 15 insertions(+), 13 deletions(-)

diff --git a/src/interfaces/libpq/fe-auth.c b/src/interfaces/libpq/fe-auth.c
index cf8af4c62e..81ec08485d 100644
--- a/src/interfaces/libpq/fe-auth.c
+++ b/src/interfaces/libpq/fe-auth.c
@@ -425,7 +425,7 @@ pg_SASL_init(PGconn *conn, int payloadlen)
 	int			initialresponselen;
 	const char *selected_mechanism;
 	PQExpBufferData mechanism_buf;
-	char	   *password;
+	char	   *password = NULL;
 	SASLStatus	status;
 
 	initPQExpBuffer(&mechanism_buf);
@@ -446,8 +446,7 @@ pg_SASL_init(PGconn *conn, int payloadlen)
 	/*
 	 * Parse the list of SASL authentication mechanisms in the
 	 * AuthenticationSASL message, and select the best mechanism that we
-	 * support.  SCRAM-SHA-256-PLUS and SCRAM-SHA-256 are the only ones
-	 * supported at the moment, listed by order of decreasing importance.
+	 * support. Mechanisms are listed by order of decreasing importance.
 	 */
 	selected_mechanism = NULL;
 	for (;;)
@@ -487,6 +486,7 @@ pg_SASL_init(PGconn *conn, int payloadlen)
 				{
 					selected_mechanism = SCRAM_SHA_256_PLUS_NAME;
 					conn->sasl = &pg_scram_mech;
+					conn->password_needed = true;
 				}
 #else
 				/*
@@ -522,6 +522,7 @@ pg_SASL_init(PGconn *conn, int payloadlen)
 		{
 			selected_mechanism = SCRAM_SHA_256_NAME;
 			conn->sasl = &pg_scram_mech;
+			conn->password_needed = true;
 		}
 	}
 
@@ -545,18 +546,19 @@ pg_SASL_init(PGconn *conn, int payloadlen)
 
 	/*
 	 * First, select the password to use for the exchange, complaining if
-	 * there isn't one.  Currently, all supported SASL mechanisms require a
-	 * password, so we can just go ahead here without further distinction.
+	 * there isn't one and the selected SASL mechanism needs it.
 	 */
-	conn->password_needed = true;
-	password = conn->connhost[conn->whichhost].password;
-	if (password == NULL)
-		password = conn->pgpass;
-	if (password == NULL || password[0] == '\0')
+	if (conn->password_needed)
 	{
-		appendPQExpBufferStr(&conn->errorMessage,
-							 PQnoPasswordSupplied);
-		goto error;
+		password = conn->connhost[conn->whichhost].password;
+		if (password == NULL)
+			password = conn->pgpass;
+		if (password == NULL || password[0] == '\0')
+		{
+			appendPQExpBufferStr(&conn->errorMessage,
+								 PQnoPasswordSupplied);
+			goto error;
+		}
 	}
 
 	Assert(conn->sasl);
-- 
2.32.1 (Apple Git-133)

From 874f5fa5398af89d237f8146349cb1cbc859e0af Mon Sep 17 00:00:00 2001
From: Daniel Gustafsson <dgustafsson@postgresql.org>
Date: Fri, 23 Feb 2024 11:09:54 +0100
Subject: [PATCH v3 1/2] Refactor SASL exchange to return tri-state status

The SASL exchange callback returned state in to output variables:
done and success.  This refactors that logic by introducing a new
return variable of type SASLStatus which makes the code easier to
read and understand, and prepares for future SASL exchanges which
operate asynchronously.

This was extracted from a larger patchset to introduce OAuthBearer
authentication and authorization.

Author: Jacob Champion <jacob.champion@enterprisedb.com>
Discussion: https://postgr.es/m/d1b467a78e0e36ed85a09adf979d04cf124a9d4b.camel@vmware.com
---
 src/interfaces/libpq/fe-auth-sasl.h  | 31 +++++++----
 src/interfaces/libpq/fe-auth-scram.c | 78 +++++++++++++---------------
 src/interfaces/libpq/fe-auth.c       | 28 +++++-----
 src/tools/pgindent/typedefs.list     |  1 +
 4 files changed, 71 insertions(+), 67 deletions(-)

diff --git a/src/interfaces/libpq/fe-auth-sasl.h b/src/interfaces/libpq/fe-auth-sasl.h
index ee5d1525b5..4eecf53a15 100644
--- a/src/interfaces/libpq/fe-auth-sasl.h
+++ b/src/interfaces/libpq/fe-auth-sasl.h
@@ -21,6 +21,17 @@
 
 #include "libpq-fe.h"
 
+/*
+ * Possible states for the SASL exchange, see the comment on exchange for an
+ * explanation of these.
+ */
+typedef enum
+{
+	SASL_COMPLETE = 0,
+	SASL_FAILED,
+	SASL_CONTINUE,
+} SASLStatus;
+
 /*
  * Frontend SASL mechanism callbacks.
  *
@@ -59,7 +70,8 @@ typedef struct pg_fe_sasl_mech
 	 * Produces a client response to a server challenge.  As a special case
 	 * for client-first SASL mechanisms, exchange() is called with a NULL
 	 * server response once at the start of the authentication exchange to
-	 * generate an initial response.
+	 * generate an initial response. Returns a SASLStatus indicating the
+	 * state and status of the exchange.
 	 *
 	 * Input parameters:
 	 *
@@ -79,22 +91,23 @@ typedef struct pg_fe_sasl_mech
 	 *
 	 *	output:	   A malloc'd buffer containing the client's response to
 	 *			   the server (can be empty), or NULL if the exchange should
-	 *			   be aborted.  (*success should be set to false in the
+	 *			   be aborted.  (The callback should return SASL_FAILED in the
 	 *			   latter case.)
 	 *
 	 *	outputlen: The length (0 or higher) of the client response buffer,
 	 *			   ignored if output is NULL.
 	 *
-	 *	done:      Set to true if the SASL exchange should not continue,
-	 *			   because the exchange is either complete or failed
+	 * Return value:
 	 *
-	 *	success:   Set to true if the SASL exchange completed successfully.
-	 *			   Ignored if *done is false.
+	 *	SASL_CONTINUE:	The output buffer is filled with a client response.
+	 *					Additional server challenge is expected
+	 *	SASL_COMPLETE:	The SASL exchange has completed successfully.
+	 *	SASL_FAILED:	The exchange has failed and the connection should be
+	 *					dropped.
 	 *--------
 	 */
-	void		(*exchange) (void *state, char *input, int inputlen,
-							 char **output, int *outputlen,
-							 bool *done, bool *success);
+	SASLStatus	(*exchange) (void *state, char *input, int inputlen,
+							 char **output, int *outputlen);
 
 	/*--------
 	 * channel_bound()
diff --git a/src/interfaces/libpq/fe-auth-scram.c b/src/interfaces/libpq/fe-auth-scram.c
index 04f0e5713d..0bb820e0d9 100644
--- a/src/interfaces/libpq/fe-auth-scram.c
+++ b/src/interfaces/libpq/fe-auth-scram.c
@@ -24,9 +24,8 @@
 /* The exported SCRAM callback mechanism. */
 static void *scram_init(PGconn *conn, const char *password,
 						const char *sasl_mechanism);
-static void scram_exchange(void *opaq, char *input, int inputlen,
-						   char **output, int *outputlen,
-						   bool *done, bool *success);
+static SASLStatus scram_exchange(void *opaq, char *input, int inputlen,
+								 char **output, int *outputlen);
 static bool scram_channel_bound(void *opaq);
 static void scram_free(void *opaq);
 
@@ -202,17 +201,14 @@ scram_free(void *opaq)
 /*
  * Exchange a SCRAM message with backend.
  */
-static void
+static SASLStatus
 scram_exchange(void *opaq, char *input, int inputlen,
-			   char **output, int *outputlen,
-			   bool *done, bool *success)
+			   char **output, int *outputlen)
 {
 	fe_scram_state *state = (fe_scram_state *) opaq;
 	PGconn	   *conn = state->conn;
 	const char *errstr = NULL;
 
-	*done = false;
-	*success = false;
 	*output = NULL;
 	*outputlen = 0;
 
@@ -225,12 +221,12 @@ scram_exchange(void *opaq, char *input, int inputlen,
 		if (inputlen == 0)
 		{
 			libpq_append_conn_error(conn, "malformed SCRAM message (empty message)");
-			goto error;
+			return SASL_FAILED;
 		}
 		if (inputlen != strlen(input))
 		{
 			libpq_append_conn_error(conn, "malformed SCRAM message (length mismatch)");
-			goto error;
+			return SASL_FAILED;
 		}
 	}
 
@@ -240,61 +236,59 @@ scram_exchange(void *opaq, char *input, int inputlen,
 			/* Begin the SCRAM handshake, by sending client nonce */
 			*output = build_client_first_message(state);
 			if (*output == NULL)
-				goto error;
+				return SASL_FAILED;
 
 			*outputlen = strlen(*output);
-			*done = false;
 			state->state = FE_SCRAM_NONCE_SENT;
-			break;
+			return SASL_CONTINUE;
 
 		case FE_SCRAM_NONCE_SENT:
 			/* Receive salt and server nonce, send response. */
 			if (!read_server_first_message(state, input))
-				goto error;
+				return SASL_FAILED;
 
 			*output = build_client_final_message(state);
 			if (*output == NULL)
-				goto error;
+				return SASL_FAILED;
 
 			*outputlen = strlen(*output);
-			*done = false;
 			state->state = FE_SCRAM_PROOF_SENT;
-			break;
+			return SASL_CONTINUE;
 
 		case FE_SCRAM_PROOF_SENT:
-			/* Receive server signature */
-			if (!read_server_final_message(state, input))
-				goto error;
-
-			/*
-			 * Verify server signature, to make sure we're talking to the
-			 * genuine server.
-			 */
-			if (!verify_server_signature(state, success, &errstr))
-			{
-				libpq_append_conn_error(conn, "could not verify server signature: %s", errstr);
-				goto error;
-			}
-
-			if (!*success)
 			{
-				libpq_append_conn_error(conn, "incorrect server signature");
+				bool		match;
+
+				/* Receive server signature */
+				if (!read_server_final_message(state, input))
+					return SASL_FAILED;
+
+				/*
+				 * Verify server signature, to make sure we're talking to the
+				 * genuine server.
+				 */
+				if (!verify_server_signature(state, &match, &errstr))
+				{
+					libpq_append_conn_error(conn, "could not verify server signature: %s", errstr);
+					return SASL_FAILED;
+				}
+
+				if (!match)
+				{
+					libpq_append_conn_error(conn, "incorrect server signature");
+				}
+				state->state = FE_SCRAM_FINISHED;
+				state->conn->client_finished_auth = true;
+				return match ? SASL_COMPLETE : SASL_FAILED;
 			}
-			*done = true;
-			state->state = FE_SCRAM_FINISHED;
-			state->conn->client_finished_auth = true;
-			break;
 
 		default:
 			/* shouldn't happen */
 			libpq_append_conn_error(conn, "invalid SCRAM exchange state");
-			goto error;
+			break;
 	}
-	return;
 
-error:
-	*done = true;
-	*success = false;
+	return SASL_FAILED;
 }
 
 /*
diff --git a/src/interfaces/libpq/fe-auth.c b/src/interfaces/libpq/fe-auth.c
index 1a8e4f6fbf..cf8af4c62e 100644
--- a/src/interfaces/libpq/fe-auth.c
+++ b/src/interfaces/libpq/fe-auth.c
@@ -423,11 +423,10 @@ pg_SASL_init(PGconn *conn, int payloadlen)
 {
 	char	   *initialresponse = NULL;
 	int			initialresponselen;
-	bool		done;
-	bool		success;
 	const char *selected_mechanism;
 	PQExpBufferData mechanism_buf;
 	char	   *password;
+	SASLStatus	status;
 
 	initPQExpBuffer(&mechanism_buf);
 
@@ -575,12 +574,11 @@ pg_SASL_init(PGconn *conn, int payloadlen)
 		goto oom_error;
 
 	/* Get the mechanism-specific Initial Client Response, if any */
-	conn->sasl->exchange(conn->sasl_state,
-						 NULL, -1,
-						 &initialresponse, &initialresponselen,
-						 &done, &success);
+	status = conn->sasl->exchange(conn->sasl_state,
+								  NULL, -1,
+								  &initialresponse, &initialresponselen);
 
-	if (done && !success)
+	if (status == SASL_FAILED)
 		goto error;
 
 	/*
@@ -629,10 +627,9 @@ pg_SASL_continue(PGconn *conn, int payloadlen, bool final)
 {
 	char	   *output;
 	int			outputlen;
-	bool		done;
-	bool		success;
 	int			res;
 	char	   *challenge;
+	SASLStatus	status;
 
 	/* Read the SASL challenge from the AuthenticationSASLContinue message. */
 	challenge = malloc(payloadlen + 1);
@@ -651,13 +648,12 @@ pg_SASL_continue(PGconn *conn, int payloadlen, bool final)
 	/* For safety and convenience, ensure the buffer is NULL-terminated. */
 	challenge[payloadlen] = '\0';
 
-	conn->sasl->exchange(conn->sasl_state,
-						 challenge, payloadlen,
-						 &output, &outputlen,
-						 &done, &success);
+	status = conn->sasl->exchange(conn->sasl_state,
+								  challenge, payloadlen,
+								  &output, &outputlen);
 	free(challenge);			/* don't need the input anymore */
 
-	if (final && !done)
+	if (final && status == SASL_CONTINUE)
 	{
 		if (outputlen != 0)
 			free(output);
@@ -670,7 +666,7 @@ pg_SASL_continue(PGconn *conn, int payloadlen, bool final)
 	 * If the exchange is not completed yet, we need to make sure that the
 	 * SASL mechanism has generated a message to send back.
 	 */
-	if (output == NULL && !done)
+	if (output == NULL && status == SASL_CONTINUE)
 	{
 		libpq_append_conn_error(conn, "no client response found after SASL exchange success");
 		return STATUS_ERROR;
@@ -692,7 +688,7 @@ pg_SASL_continue(PGconn *conn, int payloadlen, bool final)
 			return STATUS_ERROR;
 	}
 
-	if (done && !success)
+	if (status == SASL_FAILED)
 		return STATUS_ERROR;
 
 	return STATUS_OK;
diff --git a/src/tools/pgindent/typedefs.list b/src/tools/pgindent/typedefs.list
index e294f8bc4e..c87ef6107a 100644
--- a/src/tools/pgindent/typedefs.list
+++ b/src/tools/pgindent/typedefs.list
@@ -2430,6 +2430,7 @@ RuleLock
 RuleStmt
 RunningTransactions
 RunningTransactionsData
+SASLStatus
 SC_HANDLE
 SECURITY_ATTRIBUTES
 SECURITY_STATUS
-- 
2.32.1 (Apple Git-133)