Add counter key builders

#13 #29 #41
stefanprodan · Feb 17, 2019 · 481bc70 · 481bc70
1 parent 47e2722
commit 481bc70
Show file tree

Hide file tree

Showing 12 changed files with 109 additions and 43 deletions.
diff --git a/src/AspNetCoreRateLimit/Core/ClientRateLimitProcessor.cs b/src/AspNetCoreRateLimit/Core/ClientRateLimitProcessor.cs
@@ -13,8 +13,9 @@ public class ClientRateLimitProcessor : RateLimitProcessor, IRateLimitProcessor
         public ClientRateLimitProcessor(
            ClientRateLimitOptions options,
            IRateLimitCounterStore counterStore,
-           IClientPolicyStore policyStore)
-        : base(options, counterStore)
+           IClientPolicyStore policyStore,
+           IRateLimitConfiguration config)
+        : base(options, counterStore, new ClientCounterKeyBuilder(options), config)
         {
             _options = options;
             _policyStore = policyStore;
@@ -31,10 +32,5 @@ public async Task<IEnumerable<RateLimitRule>> GetMatchingRulesAsync(ClientReques
 
             return Enumerable.Empty<RateLimitRule>();
         }
-
-        protected override string GetCounterKey(ClientRequestIdentity requestIdentity, RateLimitRule rule)
-        {
-            return $"{_options.RateLimitCounterPrefix}_{requestIdentity.ClientId}_{rule.Period}";
-        }
     }
 }
diff --git a/src/AspNetCoreRateLimit/Core/IpRateLimitProcessor.cs b/src/AspNetCoreRateLimit/Core/IpRateLimitProcessor.cs
@@ -13,8 +13,9 @@ public class IpRateLimitProcessor : RateLimitProcessor, IRateLimitProcessor
         public IpRateLimitProcessor(
            IpRateLimitOptions options,
            IRateLimitCounterStore counterStore,
-           IIpPolicyStore policyStore)
-        : base(options, counterStore)
+           IIpPolicyStore policyStore,
+           IRateLimitConfiguration config)
+        : base(options, counterStore, new IpCounterKeyBuilder(options), config)
         {
             _options = options;
             _policyStore = policyStore;
@@ -51,10 +52,5 @@ public override bool IsWhitelisted(ClientRequestIdentity requestIdentity)
 
             return base.IsWhitelisted(requestIdentity);
         }
-
-        protected override string GetCounterKey(ClientRequestIdentity requestIdentity, RateLimitRule rule)
-        {
-            return $"{_options.RateLimitCounterPrefix}_{requestIdentity.ClientIp}_{rule.Period}";
-        }
     }
 }
diff --git a/src/AspNetCoreRateLimit/Core/RateLimitProcessor.cs b/src/AspNetCoreRateLimit/Core/RateLimitProcessor.cs
@@ -11,13 +11,19 @@ public abstract class RateLimitProcessor
     {
         private readonly RateLimitOptions _options;
         private readonly IRateLimitCounterStore _counterStore;
+        private readonly ICounterKeyBuilder _counterKeyBuilder;
+        private readonly IRateLimitConfiguration _config;
 
         protected RateLimitProcessor(
            RateLimitOptions options,
-           IRateLimitCounterStore counterStore)
+           IRateLimitCounterStore counterStore,
+           ICounterKeyBuilder counterKeyBuilder,
+           IRateLimitConfiguration config)
         {
             _options = options;
             _counterStore = counterStore;
+            _counterKeyBuilder = counterKeyBuilder;
+            _config = config;
         }
 
         private static readonly SemaphoreSlim Semaphore = new SemaphoreSlim(1);
@@ -47,7 +53,7 @@ public async Task<RateLimitCounter> ProcessRequestAsync(ClientRequestIdentity re
                 TotalRequests = 1
             };
 
-            var counterId = ComputeCounterKey(requestIdentity, rule);
+            var counterId = BuildCounterKey(requestIdentity, rule);
 
             // serial reads and writes
             await Semaphore.WaitAsync(cancellationToken);
