From 795e21456af5a3ce758bec41b4f0eac825a66fd3 Mon Sep 17 00:00:00 2001
From: Jelte Fennema-Nio <jelte.fennema@microsoft.com>
Date: Thu, 13 Jun 2024 19:35:43 +0200
Subject: [PATCH v2 6/8] libpq: Trace all messages received from the server

Not all messages that libpq received from the server would be sent
through our message tracing logic. This fixes that by introducing
pqFinishParsingMessage which makes it a lot harder to forget about doing
so.

The messages that we now newly send through our tracing logic are:

- CopyData (received by COPY TO STDOUT)
- Authentication requests
- NegotiateProtocolVersion
- Some ErrorResponse messages during connection startup
- ReadyForQuery when received after a FunctionCall message
---
 src/interfaces/libpq/fe-auth.c      | 21 ++++++++++++++
 src/interfaces/libpq/fe-connect.c   | 13 ++++++---
 src/interfaces/libpq/fe-protocol3.c | 45 +++++++++++++++--------------
 src/interfaces/libpq/libpq-int.h    |  1 +
 4 files changed, 54 insertions(+), 26 deletions(-)

diff --git a/src/interfaces/libpq/fe-auth.c b/src/interfaces/libpq/fe-auth.c
index f6bb58fd902..6d37f07b691 100644
--- a/src/interfaces/libpq/fe-auth.c
+++ b/src/interfaces/libpq/fe-auth.c
@@ -94,6 +94,10 @@ pg_GSS_continue(PGconn *conn, int payloadlen)
 		ginbuf.value = NULL;
 	}
 
+	/* finished parsing, trace server-to-client message */
+	if (conn->Pfdebug)
+		pqTraceOutputMessage(conn, conn->inBuffer + conn->inStart, false);
+
 	/* Only try to acquire credentials if GSS delegation isn't disabled. */
 	if (!pg_GSS_have_cred_cache(&conn->gcred))
 		conn->gcred = GSS_C_NO_CREDENTIAL;
@@ -260,6 +264,10 @@ pg_SSPI_continue(PGconn *conn, int payloadlen)
 		InBuffers[0].BufferType = SECBUFFER_TOKEN;
 	}
 
+	/* finished parsing, trace server-to-client message */
+	if (conn->Pfdebug)
+		pqTraceOutputMessage(conn, conn->inBuffer + conn->inStart, false);
+
 	OutBuffers[0].pvBuffer = NULL;
 	OutBuffers[0].BufferType = SECBUFFER_TOKEN;
 	OutBuffers[0].cbBuffer = 0;
@@ -567,6 +575,10 @@ pg_SASL_init(PGconn *conn, int payloadlen)
 		}
 	}
 
+	/* finished parsing, trace server-to-client message */
+	if (conn->Pfdebug)
+		pqTraceOutputMessage(conn, conn->inBuffer + conn->inStart, false);
+
 	Assert(conn->sasl);
 
 	/*
@@ -657,6 +669,11 @@ pg_SASL_continue(PGconn *conn, int payloadlen, bool final)
 		free(challenge);
 		return STATUS_ERROR;
 	}
+
+	/* finished parsing, trace server-to-client message */
+	if (conn->Pfdebug)
+		pqTraceOutputMessage(conn, conn->inBuffer + conn->inStart, false);
+
 	/* For safety and convenience, ensure the buffer is NULL-terminated. */
 	challenge[payloadlen] = '\0';
 
@@ -723,6 +740,10 @@ pg_password_sendauth(PGconn *conn, const char *password, AuthRequest areq)
 			return STATUS_ERROR;	/* shouldn't happen */
 	}
 
+	/* finished parsing, trace server-to-client message */
+	if (conn->Pfdebug)
+		pqTraceOutputMessage(conn, conn->inBuffer + conn->inStart, false);
+
 	/* Encrypt the password if needed. */
 
 	switch (areq)
diff --git a/src/interfaces/libpq/fe-connect.c b/src/interfaces/libpq/fe-connect.c
index 0772455c7d1..babb22c69f0 100644
--- a/src/interfaces/libpq/fe-connect.c
+++ b/src/interfaces/libpq/fe-connect.c
@@ -3776,7 +3776,7 @@ keep_going:						/* We will come back to here until there is
 						return PGRES_POLLING_READING;
 					}
 					/* OK, we read the message; mark data consumed */
-					conn->inStart = conn->inCursor;
+					pqFinishParsingMessage(conn, conn->inCursor);
 
 					/*
 					 * Before 7.2, the postmaster didn't always end its
@@ -3826,7 +3826,7 @@ keep_going:						/* We will come back to here until there is
 						goto error_return;
 					}
 					/* OK, we read the message; mark data consumed */
-					conn->inStart = conn->inCursor;
+					pqFinishParsingMessage(conn, conn->inCursor);
 
 					/*
 					 * If error is "cannot connect now", try the next host if
@@ -3855,7 +3855,7 @@ keep_going:						/* We will come back to here until there is
 						goto error_return;
 					}
 					/* OK, we read the message; mark data consumed */
-					conn->inStart = conn->inCursor;
+					pqFinishParsingMessage(conn, conn->inCursor);
 					goto error_return;
 				}
 
