From a9db2e3e018d8a6b083089e20be14743b872f937 Mon Sep 17 00:00:00 2001
From: Heikki Linnakangas <heikki.linnakangas@iki.fi>
Date: Wed, 12 Apr 2017 20:14:13 +0300
Subject: [PATCH 2/2] Improve the SASL authentication protocol.

This contains some protocol changes to SASL authentiation (which is new
in v10):

* For future-proofing, in the AuthenticationSASL message that begins SASL
  authentication, provide a list of SASL mechanisms that the server
  supports, for the client to choose from. Currently, it's always just
  SCRAM-SHA-256.

* Add a separate authentication message type for the final server->client
  SASL message, which the client doesn't need to respond to. This makes
  it unambiguous whether the client is supposed to send a response or not.
  The SASL mechanism should know that anyway, but better to be explicit.

Also, in the server, support clients that don't send an Initial Client
response in the first SASLInitialResponse message. The server is supposed
to first send an empty request in that case, to which the client will
respond with the data that usually comes in the Initial Client Response.
libpq uses the Initial Client Response field and doesn't need this, and I
would assume any other sensible implementation to use Initial Client
Response, too, but let's follow the SASL spec.

Improve the documentation on SASL authentication in protocol. Add a
section describing the SASL message flow, some details on our SCRAM-SHA-256
implementation.

Document the different kinds of PasswordMessages that the frontend sends
in different phases of SASL authentication, as well as GSS/SSPI
authentication as separate message formats. Even though they're all 'p'
messages, and the exact format depends on the context, describing them as
separate message formats makes the documentation more clear.

Discussion: https://www.postgresql.org/message-id/CAB7nPqS-aFg0iM3AQOJwKDv_0WkAedRjs1W2X8EixSz+sKBXCQ@mail.gmail.com
---
 doc/src/sgml/protocol.sgml     | 350 ++++++++++++++++++++++++++++++++++++++---
 src/backend/libpq/auth-scram.c |  27 +++-
 src/backend/libpq/auth.c       |  66 +++++++-
 src/include/libpq/pqcomm.h     |   5 +-
 src/interfaces/libpq/fe-auth.c | 159 ++++++++++++-------
 5 files changed, 521 insertions(+), 86 deletions(-)

diff --git a/doc/src/sgml/protocol.sgml b/doc/src/sgml/protocol.sgml
index 4e8bb32d33..6e271fe898 100644
--- a/doc/src/sgml/protocol.sgml
+++ b/doc/src/sgml/protocol.sgml
@@ -330,7 +330,7 @@
       <listitem>
        <para>
         The frontend must now initiate a GSSAPI negotiation. The frontend
-        will send a PasswordMessage with the first part of the GSSAPI
+        will send a GSSResponse message with the first part of the GSSAPI
         data stream in response to this. If further messages are needed,
         the server will respond with AuthenticationGSSContinue.
        </para>
@@ -342,7 +342,7 @@
       <listitem>
        <para>
         The frontend must now initiate a SSPI negotiation. The frontend
-        will send a PasswordMessage with the first part of the SSPI
+        will send a GSSResponse with the first part of the SSPI
         data stream in response to this. If further messages are needed,
         the server will respond with AuthenticationGSSContinue.
        </para>
@@ -358,7 +358,7 @@
         or a previous AuthenticationGSSContinue). If the GSSAPI
         or SSPI data in this message
         indicates more data is needed to complete the authentication,
-        the frontend must send that data as another PasswordMessage. If
+        the frontend must send that data as another GSSResponse message. If
         GSSAPI or SSPI authentication is completed by this message, the server
         will next send AuthenticationOk to indicate successful authentication
         or ErrorResponse to indicate failure.
@@ -370,27 +370,38 @@
       <term>AuthenticationSASL</term>
       <listitem>
        <para>