@@ -87,7 +93,7 @@ public async Task<RateLimitCounter> ProcessRequestAsync(ClientRequestIdentity re
         public async Task<RateLimitHeaders> GetRateLimitHeadersAsync(ClientRequestIdentity requestIdentity, RateLimitRule rule, CancellationToken cancellationToken = default)
         {
             var headers = new RateLimitHeaders();
-            var counterId = ComputeCounterKey(requestIdentity, rule);
+            var counterId = BuildCounterKey(requestIdentity, rule);
             var entry = await _counterStore.GetAsync(counterId, cancellationToken);
 
             long remaining;
@@ -111,18 +117,13 @@ public async Task<RateLimitHeaders> GetRateLimitHeadersAsync(ClientRequestIdenti
             return headers;
         }
 
-        protected abstract string GetCounterKey(ClientRequestIdentity requestIdentity, RateLimitRule rule);
-
-        protected string ComputeCounterKey(ClientRequestIdentity requestIdentity, RateLimitRule rule)
+        protected virtual string BuildCounterKey(ClientRequestIdentity requestIdentity, RateLimitRule rule)
         {
-            var key = GetCounterKey(requestIdentity, rule);
+            var key = _counterKeyBuilder.Build(requestIdentity, rule);
 
-            if (_options.EnableEndpointRateLimiting)
+            if (_options.EnableEndpointRateLimiting && _config.EndpointCounterKeyBuilder != null)
             {
-                key += $"_{requestIdentity.HttpVerb}_{requestIdentity.Path}";
-
-                // TODO: consider using the rule endpoint as key, this will allow to rate limit /api/values/1 and api/values/2 under same counter
-                //key += $"_{rule.Endpoint}";
+                key += _config.EndpointCounterKeyBuilder.Build(requestIdentity, rule);
             }
 
             var idBytes = System.Text.Encoding.UTF8.GetBytes(key);
@@ -137,7 +138,7 @@ protected string ComputeCounterKey(ClientRequestIdentity requestIdentity, RateLi
             return BitConverter.ToString(hashBytes).Replace("-", string.Empty);
         }
 
-        protected List<RateLimitRule> GetMatchingRules(ClientRequestIdentity identity, List<RateLimitRule> rules)
+        protected virtual List<RateLimitRule> GetMatchingRules(ClientRequestIdentity identity, List<RateLimitRule> rules)
         {
             var limits = new List<RateLimitRule>();
 

diff --git a/src/AspNetCoreRateLimit/CounterKeyBuilders/ClientCounterKeyBuilder.cs b/src/AspNetCoreRateLimit/CounterKeyBuilders/ClientCounterKeyBuilder.cs
@@ -0,0 +1,17 @@
+namespace AspNetCoreRateLimit
+{
+    public class ClientCounterKeyBuilder : ICounterKeyBuilder
+    {
+        private readonly ClientRateLimitOptions _options;
+
+        public ClientCounterKeyBuilder(ClientRateLimitOptions options)
+        {
+            _options = options;
+        }
+
+        public string Build(ClientRequestIdentity requestIdentity, RateLimitRule rule)
+        {
+            return $"{_options.RateLimitCounterPrefix}_{requestIdentity.ClientId}_{rule.Period}";
+        }
+    }
+}
diff --git a/src/AspNetCoreRateLimit/CounterKeyBuilders/EndpointCounterKeyBuilder.cs b/src/AspNetCoreRateLimit/CounterKeyBuilders/EndpointCounterKeyBuilder.cs
@@ -0,0 +1,11 @@
+namespace AspNetCoreRateLimit
+{
+    public class EndpointCounterKeyBuilder : ICounterKeyBuilder
+    {
+        public string Build(ClientRequestIdentity requestIdentity, RateLimitRule rule)
+        {
+            // This will allow to rate limit /api/values/1 and api/values/2 under same counter
+            return $"_{requestIdentity.HttpVerb}_{requestIdentity.Path}_{rule.Endpoint}";
+        }
+    }
+}
diff --git a/src/AspNetCoreRateLimit/CounterKeyBuilders/ICounterKeyBuilder.cs b/src/AspNetCoreRateLimit/CounterKeyBuilders/ICounterKeyBuilder.cs
@@ -0,0 +1,7 @@
+namespace AspNetCoreRateLimit
+{
+    public interface ICounterKeyBuilder
+    {
+        string Build(ClientRequestIdentity requestIdentity, RateLimitRule rule);
+    }
+}
diff --git a/src/AspNetCoreRateLimit/CounterKeyBuilders/IpCounterKeyBuilder.cs b/src/AspNetCoreRateLimit/CounterKeyBuilders/IpCounterKeyBuilder.cs
@@ -0,0 +1,17 @@
+namespace AspNetCoreRateLimit
+{
+    public class IpCounterKeyBuilder : ICounterKeyBuilder
+    {
+        private readonly IpRateLimitOptions _options;
+
+        public IpCounterKeyBuilder(IpRateLimitOptions options)
+        {
+            _options = options;
+        }
+
+        public string Build(ClientRequestIdentity requestIdentity, RateLimitRule rule)
+        {
+            return $"{_options.RateLimitCounterPrefix}_{requestIdentity.ClientIp}_{rule.Period}";
+        }
+    }
+}
diff --git a/src/AspNetCoreRateLimit/CounterKeyBuilders/PathCounterKeyBuilder.cs b/src/AspNetCoreRateLimit/CounterKeyBuilders/PathCounterKeyBuilder.cs
@@ -0,0 +1,10 @@
+namespace AspNetCoreRateLimit
+{
+    public class PathCounterKeyBuilder : ICounterKeyBuilder
+    {
+        public string Build(ClientRequestIdentity requestIdentity, RateLimitRule rule)
+        {
+            return $"_{requestIdentity.HttpVerb}_{requestIdentity.Path}";
+        }
+    }
+}
diff --git a/src/AspNetCoreRateLimit/Middleware/ClientRateLimitMiddleware.cs b/src/AspNetCoreRateLimit/Middleware/ClientRateLimitMiddleware.cs
@@ -14,7 +14,7 @@ public ClientRateLimitMiddleware(RequestDelegate next,
             IClientPolicyStore policyStore,
             IRateLimitConfiguration config,
             ILogger<ClientRateLimitMiddleware> logger)
-        : base(next, options?.Value, new ClientRateLimitProcessor(options?.Value, counterStore, policyStore), config)
+        : base(next, options?.Value, new ClientRateLimitProcessor(options?.Value, counterStore, policyStore, config), config)
         {
             _logger = logger;
         }

diff --git a/src/AspNetCoreRateLimit/Middleware/IRateLimitConfiguration.cs b/src/AspNetCoreRateLimit/Middleware/IRateLimitConfiguration.cs
@@ -4,10 +4,10 @@ namespace AspNetCoreRateLimit
 {
     public interface IRateLimitConfiguration
     {
-        //bool Enabled { get; set; }
-
         IList<IClientResolveContributor> ClientResolvers { get; }
 
         IList<IIpResolveContributor> IpResolvers { get; }
+
+        ICounterKeyBuilder EndpointCounterKeyBuilder { get; }
     }
-}
+}
diff --git a/src/AspNetCoreRateLimit/Middleware/IpRateLimitMiddleware.cs b/src/AspNetCoreRateLimit/Middleware/IpRateLimitMiddleware.cs
@@ -14,7 +14,7 @@ public IpRateLimitMiddleware(RequestDelegate next,
             IIpPolicyStore policyStore,
             IRateLimitConfiguration config,
             ILogger<IpRateLimitMiddleware> logger)
-        : base(next, options?.Value, new IpRateLimitProcessor(options?.Value, counterStore, policyStore), config)
+        : base(next, options?.Value, new IpRateLimitProcessor(options?.Value, counterStore, policyStore, config), config)
 
         {
             _logger = logger;

diff --git a/src/AspNetCoreRateLimit/Middleware/RateLimitConfiguration.cs b/src/AspNetCoreRateLimit/Middleware/RateLimitConfiguration.cs
@@ -6,33 +6,44 @@ namespace AspNetCoreRateLimit
 {
     public class RateLimitConfiguration : IRateLimitConfiguration
     {
-        //public bool Enabled { get => throw new NotImplementedException(); set => throw new NotImplementedException(); }
+        public IList<IClientResolveContributor> ClientResolvers { get; } = new List<IClientResolveContributor>();
+        public IList<IIpResolveContributor> IpResolvers { get; } = new List<IIpResolveContributor>();
 
-        public IList<IClientResolveContributor> ClientResolvers { get; }
-
-        public IList<IIpResolveContributor> IpResolvers { get; }
+        public virtual ICounterKeyBuilder EndpointCounterKeyBuilder { get; } = new PathCounterKeyBuilder();
 
         public RateLimitConfiguration(
             IHttpContextAccessor httpContextAccessor,
             IOptions<IpRateLimitOptions> ipOptions,
             IOptions<ClientRateLimitOptions> clientOptions)
         {
+            IpRateLimitOptions = ipOptions?.Value;
+            ClientRateLimitOptions = clientOptions?.Value;
+            HttpContextAccessor = httpContextAccessor;
+
             ClientResolvers = new List<IClientResolveContributor>();
+            IpResolvers = new List<IIpResolveContributor>();
+
+            RegisterResolvers();
+        }
 
-            if (!string.IsNullOrEmpty(clientOptions?.Value.ClientIdHeader))
+        protected readonly IpRateLimitOptions IpRateLimitOptions;
+        protected readonly ClientRateLimitOptions ClientRateLimitOptions;
+        protected readonly IHttpContextAccessor HttpContextAccessor;
+
+        protected virtual void RegisterResolvers()
+        {
+            if (!string.IsNullOrEmpty(ClientRateLimitOptions?.ClientIdHeader))
             {
-                ClientResolvers.Add(new ClientHeaderResolveContributor(httpContextAccessor, clientOptions.Value.ClientIdHeader));
+                ClientResolvers.Add(new ClientHeaderResolveContributor(HttpContextAccessor, ClientRateLimitOptions.ClientIdHeader));
             }
 
-            IpResolvers = new List<IIpResolveContributor>();
-
             // the contributors are resolved in the order of their collection index
-            if (!string.IsNullOrEmpty(ipOptions?.Value.RealIpHeader))
+            if (!string.IsNullOrEmpty(IpRateLimitOptions?.RealIpHeader))
             {
-                IpResolvers.Add(new IpHeaderResolveContributor(httpContextAccessor, ipOptions.Value.RealIpHeader));
+                IpResolvers.Add(new IpHeaderResolveContributor(HttpContextAccessor, IpRateLimitOptions.RealIpHeader));
             }
 
-            IpResolvers.Add(new IpConnectionResolveContributor(httpContextAccessor));
+            IpResolvers.Add(new IpConnectionResolveContributor(HttpContextAccessor));
         }
     }
 }