From 35ace4f24ca7466b112b8b5ee93f5c2b79df5e01 Mon Sep 17 00:00:00 2001
From: antiduh <antiduh@e114e3f1-24f7-e311-a78f-001bfcf5e160>
Date: Tue, 8 Jul 2014 18:56:52 +0000
Subject: [PATCH] Couple bug fixes in the client and server code. Gotta
 remember that TCP sockets may return less data than you ask for; that and the
 length of the data may be longer than the buffer you give it.

---
 TestProtocol/CustomConnection.cs | 27 ++++++++++++++++++++++-----
 TestProtocol/CustomServer.cs     | 29 ++++++++++++++++++++++-------
 2 files changed, 44 insertions(+), 12 deletions(-)

diff --git a/TestProtocol/CustomConnection.cs b/TestProtocol/CustomConnection.cs
index 274b67f..1f26baf 100644
--- a/TestProtocol/CustomConnection.cs
+++ b/TestProtocol/CustomConnection.cs
@@ -112,7 +112,10 @@ namespace TestProtocol
             byte[] readBuffer = new byte[65536];
 
             ProtocolOp operation;
-            int length;
+            int messageLength;
+            int remaining;
+            int chunkLength;
+            int position;
 
             while( this.running )
             {
@@ -132,10 +135,24 @@ namespace TestProtocol
 
                     // Read the length 
                     this.socket.Receive( readBuffer, 4, SocketFlags.None );
-                    length = ByteWriter.ReadInt32_BE( readBuffer, 0 );
+                    messageLength = ByteWriter.ReadInt32_BE( readBuffer, 0 );
+
+                    if( readBuffer.Length < messageLength )
+                    {
+                        readBuffer = new byte[messageLength];
+                    }
 
                     // Read the data
-                    this.socket.Receive( readBuffer, length, SocketFlags.None );
+                    // Keep in mind that Socket.Receive may return less data than asked for.
+                    remaining = messageLength;
+                    chunkLength = 0;
+                    position = 0;
+                    while( remaining > 0 )
+                    {
+                        chunkLength = this.socket.Receive( readBuffer, position, remaining, SocketFlags.None );
+                        remaining -= chunkLength;
+                        position += chunkLength;
+                    }
 
                 }
                 catch( SocketException e )
@@ -159,8 +176,8 @@ namespace TestProtocol
 
                 if( this.Received != null )
                 {
-                    byte[] dataCopy = new byte[length];
-                    Array.Copy( readBuffer, 0, dataCopy, 0, length );
+                    byte[] dataCopy = new byte[messageLength];
+                    Array.Copy( readBuffer, 0, dataCopy, 0, messageLength );
                     Message message = new Message( operation, dataCopy );
                  
                     try
diff --git a/TestProtocol/CustomServer.cs b/TestProtocol/CustomServer.cs
index 103407d..60a6824 100644
--- a/TestProtocol/CustomServer.cs
+++ b/TestProtocol/CustomServer.cs
@@ -149,7 +149,9 @@ namespace TestProtocol
             byte[] readBuffer = new byte[65536];
 
             ProtocolOp operation;
-            int length;
+            int messageLength;
+            int position;
+            int remaining;
 
             while( this.running )
             {
@@ -157,7 +159,7 @@ namespace TestProtocol
                 {
                     //                          |--4 bytes--|--4 bytes--|---N--|
                     // Every command is a TLV - | Operation |  Length   | Data |
-
+                    int chunkLength;
 
                     // Read the operation.
                     this.readSocket.Receive( readBuffer, 4, SocketFlags.None );
@@ -169,11 +171,24 @@ namespace TestProtocol
 
                     // Read the length 
                     this.readSocket.Receive( readBuffer, 4, SocketFlags.None );
-                    length = ByteWriter.ReadInt32_BE( readBuffer, 0 );
+                    messageLength = ByteWriter.ReadInt32_BE( readBuffer, 0 );
+
+                    if( readBuffer.Length < messageLength )
+                    {
+                        readBuffer = new byte[messageLength];
+                    }
 
                     // Read the data
-                    this.readSocket.Receive( readBuffer, length, SocketFlags.None );
-
+                    // Keep in mind that Socket.Receive may return less data than asked for.
+                    remaining = messageLength;
+                    chunkLength = 0;
+                    position = 0;
+                    while( remaining > 0 )
+                    {
+                        chunkLength = this.readSocket.Receive( readBuffer, position, remaining, SocketFlags.None );
+                        remaining -= chunkLength;
+                        position += chunkLength;
+                    }
                 }
                 catch( SocketException e )
                 {
@@ -196,8 +211,8 @@ namespace TestProtocol
 
                 if( this.Received != null )
                 {
-                    byte[] dataCopy = new byte[length];
-                    Array.Copy( readBuffer, 0, dataCopy, 0, length );
+                    byte[] dataCopy = new byte[messageLength];
+                    Array.Copy( readBuffer, 0, dataCopy, 0, messageLength );
                     Message message = new Message( operation, dataCopy );
 
                     try