v1-0004-libpq-Trace-frontend-authentication-challenges.patch

From 8cbba167f99b2ce3f5ca3010645aaebf38920429 Mon Sep 17 00:00:00 2001
From: Jelte Fennema-Nio <jelte.fennema@microsoft.com>
Date: Fri, 21 Jun 2024 00:13:05 +0200
Subject: [PATCH v1 4/8] 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.
---
 src/interfaces/libpq/fe-auth.c   | 14 ++++++++
 src/interfaces/libpq/fe-trace.c  | 62 ++++++++++++++++++++++++++++++++
 src/interfaces/libpq/libpq-int.h | 10 ++++++
 src/tools/pgindent/typedefs.list |  1 +
 4 files changed, 87 insertions(+)

diff --git a/src/interfaces/libpq/fe-auth.c b/src/interfaces/libpq/fe-auth.c
index b5b367f24d0..cb3919e0492 100644
--- a/src/interfaces/libpq/fe-auth.c
+++ b/src/interfaces/libpq/fe-auth.c
@@ -124,12 +124,15 @@ pg_GSS_continue(PGconn *conn, int payloadlen)
 		 * first or subsequent packet, just send the same kind of password
 		 * packet.
 		 */
+		conn->current_auth_response = AUTH_RESP_GSS;
 		if (pqPacketSend(conn, PqMsg_GSSResponse,
 						 goutbuf.value, goutbuf.length) != STATUS_OK)
 		{
+			conn->current_auth_response = AUTH_RESP_NONE;
 			gss_release_buffer(&lmin_s, &goutbuf);
 			return STATUS_ERROR;
 		}
+		conn->current_auth_response = AUTH_RESP_NONE;
 	}
 	gss_release_buffer(&lmin_s, &goutbuf);
 
@@ -324,12 +327,15 @@ pg_SSPI_continue(PGconn *conn, int payloadlen)
 		 */
 		if (outbuf.pBuffers[0].cbBuffer > 0)
 		{
+			conn->current_auth_response = AUTH_RESP_GSS;
 			if (pqPacketSend(conn, PqMsg_GSSResponse,
 							 outbuf.pBuffers[0].pvBuffer, outbuf.pBuffers[0].cbBuffer))
 			{
+				conn->current_auth_response = AUTH_RESP_NONE;
 				FreeContextBuffer(outbuf.pBuffers[0].pvBuffer);
 				return STATUS_ERROR;
 			}
+			conn->current_auth_response = AUTH_RESP_NONE;
 		}
 		FreeContextBuffer(outbuf.pBuffers[0].pvBuffer);
 	}
@@ -597,8 +603,11 @@ pg_SASL_init(PGconn *conn, int payloadlen)
 		if (pqPutnchar(initialresponse, initialresponselen, conn))
 			goto error;
 	}
+	conn->current_auth_response = AUTH_RESP_SASL_INITIAL;
 	if (pqPutMsgEnd(conn))
 		goto error;
+	conn->current_auth_response = AUTH_RESP_NONE;
+
 	if (pqFlush(conn))
 		goto error;
 
@@ -608,6 +617,7 @@ pg_SASL_init(PGconn *conn, int payloadlen)
 	return STATUS_OK;
 
 error:
+	conn->current_auth_response = AUTH_RESP_NONE;
 	termPQExpBuffer(&mechanism_buf);
 	free(initialresponse);
 	return STATUS_ERROR;
@@ -683,7 +693,9 @@ pg_SASL_continue(PGconn *conn, int payloadlen, bool final)
 		/*
 		 * Send the SASL response to the server.
 		 */
+		conn->current_auth_response = AUTH_RESP_SASL;
 		res = pqPacketSend(conn, PqMsg_SASLResponse, output, outputlen);
+		conn->current_auth_response = AUTH_RESP_NONE;
 		free(output);
 
 		if (res != STATUS_OK)
@@ -754,7 +766,9 @@ pg_password_sendauth(PGconn *conn, const char *password, AuthRequest areq)
 		default:
 			return STATUS_ERROR;
 	}
+	conn->current_auth_response = AUTH_RESP_PASSWORD;
 	ret = pqPacketSend(conn, PqMsg_PasswordMessage, pwd_to_send, strlen(pwd_to_send) + 1);
+	conn->current_auth_response = AUTH_RESP_NONE;
 	free(crypt_pwd);
 	return ret;
 }
diff --git a/src/interfaces/libpq/fe-trace.c b/src/interfaces/libpq/fe-trace.c
index f1e4c07cc25..23df8d0e10e 100644
--- a/src/interfaces/libpq/fe-trace.c
+++ b/src/interfaces/libpq/fe-trace.c
@@ -350,6 +350,40 @@ 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)
 {
@@ -606,6 +640,34 @@ 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);
+
+			/*
+			 * pqTraceOutput_PasswordMessage(conn->Pfdebug, message,
+			 * &logCursor);
+			 */
+			switch (conn->current_auth_response)
+			{
+				case AUTH_RESP_GSS:
+					pqTraceOutput_GSSResponse(conn->Pfdebug, message, &logCursor, length, regress);
+					break;
+				case AUTH_RESP_PASSWORD:
+					pqTraceOutput_PasswordMessage(conn->Pfdebug, message, &logCursor);
+					break;
+				case AUTH_RESP_SASL_INITIAL:
+					pqTraceOutput_SASLInitialResponse(conn->Pfdebug, message, &logCursor, regress);
+					break;
+				case AUTH_RESP_SASL:
+					pqTraceOutput_SASLResponse(conn->Pfdebug, message, &logCursor, length, regress);
+					break;
+				default:
+					fprintf(conn->Pfdebug, "UnknownAuthenticationResponse");
+					break;
+			}
+			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 f36d76bf3fe..f9b21c86ae4 100644
--- a/src/interfaces/libpq/libpq-int.h
+++ b/src/interfaces/libpq/libpq-int.h
@@ -331,6 +331,15 @@ typedef enum
 	PGQUERY_CLOSE				/* Close Statement or Portal */
 } PGQueryClass;
 
+typedef enum
+{
+	AUTH_RESP_NONE = 0,
+	AUTH_RESP_GSS,
+	AUTH_RESP_PASSWORD,
+	AUTH_RESP_SASL_INITIAL,
+	AUTH_RESP_SASL,
+} AuthResponseType;
+
 /*
  * An entry in the pending command queue.
  */
@@ -490,6 +499,7 @@ struct pg_conn
 										 * codes */
 	bool		client_finished_auth;	/* have we finished our half of the
 										 * authentication exchange? */
+	AuthResponseType current_auth_response;
 
 
 	/* Transient state needed while establishing connection */
diff --git a/src/tools/pgindent/typedefs.list b/src/tools/pgindent/typedefs.list
index 61ad417cde6..06ac0574bfd 100644
--- a/src/tools/pgindent/typedefs.list
+++ b/src/tools/pgindent/typedefs.list
@@ -167,6 +167,7 @@ AttrMissing
 AttrNumber
 AttributeOpts
 AuthRequest
+AuthResponseType
 AuthToken
 AutoPrewarmSharedState
 AutoVacOpts
-- 
2.34.1