-        The frontend must now initiate a SASL negotiation, using the SASL
-        mechanism specified in the message. The frontend will send a
-        PasswordMessage with the first part of the SASL data stream in
-        response to this. If further messages are needed, the server will
-        respond with AuthenticationSASLContinue.
+        The frontend must now initiate a SASL negotiation, using one of the
+        SASL mechanisms listed in the message. The frontend will send a
+        SASLInitialResponse with the name of the selected mechanism, and the
+        first part of the SASL data stream in response to this. If further
+        messages are needed, the server will respond with
+        AuthenticationSASLContinue. See <xref linkend="sasl-authentication">
+        for details.
        </para>
       </listitem>
-
      </varlistentry>
+
      <varlistentry>
       <term>AuthenticationSASLContinue</term>
       <listitem>
        <para>
-        This message contains the response data from the previous step
-        of SASL negotiation (AuthenticationSASL, or a previous
-        AuthenticationSASLContinue). If the SASL data in this message
-        indicates more data is needed to complete the authentication,
-        the frontend must send that data as another PasswordMessage. If
-        SASL authentication is completed by this message, the server
-        will next send AuthenticationOk to indicate successful authentication
-        or ErrorResponse to indicate failure.
+        This message contains challenge data from the previous step of SASL
+        negotiation (AuthenticationSASL, or a previous
+        AuthenticationSASLContinue). The frontend must respond with a
+        SASLResponse message.
+       </para>
+      </listitem>
+     </varlistentry>
+
+     <varlistentry>
+      <term>AuthenticationSASLFinal</term>
+      <listitem>
+       <para>
+        SASL authentication has completed with additional mechanism-specific
+        data for the client. The server will next send AuthenticationOk to
+        indicate successful authentication, or an ErrorResponse to indicate
+        failure. This message is sent only if the SASL mechanism specifies
+        additional data to be sent from server to client at completion.
        </para>
       </listitem>
      </varlistentry>
@@ -1326,6 +1337,141 @@
   </sect2>
  </sect1>
 
