v6-0002-WIP-use-secure_read-to-peek.patch

From a0105faf49ade2c695523b88c927c402aae6eb53 Mon Sep 17 00:00:00 2001
From: Thomas Munro <thomas.munro@gmail.com>
Date: Tue, 23 Mar 2021 22:15:23 +1300
Subject: [PATCH v6 2/2] WIP -- use secure_read() to peek

XXX Just a sketch, probably has bugs
---
 doc/src/sgml/config.sgml    |  4 +-
 src/backend/libpq/pqcomm.c  | 84 ++++++++++++++++++-------------------
 src/backend/tcop/postgres.c | 11 ++++-
 src/include/libpq/libpq.h   |  2 +-
 4 files changed, 53 insertions(+), 48 deletions(-)

diff --git a/doc/src/sgml/config.sgml b/doc/src/sgml/config.sgml
index 5cd0d38dbf..e522be460c 100644
--- a/doc/src/sgml/config.sgml
+++ b/doc/src/sgml/config.sgml
@@ -1008,8 +1008,8 @@ include_dir 'conf.d'
        <para>
         Sets the time interval between checks that the client is still
         connected, while running queries.  The check is performed by testing
-        whether one byte could be read from the socket with
-        <symbol>MSG_PEEK</symbol>.  If the kernel reports that the connection
+        whether a part of the next message can be read from the client.
+        If the kernel reports that the connection
         has been closed or lost, a long running query can abort immediately,
         rather than discovering the problem when it eventually tries to send
         the response.
diff --git a/src/backend/libpq/pqcomm.c b/src/backend/libpq/pqcomm.c
index 74b309a1ef..a763a1b535 100644
--- a/src/backend/libpq/pqcomm.c
+++ b/src/backend/libpq/pqcomm.c
@@ -921,13 +921,13 @@ socket_set_nonblocking(bool nonblocking)
 }
 
 /* --------------------------------
- *		pq_recvbuf - load some bytes into the input buffer
+ *		pq_recvbuf_ext - load some bytes into the input buffer
  *
  *		returns 0 if OK, EOF if trouble
  * --------------------------------
  */
 static int
-pq_recvbuf(void)
+pq_recvbuf_ext(bool nonblocking)
 {
 	if (PqRecvPointer > 0)
 	{
@@ -943,8 +943,7 @@ pq_recvbuf(void)
 			PqRecvLength = PqRecvPointer = 0;
 	}
 
-	/* Ensure that we're in blocking mode */
-	socket_set_nonblocking(false);
+	socket_set_nonblocking(nonblocking);
 
 	/* Can fill buffer from PqRecvLength and upwards */
 	for (;;)
@@ -956,6 +955,9 @@ pq_recvbuf(void)
 
 		if (r < 0)
 		{
+			if (nonblocking && (errno == EAGAIN || errno == EWOULDBLOCK))
+				return 0;
+
 			if (errno == EINTR)
 				continue;		/* Ok if interrupted */
 
@@ -983,6 +985,13 @@ pq_recvbuf(void)
 	}
 }
 
+static int
+pq_recvbuf(void)
+{
+	return pq_recvbuf_ext(false);
+}
+
+
 /* --------------------------------
  *		pq_getbyte	- get a single byte from connection, or return EOF
  * --------------------------------
@@ -1924,52 +1933,39 @@ pq_settcpusertimeout(int timeout, Port *port)
 	return STATUS_OK;
 }
 
-/* --------------------------------
- *	pq_check_client_connection - is the client still connected?
- * --------------------------------
+/*
+ * Proactively check if the client connection has gone away.  Return 0 if still
+ * connected or currently reading a message already, EOF if disconnected, and 1
+ * if at least one byte is available to read.  If 1 is returned, the first
+ * byte of the next message is written to *c without consuming it, but we can't
+ * find out if the client has disconnected until we consume more data.
  */
-bool
-pq_check_client_connection(void)
+int
+pq_check_client_connection(unsigned char *c)
 {
-	bool		connected = true;
-	char		nextbyte;
-	int			r;
-
-retry:
-#ifdef WIN32
-	pgwin32_noblock = 1;
-#endif
-	r = recv(MyProcPort->sock, &nextbyte, 1, MSG_PEEK);
-#ifdef WIN32
-	pgwin32_noblock = 0;
-#endif
+	/* We're already in the middle of a message, so assume we are connected. */
+	if (PqCommReadingMsg)
+		return 0;
 
-	if (r == 0)
-	{
-		/* EOF detected. */
-		connected = false;
-	}
-	else if (r > 0)
-	{
-		/* Data available to read.  Connection looks good. */
-	}
-	else if (errno == EINTR)
-	{
-		/* Interrupted by a signal, so retry. */
-		goto retry;
-	}
-	else if (errno == EAGAIN || errno == EWOULDBLOCK)
+	/* Do we have a byte already in our read buffer? */
+	if (PqRecvPointer < PqRecvLength)
 	{
-		/* No data available to read.  Connection looks good. */
+		*c = PqRecvBuffer[PqRecvPointer];
+		return 1;
 	}
-	else
+
+	/* Try to read at least one byte from secure_read() without blocking. */
+	if (pq_recvbuf_ext(true))
+		return EOF;
+
+	/* Now do we have a byte in our read buffer? */
+	if (PqRecvPointer < PqRecvLength)
 	{
-		/* Got some other error.  We'd better log the reason. */
-		ereport(COMMERROR,
-				(errcode(ERRCODE_CONNECTION_EXCEPTION),
-				 errmsg("could not check client connection: %m")));
-		connected = false;
+		*c = PqRecvBuffer[PqRecvPointer];
+		return 1;
 	}
 
-	return connected;
+	/* No data, but we're still connected. */
+
+	return 0;
 }
diff --git a/src/backend/tcop/postgres.c b/src/backend/tcop/postgres.c
index 44fb7b7c30..23d1ee84da 100644
--- a/src/backend/tcop/postgres.c
+++ b/src/backend/tcop/postgres.c
@@ -3170,7 +3170,16 @@ ProcessInterrupts(void)
 		 */
 		if (!DoingCommandRead && client_connection_check_interval > 0)
 		{
-			if (!pq_check_client_connection())
+			unsigned char peekbyte;
+			int		r;
+
+			/*
+			 * Try to peek ahead to see if we've been disconnected.  If we see
+			 * a pipelined 'X' message, we'll treat that as a disconnection
+			 * too, but we don't try to look further ahead than that.
+			 */
+			r = pq_check_client_connection(&peekbyte);
+			if (r == EOF || (r == 1 && peekbyte == 'X'))
 				ClientConnectionLost = true;
 			enable_timeout_after(CLIENT_CONNECTION_CHECK_TIMEOUT,
 								 client_connection_check_interval);
diff --git a/src/include/libpq/libpq.h b/src/include/libpq/libpq.h
index 3bd97c4e93..f6cd3ee067 100644
--- a/src/include/libpq/libpq.h
+++ b/src/include/libpq/libpq.h
@@ -71,7 +71,7 @@ extern int	pq_getbyte(void);
 extern int	pq_peekbyte(void);
 extern int	pq_getbyte_if_available(unsigned char *c);
 extern int	pq_putmessage_v2(char msgtype, const char *s, size_t len);
-extern bool pq_check_client_connection(void);
+extern int	pq_check_client_connection(unsigned char *c);
 
 /*
  * prototypes for functions in be-secure.c
-- 
2.30.1