From f34fb8a816a74301c402da32897deab98d1ac069 Mon Sep 17 00:00:00 2001
From: Jelte Fennema-Nio <jelte.fennema@microsoft.com>
Date: Thu, 20 Jun 2024 17:40:08 +0200
Subject: [PATCH v3 5/8] libpq: Trace responses to SSLRequest and GSSENCRequest

Since these are single bytes instead of v2 or v3 messages they need
custom tracing logic.  These "messages" don't even have official names
in the protocol specification, so I (Jelte) called them SSLResponse and
GSSENCResponse here.

Author: Jelte Fennema-Nio <postgres@jeltef.nl>
Discussion: https://postgr.es/m/CAGECzQSoPHtZ4xe0raJ6FYSEiPPS+YWXBhOGo+Y1YecLgknF3g@mail.gmail.com
---
 src/interfaces/libpq/fe-connect.c | 16 ++++++++++++++++
 src/interfaces/libpq/fe-trace.c   | 20 ++++++++++++++++++++
 src/interfaces/libpq/libpq-int.h  |  2 ++
 3 files changed, 38 insertions(+)

diff --git a/src/interfaces/libpq/fe-connect.c b/src/interfaces/libpq/fe-connect.c
index 360d9a4547..f7ece66b69 100644
--- a/src/interfaces/libpq/fe-connect.c
+++ b/src/interfaces/libpq/fe-connect.c
@@ -3493,11 +3493,17 @@ keep_going:						/* We will come back to here until there is
 					}
 					if (SSLok == 'S')
 					{
+						if (conn->Pfdebug)
+							pqTraceOutputCharResponse(conn, "SSLResponse",
+													  SSLok);
 						/* mark byte consumed */
 						conn->inStart = conn->inCursor;
 					}
 					else if (SSLok == 'N')
 					{
+						if (conn->Pfdebug)
+							pqTraceOutputCharResponse(conn, "SSLResponse",
+													  SSLok);
 						/* mark byte consumed */
 						conn->inStart = conn->inCursor;
 
@@ -3635,6 +3641,11 @@ keep_going:						/* We will come back to here until there is
 
 					if (gss_ok == 'N')
 					{
+						if (conn->Pfdebug)
+							pqTraceOutputCharResponse(conn,
+													  "GSSENCResponse",
+													  gss_ok);
+
 						/*
 						 * The connection is still valid, so if it's OK to
 						 * continue without GSS, we can proceed using this
@@ -3648,6 +3659,11 @@ keep_going:						/* We will come back to here until there is
 												gss_ok);
 						goto error_return;
 					}
+
+					if (conn->Pfdebug)
+						pqTraceOutputCharResponse(conn,
+												  "GSSENCResponse",
+												  gss_ok);
 				}
 
 				/* Begin or continue GSSAPI negotiation */
diff --git a/src/interfaces/libpq/fe-trace.c b/src/interfaces/libpq/fe-trace.c
index 367b322b99..3bb62397ec 100644
--- a/src/interfaces/libpq/fe-trace.c
+++ b/src/interfaces/libpq/fe-trace.c
@@ -840,3 +840,23 @@ pqTraceOutputNoTypeByteMessage(PGconn *conn, const char *message)
 
 	fputc('\n', conn->Pfdebug);
 }
+
+/*
+ * Trace a single-byte backend response received for a specific request
+ * type the frontend previously sent.  Only useful for the simplest of
+ * FE/BE interaction workflows such as authentication.
+ */
+void
+pqTraceOutputCharResponse(PGconn *conn, const char *requestType,
+						  char response)
+{
+	if ((conn->traceFlags & PQTRACE_SUPPRESS_TIMESTAMPS) == 0)
+	{
+		char		timestr[128];
+
+		pqTraceFormatTimestamp(timestr, sizeof(timestr));
+		fprintf(conn->Pfdebug, "%s\t", timestr);
+	}
+
+	fprintf(conn->Pfdebug, "B\t1\t%s\t %c\n", requestType, response);
+}
diff --git a/src/interfaces/libpq/libpq-int.h b/src/interfaces/libpq/libpq-int.h
index 03e4da40ba..bca3284041 100644
--- a/src/interfaces/libpq/libpq-int.h
+++ b/src/interfaces/libpq/libpq-int.h
@@ -889,6 +889,8 @@ extern ssize_t pg_GSS_read(PGconn *conn, void *ptr, size_t len);
 extern void pqTraceOutputMessage(PGconn *conn, const char *message,
 								 bool toServer);
 extern void pqTraceOutputNoTypeByteMessage(PGconn *conn, const char *message);
+extern void pqTraceOutputCharResponse(PGconn *conn, const char *requestType,
+									  char response);
 
 /* === miscellaneous macros === */
 
-- 
2.39.2