+<sect1 id="sasl-authentication">
+<title>SASL Authentication</title>
+
+<para>
+<firstterm>SASL</> is a framework for authentication in connection-oriented
+protocols. At the moment, <productname>PostgreSQL</> implements only one SASL
+authentication mechanism, SCRAM-SHA-256, but more might be added in the
+future. The below steps illustrate how SASL authentication is performed in
+general, while the next subsection gives more details on SCRAM-SHA-256.
+</para>
+
+<procedure>
+<title>SASL Authentication Message Flow</title>
+
+<step id="sasl-auth-begin">
+<para>  
+  To begin a SASL authentication exchange, the server an AuthenticationSASL
+  message. It includes a list of SASL authentication mechanisms that the
+  server can accept, in the server's preferred order.
+</para>
+</step>
+
+<step id="sasl-auth-initial-response">
+<para>
+  The client selects one of the supported mechanisms from the list, and sends
+  a SASLInitialResponse message to the server. The message includes the name
+  of the selected mechanism, and an optional Initial Client Response, if the
+  selected mechanism uses that.
+</para>
+</step>
+
+<step id="sasl-auth-continue">
+<para>
+  One or more server-challenge and client-response message will follow. Each
+  server-challenge is sent in an AuthenticationSASLContinue message, followed
+  by a response from client in an SASLResponse message. The particulars of
+  the messages are mechanism specific.
+</para>
+</step>
+
+<step id="sasl-auth-end">
+<para>
+  Finally, when the authentication exchange is completed successfully, the
+  server sends an AuthenticationSASLFinal message, followed
+  immediately by an AuthenticationOk message. The AuthenticationSASLFinal
+  contains additional server-to-client data, whose content is particular to the
+  selected authentication mechanism. If the authentication mechanism doesn't
+  use additional data that's sent at completion, the AuthenticationSASLFinal
+  message is not sent.
+</para>
+</step>
+</procedure>
+
+<para>
+On error, the server can abort the authentication at any stage, and send an
+ErrorMessage.
+</para>
+
+ <sect2 id="sasl-scram-sha256">
+  <title>SCRAM-SHA-256 authentication</title>
+
+  <para>
+    <firstterm>SCRAM-SHA-256</> (called just <firstterm>SCRAM</> from now on) is
+    the only implemented SASL mechanism, at the moment. It is described in detail
+    in RFC 7677 and RFC 5741. 
+  </para>
+
+  <para>
+When SCRAM-SHA-256 is used in PostgreSQL, the server will ignore the username
+that the client sends in the <structname>client-first-message</>. The username
+that was already sent in the startup message is used instead.
+<productname>PostgreSQL</> supports multiple character encodings, while SCRAM
+dictates UTF-8 to be used for the username, so it might be impossible to
+represent the PostgreSQL username in UTF-8. To avoid confusion, the client
+should use <literal>pg_same_as_startup_message</literal> as the username in the
+<structname>client-first-message</>.
+  </para>
+
+  <para>
+The SCRAM specification dictates that the password is also in UTF-8, and is
+processed with the <firstterm>SASLprep</> algorithm.
+<productname>PostgreSQL</>, however, does not require UTF-8 to be used for
+the password. When a user's password is set, it is processed with SASLprep
+as if it was in UTF-8, regardless of the actual encoding used. However, if
+it is not a legal UTF-8 byte sequence, or it contains UTF-8 byte sequences
+that are prohibited by the SASLprep algorithm, the raw password will be used
+without SASLprep processing, instead of throwing an error. This allows the
+password to be normalized when it is in UTF-8, but still allows a non-UTF-8
+password to be used, and doesn't require the system to know which encoding
+the password is in.
+  </para>
+
+  <para>
+<firstterm>Channel binding</> has not been implemented yet.
+  </para>
+
+<procedure>
+<title>Example</title>
+  <step id="scram-begin">
+<para>
+  The server sends an AuthenticationSASL message. It includes a list of
+  SASL authentication mechanisms that the server can accept.
+</para>
+</step>
+<step id="scram-client-first">
+<para>
+  The client responds by sending a SASLInitialResponse message, which
+  indicates the chosen mechanism, <literal>SCRAM-SHA-256</>. In the Initial
+  Client response field, the message contains the SCRAM
+  <structname>client-first-message</>.
+</para>
+</step>
+<step id="scram-server-first">
+<para>
+  Server sends an AuthenticationSASLContinue message, with a SCRAM
+  <structname>server-first message</> as the content.
+</para>
+</step>
+<step id="scram-client-final">
+<para>
+  Client sends a SASLResponse message, with SCRAM
+  <structname>client-final-message</> as the content.
+</para>
+</step>
+<step id="scram-server-final">
+<para>
+  Server sends an AuthenticationSASLFinal message, with the SCRAM
+  <structname>server-final-message</>, followed immediately by
+  an AuthenticationOk message.
+</para>
+</step>
+</procedure>
+</sect2>
+
+
 <sect1 id="protocol-replication">
 <title>Streaming Replication Protocol</title>
 
@@ -2894,6 +3040,12 @@ AuthenticationSASL (B)
 </para>
 </listitem>
 </varlistentry>
+</variablelist>
+The message body is a list of SASL authentication mechanisms, in the
+server's order of preference. A zero byte is required as terminator after
+the last authentication mechanism name. For each mechanism, there is the
+following:
+<variablelist>
 <varlistentry>
 <term>
         String
@@ -4313,6 +4465,50 @@ FunctionCallResponse (B)
 </listitem>
 </varlistentry>
 
+<varlistentry>
+<term>
+GSSResponse (F)
+</term>
+<listitem>
+<para>
+
+<variablelist>
+<varlistentry>
+<term>
+        Byte1('p')
+</term>
+<listitem>
+<para>
+                Identifies the message as a GSSAPI or SSPI response. Note that
+                this is also used for SASL and password response messages.
+                The exact message type can be deduced from the context.
+</para>
+</listitem>
+</varlistentry>
+<varlistentry>
+<term>
+        Int32
+</term>
+<listitem>
+<para>
+                Length of message contents in bytes, including self.
+</para>
+</listitem>
+</varlistentry>
+<varlistentry>
+<term>
+        Byte<replaceable>n</replaceable>
+</term>
+<listitem>
+<para>
+                GSSAPI/SSPI specific message data.
+</para>
+</listitem>
+</varlistentry>
+</variablelist>
+</para>
+</listitem>
+</varlistentry>
 
 <varlistentry>
 <term>
