ssl-write-retry.patch

diff --git a/src/interfaces/libpq/fe-misc.c b/src/interfaces/libpq/fe-misc.c
index 17dde4a..b5e62c9 100644
--- a/src/interfaces/libpq/fe-misc.c
+++ b/src/interfaces/libpq/fe-misc.c
@@ -781,6 +781,8 @@ pqSendSome(PGconn *conn, int len)
 	char	   *ptr = conn->outBuffer;
 	int			remaining = conn->outCount;
 	int			result = 0;
+	int			sent = 0;
+	bool		usingSSLWriteBuffer = false;
 
 	if (conn->sock < 0)
 	{
@@ -789,10 +791,28 @@ pqSendSome(PGconn *conn, int len)
 		return -1;
 	}
 
+#ifdef USE_SSL
+	if (conn->sslRetryPos)
+	{
+		/* 
+		 * We have some leftovers from previous SSL write attempts, send those
+		 * instead of the regular output buffer.
+		 */
+		len = conn->sslRetryBytes;
+		ptr = conn->sslRetryPos;
+		remaining = conn->sslRetryBufSize;
+
+		if (conn->Pfdebug)
+			fprintf(conn->Pfdebug, "pqSendSome: using SSL write retry buffer: %p retry: %d total: %d\n",
+				ptr, len, remaining);
+
+		usingSSLWriteBuffer = true;
+	}
+#endif
+
 	/* while there's still data to send */
 	while (len > 0)
 	{
-		int			sent;
 		char		sebuf[256];
 
 #ifndef WIN32
@@ -807,6 +827,10 @@ pqSendSome(PGconn *conn, int len)
 		sent = pqsecure_write(conn, ptr, Min(len, 65536));
 #endif
 
+		if (conn->Pfdebug)
+			fprintf(conn->Pfdebug, "pqSendSome: write buf: %p len: %d sent: %d\n",
+					ptr, len, sent);
+
 		if (sent < 0)
 		{
 			/*
@@ -857,6 +881,35 @@ pqSendSome(PGconn *conn, int len)
 					return -1;
 			}
 		}
+#ifdef USE_SSL
+		else if (sent == 0 && conn->ssl)
+		{
+			/*
+			 * With non-blocking SSL connections we need to ensure that the
+			 * buffer passed to the pqsecure_write() retry is the the exact
+			 * same buffer as in previous write -- see SSL_write(3SSL) for more
+			 * on this. For this we need to keep the the original buffer and
+			 * remember its length. Also a new outBuffer is needed, but the
+			 * actual allocation is deferred to the end of the function.
+			 *
+			 * For non-SSL connections none of this is needed.
+			 */
+			if (!conn->sslRetryPos)
+			{
+				if (conn->Pfdebug)
+					fprintf(conn->Pfdebug, "pqSendSome: SSL retry setup: buf: %p len: %d remain: %d\n",
+						conn->sslRetryBuf, len, remaining);
+
+				conn->sslRetryBuf = conn->outBuffer;
+				conn->sslRetryPos = ptr;
+				conn->sslRetryBytes = len;
+				conn->sslRetryBufSize = remaining;
+				conn->outBuffer = NULL;
+			} 
+
+			usingSSLWriteBuffer = true;
+		}
+#endif
 		else
 		{
 			ptr += sent;
@@ -903,10 +956,73 @@ pqSendSome(PGconn *conn, int len)
 		}
 	}
 
-	/* shift the remaining contents of the buffer */
-	if (remaining > 0)
-		memmove(conn->outBuffer, ptr, remaining);
-	conn->outCount = remaining;
+#ifdef USE_SSL
+	if (conn->sslRetryPos && sent > 0)
+	{
+		/*
+		 * A SSL write was successfully retried, advance the retry buffer
+		 * position to the rest of the buffer. Free the buffer if all of its
+		 * contents are delivered.
+		 */
+		conn->sslRetryPos = ptr;
+		conn->sslRetryBytes = remaining;
+		conn->sslRetryBufSize = remaining;
+
+		if (conn->sslRetryBufSize <= 0)
+		{
+			if (conn->Pfdebug)
+				fprintf(conn->Pfdebug, "SSL retry clean-up: freed buf=%p\n",
+					conn->sslRetryBuf);
+
+			free(conn->sslRetryBuf);
+			conn->sslRetryBytes = conn->sslRetryBufSize = 0;
+			conn->sslRetryBuf = conn->sslRetryPos = NULL;
+		}
+
+		if (conn->sslRetryBytes || (conn->outCount - remaining) > 0)
+		{
+			/*
+			 * We still have some data left in the SSL retry buffers or the 
+			 * outBuffer. Return 1 to indicate that further retries are needed
+			 * to flush the entire output buffer.
+			 */
+			result = 1;
+		}
+	}
+
+	if (!conn->outBuffer)
+	{
+		/*
+		 * For non-blocking sockets the outBuffer is set to NULL whenever a SSL
+		 * write needs to be retried. Just allocate a new buffer and move on.
+		 */
+		conn->outBufSize = 16 * 1024;
+		conn->outCount = 0;
+		conn->outBuffer = malloc(conn->outBufSize);
+		if (!conn->outBuffer)
+		{
+			printfPQExpBuffer(&conn->errorMessage,
+							  libpq_gettext("cannot allocate send buffer\n"));
+			return -1;
+		}
+	}
+#endif
+
+	if (!usingSSLWriteBuffer)
+	{
+		/*
+		 * Shift the remaining contents of the buffer if were using
+		 * the primary outBuffer. The SSL write retry buffer should already
+		 * have been updated above.
+		 */
+		if (remaining > 0)
+			memmove(conn->outBuffer, ptr, remaining);
+		conn->outCount = remaining;
+	}
+
+	if (conn->Pfdebug)
+		fprintf(conn->Pfdebug, "pqSendSome: in the end sslRetryCount: %d outCount: %d\n",
+		conn->sslRetryBytes, conn->outCount);
 
 	return result;
 }
@@ -924,7 +1040,7 @@ pqFlush(PGconn *conn)
 	if (conn->Pfdebug)
 		fflush(conn->Pfdebug);
 
-	if (conn->outCount > 0)
+	if (conn->outCount > 0 || conn->sslRetryBufSize)
 		return pqSendSome(conn, conn->outCount);
 
 	return 0;
diff --git a/src/interfaces/libpq/libpq-int.h b/src/interfaces/libpq/libpq-int.h
index d56ef5d..18260a3 100644
--- a/src/interfaces/libpq/libpq-int.h
+++ b/src/interfaces/libpq/libpq-int.h
@@ -407,6 +407,11 @@ struct pg_conn
 	X509	   *peer;			/* X509 cert of server */
 	char		peer_dn[256 + 1];		/* peer distinguished name */
 	char		peer_cn[SM_USER + 1];	/* peer common name */
+
+	char	   *sslRetryBuf;	/* buffer for retrying SSL write operations */
+	char	   *sslRetryPos;	/* last buffer address passed SSL write */
+	int			sslRetryBytes;	/* last buffer length passed SSL write */
+	int			sslRetryBufSize;	/* total retry buffer size */
 #ifdef USE_SSL_ENGINE
 	ENGINE	   *engine;			/* SSL engine, if any */
 #else