Implement Keep Alive packets for TCP

Hazel.UnitTests/TcpConnectionTests.cs

@@ -1,4 +1,5 @@
 using System;
+using System.Diagnostics;
 using Microsoft.VisualStudio.TestTools.UnitTesting;
 using System.Net;
 using System.Threading;
@@ -136,5 +137,66 @@ public void ServerDisconnectTest()
                 TestHelper.RunServerDisconnectTest(listener, connection);
             }
         }
+
+        /// <summary>
+        ///     Tests the keepalive functionality from the client,
+        /// </summary>
+        [TestMethod]
+        public void KeepAliveClientTest()
+        {
+            using (TcpConnectionListener listener = new TcpConnectionListener(new NetworkEndPoint(IPAddress.Any, 4296)))
+            using (TcpConnection connection = new TcpConnection(new NetworkEndPoint(IPAddress.Loopback, 4296)))
+            {
+                listener.Start();
+
+                connection.Connect();
+                connection.KeepAliveInterval = 100;
+
+                System.Threading.Thread.Sleep(1050);    //Enough time for ~10 keep alive packets
+
+                Assert.IsTrue(
+                    connection.Statistics.TotalBytesSent >= 30 &&
+                    connection.Statistics.TotalBytesSent <= 50,
+                    "Sent: " + connection.Statistics.TotalBytesSent
+                );
+            }
+        }
+
+        /// <summary>
+        ///     Tests the keepalive functionality from the client,
+        /// </summary>
+        [TestMethod]
+        public void KeepAliveServerTest()
+        {
+            ManualResetEvent mutex = new ManualResetEvent(false);
+            TcpConnection listenerConnectionToClient = null;
+
+            using (TcpConnectionListener listener = new TcpConnectionListener(new NetworkEndPoint(IPAddress.Any, 4296)))
+            using (TcpConnection connection = new TcpConnection(new NetworkEndPoint(IPAddress.Loopback, 4296)))
+            {                
+                listener.NewConnection += delegate (object sender, NewConnectionEventArgs args)
+                {
+                    listenerConnectionToClient = (TcpConnection)args.Connection;
+                    listenerConnectionToClient.KeepAliveInterval = 100;
+
+                    Thread.Sleep(1050);    //Enough time for ~10 keep alive packets
+
+                    mutex.Set();
+                };
+
+                listener.Start();
+
+                connection.Connect();
+
+                mutex.WaitOne();
+
+                Assert.IsNotNull(listenerConnectionToClient);
+                Assert.IsTrue(
+                    listenerConnectionToClient.Statistics.TotalBytesSent >= 30 &&
+                    listenerConnectionToClient.Statistics.TotalBytesSent <= 50,
+                    "Sent: " + listenerConnectionToClient.Statistics.TotalBytesSent
+                );
+            }
+        }
     }
 }

Hazel.UnitTests/UdpConnectionTests.cs

@@ -194,22 +194,18 @@ public void KeepAliveClientTest()
         public void KeepAliveServerTest()
         {
             ManualResetEvent mutex = new ManualResetEvent(false);
+            UdpConnection listenerConnectionToClient = null;
 
             using (UdpConnectionListener listener = new UdpConnectionListener(new NetworkEndPoint(IPAddress.Any, 4296)))
             using (UdpConnection connection = new UdpClientConnection(new NetworkEndPoint(IPAddress.Loopback, 4296)))
             {
                 listener.NewConnection += delegate(object sender, NewConnectionEventArgs args)
                 {
-                    ((UdpConnection)args.Connection).KeepAliveInterval = 100;
+                    listenerConnectionToClient = (UdpConnection) args.Connection;
+                    listenerConnectionToClient.KeepAliveInterval = 100;
 
                     Thread.Sleep(1050);    //Enough time for ~10 keep alive packets
 
-                    Assert.IsTrue(
-                        args.Connection.Statistics.TotalBytesSent >= 30 &&
-                        args.Connection.Statistics.TotalBytesSent <= 50,
-                        "Sent: " + args.Connection.Statistics.TotalBytesSent
-                    );
-
                     mutex.Set();
                 };
 
@@ -218,6 +214,13 @@ public void KeepAliveServerTest()
                 connection.Connect();
 
                 mutex.WaitOne();
+
+                Assert.IsNotNull(listenerConnectionToClient);
+                Assert.IsTrue(
+                    listenerConnectionToClient.Statistics.TotalBytesSent >= 30 &&
+                    listenerConnectionToClient.Statistics.TotalBytesSent <= 50,
+                    "Sent: " + listenerConnectionToClient.Statistics.TotalBytesSent
+                );
             }
         }
 

Hazel/Udp/UdpConnection.KeepAlive.cs → Hazel/Connection.KeepAlive.cs

@@ -3,12 +3,13 @@
 using System.Diagnostics;
 using System.Linq;
 using System.Text;
+using System.Net.Sockets;
+using System.Net;
 using System.Threading;
 
