Add a configurable endpoint rate-limiter

Refresh.Common/RefreshContext.cs

@@ -13,4 +13,5 @@ public enum RefreshContext
     Presence,
     Database,
     CacheService,
+    RateLimit,
 }

Refresh.Core/Configuration/ConfigRateLimitBucket.cs

@@ -0,0 +1,15 @@
+namespace Refresh.Core.Configuration;
+
+public class ConfigRateLimitBucket
+{
+    public int TimeWindowSeconds { get; set; }
+    public int MaxRequestAmount { get; set; }
+    public int BlockDurationSeconds { get; set; }
+
+    public ConfigRateLimitBucket(int windowDurationSeconds, int maxRequestAmount, int blockDurationSeconds)
+    {
+        this.TimeWindowSeconds = windowDurationSeconds;
+        this.MaxRequestAmount = maxRequestAmount;
+        this.BlockDurationSeconds = blockDurationSeconds;
+    }
+}

Refresh.Core/Configuration/ConfigStore.cs

@@ -14,6 +14,7 @@ public class ConfigStore
     public RichPresenceConfig RichPresence { get; }
 
     public DryArchiveConfig DryArchive { get; }
+    public EndpointRateLimitConfig EndpointRateLimits { get; }
 
     private static readonly Lock ConfigLock = new();
     public ConfigStore(Logger logger)
@@ -28,6 +29,7 @@ public ConfigStore(Logger logger)
             this.RichPresence = Config.LoadFromJsonFile<RichPresenceConfig>("rpc.json", logger);
 
             this.DryArchive = Config.LoadFromJsonFile<DryArchiveConfig>("dry.json", logger);
+            this.EndpointRateLimits = Config.LoadFromJsonFile<EndpointRateLimitConfig>("endpointRateLimits.json", logger);
         }
     }
 
@@ -41,6 +43,7 @@ public ConfigStore()
         this.RichPresence = new RichPresenceConfig();
 
         this.DryArchive = new DryArchiveConfig();
+        this.EndpointRateLimits = new EndpointRateLimitConfig();
     }
 
     public void AddToBunkum(BunkumServer server)
@@ -51,5 +54,6 @@ public void AddToBunkum(BunkumServer server)
         server.AddConfig(this.Integration);
         server.AddConfig(this.RichPresence);
         server.AddConfig(this.DryArchive);
+        server.AddConfig(this.EndpointRateLimits);
     }
 }

Refresh.Core/Configuration/EndpointRateLimitConfig.cs

@@ -0,0 +1,46 @@
+using Bunkum.Core.Configuration;
+using Refresh.Core.RateLimits.EndpointRateLimiting;
+
+namespace Refresh.Core.Configuration;
+
+public class EndpointRateLimitConfig : Config
+{
+    public override int CurrentConfigVersion => 1;
+    public override int Version { get; set; }
+
+    protected override void Migrate(int oldVer, dynamic oldConfig)
+    {
+        // initialize
+        if (oldVer < 1)
+        {
+            this.AddMissingBucketsFromDefaults();
+        }
+
+        // think of how exactly to overwrite in the future
+    }
+
+    public void AddMissingBucketsFromDefaults()
+    {
+        foreach (KeyValuePair<EndpointBucketId, ConfigRateLimitBucket> defaultPair in EndpointBucketDefaults.Buckets)
+        {
+            string bucketName = defaultPair.Key.ToString();
+            ConfigRateLimitBucket bucket = defaultPair.Value;
+            this.Buckets.TryAdd(bucketName, bucket);
+        }
+    }
+
+    /// <summary>
+    /// If a bucket's default values are updated in a new server release, this will determine whether the bucket's configured values,
+    /// which might or might not have been changed by the server owner, will be overwritten with the new default or not.
+    /// 
+    /// Although this option does nothing for now, it already exists so owners can already decide to opt out of this ahead of time.
+    /// </summary>
+    public bool OverwriteBucketValuesIfDefaultsAreUpdated { get; set; } = true;
+
+    /// <summary>
+    /// Whether we should print the IDs of all buckets that are missing from the config, but exist in the default map,
+    /// while the <see cref="EndpointRateLimiter"/> is being initialized. These logs will be printed as warnings if enabled.
+    /// </summary>
+    public bool PrintMissingBuckets { get; set; } = false;
+    public Dictionary<string, ConfigRateLimitBucket> Buckets { get; set; } = new();
+}

Refresh.Core/RateLimits/EndpointRateLimiting/Client/IClientBucketBaseData.cs

@@ -0,0 +1,7 @@
+namespace Refresh.Core.RateLimits.EndpointRateLimiting.Client;
+
+public interface IClientBucketBaseData
+{
+    public List<int> RequestTimes { get; init; }
+    public int LimitedUntil { get; set; }
+}

Refresh.Core/RateLimits/EndpointRateLimiting/Client/TrackedClientBucketData.cs

