From a236409034c1089a5da811b9df1c19d22ee91318 Mon Sep 17 00:00:00 2001
From: "Tristan B. Kildaire" <deavmi@localhost.localdomain>
Date: Fri, 4 Dec 2020 18:55:25 +0200
Subject: [PATCH] Full wait doesn't work

---
 source/bmessage.d | 89 +++++++++++++++++++++++++++++++++++++++++++----
 1 file changed, 83 insertions(+), 6 deletions(-)

diff --git a/source/bmessage.d b/source/bmessage.d
index 718db2a..734fd0c 100644
--- a/source/bmessage.d
+++ b/source/bmessage.d
@@ -7,9 +7,36 @@ public bool receiveMessage(Socket originator, ref byte[] receiveMessage)
 	/* Construct a buffer to receive into */
 	byte[] receiveBuffer;
 
-	/* Read 4 bytes of length */
-	receiveBuffer.length = 4;
-	originator.receive(receiveBuffer, cast(SocketFlags)MSG_WAITALL);
+	/* The current byte */
+	uint currentByte = 0;
+
+	/* The amount of bytes received */
+	long bytesReceived;
+
+	/* Loop consume the next 4 bytes */
+	while(currentByte < 4)
+	{
+		/* Temporary buffer */
+		byte[4] tempBuffer;
+
+		/* Read at-most 4 bytes */
+		bytesReceived = originator.receive(tempBuffer);
+
+		/* If there was an error reading from the socket */
+		if(!(bytesReceived > 0))
+		{
+			return false;
+		}
+		/* If there is no error reading from the socket */
+		else
+		{
+			/* Add the read bytes to the *real* buffer */
+			receiveBuffer ~= tempBuffer[0..bytesReceived];
+
+			/* Increment the byte counter */
+			currentByte += bytesReceived;
+		}
+	}
 
 	/* Response message length */
 	int messageLength;
@@ -37,10 +64,60 @@ public bool receiveMessage(Socket originator, ref byte[] receiveMessage)
 	/* Reset buffer */
 	receiveBuffer.length = cast(uint)messageLength;
 
-	/* Read the full message */
-	originator.receive(receiveBuffer, cast(SocketFlags)MSG_WAITALL);
+	/* Reset the byte counter */
+	currentByte = 0;
 
-	// writeln("Message ", fullMessage);
+	while(currentByte < messageLength)
+	{
+		/**
+		 * Receive 20 bytes (at most) at a time and don't dequeue from
+		 * the kernel's TCP stack's buffer.
+		 */
+		byte[20] tempBuffer;
+		bytesReceived = originator.receive(tempBuffer, SocketFlags.PEEK);
+
+		/* Check for an error whilst receiving */
+		if(!(bytesReceived > 0))
+		{
+			return false;
+		}
+		else
+		{
+			if(cast(uint)bytesReceived+currentByte > messageLength)
+			{
+				byte[] remainingBytes;
+				remainingBytes.length = messageLength-currentByte;
+
+				/* Receive the remaining bytes */
+				originator.receive(remainingBytes);
+
+				/* Increment counter of received bytes */
+				currentByte += remainingBytes.length;
+
+				/* Append the received bytes to the FULL message buffer */
+				fullMessage ~= remainingBytes;
+
+				// writeln("Received ", currentByte, "/", cast(uint)messageLength, " bytes");
+			}
+			else
+			{
+				/* Increment counter of received bytes */
+				currentByte += bytesReceived;
+	
+				/* Append the received bytes to the FULL message buffer */
+				fullMessage ~= tempBuffer[0..bytesReceived];
+
+				/* TODO: Bug when over send, we must not allow this */
+				// writeln("Received ", currentByte, "/", cast(uint)messageLength, " bytes");	
+
+				/* Dequeue the received bytes */
+				originator.receive(tempBuffer);
+			}
+		}
+	}
+
+	/* Set the message in `receiveMessage */
+	receiveMessage = fullMessage;
 
 	return true;
 }