-
-namespace Hazel.Udp
+namespace Hazel
 {
-    partial class UdpConnection
+    partial class Connection
     {
         /// <summary>
         ///     The interval from data being received or transmitted to a keepalive packet being sent in milliseconds.
@@ -33,7 +34,7 @@ public int KeepAliveInterval
             set
             {
                 keepAliveInterval = value;
-                
+
                 //Update timer
                 ResetKeepAliveTimer();
             }
@@ -55,18 +56,20 @@ public int KeepAliveInterval
         /// </summary>
         bool keepAliveTimerDisposed;
 
+        readonly byte[] keepAlivePacket = new byte[1];
+
         /// <summary>
         ///     Starts the keepalive timer.
         /// </summary>
-        void InitializeKeepAliveTimer()
+        protected void InitializeKeepAliveTimer()
         {
             lock (keepAliveTimerLock)
             {
                 keepAliveTimer = new Timer(
                     (o) =>
                     {
                         Trace.WriteLine("Keepalive packet sent.");
-                        SendHello(null, null);
+                        SendBytes(keepAlivePacket, SendOption.KeepAlive);
                     },
                     null,
                     keepAliveInterval,
@@ -78,20 +81,23 @@ void InitializeKeepAliveTimer()
         /// <summary>
         ///     Resets the keepalive timer to zero.
         /// </summary>
-        void ResetKeepAliveTimer()
+        protected void ResetKeepAliveTimer()
         {
             lock (keepAliveTimerLock)
-                keepAliveTimer.Change(keepAliveInterval, keepAliveInterval);
+            {                
+                if(keepAliveTimer != null)
+                    keepAliveTimer.Change(keepAliveInterval, keepAliveInterval);
+            }
         }
 
         /// <summary>
         ///     Disposes of the keep alive timer.
         /// </summary>
-        void DisposeKeepAliveTimer()
+        protected void DisposeKeepAliveTimer()
         {
-            lock(keepAliveTimerLock)
+            lock (keepAliveTimerLock)
             {
-                if (!keepAliveTimerDisposed)
+                if (!keepAliveTimerDisposed && keepAliveTimer != null)
                     keepAliveTimer.Dispose();
                 keepAliveTimerDisposed = true;
             }

Hazel/Connection.cs

@@ -32,7 +32,7 @@ namespace Hazel
     ///     </para>
     /// </remarks>
     /// <threadsafety static="true" instance="true"/>
-    public abstract class Connection : IDisposable
+    public abstract partial class Connection : IDisposable
     {
         /// <summary>
         ///     Called when a message has been received.
@@ -261,6 +261,7 @@ protected virtual void Dispose(bool disposing)
         {
             if (disposing)
             {
+                DisposeKeepAliveTimer();
             }
         }
     }

Hazel/Hazel.csproj

@@ -49,6 +49,7 @@
     <Reference Include="System.Xml" />
   </ItemGroup>
   <ItemGroup>
+    <Compile Include="Connection.KeepAlive.cs" />
     <Compile Include="Connection.cs" />
     <Compile Include="ConnectionEndPoint.cs" />
     <Compile Include="ConnectionListener.cs" />
@@ -75,7 +76,6 @@
       <SubType>Code</SubType>
     </Compile>
     <Compile Include="Udp\UdpConnection.Fragmented.cs" />
-    <Compile Include="Udp\UdpConnection.KeepAlive.cs" />
     <Compile Include="Udp\UdpConnection.Reliable.cs" />
     <Compile Include="Udp\UdpConnectionListener.cs" />
     <Compile Include="Udp\UdpServerConnection.cs" />

Hazel/SendOption.cs

@@ -42,6 +42,15 @@ public enum SendOption : byte
         ///     guaranteed to arrive and to arrive only once but the sending process may require more memory, processing,
         ///     a larger number protocol bytes and may be slower than sending unreliably.
         /// </remarks>
-        FragmentedReliable = 2
+        FragmentedReliable = 2,
+
+        /// <summary>
+        ///     Requests data be treated as keep alive packet.
+        /// </summary>
+        /// <remarks>
+        ///     This packet type is used soly to keep connection working.
+        ///     Data inside keep alive packet will be ignored.
+        /// </remarks>
+        KeepAlive = 3
     }
 }

Hazel/Tcp/TcpConnection.cs

@@ -42,6 +42,8 @@ internal TcpConnection(Socket socket)
                 this.socket.NoDelay = true;
 
                 State = ConnectionState.Connected;
+
+                InitializeKeepAliveTimer();
             }
         }
 
@@ -69,7 +71,7 @@ public TcpConnection(NetworkEndPoint remoteEndPoint)
                         throw new HazelException("IPV6 not supported!");
 
                     socket = new Socket(AddressFamily.InterNetworkV6, SocketType.Stream, System.Net.Sockets.ProtocolType.Tcp);
-                    socket.SetSocketOption(SocketOptionLevel.IPv6, (SocketOptionName)27, false);
+                    socket.SetSocketOption(SocketOptionLevel.IPv6, (SocketOptionName)27, false);                    
                 }
 
                 socket.NoDelay = true;
@@ -123,28 +125,33 @@ public override void Connect(byte[] bytes = null, int timeout = 5000)
                 State = ConnectionState.Connected;
 
                 SendBytes(actualBytes);
+
+                InitializeKeepAliveTimer();
             }
         }
 
         /// <inheritdoc/>
         /// <remarks>
         ///     <include file="DocInclude/common.xml" path="docs/item[@name='Connection_SendBytes_General']/*" />
         ///     <para>