@@ -4726,10 +4922,8 @@ PasswordMessage (F)
 <listitem>
 <para>
                 Identifies the message as a password response. Note that
-                this is also used for GSSAPI, SSPI and SASL response messages
-                (which is really a design error, since the contained data
-                is not a null-terminated string in that case, but can be
-                arbitrary binary data).
+                this is also used for GSSAPI, SSPI and SASL response messages.
+                The exact message type can be deduced from the context.
 </para>
 </listitem>
 </varlistentry>
@@ -4761,6 +4955,120 @@ PasswordMessage (F)
 
 <varlistentry>
 <term>
+SASLInitialresponse (F)
+</term>
+<listitem>
+<para>
+
+<variablelist>
+<varlistentry>
+<term>
+        Byte1('p')
+</term>
+<listitem>
+<para>
+                Identifies the message as an initial SASL response. Note that
+                this is also used for GSSAPI, SSPI and password response messages.
+                The exact message type is deduced from the context.
+</para>
+</listitem>
+</varlistentry>
+<varlistentry>
+<term>
+        Int32
+</term>
+<listitem>
+<para>
+                Length of message contents in bytes, including self.
+</para>
+</listitem>
+</varlistentry>
+<varlistentry>
+<term>
+        String
+</term>
+<listitem>
+<para>
+                Name of the SASL authentication mechanism that the client
+                selected.
+</para>
+</listitem>
+</varlistentry>
+<varlistentry>
+<term>
+        Int32
+</term>
+<listitem>
+<para>
+                Length of SASL mechanism specific "Initial Client Response" that
+                follows, or -1 if there is no Initial Response.
+</para>
+</listitem>
+</varlistentry>
+<varlistentry>
+<term>
+        Byte<replaceable>n</replaceable>
+</term>
+<listitem>
+<para>
+                SASL mechanism specific "Initial Response".
+</para>
+</listitem>
+</varlistentry>
+</variablelist>
+</para>
+</listitem>
+</varlistentry>
+
+
+<varlistentry>
+<term>
+SASLResponse (F)
+</term>
+<listitem>
+<para>
+
+<variablelist>
+<varlistentry>
+<term>
+        Byte1('p')
+</term>
+<listitem>
+<para>
+                Identifies the message as a SASL response. Note that
+                this is also used for GSSAPI, SSPI and password response messages.
+                The exact message type can be deduced from the context.
+</para>
+</listitem>
+</varlistentry>
+<varlistentry>
+<term>
+        Int32
+</term>
+<listitem>
+<para>
+                Length of message contents in bytes, including self.
+</para>
+</listitem>
+</varlistentry>
+<varlistentry>
+<term>
+        Byte<replaceable>n</replaceable>
+</term>
+<listitem>
+<para>
+                SASL mechanism specific message data.
+</para>
+</listitem>
+</varlistentry>
+</variablelist>
+</para>
+</listitem>
+</varlistentry>
+
+
+<varlistentry>
+<term>
 PortalSuspended (B)
 </term>
 <listitem>
diff --git a/src/backend/libpq/auth-scram.c b/src/backend/libpq/auth-scram.c
index a47c48d980..59e997f96d 100644
--- a/src/backend/libpq/auth-scram.c
+++ b/src/backend/libpq/auth-scram.c
@@ -254,8 +254,16 @@ pg_be_scram_init(const char *username, const char *shadow_pass)
 /*
  * Continue a SCRAM authentication exchange.
  *
- * The next message to send to client is saved in "output", for a length
- * of "outputlen".  In the case of an error, optionally store a palloc'd
+ * 'input' is the SCRAM payload sent by the client.  On the first call,
+ * 'input' contains the "Initial Client Response" that the client sent as
+ * part of the SASLInitialResponse message, or NULL if no Initial Client
+ * Response was given.  (The SASL specification distinguishes between an
+ * empty response and non-existing one.)  On subsequent calls, 'input'
+ * cannot be NULL.  For convenience in this function, the caller must
+ * ensure that there is a null terminator at input[inputlen].
+ *
+ * The next message to send to client is saved in 'output', for a length
+ * of 'outputlen'.  In the case of an error, optionally store a palloc'd
  * string at *logdetail that will be sent to the postmaster log (but not
  * the client).
  */