@@ -3880,7 +3880,12 @@ keep_going:						/* We will come back to here until there is
 				 */
 				res = pg_fe_sendauth(areq, msgLength, conn);
 
-				/* OK, we have processed the message; mark data consumed */
+				/*
+				 * OK, we have processed the message; mark data consumed We
+				 * don't call pqFinishParsingMessage here because we already
+				 * trace this message inside pg_fe_sendauth (since we need to
+				 * trace it before sending the response.
+				 */
 				conn->inStart = conn->inCursor;
 
 				if (res != STATUS_OK)
diff --git a/src/interfaces/libpq/fe-protocol3.c b/src/interfaces/libpq/fe-protocol3.c
index 3170d484f02..f165009d3bb 100644
--- a/src/interfaces/libpq/fe-protocol3.c
+++ b/src/interfaces/libpq/fe-protocol3.c
@@ -57,6 +57,22 @@ static int	build_startup_packet(const PGconn *conn, char *packet,
 								 const PQEnvironmentOption *options);
 
 
+/*
+ * pqFinishParsingMessage is a simple helper function to advance inStart after
+ * a server-to-client message has successfully been parsed. It's main purpose
+ * is making it hard forget to call pqTraceOutputMessage when doing so.
+ */
+void
+pqFinishParsingMessage(PGconn *conn, int newInStart)
+{
+	/* trace server-to-client message */
+	if (conn->Pfdebug)
+		pqTraceOutputMessage(conn, conn->inBuffer + conn->inStart, false);
+
+	/* Mark message as done */
+	conn->inStart = newInStart;
+}
+
 /*
  * parseInput: if appropriate, parse input data from backend
  * until input is exhausted or a stopping state is reached.
@@ -454,12 +470,8 @@ pqParseInput3(PGconn *conn)
 		/* Successfully consumed this message */
 		if (conn->inCursor == conn->inStart + 5 + msgLength)
 		{
-			/* trace server-to-client message */
-			if (conn->Pfdebug)
-				pqTraceOutputMessage(conn, conn->inBuffer + conn->inStart, false);
-
 			/* Normal case: parsing agrees with specified length */
-			conn->inStart = conn->inCursor;
+			pqFinishParsingMessage(conn, conn->inCursor);
 		}
 		else
 		{
@@ -1728,12 +1740,7 @@ getCopyDataMessage(PGconn *conn)
 				return -1;
 		}
 
-		/* trace server-to-client message */
-		if (conn->Pfdebug)
-			pqTraceOutputMessage(conn, conn->inBuffer + conn->inStart, false);
-
-		/* Drop the processed message and loop around for another */
-		conn->inStart = conn->inCursor;
+		pqFinishParsingMessage(conn, conn->inCursor);
 	}
 }
 
@@ -1791,13 +1798,13 @@ pqGetCopyData3(PGconn *conn, char **buffer, int async)
 			(*buffer)[msgLength] = '\0';	/* Add terminating null */
 
 			/* Mark message consumed */
-			conn->inStart = conn->inCursor + msgLength;
+			pqFinishParsingMessage(conn, conn->inCursor + msgLength);
 
 			return msgLength;
 		}
 
 		/* Empty, so drop it and loop around for another */
-		conn->inStart = conn->inCursor;
+		pqFinishParsingMessage(conn, conn->inCursor);
 	}
 }
 
@@ -2168,8 +2175,8 @@ pqFunctionCall3(PGconn *conn, Oid fnid,
 			case 'Z':			/* backend is ready for new query */
 				if (getReadyForQuery(conn))
 					continue;
-				/* consume the message and exit */
-				conn->inStart += 5 + msgLength;
+
+				pqFinishParsingMessage(conn, conn->inStart + 5 + msgLength);
 
 				/*
 				 * If we already have a result object (probably an error), use
@@ -2208,13 +2215,7 @@ pqFunctionCall3(PGconn *conn, Oid fnid,
 				return pqPrepareAsyncResult(conn);
 		}
 
-		/* trace server-to-client message */
-		if (conn->Pfdebug)
-			pqTraceOutputMessage(conn, conn->inBuffer + conn->inStart, false);
-
-		/* Completed this message, keep going */
-		/* trust the specified message length as what to skip */
-		conn->inStart += 5 + msgLength;
+		pqFinishParsingMessage(conn, conn->inStart + 5 + msgLength);
 		needInput = false;
 	}
 
diff --git a/src/interfaces/libpq/libpq-int.h b/src/interfaces/libpq/libpq-int.h
index 82a568f8040..74af7d36676 100644
--- a/src/interfaces/libpq/libpq-int.h
+++ b/src/interfaces/libpq/libpq-int.h
@@ -726,6 +726,7 @@ extern int	PQsendQueryContinue(PGconn *conn, const char *query);
 
 /* === in fe-protocol3.c === */
 
+extern void pqFinishParsingMessage(PGconn *conn, int newInStart);
 extern char *pqBuildStartupPacket3(PGconn *conn, int *packetlen,
 								   const PQEnvironmentOption *options);
 extern void pqParseInput3(PGconn *conn);
-- 
2.34.1

