Add a loop around getting the packet size

Believe it or not, we can get short reads here.

1 files changed, 11 insertions, 10 deletions

diff --git a/sockutil.c b/sockutil.c
index 73ab304..05aa0bc 100644
--- a/sockutil.c
+++ b/sockutil.c
@@ -106,26 +106,27 @@ recv_packet(int fd)
 {
 	char buf[QARSH_MAX_PACKET_SIZE];
 	uint32_t packetsize;
+	char *psbuf = &packetsize;
 	int bufused = 0;
 	int ret = 0;
 
-	if ((ret = read(fd, &packetsize, sizeof packetsize)) < 0) {
-		if (errno != ECONNRESET) {
-		       	fprintf(stderr, "Read error while reading packet size: %s\n", strerror(errno));
+	do {
+		if ((ret = read(fd, psbuf+bufused, sizeof packetsize - bufused)) < 0) {
+			fprintf(stderr, "Read error while reading packet size: %s\n", strerror(errno));
+			return NULL;
+		} else if (ret == 0) {
+			return NULL;
 		}
-		return NULL;
-	} else if (ret == 0) {
-		return NULL;
-	} else if (ret != sizeof packetsize) {
-		fprintf(stderr, "Did not read packet size");
-		return NULL;
-	}
+		bufused += ret;
+	} while (bufused < sizeof packetsize);
+
 	packetsize = ntohl(packetsize);
 	if (packetsize > QARSH_MAX_PACKET_SIZE) {
 		fprintf(stderr, "Packet size too large, %d > %d\n", packetsize, QARSH_MAX_PACKET_SIZE);
 		return NULL;
 	}
 	/* Keep reading until we get the whole packet and nothing but the packet, so help me socket */
+	bufused = 0;
 	do {
 		if ((ret = read(fd, buf+bufused, packetsize-bufused)) < 0) {
 			fprintf(stderr, "Read error while reading packet data: %s\n", strerror(errno));