@@ -269,6 +277,21 @@ pg_be_scram_exchange(void *opaq, char *input, int inputlen,
 	*output = NULL;
 
 	/*
+	 * If the client didn't send an "Initial Client Response" in the
+	 * SASLInitialResponse message, send an empty challenge, to which the
+	 * client will respond with the same data that usually comes in the
+	 * Initial Client Response.
+	 */
+	if (input == NULL)
+	{
+		Assert(state->state == SCRAM_AUTH_INIT);
+
+		*output = pstrdup("");
+		*outputlen = 0;
+		return SASL_EXCHANGE_CONTINUE;
+	}
+
+	/*
 	 * Check that the input length agrees with the string length of the input.
 	 * We can ignore inputlen after this.
 	 */
diff --git a/src/backend/libpq/auth.c b/src/backend/libpq/auth.c
index 681b93855f..dbb9e8dcc1 100644
--- a/src/backend/libpq/auth.c
+++ b/src/backend/libpq/auth.c
@@ -620,10 +620,10 @@ sendAuthRequest(Port *port, AuthRequest areq, char *extradata, int extralen)
 	pq_endmessage(&buf);
 
 	/*
-	 * Flush message so client will see it, except for AUTH_REQ_OK, which need
-	 * not be sent until we are ready for queries.
+	 * Flush message so client will see it, except for AUTH_REQ_OK and
+	 * AUTH_REQ_SASL_FIN, which need not be sent until we are ready for queries.
 	 */
-	if (areq != AUTH_REQ_OK)
+	if (areq != AUTH_REQ_OK && areq != AUTH_REQ_SASL_FIN)
 		pq_flush();
 
 	CHECK_FOR_INTERRUPTS();
@@ -850,7 +850,10 @@ CheckSCRAMAuth(Port *port, char *shadow_pass, char **logdetail)
 	void	   *scram_opaq;
 	char	   *output = NULL;
 	int			outputlen = 0;
+	char	   *input;
+	int			inputlen;
 	int			result;
+	bool		initial;
 
 	/*
 	 * SASL auth is not supported for protocol versions before 3, because it
@@ -866,10 +869,13 @@ CheckSCRAMAuth(Port *port, char *shadow_pass, char **logdetail)
 				 errmsg("SASL authentication is not supported in protocol version 2")));
 
 	/*
-	 * Send first the authentication request to user.
+	 * Send the SASL authentication request to user.  It includes the list of
+	 * authentication mechanisms (which is trivial, because we only support
+	 * SCRAM-SHA-256 at the moment).  The extra "\0" is for an empty string to
+	 * terminate the list.
 	 */
-	sendAuthRequest(port, AUTH_REQ_SASL, SCRAM_SHA256_NAME,
-					strlen(SCRAM_SHA256_NAME) + 1);
+	sendAuthRequest(port, AUTH_REQ_SASL, SCRAM_SHA256_NAME "\0",
+					strlen(SCRAM_SHA256_NAME) + 2);
 
 	/*
 	 * Initialize the status tracker for message exchanges.
@@ -890,6 +896,7 @@ CheckSCRAMAuth(Port *port, char *shadow_pass, char **logdetail)
 	 * from the client.  All messages from client to server are password
 	 * packets (type 'p').
 	 */
+	initial = true;
 	do
 	{
 		pq_startmsgread();
@@ -921,10 +928,50 @@ CheckSCRAMAuth(Port *port, char *shadow_pass, char **logdetail)
 		elog(DEBUG4, "Processing received SASL response of length %d", buf.len);
 
 		/*
+		 * The first SASLInitialResponse message is different from the
+		 * others.  It indicates which SASL mechanism the selected, and
+		 * contains an optional Initial Client Response payload.  The
+		 * subsequent SASLResponse messages contain just the SASL payload.
+		 */
+		if (initial)
+		{
+			const char *selected_mech;
+
+			/*
+			 * We only support SCRAM-SHA-256 at the moment, so anything else
+			 * is an error.
+			 */
+			selected_mech = pq_getmsgrawstring(&buf);
+			if (strcmp(selected_mech, SCRAM_SHA256_NAME) != 0)
+				ereport(COMMERROR,
+						(errcode(ERRCODE_PROTOCOL_VIOLATION),
+						 errmsg("client selected an invalid SASL authentication mechanism")));
+
+			inputlen = pq_getmsgint(&buf, 4);
+			if (inputlen == -1)
+				input = NULL;
+			else
+				input = (char *) pq_getmsgbytes(&buf, inputlen);
+
+			initial = false;
+		}
+		else
+		{
+			inputlen = buf.len;
+			input = (char *) pq_getmsgbytes(&buf, buf.len);
+		}
+		/*
+		 * The StringInfo guarantees that there's a \0 byte after the
+		 * response.
+		 */
+		Assert(input[inputlen] == '\0');
+		pq_getmsgend(&buf);
+
+		/*
 		 * we pass 'logdetail' as NULL when doing a mock authentication,
 		 * because we should already have a better error message in that case
 		 */
-		result = pg_be_scram_exchange(scram_opaq, buf.data, buf.len,
+		result = pg_be_scram_exchange(scram_opaq, input, inputlen,
 									  &output, &outputlen,
 									  logdetail);
 
@@ -938,7 +985,10 @@ CheckSCRAMAuth(Port *port, char *shadow_pass, char **logdetail)
 			 */
 			elog(DEBUG4, "sending SASL challenge of length %u", outputlen);
 
-			sendAuthRequest(port, AUTH_REQ_SASL_CONT, output, outputlen);
+			if (result == SASL_EXCHANGE_SUCCESS)
+				sendAuthRequest(port, AUTH_REQ_SASL_FIN, output, outputlen);
+			else
+				sendAuthRequest(port, AUTH_REQ_SASL_CONT, output, outputlen);
 
 			pfree(output);
 		}
diff --git a/src/include/libpq/pqcomm.h b/src/include/libpq/pqcomm.h
index 5441aaa93a..b6de569c5c 100644
--- a/src/include/libpq/pqcomm.h
+++ b/src/include/libpq/pqcomm.h
@@ -172,8 +172,9 @@ extern bool Db_user_namespace;
 #define AUTH_REQ_GSS		7	/* GSSAPI without wrap() */
 #define AUTH_REQ_GSS_CONT	8	/* Continue GSS exchanges */
 #define AUTH_REQ_SSPI		9	/* SSPI negotiate without wrap() */
-#define AUTH_REQ_SASL	   10	/* SASL */
-#define AUTH_REQ_SASL_CONT 11	/* continue SASL exchange */
+#define AUTH_REQ_SASL	   10	/* Begin SASL authentication */
+#define AUTH_REQ_SASL_CONT 11	/* Continue SASL authentication */
+#define AUTH_REQ_SASL_FIN  12	/* Final SASL message */
 
 typedef uint32 AuthRequest;
 
diff --git a/src/interfaces/libpq/fe-auth.c b/src/interfaces/libpq/fe-auth.c
index 06a7140a2b..7ff01cdc52 100644
--- a/src/interfaces/libpq/fe-auth.c
+++ b/src/interfaces/libpq/fe-auth.c
@@ -471,88 +471,128 @@ pg_SSPI_startup(PGconn *conn, int use_negotiate, int payloadLen)
 static int
 pg_SASL_init(PGconn *conn, int payloadLen)
 {
-	char		auth_mechanism[21];
-	char	   *initialresponse;
+	char	   *initialresponse = NULL;
 	int			initialresponselen;
 	bool		done;
 	bool		success;
-	int			res;
+	const char *selected_mechanism;
+	PQExpBufferData mechanism_buf;
 
-	/*
-	 * Read the authentication mechanism the server told us to use.
-	 */
-	if (payloadLen > sizeof(auth_mechanism) - 1)
-		printfPQExpBuffer(&conn->errorMessage,
-						  libpq_gettext("SASL authentication mechanism not supported\n"));
-	if (pqGetnchar(auth_mechanism, payloadLen, conn))
+	initPQExpBuffer(&mechanism_buf);
+
+	if (conn->sasl_state)
 	{
 		printfPQExpBuffer(&conn->errorMessage,
-						  "fe_sendauth: invalid authentication request from server: invalid authentication mechanism\n");
-
-		return STATUS_ERROR;
+						  libpq_gettext("duplicate SASL authentication request\n"));
+		goto error;
 	}
-	auth_mechanism[payloadLen] = '\0';
 
 	/*
-	 * Check the authentication mechanism (only SCRAM-SHA-256 is supported at
-	 * the moment.)
+	 * Parse the list of SASL authentication mechanisms in the
+	 * AuthenticationSASL message, and select the best mechanism that we
+	 * support.  (Only SCRAM-SHA-256 is supported at the moment.)
 	 */
-	if (strcmp(auth_mechanism, SCRAM_SHA256_NAME) == 0)
+	for (;;)
 	{
-		char	   *password;
-
-		conn->password_needed = true;
-		password = conn->connhost[conn->whichhost].password;
-		if (password == NULL)
-			password = conn->pgpass;
-		if (password == NULL || password[0] == '\0')
+		if (pqGets(&mechanism_buf, conn))
 		{
 			printfPQExpBuffer(&conn->errorMessage,
-							  PQnoPasswordSupplied);
-			return STATUS_ERROR;
+							  "fe_sendauth: invalid authentication request from server: invalid list of authentication mechanisms\n");
+			goto error;
 		}
+		if (PQExpBufferDataBroken(mechanism_buf))
+			goto oom_error;
+
+		/* An empty string indicates end of list */
+		if (mechanism_buf.data[0] == '\0')
+			break;
+
+		/*
+		 * If we have already selected a mechanism, just skip through the
+		 * rest of the list.
+		 */
+		if (selected_mechanism)
+			continue;
 
-		conn->sasl_state = pg_fe_scram_init(conn->pguser, password);
-		if (!conn->sasl_state)
+		/*
+		 * Do we support this mechanism?
+		 */
+		if (strcmp(mechanism_buf.data, SCRAM_SHA256_NAME) == 0)
 		{
-			printfPQExpBuffer(&conn->errorMessage,
-							  libpq_gettext("out of memory\n"));
-			return STATUS_ERROR;
+			char	   *password;
+
+			conn->password_needed = true;
+			password = conn->connhost[conn->whichhost].password;
+			if (password == NULL)
+				password = conn->pgpass;
+			if (password == NULL || password[0] == '\0')
+			{
+				printfPQExpBuffer(&conn->errorMessage,
+								  PQnoPasswordSupplied);
+				goto error;
+			}
+
+			conn->sasl_state = pg_fe_scram_init(conn->pguser, password);
+			if (!conn->sasl_state)
+				goto oom_error;
+			selected_mechanism = SCRAM_SHA256_NAME;
 		}
 	}
-	else
+
+	if (!selected_mechanism)
 	{
 		printfPQExpBuffer(&conn->errorMessage,
-		   libpq_gettext("SASL authentication mechanism %s not supported\n"),
-						  auth_mechanism);
-		return STATUS_ERROR;
+						  libpq_gettext("none of the server's SASL authentication mechanisms are supported\n"));
+		goto error;
 	}
 
