From 0ea4b820022ca06d604c8773b8d56dd5cc25be2d Mon Sep 17 00:00:00 2001
From: JamJar00 <jamster.30@btinternet.com>
Date: Tue, 27 Sep 2016 01:02:05 +0100
Subject: [PATCH] Fixed bugs in TCP exception handling and improved some
 exception messages.

Previously most exceptions from Socket.BeginReceive in TCP were unhandled when they should have marked a disconnection, this is fixed and now more obvious.

Also TcpConnection.StartListening is no longer needed, simply bootstrap it using StartWaitingForHeader.
---
 Hazel/Tcp/TcpConnection.cs         | 44 ++++++++++++++----------------
 Hazel/Tcp/TcpConnectionListener.cs | 10 ++++++-
 2 files changed, 30 insertions(+), 24 deletions(-)

diff --git a/Hazel/Tcp/TcpConnection.cs b/Hazel/Tcp/TcpConnection.cs
index 2507870..b2fa616 100644
--- a/Hazel/Tcp/TcpConnection.cs
+++ b/Hazel/Tcp/TcpConnection.cs
@@ -76,22 +76,6 @@ namespace Hazel.Tcp
             }
         }
 
-        /// <summary>
-        ///     Internal call to start listening once this socket has been constructed and is ready.
-        /// </summary>
-        internal void StartListening()
-        {
-            //Start receiving data
-            try
-            {
-                StartWaitingForHeader();
-            }
-            catch (SocketException e)
-            {
-                throw new HazelException("A Socket exception occured while initiating a receive operation.", e);
-            }
-        }
-
         /// <inheritdoc />
         public override void Connect()
         {
@@ -110,10 +94,17 @@ namespace Hazel.Tcp
                 }
 
                 //Start receiving data
-                StartListening();
+                try
+                {
+                    StartListening();
+                }
+                catch (SocketException e)
+                {
+                    throw new HazelException("A Socket exception occured while initiating the first receive operation.", e);
+                }
 
-                //Set connected
-                State = ConnectionState.Connected;
+            //Set connected
+            State = ConnectionState.Connected;
             }
         }
 
@@ -167,7 +158,7 @@ namespace Hazel.Tcp
             }
             catch (SocketException e)
             {
-                HandleDisconnect(new HazelException("A Socket exception occured while initiating a receive operation.", e));
+                HandleDisconnect(new HazelException("A Socket exception occured while initiating a body receive operation.", e));
             }
         }
 
@@ -178,7 +169,14 @@ namespace Hazel.Tcp
         void BodyReadCallback(byte[] bytes)
         {
             //Begin receiving from the start
-            StartWaitingForHeader();
+            try
+            {
+                StartWaitingForHeader();
+            }
+            catch (SocketException e)
+            {
+                HandleDisconnect(new HazelException("A Socket exception occured while initiating a body receive operation.", e));
+            }
 
             Statistics.LogReceive(bytes.Length, bytes.Length + 4);
 
@@ -189,7 +187,7 @@ namespace Hazel.Tcp
         /// <summary>
         ///     Starts this connections waiting for the header.
         /// </summary>
-        void StartWaitingForHeader()
+        internal void StartWaitingForHeader()
         {
             StartWaitingForBytes(4, HeaderReadCallback);
         }
@@ -267,7 +265,7 @@ namespace Hazel.Tcp
                 }
                 catch (SocketException e)
                 {
-                    HandleDisconnect(new HazelException("A Socket exception occured while initiating a receive operation.", e));
+                    HandleDisconnect(new HazelException("A Socket exception occured while initiating a chunk receive operation.", e));
                     return;
                 }
             }
diff --git a/Hazel/Tcp/TcpConnectionListener.cs b/Hazel/Tcp/TcpConnectionListener.cs
index 78b7fa0..9e174b9 100644
--- a/Hazel/Tcp/TcpConnectionListener.cs
+++ b/Hazel/Tcp/TcpConnectionListener.cs
@@ -101,7 +101,15 @@ namespace Hazel.Tcp
                 //Invoke
                 InvokeNewConnection(tcpConnection);
 
-                tcpConnection.StartListening();
+                try
+                {
+                    tcpConnection.StartWaitingForHeader();
+                }
+                catch (SocketException)
+                {
+                    //Receive operation couldn't be started, there's nothing we can do to save it
+                    return;
+                }
             }
         }
 
-- 
2.39.5