WIP: Do the same block-raw-read dance in backend.ssl-reneg

This is more complicated than in the client, because the callers of
secure_write() will not automatically do a read, if the write fails.
Need to expose pq_recvbuf() and call it from secure_write(), which is
modularity violation.

4 files changed, 44 insertions, 19 deletions

diff --git a/src/backend/libpq/be-secure-openssl.c b/src/backend/libpq/be-secure-openssl.c
index b06f987b3f..b9bcf4eb78 100644
--- a/src/backend/libpq/be-secure-openssl.c
+++ b/src/backend/libpq/be-secure-openssl.c
@@ -93,6 +93,9 @@ static const char *SSLerrmessage(void);
 /* are we in the middle of a renegotiation? */
 static bool in_ssl_renegotiation = false;
 
+/* kill-switch to have my_sock_read pretend there's no data */
+static bool my_block_raw_read = false;
+
 static SSL_CTX *SSL_context = NULL;
 
 /* ------------------------------------------------------------ */
@@ -606,7 +609,9 @@ be_tls_write(Port *port, void *ptr, size_t len, int *waitfor)
 	}
 
 	errno = 0;
+	my_block_raw_read = true;
 	n = SSL_write(port->ssl, ptr, len);
+	my_block_raw_read = false;
 	err = SSL_get_error(port->ssl, n);
 	switch (err)
 	{
@@ -699,19 +704,28 @@ static BIO_METHOD my_bio_methods;
 static int
 my_sock_read(BIO *h, char *buf, int size)
 {
-	int			res = 0;
+	int			res;
 
-	if (buf != NULL)
+	if (buf == NULL)
+		return 0;	/* XXX: can this happen? */
+
+	/* If the kill-switch is set, pretend that there is no data. */
+	if (my_block_raw_read)
 	{
-		res = secure_raw_read(((Port *)h->ptr), buf, size);
+		errno = EWOULDBLOCK;
 		BIO_clear_retry_flags(h);
-		if (res <= 0)
+		BIO_set_retry_read(h);
+		return -1;
+	}
+
+	res = secure_raw_read(((Port *)h->ptr), buf, size);
+	BIO_clear_retry_flags(h);
+	if (res <= 0)
+	{
+		/* If we were interrupted, tell caller to retry */
+		if (errno == EINTR || errno == EWOULDBLOCK || errno == EAGAIN)
 		{
-			/* If we were interrupted, tell caller to retry */
-			if (errno == EINTR || errno == EWOULDBLOCK || errno == EAGAIN)
-			{
-				BIO_set_retry_read(h);
-			}
+			BIO_set_retry_read(h);
 		}
 	}
 
diff --git a/src/backend/libpq/be-secure.c b/src/backend/libpq/be-secure.c
index 4e7acbe080..e7bc52bb35 100644
--- a/src/backend/libpq/be-secure.c
+++ b/src/backend/libpq/be-secure.c
@@ -245,6 +245,9 @@ retry:
 			 * for the socket to become ready again.
 			 */
 		}
+		if (waitfor == WL_SOCKET_READABLE)
+			pq_recvbuf(false);
+
 		goto retry;
 	}
 
diff --git a/src/backend/libpq/pqcomm.c b/src/backend/libpq/pqcomm.c
index 34efac4865..bedc19f6da 100644
--- a/src/backend/libpq/pqcomm.c
+++ b/src/backend/libpq/pqcomm.c
@@ -843,11 +843,11 @@ socket_set_nonblocking(bool nonblocking)
 /* --------------------------------
  *		pq_recvbuf - load some bytes into the input buffer
  *
- *		returns 0 if OK, EOF if trouble
+ *		returns 0 if OK, EOF if	trouble
  * --------------------------------
  */
-static int
-pq_recvbuf(void)
+int
+pq_recvbuf(bool block)
 {
 	if (PqRecvPointer > 0)
 	{
@@ -863,8 +863,8 @@ pq_recvbuf(void)
 			PqRecvLength = PqRecvPointer = 0;
 	}
 
-	/* Ensure that we're in blocking mode */
-	socket_set_nonblocking(false);
+	/* Ensure that we're in blocking mode (or not) */
+	socket_set_nonblocking(!block);
 
 	/* Can fill buffer from PqRecvLength and upwards */
 	for (;;)
@@ -879,6 +879,12 @@ pq_recvbuf(void)
 			if (errno == EINTR)
 				continue;		/* Ok if interrupted */
 
+			if (!block)
+			{
+				if (errno == EWOULDBLOCK || errno == EAGAIN)
+					return 0;
+			}
+
 			/*
 			 * Careful: an ereport() that tries to write to the client would
 			 * cause recursion to here, leading to stack overflow and core
@@ -914,7 +920,7 @@ pq_getbyte(void)
 
 	while (PqRecvPointer >= PqRecvLength)
 	{
-		if (pq_recvbuf())		/* If nothing in buffer, then recv some */
+		if (pq_recvbuf(true))		/* If nothing in buffer, then recv some */
 			return EOF;			/* Failed to recv data */
 	}
 	return (unsigned char) PqRecvBuffer[PqRecvPointer++];
@@ -933,7 +939,7 @@ pq_peekbyte(void)
 
 	while (PqRecvPointer >= PqRecvLength)
 	{
-		if (pq_recvbuf())		/* If nothing in buffer, then recv some */
+		if (pq_recvbuf(true))	/* If nothing in buffer, then recv some */
 			return EOF;			/* Failed to recv data */
 	}
 	return (unsigned char) PqRecvBuffer[PqRecvPointer];
@@ -1012,7 +1018,7 @@ pq_getbytes(char *s, size_t len)
 	{
 		while (PqRecvPointer >= PqRecvLength)
 		{
-			if (pq_recvbuf())	/* If nothing in buffer, then recv some */
+			if (pq_recvbuf(true))	/* If nothing in buffer, then recv some */
 				return EOF;		/* Failed to recv data */
 		}
 		amount = PqRecvLength - PqRecvPointer;
@@ -1046,7 +1052,7 @@ pq_discardbytes(size_t len)
 	{
 		while (PqRecvPointer >= PqRecvLength)
 		{
-			if (pq_recvbuf())	/* If nothing in buffer, then recv some */
+			if (pq_recvbuf(true))	/* If nothing in buffer, then recv some */
 				return EOF;		/* Failed to recv data */
 		}
 		amount = PqRecvLength - PqRecvPointer;
@@ -1087,7 +1093,7 @@ pq_getstring(StringInfo s)
 	{
 		while (PqRecvPointer >= PqRecvLength)
 		{
-			if (pq_recvbuf())	/* If nothing in buffer, then recv some */
+			if (pq_recvbuf(true))	/* If nothing in buffer, then recv some */
 				return EOF;		/* Failed to recv data */
 		}
 
diff --git a/src/include/libpq/libpq.h b/src/include/libpq/libpq.h
index af4ba2ab07..152e30e409 100644
--- a/src/include/libpq/libpq.h
+++ b/src/include/libpq/libpq.h
@@ -1,3 +1,4 @@
+
 /*-------------------------------------------------------------------------
  *
  * libpq.h
@@ -86,6 +87,7 @@ extern int	pq_getbyte(void);
 extern int	pq_peekbyte(void);
 extern int	pq_getbyte_if_available(unsigned char *c);
 extern int	pq_putbytes(const char *s, size_t len);
+extern int	pq_recvbuf(bool block);
 
 /*
  * prototypes for functions in be-secure.c