@@ -0,0 +1,15 @@
+namespace Refresh.Core.RateLimits.EndpointRateLimiting.Client;
+
+public class TrackedClientBucketData<TClientIdType> : IClientBucketBaseData
+{
+    public List<int> RequestTimes { get; init; } = new(25);
+    public int LimitedUntil { get; set; }
+    public TClientIdType ClientId { get; init; }
+    public EndpointBucketId Bucket { get; init; }
+
+    public TrackedClientBucketData(TClientIdType clientId, EndpointBucketId bucket)
+    {
+        this.ClientId = clientId;
+        this.Bucket = bucket;
+    }
+}

Refresh.Core/RateLimits/EndpointRateLimiting/EndpointBucketDefaults.cs

@@ -0,0 +1,20 @@
+using System.Collections.Frozen;
+using Refresh.Core.Configuration;
+
+namespace Refresh.Core.RateLimits.EndpointRateLimiting;
+
+public static class EndpointBucketDefaults
+{
+    public static readonly FrozenDictionary<EndpointBucketId, ConfigRateLimitBucket> Buckets = new Dictionary<EndpointBucketId, ConfigRateLimitBucket>()
+    {
+        #region Misc
+        {EndpointBucketId.Default, new(90, 300, 45)},
+        #endregion
+
+        #region Levels
+        // game sometimes requests many levels in bursts
+        {EndpointBucketId.GameGetSingleLevel, new(240, 200, 180)},
+        {EndpointBucketId.ApiGetSingleLevel, new(240, 50, 180)},
+        #endregion
+    }.ToFrozenDictionary();
+}

Refresh.Core/RateLimits/EndpointRateLimiting/EndpointBucketId.cs

@@ -0,0 +1,16 @@
+namespace Refresh.Core.RateLimits.EndpointRateLimiting;
+
+// TODO add IDs for all API/game buckets here
+// Generally, fetch endpoints should use separate buckets depending on whether they are game/API endpoints,
+// while upload/modification/deletion endpoints should share buckets.
+public enum EndpointBucketId
+{
+    #region Misc
+    Default,
+    #endregion
+
+    #region Levels
+    GameGetSingleLevel,
+    ApiGetSingleLevel,
+    #endregion
+}

Refresh.Core/RateLimits/EndpointRateLimiting/EndpointRateLimitAttribute.cs

@@ -0,0 +1,12 @@
+namespace Refresh.Core.RateLimits.EndpointRateLimiting;
+
+[AttributeUsage(AttributeTargets.Method)]
+public class EndpointRateLimitAttribute : Attribute
+{
+    public readonly EndpointBucketId MainBucket;
+
+    public EndpointRateLimitAttribute(EndpointBucketId bucket)
+    {
+        this.MainBucket = bucket;
+    }
+}

Refresh.Core/RateLimits/EndpointRateLimiting/EndpointRateLimiter.cs

