From: AeonLucid <aeonlucid@outlook.com>
Date: Sat, 24 Oct 2020 00:13:56 +0000 (+0200)
Subject: Basic Hazel connection rate limiting
X-Git-Tag: v1.2.2~96^2~29
X-Git-Url: https://git.deb.at/?a=commitdiff_plain;h=553c8b9b1b0e0a96beae9a8b830e8792e94f5a1d;p=rhonda%2Fimpostor.git

Basic Hazel connection rate limiting
---

diff --git a/src/Impostor.Hazel/Udp/UdpConnectionListener.cs b/src/Impostor.Hazel/Udp/UdpConnectionListener.cs
index 1edd56a..a2b6180 100644
--- 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
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