Basic Hazel connection rate limiting

diff --git a/src/Impostor.Hazel/Udp/UdpConnectionListener.cs b/src/Impostor.Hazel/Udp/UdpConnectionListener.cs
index 1edd56a76a439059191dd25dfd7eda5bb03701ad..a2b618032539ef914d9cd029c39b053c637a9b53 100644 (file)
--- a/src/Impostor.Hazel/Udp/UdpConnectionListener.cs
+++ b/src/Impostor.Hazel/Udp/UdpConnectionListener.cs
@@ -31,6 +31,7 @@ namespace Impostor.Hazel.Udp
         private readonly Timer _reliablePacketTimer;
         private readonly ConcurrentDictionary<EndPoint, UdpServerConnection> _allConnections;
         private readonly CancellationTokenSource _stoppingCts;
+        private readonly UdpConnectionRateLimit _connectionRateLimit;
         private Task _executingTask;
 
         /// <summary>
@@ -58,6 +59,8 @@ namespace Impostor.Hazel.Udp
             {
                 _socket.Dispose();
             });
+
+            _connectionRateLimit = new UdpConnectionRateLimit();
         }
 
         public int ConnectionCount => this._allConnections.Count;
@@ -154,6 +157,13 @@ namespace Impostor.Hazel.Udp
                             continue;
                         }
 
+                        // Check rateLimit.
+                        if (!_connectionRateLimit.IsAllowed(data.RemoteEndPoint.Address))
+                        {
+                            Logger.Warning("Ratelimited connection attempt from {0}.", data.RemoteEndPoint);
+                            continue;
+                        }
+
                         // Create new client
                         client = new UdpServerConnection(this, data.RemoteEndPoint, IPMode);
 
@@ -274,6 +284,8 @@ namespace Impostor.Hazel.Udp
 
             await _reliablePacketTimer.DisposeAsync();
 
+            _connectionRateLimit.Dispose();
+
             await base.DisposeAsync();
         }
     }

diff --git a/src/Impostor.Hazel/Udp/UdpConnectionRateLimit.cs b/src/Impostor.Hazel/Udp/UdpConnectionRateLimit.cs
new file mode 100644 (file)
index 0000000..64881d3
--- /dev/null
+++ b/src/Impostor.Hazel/Udp/UdpConnectionRateLimit.cs
@@ -0,0 +1,75 @@
+using System;
+using System.Collections.Concurrent;
+using System.Net;
+using System.Threading;
+using Serilog;
+
+namespace Impostor.Hazel.Udp
+{
+    public class UdpConnectionRateLimit : IDisposable
+    {
+        private static readonly ILogger Logger = Log.ForContext<UdpConnectionRateLimit>();
+
+        // Allow burst to 5 connections.
+        // Decrease by 1 every second.
+        private const int MaxConnections = 5;
+        private const int FalloffMs = 1000;
+
+        private readonly ConcurrentDictionary<IPAddress, int> _connectionCount;
+        private readonly Timer _timer;
+        private bool _isDisposed;
+
+        public UdpConnectionRateLimit()
+        {
+            _connectionCount = new ConcurrentDictionary<IPAddress, int>();
+            _timer = new Timer(UpdateRateLimit, null, FalloffMs, Timeout.Infinite);
+        }
+
+        private void UpdateRateLimit(object state)
+        {
+            try
+            {
+                foreach (var pair in _connectionCount)
+                {
+                    var count = pair.Value - 1;
+                    if (count > 0)
+                    {
+                        _connectionCount.TryUpdate(pair.Key, count, pair.Value);
+                    }
+                    else
+                    {
+                        _connectionCount.TryRemove(pair);
+                    }
+                }
+            }
+            catch (Exception e)
+            {
+                Logger.Error(e, "Exception caught in UpdateRateLimit.");
+            }
+            finally
+            {
+                if (!_isDisposed)
+                {
+                    _timer.Change(FalloffMs, Timeout.Infinite);
+                }
+            }
+        }
+
+        public bool IsAllowed(IPAddress key)
+        {
+            if (_connectionCount.TryGetValue(key, out var value) && value >= MaxConnections)
+            {
+                return false;
+            }
+
+            _connectionCount.AddOrUpdate(key, _ => 1, (_, i) => i + 1);
+            return true;
+        }
+
+        public void Dispose()
+        {
+            _isDisposed = true;
+            _timer.Dispose();
+        }
+    }
+}
\ No newline at end of file