@@ -0,0 +1,157 @@
+using System.Collections.Frozen;
+using System.Net;
+using System.Reflection;
+using Bunkum.Listener.Request;
+using MongoDB.Bson;
+using NotEnoughLogs;
+using Refresh.Common;
+using Refresh.Common.Time;
+using Refresh.Core.Configuration;
+using Refresh.Core.RateLimits.EndpointRateLimiting.Client;
+using Refresh.Database.Models.Users;
+
+namespace Refresh.Core.RateLimits.EndpointRateLimiting;
+
+public class EndpointRateLimiter
+{
+    private readonly Logger _logger;
+    private readonly IDateTimeProvider _timeProvider;
+    protected FrozenDictionary<EndpointBucketId, ConfigRateLimitBucket> Buckets;
+
+    protected List<TrackedClientBucketData<ObjectId>> UserInfos = new(25);
+    protected List<TrackedClientBucketData<IPAddress>> RemoteEndpointInfos = new(25);
+
+    public EndpointRateLimiter(IDateTimeProvider timeProvider, Logger logger, EndpointRateLimitConfig config)
+    {
+        this._timeProvider = timeProvider;
+        this._logger = logger;
+
+        // Copy the buckets over, converting the string bucket names to their corresponding enum values.
+        Dictionary<EndpointBucketId, ConfigRateLimitBucket> validBuckets = new();
+
+        foreach (KeyValuePair<string, ConfigRateLimitBucket> bucket in config.Buckets)
+        {
+            bool parsed = Enum.TryParse(bucket.Key, true, out EndpointBucketId nameParsed);
+            if (!parsed)
+            {
+                this._logger.LogDebug(RefreshContext.RateLimit, $"Bucket name '{bucket.Key}' found in rate-limit config is unknown (does not map to a valid {nameof(EndpointBucketId)} enum value), its bucket will be ignored.");
+                continue;
+            }
+
+            validBuckets.Add(nameParsed, bucket.Value);
+        }
+
+        // check for any buckets missing from the config, and insert default buckets in their place.
+        // this way, instead of logging the missing bucket every single time it's looked up during a request,
+        // we instead just print it once here.
+        foreach (KeyValuePair<EndpointBucketId, ConfigRateLimitBucket> defaultPair in EndpointBucketDefaults.Buckets)
+        {
+            bool existsInConfig = validBuckets.ContainsKey(defaultPair.Key);
+            if (existsInConfig) continue;
+
+            if (config.PrintMissingBuckets)
+            {
+                logger.LogWarning(RefreshContext.RateLimit, $"Bucket {defaultPair.Key} is missing from your config, we will use its hardcoded defaults instead.");
+            }
+
+            validBuckets.Add(defaultPair.Key, defaultPair.Value);
+        }
+
+        this.Buckets = validBuckets.ToFrozenDictionary();
+    }
+
+    private LoadedBucketData GetBucketNameAndData(ListenerContext context, MethodInfo? method)
+    {
+        EndpointRateLimitAttribute? attribute = method?.GetCustomAttribute<EndpointRateLimitAttribute>();
+
+        EndpointBucketId bucketName = EndpointBucketId.Default;
+        if (attribute != null) bucketName = attribute.MainBucket;
+
+        ConfigRateLimitBucket? bucketData = this.Buckets.GetValueOrDefault(bucketName);
+
+        if (bucketData == null)
+        {
+            // Don't look this bucket up in the defaults, because we've already merged with the defaults in the constructor above,
+            // so all buckets missing from the config should already have their default versions in the map we use here.
+            throw new NotImplementedException($"Could not find bucket '{bucketName}' in neither the config file nor the hardcoded defaults! You should open an issue about this.");
+        }
+
+        return new LoadedBucketData(bucketName, bucketData);
+    }
+
+    public bool UserViolatesRateLimit(ListenerContext context, MethodInfo method, GameUser user)
+    {
+        LoadedBucketData bucketData = this.GetBucketNameAndData(context, method);
+
+        lock (this.UserInfos)
+        {
+            TrackedClientBucketData<ObjectId>? info = this.UserInfos
+                .FirstOrDefault(i => user.UserId.Equals(i.ClientId) && i.Bucket == bucketData.Id);
+
+            if (info == null)
+            {
+                info = new TrackedClientBucketData<ObjectId>(user.UserId, bucketData.Id);
+                this.UserInfos.Add(info);
+            }
+
+            lock (info)
+            {
+                return this.ViolatesRateLimit(context, bucketData, info, user);
+            }
+        }
+    }
+
+    public bool RemoteEndpointViolatesRateLimit(ListenerContext context, MethodInfo method)
+    {
+        IPAddress ipAddress = context.RemoteEndpoint.Address;
+
+        LoadedBucketData bucketData = this.GetBucketNameAndData(context, method);
+
+        lock (this.RemoteEndpointInfos)
+        {
+            TrackedClientBucketData<IPAddress>? info = this.RemoteEndpointInfos
+                .FirstOrDefault(i => ipAddress.Equals(i.ClientId) && i.Bucket == bucketData.Id);
+
+            if (info == null)
+            {
+                info = new TrackedClientBucketData<IPAddress>(ipAddress, bucketData.Id);
+                this.RemoteEndpointInfos.Add(info);
+            }
+
+            lock (info)
+            {
+                return this.ViolatesRateLimit(context, bucketData, info, null);
+            }
+        }
+    }
+
+    public bool ViolatesRateLimit(ListenerContext context, LoadedBucketData bucket, IClientBucketBaseData info, GameUser? user)
+    {
+        int now = (int)this._timeProvider.TimestampSeconds;
+
+        this._logger.LogTrace(RefreshContext.RateLimit, $"{this.GetType().Name}.{nameof(this.ViolatesRateLimit)}() - Request times count: {info.RequestTimes.Count}, limited until: {info.LimitedUntil}.");
+
+        if (info.LimitedUntil != 0)
+        {
+            // TODO also track requests received while the client is already rate-limited, to increase their block duration as punishment
+            if (info.LimitedUntil > now) return true;
+            info.LimitedUntil = 0;
+
+            // TODO don't clear all tracked requests once the block duration is over, only ever clear expired ones
+            info.RequestTimes.Clear();
+        }
+
+        info.RequestTimes.RemoveAll(r => r <= now - bucket.Data.TimeWindowSeconds);
+
+        if (info.RequestTimes.Count + 1 > bucket.Data.MaxRequestAmount)
+        {
+            info.LimitedUntil = now + bucket.Data.BlockDurationSeconds;
+            context.ResponseHeaders.TryAdd("Retry-After", bucket.Data.BlockDurationSeconds.ToString());
+
+            return true;
+        }
+
+        info.RequestTimes.Add(now);
+        return false;
+    }
+}

Refresh.Core/RateLimits/EndpointRateLimiting/LoadedBucketData.cs

@@ -0,0 +1,5 @@
+using Refresh.Core.Configuration;
+
+namespace Refresh.Core.RateLimits.EndpointRateLimiting;
+
+public record LoadedBucketData(EndpointBucketId Id, ConfigRateLimitBucket Data);