From a8b72ec7a989aee3fd360c71b9f63b6111709452 Mon Sep 17 00:00:00 2001
From: Forest <forest@innersloth.com>
Date: Wed, 29 Apr 2020 15:46:21 -0700
Subject: [PATCH] Fix a severe issue with redundant acks Fix a potential issue
 in ObjectPool on Unity IL2CPP Improve naming of thread limited classes.

---
 ... => ThreadLimitedUdpConnectionListener.cs} | 29 +++++++------
 ...cs => ThreadLimitedUdpServerConnection.cs} |  8 ++--
 Hazel/Hazel.csproj                            |  8 ++--
 Hazel/ObjectPool.cs                           | 31 +++++++++++++-
 Hazel/Udp/UdpConnection.KeepAlive.cs          | 42 ++++++++++---------
 Hazel/Udp/UdpConnection.Reliable.cs           | 39 +++++++++--------
 Hazel/Udp/UdpConnectionListener.cs            |  6 +--
 Hazel/Udp/UnityUdpClientConnection.cs         |  9 +++-
 8 files changed, 107 insertions(+), 65 deletions(-)
 rename Hazel/FewerThreads/{UdpConnectionListener2.cs => ThreadLimitedUdpConnectionListener.cs} (89%)
 rename Hazel/FewerThreads/{UdpServerConnection2.cs => ThreadLimitedUdpServerConnection.cs} (89%)

diff --git a/Hazel/FewerThreads/UdpConnectionListener2.cs b/Hazel/FewerThreads/ThreadLimitedUdpConnectionListener.cs
similarity index 89%
rename from Hazel/FewerThreads/UdpConnectionListener2.cs
rename to Hazel/FewerThreads/ThreadLimitedUdpConnectionListener.cs
index d4e8876..df5e000 100644
--- a/Hazel/FewerThreads/UdpConnectionListener2.cs
+++ b/Hazel/FewerThreads/ThreadLimitedUdpConnectionListener.cs
@@ -11,7 +11,7 @@ namespace Hazel.Udp.FewerThreads
     ///     Listens for new UDP connections and creates UdpConnections for them.
     /// </summary>
     /// <inheritdoc />
