From f60b30160f074f2991af3fab636b0394d1e271c9 Mon Sep 17 00:00:00 2001
From: JamJar00 <jamster.30@btinternet.com>
Date: Wed, 30 Nov 2016 00:25:15 +0000
Subject: [PATCH] Added timeouts on Connect

---
 Hazel/Connection.cs              |  8 +++++---
 Hazel/Tcp/TcpConnection.cs       |  8 ++++++--
 Hazel/Udp/UdpClientConnection.cs | 13 +++++++++++--
 Hazel/Udp/UdpServerConnection.cs |  2 +-
 4 files changed, 23 insertions(+), 8 deletions(-)

diff --git a/Hazel/Connection.cs b/Hazel/Connection.cs
index 3f12dc0..45058c1 100644
--- a/Hazel/Connection.cs
+++ b/Hazel/Connection.cs
@@ -162,12 +162,13 @@ namespace Hazel
         ///     Connects the connection to a server and begins listening.
         /// </summary>
         /// <param name="bytes">The bytes of data to send in the handshake.</param>
+        /// <param name="timeout">The number of milliseconds to wait before giving up on the connect attempt.</param>
         /// <remarks>
         ///     Calling Connect makes the connection attempt to connect to the end point that's specified in the 
         ///     constructor. This method will block until the connection attempt completes and will throw a 
         ///     <see cref="HazelException"/> if there is a problem connecting.
         /// </remarks>
-        public abstract void Connect(byte[] bytes = null);
+        public abstract void Connect(byte[] bytes = null, int timeout = 5000);
 
         /// <summary>
         ///     Invokes the DataReceived event.
@@ -213,14 +214,15 @@ namespace Hazel
         /// <summary>
         ///     Blocks until the Connection is connected.
         /// </summary>
+        /// <param name="timeout">The number of milliseconds to wait before timing out.</param>
         /// <remarks>
         ///     This is a helper method for waiting until the connection is connected. It will block until the 
         ///     <see cref="State"/> property is set to <see cref="ConnectionState.Connected"/> allowing the main thread to 
         ///     wait until specific data is received etc. before returning to the user's code.
         /// </remarks>
-        protected void WaitOnConnect()
+        protected bool WaitOnConnect(int timeout)
         {
-            connectWaitLock.WaitOne();
+            return connectWaitLock.WaitOne(timeout);
         }
 
         /// <summary>
diff --git a/Hazel/Tcp/TcpConnection.cs b/Hazel/Tcp/TcpConnection.cs
index d0c2a8f..c641907 100644
--- a/Hazel/Tcp/TcpConnection.cs
+++ b/Hazel/Tcp/TcpConnection.cs
@@ -77,7 +77,7 @@ namespace Hazel.Tcp
         }
 
         /// <inheritdoc />
-        public override void Connect(byte[] bytes = null)
+        public override void Connect(byte[] bytes = null, int timeout = 5000)
         {
             lock(socketLock)
             {
@@ -86,7 +86,11 @@ namespace Hazel.Tcp
 
                 try
                 {
-                    socket.Connect(RemoteEndPoint);
+                    IAsyncResult result = socket.BeginConnect(RemoteEndPoint, null, null);
+
+                    result.AsyncWaitHandle.WaitOne(timeout);
+
+                    socket.EndConnect(result);
                 }
                 catch (Exception e)
                 {
diff --git a/Hazel/Udp/UdpClientConnection.cs b/Hazel/Udp/UdpClientConnection.cs
index afcc124..4f2e33f 100644
--- a/Hazel/Udp/UdpClientConnection.cs
+++ b/Hazel/Udp/UdpClientConnection.cs
@@ -102,7 +102,7 @@ namespace Hazel.Udp
         }
 
         /// <inheritdoc />
-        public override void Connect(byte[] bytes = null)
+        public override void Connect(byte[] bytes = null, int timeout = 5000)
         {
             lock(socketLock)
             {
@@ -121,6 +121,7 @@ namespace Hazel.Udp
                 }
                 catch (SocketException e)
                 {
+                    State = ConnectionState.NotConnected;
                     throw new HazelException("A socket exception occured while binding to the port.", e);
                 }
 
@@ -137,6 +138,7 @@ namespace Hazel.Udp
                 }
                 catch (SocketException e)
                 {
+                    Dispose();
                     throw new HazelException("A Socket exception occured while initiating a receive operation.", e);
                 }
             }
@@ -146,7 +148,14 @@ namespace Hazel.Udp
             SendHello(bytes, () => { lock (socketLock) State = ConnectionState.Connected; });
 
             //Wait till hello packet is acknowledged and the state is set to Connected
-            WaitOnConnect();
+            bool timedOut = !WaitOnConnect(timeout);
+
+            //If we timed out raise an exception
+            if (timedOut)
+            {
+                Dispose();
+                throw new HazelException("Connection attempt timed out.");
+            }
         }
 
         /// <summary>
diff --git a/Hazel/Udp/UdpServerConnection.cs b/Hazel/Udp/UdpServerConnection.cs
index 98a2047..ff19df4 100644
--- a/Hazel/Udp/UdpServerConnection.cs
+++ b/Hazel/Udp/UdpServerConnection.cs
@@ -60,7 +60,7 @@ namespace Hazel.Udp
         /// <remarks>
         ///     This will always throw a HazelException.
         /// </remarks>
-        public override void Connect(byte[] bytes)
+        public override void Connect(byte[] bytes = null, int timeout = 5000)
         {
             throw new HazelException("Cannot manually connect a UdpServerConnection, did you mean to use UdpClientConnection?");
         }
-- 
2.39.5