-	/* Send the initial client response */
+	/* Get the mechanism-specific Initial Client Response, if any */
 	pg_fe_scram_exchange(conn->sasl_state,
 						 NULL, -1,
 						 &initialresponse, &initialresponselen,
 						 &done, &success, &conn->errorMessage);
 
-	if (initialresponselen != 0)
-	{
-		res = pqPacketSend(conn, 'p', initialresponse, initialresponselen);
-		free(initialresponse);
-
-		if (res != STATUS_OK)
-			return STATUS_ERROR;
-	}
-
 	if (done && !success)
+		goto error;
+
+	/*
+	 * Build a SASLInitialResponse message, and send it.
+	 */
+	if (pqPutMsgStart('p', true, conn))
+		goto error;
+	if (pqPuts(selected_mechanism, conn))
+		goto error;
+	if (initialresponselen >= 0)
 	{
-		/* Use error message, if set already */
-		if (conn->errorMessage.len == 0)
-			printfPQExpBuffer(&conn->errorMessage,
-							  "fe_sendauth: error in SASL authentication\n");
-		return STATUS_ERROR;
+		if (pqPutInt(initialresponselen, 4, conn))
+			goto error;
+		if (pqPutnchar(initialresponse, initialresponselen, conn))
+			goto error;
 	}
