From 409f114d9ee8aa556059163b612cb8cb45c9aed3 Mon Sep 17 00:00:00 2001
From: Nikos Mavrogiannopoulos <nmav@gnutls.org>
Date: Sat, 18 Mar 2017 09:33:32 +0100
Subject: [PATCH] cstp_recv_nb: improve operation under receiving from UNIX
 socket

That is, ensure that all possible packet size combinations are
correctly received.

Signed-off-by: Nikos Mavrogiannopoulos <nmav@gnutls.org>
---
 src/tlslib.c | 86 ++++++++++++++++++++++++++++++++++------------------
 1 file changed, 57 insertions(+), 29 deletions(-)

diff --git a/src/tlslib.c b/src/tlslib.c
index 20d7f084..a5203767 100644
--- a/src/tlslib.c
+++ b/src/tlslib.c
@@ -190,6 +190,33 @@ ssize_t cstp_recv(worker_st *ws, void *data, size_t data_size)
 	return ret;
 }
 
+static
+int recv_remaining(int fd, uint8_t *p, int left)
+{
+	int counter = 100; /* allow 10 seconds for a full packet */
+	unsigned total = 0;
+	int ret;
+
+	while(left > 0) {
+		ret = recv(fd, p, left, 0);
+		if (ret == -1 && counter > 0 && (errno == EINTR || errno == EAGAIN)) {
+			counter--;
+			ms_sleep(100);
+			continue;
+		}
+		if (ret == 0)
+			ret = GNUTLS_E_PREMATURE_TERMINATION;
+		if (ret < 0)
+			break;
+
+		left -= ret;
+		p += ret;
+		total += ret;
+	}
+
+	return total;
+}
+
 ssize_t cstp_recv_nb(worker_st *ws, void *data, size_t data_size)
 {
 	int ret;
@@ -203,45 +230,46 @@ ssize_t cstp_recv_nb(worker_st *ws, void *data, size_t data_size)
 		 * incomplete CSTP packet. In that case we attempt to read
 		 * a full CSTP packet.
 		 */
-		int counter = 100; /* allow 10 seconds for a full packet */
 		unsigned pktlen;
 		unsigned total = 0;
 		int left = data_size;
 		uint8_t *p = data;
 
  		ret = recv(ws->conn_fd, p, left, 0);
-		if (ret > 6) {
+ 		if (ret <= 0)	
+ 			return ret;
+
+		if (ret < 8) {
+			/* ensure we have the CSTP header */
+			p += ret;
+			left = 8 - ret;
+
+			ret = recv_remaining(ws->conn_fd, p, left);
+			if (ret <= 0)
+				return ret;
+
+			total = 8;
+			p += 8;
+		} else {
 			total += ret;
-			/* get the actual length from headers */
-			pktlen = (p[4] << 8) + p[5];
-			if (pktlen+8 > data_size) {
-				oclog(ws, LOG_ERR, "error in CSTP packet length");
-				return GNUTLS_E_UNEXPECTED_PACKET_LENGTH;
-			}
-
-			p = p + ret;
-			left = pktlen+8 - ret;
-
-			while(left > 0) {
-				ret = recv(ws->conn_fd, p, left, 0);
-				if (ret == -1 && counter > 0 && (errno == EINTR || errno == EAGAIN)) {
-					counter--;
-					ms_sleep(100);
-					continue;
-				}
-				if (ret == 0)
-					ret = GNUTLS_E_PREMATURE_TERMINATION;
-				if (ret < 0)
-					break;
-
-				left -= ret;
-				p += ret;
-				total += ret;
-			}
+			p += ret;
 		}
 
-		ret = total;
+		/* get the actual length from headers */
+		pktlen = (p[4] << 8) + p[5];
+		if (pktlen+8 > data_size) {
+			oclog(ws, LOG_ERR, "error in CSTP packet length");
+			return GNUTLS_E_UNEXPECTED_PACKET_LENGTH;
+		}
 
+		left = pktlen+8 - total;
+
+		ret = recv_remaining(ws->conn_fd, p, left);
+		if (ret <= 0)
+			return ret;
+
+		total += ret;
+		ret = total;
 	}
 
 	return ret;