-        ///         The sendOption parameter is ignored by the TcpConnection as TCP only supports FragmentedReliable 
-        ///         communication, specifying anything else will have no effect.
+        ///         The sendOption parameter can only be set as FragmentedReliable or KeepAlive.
+        ///         TCP only supports FragmentedReliable communication, specifying anything else will have no effect.
+        ///         KeepAlive packets are still send as FragmentedReliable.
         ///     </para>
         /// </remarks>
         public override void SendBytes(byte[] bytes, SendOption sendOption = SendOption.FragmentedReliable)
         {
             //Get bytes for length
-            byte[] fullBytes = AppendLengthHeader(bytes);
+            byte[] fullBytes = AppendLengthHeader(bytes, sendOption);
 
             //Write the bytes to the socket
             lock (socketLock)
             {
                 if (State != ConnectionState.Connected)
                     throw new InvalidOperationException("Could not send data as this Connection is not connected. Did you disconnect?");
 
+                ResetKeepAliveTimer();
+
                 try
                 {
                     socket.BeginSend(fullBytes, 0, fullBytes.Length, SocketFlags.None, null, null);
@@ -170,6 +177,22 @@ void HeaderReadCallback(byte[] bytes, Action<byte[]> callback)
             //Get length 
             int length = GetLengthFromBytes(bytes);
 
+            //Check if there is no body -> [KeepAlive] packet, if so then log it and wait for next message
+            if (length == 0)
+            {
+                try
+                {
+                    StartWaitingForHeader(callback);
+                }
+                catch (Exception e)
+                {
+                    HandleDisconnect(new HazelException("An exception occured while initiating a header receive operation.", e));
+                }
+
+                Statistics.LogHelloReceive(bytes.Length + 4);
+                return;
+            }
+
             //Begin receiving the body
             try
             {
@@ -303,6 +326,8 @@ void ChunkReadCallback(IAsyncResult result)
                 return;
             }
 
+            ResetKeepAliveTimer();
+
             StateObject state = (StateObject)result.AsyncState;
 
             state.totalBytesReceived += bytesReceived;      //TODO threading issues on state?
@@ -362,16 +387,28 @@ void HandleDisconnect(HazelException e = null)
         ///     Appends the length header to the bytes.
         /// </summary>
         /// <param name="bytes">The source bytes.</param>
+        /// <param name="sendOption">Allows for marking packet as KeepAlive. (Zero body length)</param>
         /// <returns>The new bytes.</returns>
-        static byte[] AppendLengthHeader(byte[] bytes)
+        static byte[] AppendLengthHeader(byte[] bytes, SendOption sendOption)
         {
             byte[] fullBytes = new byte[bytes.Length + 4];
 
-            //Append length
-            fullBytes[0] = (byte)(((uint)bytes.Length >> 24) & 0xFF);
-            fullBytes[1] = (byte)(((uint)bytes.Length >> 16) & 0xFF);
-            fullBytes[2] = (byte)(((uint)bytes.Length >> 8) & 0xFF);
-            fullBytes[3] = (byte)(uint)bytes.Length;
+            if (sendOption == SendOption.KeepAlive)
+            {
+                //Append length as 0
+                fullBytes[0] = 0;
+                fullBytes[1] = 0;
+                fullBytes[2] = 0;
+                fullBytes[3] = 0;
+            }
+            else
+            {
+                //Append length
+                fullBytes[0] = (byte) (((uint) bytes.Length >> 24) & 0xFF);
+                fullBytes[1] = (byte) (((uint) bytes.Length >> 16) & 0xFF);
+                fullBytes[2] = (byte) (((uint) bytes.Length >> 8) & 0xFF);
+                fullBytes[3] = (byte) (uint) bytes.Length;                
+            }
 
             //Add rest of bytes
             Buffer.BlockCopy(bytes, 0, fullBytes, 4, bytes.Length);

Hazel/Udp/UdpClientConnection.cs

@@ -160,6 +160,8 @@ public override void Connect(byte[] bytes = null, int timeout = 5000)
                 Dispose();
                 throw new HazelException("Connection attempt timed out.");
             }
+
+            InitializeKeepAliveTimer();
         }
 
         /// <summary>

Hazel/Udp/UdpConnection.Fragmented.cs

@@ -10,7 +10,7 @@ partial class UdpConnection
         /// <summary>
         ///     The amount of data that can be put into a fragment.
         /// </summary>
-        public int FragmentSize { get { return _fragmentSize; } }
+        public int FragmentSize { get { return fragmentSize; } }
         int fragmentSize = 65507 - 1 - 2 - 2 - 2;
 
         /// <summary>