+	if (pqPutMsgEnd(conn))
+		goto error;
+	if (pqFlush(conn))
+		goto error;
+
+	termPQExpBuffer(&mechanism_buf);
+	if (initialresponse)
+		free(initialresponse);
 
 	return STATUS_OK;
+
+error:
+	termPQExpBuffer(&mechanism_buf);
+	if (initialresponse)
+		free(initialresponse);
+	return STATUS_ERROR;
+
+oom_error:
+	termPQExpBuffer(&mechanism_buf);
+	if (initialresponse)
+		free(initialresponse);
+	printfPQExpBuffer(&conn->errorMessage,
+					  libpq_gettext("out of memory\n"));
+	return STATUS_ERROR;
 }
 
 /*
@@ -561,7 +601,7 @@ pg_SASL_init(PGconn *conn, int payloadLen)
  * the protocol.
  */
 static int
-pg_SASL_exchange(PGconn *conn, int payloadLen)
+pg_SASL_exchange(PGconn *conn, int payloadLen, bool final)
 {
 	char	   *output;
 	int			outputlen;
@@ -594,9 +634,20 @@ pg_SASL_exchange(PGconn *conn, int payloadLen)
 						 &done, &success, &conn->errorMessage);
 	free(challenge);	/* don't need the input anymore */
 
