From 4a26851519266d31624ad9e332bd89483d7280a2 Mon Sep 17 00:00:00 2001
From: Alvaro Herrera <alvherre@alvh.no-ip.org>
Date: Fri, 9 Aug 2024 19:05:08 -0400
Subject: [PATCH] libpq: Trace frontend authentication challenges

If tracing was enabled during connection startup, these messages would
previously be listed in the trace output as something like this:

F	54	Unknown message: 70
mismatched message length: consumed 4, expected 54

With this commit their type and contents are now correctly listed.

Author: Jelte Fennema-Nio <postgres@jeltef.nl>
Reviewed-by: Michael Paquier <michael@paquier.xyz>
Discussion: https://postgr.es/m/CAGECzQSoPHtZ4xe0raJ6FYSEiPPS+YWXBhOGo+Y1YecLgknF3g@mail.gmail.com
---
 src/interfaces/libpq/fe-auth.c   |  6 +++
 src/interfaces/libpq/fe-trace.c  | 69 ++++++++++++++++++++++++++++++++
 src/interfaces/libpq/libpq-int.h | 16 +++++++-
 3 files changed, 90 insertions(+), 1 deletion(-)

diff --git a/src/interfaces/libpq/fe-auth.c b/src/interfaces/libpq/fe-auth.c
index 3b25d8afda..4da07c1f98 100644
--- a/src/interfaces/libpq/fe-auth.c
+++ b/src/interfaces/libpq/fe-auth.c
@@ -124,6 +124,7 @@ pg_GSS_continue(PGconn *conn, int payloadlen)
 		 * first or subsequent packet, just send the same kind of password
 		 * packet.
 		 */
+		conn->current_auth_response = AUTH_RESPONSE_GSS;
 		if (pqPacketSend(conn, PqMsg_GSSResponse,
 						 goutbuf.value, goutbuf.length) != STATUS_OK)
 		{
@@ -324,6 +325,7 @@ pg_SSPI_continue(PGconn *conn, int payloadlen)
 		 */
 		if (outbuf.pBuffers[0].cbBuffer > 0)
 		{
+			conn->current_auth_response = AUTH_RESPONSE_GSS;
 			if (pqPacketSend(conn, PqMsg_GSSResponse,
 							 outbuf.pBuffers[0].pvBuffer, outbuf.pBuffers[0].cbBuffer))
 			{
@@ -597,8 +599,10 @@ pg_SASL_init(PGconn *conn, int payloadlen)
 		if (pqPutnchar(initialresponse, initialresponselen, conn))
 			goto error;
 	}
+	conn->current_auth_response = AUTH_RESPONSE_SASL_INITIAL;
 	if (pqPutMsgEnd(conn))
 		goto error;
+
 	if (pqFlush(conn))
 		goto error;
 
@@ -683,6 +687,7 @@ pg_SASL_continue(PGconn *conn, int payloadlen, bool final)
 		/*
 		 * Send the SASL response to the server.
 		 */
+		conn->current_auth_response = AUTH_RESPONSE_SASL;
 		res = pqPacketSend(conn, PqMsg_SASLResponse, output, outputlen);
 		free(output);
 
@@ -754,6 +759,7 @@ pg_password_sendauth(PGconn *conn, const char *password, AuthRequest areq)
 		default:
 			return STATUS_ERROR;
 	}
+	conn->current_auth_response = AUTH_RESPONSE_PASSWORD;
 	ret = pqPacketSend(conn, PqMsg_PasswordMessage,
 					   pwd_to_send, strlen(pwd_to_send) + 1);
 	free(crypt_pwd);
diff --git a/src/interfaces/libpq/fe-trace.c b/src/interfaces/libpq/fe-trace.c
index c34047a6e4..367b322b99 100644
--- a/src/interfaces/libpq/fe-trace.c
+++ b/src/interfaces/libpq/fe-trace.c
@@ -354,6 +354,42 @@ pqTraceOutput_CopyFail(FILE *f, const char *message, int *cursor)
 	pqTraceOutputString(f, message, cursor, false);
 }
 
