Skip to content

Commit 1e9871e

Browse files
Toastbrot236Toastbrot236
committed
Port RateLimiter stuff from Bunkum, minor rate-limiting logic improvements, lookup bucket data by name
1 parent 9b77c23 commit 1e9871e

7 files changed

Lines changed: 206 additions & 0 deletions

File tree

‎Refresh.Common/RefreshContext.cs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,4 +13,5 @@ public enum RefreshContext
1313
Presence,
1414
Database,
1515
CacheService,
16+
RateLimit,
1617
}

‎Refresh.Core/RateLimits/BasicRateLimitAttribute.cs‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,12 @@
1+
namespace Refresh.Core.RateLimits;
2+
3+
[AttributeUsage(AttributeTargets.Method)]
4+
public class BasicRateLimitAttribute : Attribute
5+
{
6+
public readonly string Bucket;
7+
8+
public BasicRateLimitAttribute(string bucket)
9+
{
10+
this.Bucket = bucket;
11+
}
12+
}

‎Refresh.Core/RateLimits/BasicRateLimitBucket.cs‎

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,22 @@
1+
using Refresh.Core.Configuration.Structs;
2+
3+
namespace Refresh.Core.RateLimits;
4+
5+
public class BasicRateLimitBucket
6+
{
7+
public int TimeWindowSeconds { get; set; }
8+
public int MaxRequestCount { get; set; }
9+
public int BlockDurationSeconds { get; set; }
10+
public string BucketName { get; set; } = "";
11+
12+
public static BasicRateLimitBucket FromOld(ConfigRateLimitBucket old, string bucketName)
13+
{
14+
return new()
15+
{
16+
TimeWindowSeconds = old.TimeWindowSeconds,
17+
MaxRequestCount = old.MaxRequestCount,
18+
BlockDurationSeconds = old.BlockDurationSeconds,
19+
BucketName = bucketName,
20+
};
21+
}
22+
}

‎Refresh.Core/RateLimits/BasicRateLimiter.cs‎

Lines changed: 126 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,126 @@
1+
using System.Collections.Frozen;
2+
using System.Net;
3+
using System.Reflection;
4+
using Bunkum.Core.RateLimit;
5+
using Bunkum.Listener.Request;
6+
using NotEnoughLogs;
7+
using Refresh.Common;
8+
using Refresh.Common.Time;
9+
using Refresh.Core.Configuration;
10+
using Refresh.Core.Configuration.Structs;
11+
using Refresh.Core.RateLimits.Info;
12+
13+
namespace Refresh.Core.RateLimits;
14+
15+
public class BasicRateLimiter : IRateLimiter
16+
{
17+
private readonly Logger _logger;
18+
private readonly IDateTimeProvider _timeProvider;
19+
private readonly FrozenDictionary<string, ConfigRateLimitBucket> _buckets;
20+
21+
public BasicRateLimiter(IDateTimeProvider timeProvider, Logger logger, Dictionary<string, ConfigRateLimitBucket> buckets)
22+
{
23+
this._timeProvider = timeProvider;
24+
this._logger = logger;
25+
this._buckets = buckets.ToFrozenDictionary();
26+
}
27+
28+
private readonly List<RateLimitUserInfo> _userInfos = new(25);
29+
private readonly List<RateLimitRemoteEndpointInfo> _remoteEndpointInfos = new(25);
30+
31+
private BasicRateLimitBucket GetBucket(MethodInfo? method)
32+
{
33+
string bucketName = method?.GetCustomAttribute<BasicRateLimitAttribute>()?.Bucket ?? "global";
34+
ConfigRateLimitBucket? bucketData = this._buckets.GetValueOrDefault(bucketName);
35+
36+
if (bucketData == null)
37+
{
38+
this._logger.LogDebug(RefreshContext.RateLimit, $"Could not find bucket '{bucketName}' in config, falling back to hardcoded defaults.");
39+
bucketData = BasicEndpointRateLimitConfig.DefaultBuckets.GetValueOrDefault(bucketName);
40+
41+
if (bucketData == null)
42+
{
43+
throw new NotImplementedException($"Could not find bucket '{bucketName}' in neither the config file nor the hardcoded defaults!");
44+
}
45+
}
46+
47+
return BasicRateLimitBucket.FromOld(bucketData.Value, bucketName);
48+
}
49+
50+
public bool UserViolatesRateLimit(ListenerContext context, MethodInfo? method, IRateLimitUser user)
51+
{
52+
BasicRateLimitBucket bucket = this.GetBucket(method);
53+
54+
lock (this._remoteEndpointInfos)
55+
{
56+
RateLimitUserInfo? info = this._userInfos
57+
.FirstOrDefault(i =>
58+
user.RateLimitUserIdIsEqual(i.User.RateLimitUserId) && i.Bucket == bucket.BucketName);
59+
60+
if (info == null)
61+
{
62+
info = new RateLimitUserInfo(user, bucket.BucketName);
63+
lock (this._userInfos)
64+
{
65+
this._userInfos.Add(info);
66+
}
67+
}
68+
69+
lock (info)
70+
{
71+
return this.ViolatesRateLimit(context, info, bucket);
72+
}
73+
}
74+
}
75+
76+
public bool RemoteEndpointViolatesRateLimit(ListenerContext context, MethodInfo? method)
77+
{
78+
IPAddress ipAddress = context.RemoteEndpoint.Address;
79+
80+
BasicRateLimitBucket bucket = this.GetBucket(method);
81+
82+
lock (this._remoteEndpointInfos)
83+
{
84+
RateLimitRemoteEndpointInfo? info = this._remoteEndpointInfos
85+
.FirstOrDefault(i => ipAddress.Equals(i.IpAddress) && i.Bucket == bucket.BucketName);
86+
87+
if (info == null)
88+
{
89+
info = new RateLimitRemoteEndpointInfo(ipAddress, bucket.BucketName);
90+
91+
this._remoteEndpointInfos.Add(info);
92+
}
93+
94+
lock (info)
95+
{
96+
return this.ViolatesRateLimit(context, info, bucket);
97+
}
98+
}
99+
}
100+
101+
private bool ViolatesRateLimit(ListenerContext context, IRateLimitInfo info, BasicRateLimitBucket bucket)
102+
{
103+
int now = (int)this._timeProvider.TimestampSeconds;
104+
105+
// Always clear all expired and add this request as a new one,
106+
// to make the block duration longer the more the client continues to spam
107+
info.RequestTimes.RemoveAll(r => r <= now - bucket.TimeWindowSeconds);
108+
info.RequestTimes.Add(now);
109+
110+
// Repeat block
111+
// If a rate-limit is already triggered, block regardless of whether MaxRequestCount has been reached
112+
if (info.LimitedUntil != 0 && info.LimitedUntil > now)
113+
return true;
114+
115+
// Initial block
116+
if (info.RequestTimes.Count > bucket.MaxRequestCount)
117+
{
118+
info.LimitedUntil = now + bucket.BlockDurationSeconds;
119+
context.ResponseHeaders.TryAdd("Retry-After", bucket.BlockDurationSeconds.ToString()); // TODO also include overshot time
120+
121+
return true;
122+
}
123+
124+
return false;
125+
}
126+
}