-	/* Send the SASL response to the server, if any. */
+	if (final && !done)
+	{
+		if (outputlen != 0)
+			free(output);
+
+		printfPQExpBuffer(&conn->errorMessage,
+						  libpq_gettext("AuthenticationSASLFinal received from server, but SASL authentication was not completed\n"));
+		return STATUS_ERROR;
+	}
 	if (outputlen != 0)
 	{
+		/*
+		 * Send the SASL response to the server.
+		 */
 		res = pqPacketSend(conn, 'p', output, outputlen);
 		free(output);
 
@@ -914,13 +965,15 @@ pg_fe_sendauth(AuthRequest areq, int payloadLen, PGconn *conn)
 			break;
 
 		case AUTH_REQ_SASL_CONT:
+		case AUTH_REQ_SASL_FIN:
 			if (conn->sasl_state == NULL)
 			{
 				printfPQExpBuffer(&conn->errorMessage,
 								  "fe_sendauth: invalid authentication request from server: AUTH_REQ_SASL_CONT without AUTH_REQ_SASL\n");
 				return STATUS_ERROR;
 			}
-			if (pg_SASL_exchange(conn, payloadLen) != STATUS_OK)
+			if (pg_SASL_exchange(conn, payloadLen,
+								 (areq == AUTH_REQ_SASL_FIN)) != STATUS_OK)
 			{
 				/* Use error message, if set already */
 				if (conn->errorMessage.len == 0)
-- 
2.11.0