+static void
+pqTraceOutput_GSSResponse(FILE *f, const char *message, int *cursor,
+						  int length, bool regress)
+{
+	fprintf(f, "GSSResponse\t");
+	pqTraceOutputNchar(f, length - *cursor + 1, message, cursor, regress);
+}
+
+static void
+pqTraceOutput_PasswordMessage(FILE *f, const char *message, int *cursor)
+{
+	fprintf(f, "PasswordMessage\t");
+	pqTraceOutputString(f, message, cursor, false);
+}
+
+static void
+pqTraceOutput_SASLInitialResponse(FILE *f, const char *message, int *cursor,
+								  bool regress)
+{
+	int			initialResponse;
+
+	fprintf(f, "SASLInitialResponse\t");
+	pqTraceOutputString(f, message, cursor, false);
+	initialResponse = pqTraceOutputInt32(f, message, cursor, false);
+	if (initialResponse != -1)
+		pqTraceOutputNchar(f, initialResponse, message, cursor, regress);
+}
+
+static void
+pqTraceOutput_SASLResponse(FILE *f, const char *message, int *cursor,
+						   int length, bool regress)
+{
+	fprintf(f, "SASLResponse\t");
+	pqTraceOutputNchar(f, length - *cursor + 1, message, cursor, regress);
+}
+
 static void
 pqTraceOutput_FunctionCall(FILE *f, const char *message, int *cursor, bool regress)
 {
@@ -610,6 +646,39 @@ pqTraceOutputMessage(PGconn *conn, const char *message, bool toServer)
 		case PqMsg_CopyFail:
 			pqTraceOutput_CopyFail(conn->Pfdebug, message, &logCursor);
 			break;
+		case PqMsg_GSSResponse:
+			Assert(PqMsg_GSSResponse == PqMsg_PasswordMessage);
+			Assert(PqMsg_GSSResponse == PqMsg_SASLInitialResponse);
+			Assert(PqMsg_GSSResponse == PqMsg_SASLResponse);
+
+			/*
+			 * These messages share a common type byte, so we discriminate by
+			 * having the code store the auth type separately.
+			 */
+			switch (conn->current_auth_response)
+			{
+				case AUTH_RESPONSE_GSS:
+					pqTraceOutput_GSSResponse(conn->Pfdebug, message,
+											  &logCursor, length, regress);
+					break;
+				case AUTH_RESPONSE_PASSWORD:
+					pqTraceOutput_PasswordMessage(conn->Pfdebug, message,
+												  &logCursor);
+					break;
+				case AUTH_RESPONSE_SASL_INITIAL:
+					pqTraceOutput_SASLInitialResponse(conn->Pfdebug, message,
+													  &logCursor, regress);
+					break;
+				case AUTH_RESPONSE_SASL:
+					pqTraceOutput_SASLResponse(conn->Pfdebug, message,
+											   &logCursor, length, regress);
+					break;
+				default:
+					fprintf(conn->Pfdebug, "UnknownAuthenticationResponse");
+					break;
+			}
+			conn->current_auth_response = '\0';
+			break;
 		case PqMsg_FunctionCall:
 			pqTraceOutput_FunctionCall(conn->Pfdebug, message, &logCursor, regress);
 			break;
diff --git a/src/interfaces/libpq/libpq-int.h b/src/interfaces/libpq/libpq-int.h
index f36d76bf3f..03e4da40ba 100644
--- a/src/interfaces/libpq/libpq-int.h
+++ b/src/interfaces/libpq/libpq-int.h
@@ -331,6 +331,18 @@ typedef enum
 	PGQUERY_CLOSE				/* Close Statement or Portal */
 } PGQueryClass;
 
+
+/*
+ * valid values for pg_conn->current_auth_response.  These are just for
+ * libpq internal use: since authentication response types all use the
+ * protocol byte 'p', fe-trace.c needs a way to distinguish them in order
+ * to print them correctly.
+ */
+#define AUTH_RESPONSE_GSS			'G'
+#define AUTH_RESPONSE_PASSWORD		'P'
+#define AUTH_RESPONSE_SASL_INITIAL	'I'
+#define AUTH_RESPONSE_SASL			'S'
+
 /*
  * An entry in the pending command queue.
  */
@@ -490,7 +502,9 @@ struct pg_conn
 										 * codes */
 	bool		client_finished_auth;	/* have we finished our half of the
 										 * authentication exchange? */
-
+	char		current_auth_response;	/* used by pqTraceOutputMessage to
+										 * know which auth response we're
+										 * sending */
 
 	/* Transient state needed while establishing connection */
 	PGTargetServerType target_server_type;	/* desired session properties */
-- 
2.39.2