‎Refresh.Core/RateLimits/Info/IRateLimitInfo.cs‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,9 @@
1+
namespace Refresh.Core.RateLimits.Info;
2+
3+
// For some reason, this interface is internal in Bunkum.Core, so we have to copy it.
4+
public interface IRateLimitInfo
5+
{
6+
internal List<int> RequestTimes { get; init; }
7+
internal int LimitedUntil { get; set; }
8+
public string Bucket { get; init; }
9+
}

‎Refresh.Core/RateLimits/Info/RateLimitRemoteEndpointInfo.cs‎

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,18 @@
1+
using System.Net;
2+
3+
namespace Refresh.Core.RateLimits.Info;
4+
5+
// For some reason, this class, unlike the others, is NOT internal in Bunkum.Core, but we're copying it anyway for consistency.
6+
public class RateLimitRemoteEndpointInfo : IRateLimitInfo
7+
{
8+
public RateLimitRemoteEndpointInfo(IPAddress ipAddress, string bucket)
9+
{
10+
this.IpAddress = ipAddress;
11+
this.Bucket = bucket;
12+
}
13+
14+
internal IPAddress IpAddress { get; init; }
15+
public List<int> RequestTimes { get; init; } = new(25);
16+
public int LimitedUntil { get; set; }
17+
public string Bucket { get; init; }
18+
}

‎Refresh.Core/RateLimits/Info/RateLimitUserInfo.cs‎

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,18 @@
1+
using Bunkum.Core.RateLimit;
2+
3+
namespace Refresh.Core.RateLimits.Info;
4+
5+
// For some reason, this class is internal in Bunkum.Core, so we have to copy it.
6+
public class RateLimitUserInfo : IRateLimitInfo
7+
{
8+
internal RateLimitUserInfo(IRateLimitUser user, string bucket)
9+
{
10+
this.User = user;
11+
this.Bucket = bucket;
12+
}
13+
14+
internal IRateLimitUser User { get; init; }
15+
public string Bucket { get; init; }
16+
public List<int> RequestTimes { get; init; } = new(25);
17+
public int LimitedUntil { get; set; }
18+
}

0 commit comments

Comments
 (0)