-    public class UdpConnectionListener2 : IDisposable
+    public class ThreadLimitedUdpConnectionListener : IDisposable
     {
         private struct SendMessageInfo
         {
@@ -49,7 +49,7 @@ namespace Hazel.Udp.FewerThreads
         private Thread sendThread;
         private HazelThreadPool processThreads;
 
-        private ConcurrentDictionary<EndPoint, UdpServerConnection2> allConnections = new ConcurrentDictionary<EndPoint, UdpServerConnection2>();
+        private ConcurrentDictionary<EndPoint, ThreadLimitedUdpServerConnection> allConnections = new ConcurrentDictionary<EndPoint, ThreadLimitedUdpServerConnection>();
 
         private Queue<ReceiveMessageInfo> receiveQueue = new Queue<ReceiveMessageInfo>();
         private Queue<SendMessageInfo> sendQueue = new Queue<SendMessageInfo>();
@@ -60,7 +60,7 @@ namespace Hazel.Udp.FewerThreads
 
         private bool isActive;
 
-        public UdpConnectionListener2(IPEndPoint endPoint, ILogger logger, IPMode ipMode = IPMode.IPv4)
+        public ThreadLimitedUdpConnectionListener(int numWorkers, IPEndPoint endPoint, ILogger logger, IPMode ipMode = IPMode.IPv4)
         {
             this.Logger = logger;
             this.EndPoint = endPoint;
@@ -75,10 +75,10 @@ namespace Hazel.Udp.FewerThreads
             this.reliablePacketThread = new Thread(ManageReliablePackets);
             this.sendThread = new Thread(SendLoop);
             this.receiveThread = new Thread(ReceiveLoop);
-            this.processThreads = new HazelThreadPool(4, ProcessingLoop);
+            this.processThreads = new HazelThreadPool(numWorkers, ProcessingLoop);
         }
 
-        ~UdpConnectionListener2()
+        ~ThreadLimitedUdpConnectionListener()
         {
             this.Dispose(false);
         }
@@ -121,7 +121,7 @@ namespace Hazel.Udp.FewerThreads
             {
                 if (this.socket.Poll(Timeout.Infinite, SelectMode.SelectRead))
                 {
-                    EndPoint remoteEP = new IPEndPoint(IPMode == IPMode.IPv4 ? IPAddress.Any : IPAddress.IPv6Any, this.EndPoint.Port);
+                    EndPoint remoteEP = new IPEndPoint(this.EndPoint.Address, this.EndPoint.Port);
                     MessageReader message = MessageReader.GetSized(BufferSize);
                     try
                     {
@@ -215,7 +215,7 @@ namespace Hazel.Udp.FewerThreads
 
             // If we're aware of this connection use the one already
             // If this is a new client then connect with them!
-            UdpServerConnection2 connection;
+            ThreadLimitedUdpServerConnection connection;
             if (!this.allConnections.TryGetValue(remoteEndPoint, out connection))
             {
                 lock (this.allConnections)
@@ -244,7 +244,7 @@ namespace Hazel.Udp.FewerThreads
                         }
 
                         aware = false;
-                        connection = new UdpServerConnection2(this, (IPEndPoint)remoteEndPoint, this.IPMode);
+                        connection = new ThreadLimitedUdpServerConnection(this, (IPEndPoint)remoteEndPoint, this.IPMode);
                         if (!this.allConnections.TryAdd(remoteEndPoint, connection))
                         {
                             throw new HazelException("Failed to add a connection. This should never happen.");
@@ -253,10 +253,9 @@ namespace Hazel.Udp.FewerThreads
                 }
             }
 
-            //Inform the connection of the buffer (new connections need to send an ack back to client)
-            connection.HandleReceive(message, bytesReceived);
-            
-            //If it's a new connection invoke the NewConnection event.
+            // If it's a new connection invoke the NewConnection event.
+            // This needs to happen before handling the message because in localhost scenarios, the ACK and
+            // subsequent messages can happen before the NewConnection event sets up OnDataRecieved handlers
             if (!aware)
             {
                 // Skip header and hello byte;
@@ -265,7 +264,11 @@ namespace Hazel.Udp.FewerThreads
                 message.Position = 0;
                 this.NewConnection?.Invoke(new NewConnectionEventArgs(message, connection));
             }
-            else if (isHello)
+
+            // Inform the connection of the buffer (new connections need to send an ack back to client)
+            connection.HandleReceive(message, bytesReceived);
+
+            if (isHello)
             {
                 message.Recycle();
             }
diff --git a/Hazel/FewerThreads/UdpServerConnection2.cs b/Hazel/FewerThreads/ThreadLimitedUdpServerConnection.cs
similarity index 89%
rename from Hazel/FewerThreads/UdpServerConnection2.cs
rename to Hazel/FewerThreads/ThreadLimitedUdpServerConnection.cs
index c5e6795..ee5d2cd 100644
--- a/Hazel/FewerThreads/UdpServerConnection2.cs
+++ b/Hazel/FewerThreads/ThreadLimitedUdpServerConnection.cs
@@ -4,10 +4,10 @@ using System.Net;
 namespace Hazel.Udp.FewerThreads
 {
     /// <summary>
-    ///     Represents a servers's connection to a client that uses the UDP protocol.
+    /// Represents a servers's connection to a client that uses the UDP protocol.
     /// </summary>
     /// <inheritdoc/>
-    internal sealed class UdpServerConnection2 : UdpConnection
+    internal sealed class ThreadLimitedUdpServerConnection : UdpConnection
     {
         /// <summary>
         ///     The connection listener that we use the socket of.
@@ -16,7 +16,7 @@ namespace Hazel.Udp.FewerThreads
         ///     Udp server connections utilize the same socket in the listener for sends/receives, this is the listener that 
         ///     created this connection and is hence the listener this conenction sends and receives via.
         /// </remarks>
-        public UdpConnectionListener2 Listener { get; private set; }
+        public ThreadLimitedUdpConnectionListener Listener { get; private set; }
 
         /// <summary>
         ///     Creates a UdpConnection for the virtual connection to the endpoint.
@@ -24,7 +24,7 @@ namespace Hazel.Udp.FewerThreads
         /// <param name="listener">The listener that created this connection.</param>
         /// <param name="endPoint">The endpoint that we are connected to.</param>
         /// <param name="IPMode">The IPMode we are connected using.</param>
-        internal UdpServerConnection2(UdpConnectionListener2 listener, IPEndPoint endPoint, IPMode IPMode)
+        internal ThreadLimitedUdpServerConnection(ThreadLimitedUdpConnectionListener listener, IPEndPoint endPoint, IPMode IPMode)
             : base()
         {
             this.Listener = listener;
diff --git a/Hazel/Hazel.csproj b/Hazel/Hazel.csproj
index e249199..f6b10e6 100644
--- a/Hazel/Hazel.csproj
+++ b/Hazel/Hazel.csproj
@@ -28,7 +28,7 @@
     <LangVersion>7.3</LangVersion>
   </PropertyGroup>
   <PropertyGroup Condition=" '$(Configuration)|$(Platform)' == 'Release|AnyCPU' ">
-    <DebugType>pdbonly</DebugType>
+    <DebugType>portable</DebugType>
     <Optimize>true</Optimize>
     <OutputPath>bin\Release\</OutputPath>
     <DefineConstants>TRACE</DefineConstants>
@@ -39,6 +39,7 @@
     <AllowUnsafeBlocks>true</AllowUnsafeBlocks>
     <Prefer32Bit>false</Prefer32Bit>
     <LangVersion>7.3</LangVersion>
+    <DebugSymbols>true</DebugSymbols>
   </PropertyGroup>
   <PropertyGroup>
     <SignAssembly>true</SignAssembly>
@@ -73,8 +74,8 @@
     <Compile Include="DataReceivedEventArgs.cs" />
     <Compile Include="DisconnectedEventArgs.cs" />
     <Compile Include="FewerThreads\HazelThreadPool.cs" />
-    <Compile Include="FewerThreads\UdpConnectionListener2.cs" />
-    <Compile Include="FewerThreads\UdpServerConnection2.cs" />
+    <Compile Include="FewerThreads\ThreadLimitedUdpConnectionListener.cs" />
+    <Compile Include="FewerThreads\ThreadLimitedUdpServerConnection.cs" />
     <Compile Include="HazelException.cs" />
     <Compile Include="IPMode.cs" />
     <Compile Include="IRecyclable.cs" />
@@ -103,7 +104,6 @@
     <Compile Include="UPnP\NetUtility.cs" />
     <Compile Include="UPnP\UPnPHelper.cs" />
   </ItemGroup>
-  <ItemGroup />
   <Import Project="$(MSBuildToolsPath)\Microsoft.CSharp.targets" />
   <!-- To modify your build process, add your task inside one of the targets below and uncomment it. 
        Other similar extension points exist, see Microsoft.Common.targets.
diff --git a/Hazel/ObjectPool.cs b/Hazel/ObjectPool.cs
index 163301d..db3e3b4 100644
--- a/Hazel/ObjectPool.cs
+++ b/Hazel/ObjectPool.cs
@@ -1,5 +1,6 @@
 using System;
 using System.Collections.Concurrent;
+using System.Collections.Generic;
 using System.Threading;
 
 namespace Hazel
@@ -18,8 +19,11 @@ namespace Hazel
         public int NumberNotInUse { get { return this.pool.Count; } }
         public int Size { get { return this.NumberInUse + this.NumberNotInUse; } }
 
-        // Available objects
+#if HAZEL_BAG
         private readonly ConcurrentBag<T> pool = new ConcurrentBag<T>();
+#else
+        private readonly List<T> pool = new List<T>();
+#endif
 
         // Unavailable objects
         private readonly ConcurrentDictionary<T, bool> inuse = new ConcurrentDictionary<T, bool>();
@@ -44,11 +48,29 @@ namespace Hazel
         /// <returns>An instance of T.</returns>
         internal T GetObject()
         {
+#if HAZEL_BAG
             if (!pool.TryTake(out T item))
             {
                 Interlocked.Increment(ref numberCreated);
                 item = objectFactory.Invoke();
             }
+#else
+            T item;
+            lock (this.pool)
+            {
+                if (this.pool.Count > 0)
+                {
+                    var idx = this.pool.Count - 1;
+                    item = this.pool[idx];
+                    this.pool.RemoveAt(idx);
+                }
+                else
+                {
+                    Interlocked.Increment(ref numberCreated);
+                    item = objectFactory.Invoke();
+                }
+            }
+#endif
 
             if (!inuse.TryAdd(item, true))
             {
@@ -66,7 +88,14 @@ namespace Hazel
         {
             if (inuse.TryRemove(item, out bool b))
             {
+#if HAZEL_BAG
                 pool.Add(item);
+#else
+                lock (this.pool)
+                {
+                    pool.Add(item);
+                }
+#endif
             }
             else
             {
diff --git a/Hazel/Udp/UdpConnection.KeepAlive.cs b/Hazel/Udp/UdpConnection.KeepAlive.cs
index 8d55f05..9ae7697 100644
--- a/Hazel/Udp/UdpConnection.KeepAlive.cs
+++ b/Hazel/Udp/UdpConnection.KeepAlive.cs
@@ -74,32 +74,34 @@ namespace Hazel.Udp
         protected void InitializeKeepAliveTimer()
         {
             keepAliveTimer = new Timer(
-                (o) =>
-                {
-                    if (this.State != ConnectionState.Connected) return;
-
-                    if (this.pingsSinceAck >= this.MissingPingsUntilDisconnect)
-                    {
-                        this.DisposeKeepAliveTimer();
-                        this.DisconnectInternal(HazelInternalErrors.PingsWithoutResponse, $"Sent {this.pingsSinceAck} pings that remote has not responded to.");
-                        return;
-                    }
-
-                    try
-                    {
-                        this.pingsSinceAck++;
-                        SendPing();
-                    }
-                    catch
-                    {
-                    }
-                },
+                HandleKeepAlive,
                 null,
                 keepAliveInterval,
                 keepAliveInterval
             );
         }
 
+        private void HandleKeepAlive(object state)
+        {
+            if (this.State != ConnectionState.Connected) return;
+
+            if (this.pingsSinceAck >= this.MissingPingsUntilDisconnect)
+            {
+                this.DisposeKeepAliveTimer();
+                this.DisconnectInternal(HazelInternalErrors.PingsWithoutResponse, $"Sent {this.pingsSinceAck} pings that remote has not responded to.");
+                return;
+            }
+
+            try
+            {
+                this.pingsSinceAck++;
+                SendPing();
+            }
+            catch
+            {
+            }
+        }
+
         // Pings are special, quasi-reliable packets. 
         // We send them to trigger responses that validate our connection is alive
         // An unacked ping should never be the sole cause of a disconnect.
diff --git a/Hazel/Udp/UdpConnection.Reliable.cs b/Hazel/Udp/UdpConnection.Reliable.cs
index b15f883..f09c432 100644
--- a/Hazel/Udp/UdpConnection.Reliable.cs
+++ b/Hazel/Udp/UdpConnection.Reliable.cs
@@ -49,7 +49,7 @@ namespace Hazel.Udp
         internal ConcurrentDictionary<ushort, Packet> reliableDataPacketsSent = new ConcurrentDictionary<ushort, Packet>();
 
         /// <summary>
-        ///     The last packets that were received.
+        ///     Packet ids that have not been received, but are expected. 
         /// </summary>
         private HashSet<ushort> reliableDataPacketsMissing = new HashSet<ushort>();
 
@@ -367,10 +367,21 @@ namespace Hazel.Udp
                 //If it's new or we've not received anything yet
                 if (isNew)
                 {
-                    //Mark items between the most recent receive and the id received as missing
-                    for (ushort i = (ushort)(reliableReceiveLast + 1); i < id; i++)
+                    // Mark items between the most recent receive and the id received as missing
+                    if (id > reliableReceiveLast)
                     {
-                        reliableDataPacketsMissing.Add(i);
+                        for (ushort i = (ushort)(reliableReceiveLast + 1); i < id; i++)
+                        {
+                            reliableDataPacketsMissing.Add(i);
+                        }
+                    }
+                    else
+                    {
+                        int cnt = (ushort.MaxValue - reliableReceiveLast) + id;
+                        for (ushort i = 1; i < cnt; ++i)
+                        {
+                            reliableDataPacketsMissing.Add((ushort)(i + reliableReceiveLast));
+                        }
                     }
 
                     //Update the most recently received
@@ -399,7 +410,6 @@ namespace Hazel.Udp
         {
             this.pingsSinceAck = 0;
 
-            // Get ID
             ushort id = (ushort)((bytes[1] << 8) + bytes[2]);
             AcknowledgeMessageId(id);
 
@@ -415,13 +425,9 @@ namespace Hazel.Udp
 
                     recentPackets >>= 1;
                 }
-
-                Statistics.LogReliableReceive(bytesReceived - 4, bytesReceived);
-            }
-            else
-            {
-                Statistics.LogReliableReceive(bytesReceived - 3, bytesReceived);
             }
+
+            Statistics.LogReliableReceive(0, bytesReceived);
         }
 
         private void AcknowledgeMessageId(ushort id)
@@ -459,16 +465,15 @@ namespace Hazel.Udp
         /// <param name="byte2">The second identification byte.</param>
         private void SendAck(ushort id)
         {
-            const byte Found = 1;
-            const byte Missing = 1;
-
             byte recentPackets = 0;
             lock (this.reliableDataPacketsMissing)
             {
                 for (int i = 1; i <= 8; ++i)
                 {
-                    recentPackets |= this.reliableDataPacketsMissing.Contains((ushort)(id - i)) ? Found : Missing;
-                    recentPackets <<= 1;
+                    if (!this.reliableDataPacketsMissing.Contains((ushort)(id - i)))
+                    {
+                        recentPackets |= (byte)(1 << (i - 1));
+                    }
                 }
             }
 
@@ -480,8 +485,6 @@ namespace Hazel.Udp
                 recentPackets
             };
 
-            // Always reply with acknowledgement in order to stop the sender repeatedly sending it
-            // TODO: group acks together
             try
             {
                 WriteBytesToConnection(bytes, bytes.Length);
diff --git a/Hazel/Udp/UdpConnectionListener.cs b/Hazel/Udp/UdpConnectionListener.cs
index 5570ae6..59089bf 100644
--- a/Hazel/Udp/UdpConnectionListener.cs
+++ b/Hazel/Udp/UdpConnectionListener.cs
@@ -42,7 +42,7 @@ namespace Hazel.Udp
             this.IPMode = ipMode;
 
             this.socket = UdpConnection.CreateSocket(this.IPMode);
-
+            
             socket.ReceiveBufferSize = SendReceiveBufferSize;
             socket.SendBufferSize = SendReceiveBufferSize;
             
@@ -119,7 +119,7 @@ namespace Hazel.Udp
         {
             var message = (MessageReader)result.AsyncState;
             int bytesReceived;
-            EndPoint remoteEndPoint = new IPEndPoint(IPMode == IPMode.IPv4 ? IPAddress.Any : IPAddress.IPv6Any, 0);
+            EndPoint remoteEndPoint = new IPEndPoint(this.EndPoint.Address, this.EndPoint.Port);
 
             //End the receive operation
             try
@@ -225,7 +225,7 @@ namespace Hazel.Udp
                 InvokeNewConnection(message, connection);
             }
 
-            //Inform the connection of the buffer (new connections need to send an ack back to client)
+            // Inform the connection of the buffer (new connections need to send an ack back to client)
             connection.HandleReceive(message, bytesReceived);
 
             if (aware && isHello)
diff --git a/Hazel/Udp/UnityUdpClientConnection.cs b/Hazel/Udp/UnityUdpClientConnection.cs
index 4f4640f..2e217d8 100644
--- a/Hazel/Udp/UnityUdpClientConnection.cs
+++ b/Hazel/Udp/UnityUdpClientConnection.cs
@@ -15,6 +15,8 @@ namespace Hazel.Udp
     {
         private Socket socket;
 
+        public EndPoint LocalEndpoint { get { return this.socket.LocalEndPoint; } }
+
         public UnityUdpClientConnection(IPEndPoint remoteEndPoint, IPMode ipMode = IPMode.IPv4)
             : base()
         {
@@ -23,6 +25,7 @@ namespace Hazel.Udp
             this.IPMode = ipMode;
 
             this.socket = CreateSocket(ipMode);
+            this.socket.SetSocketOption(SocketOptionLevel.Socket, SocketOptionName.ExclusiveAddressUse, true);
         }
         
         ~UnityUdpClientConnection()
@@ -135,7 +138,8 @@ namespace Hazel.Udp
             var msg = MessageReader.GetSized(ushort.MaxValue);
             try
             {
-                socket.BeginReceive(msg.Buffer, 0, msg.Buffer.Length, SocketFlags.None, ReadCallback, msg);
+                var ep = this.RemoteEndPoint;
+                socket.BeginReceiveFrom(msg.Buffer, 0, msg.Buffer.Length, SocketFlags.None, ref ep, ReadCallback, msg);
             }
             catch
             {
@@ -154,7 +158,8 @@ namespace Hazel.Udp
 
             try
             {
-                msg.Length = socket.EndReceive(result);
+                var ep = this.RemoteEndPoint;
+                msg.Length = socket.EndReceiveFrom(result, ref ep);
             }
             catch (SocketException e)
             {
-- 
2.39.5