1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
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();
}
}
}
|
