From 5109c164e3f640d4d96ef488fff671cb1af38e23 Mon Sep 17 00:00:00 2001
From: Elliot Coyle <elliotcoyle@ekmsystems.co.uk>
Date: Wed, 13 Nov 2019 10:06:00 +0000
Subject: [PATCH] Switching to use DI to push the strategy into the middleware

---
 .../RateLimiterMiddleware.cs                  |  4 +--
 .../RateLimiterMiddlewareExtensions.cs        | 31 ++++++++++++++----
 KnowYourLimits.sln                            |  7 ++++
 README.md                                     | 32 +++++++++++++------
 4 files changed, 55 insertions(+), 19 deletions(-)

diff --git a/KnowYourLimits.AspNetCore/RateLimiterMiddleware.cs b/KnowYourLimits.AspNetCore/RateLimiterMiddleware.cs
index 17158b6..203900d 100644
--- a/KnowYourLimits.AspNetCore/RateLimiterMiddleware.cs
+++ b/KnowYourLimits.AspNetCore/RateLimiterMiddleware.cs
@@ -12,9 +12,9 @@ public class RateLimiterMiddleware<TClientIdentity>
         where TClientIdentity : IClientIdentity, new()
     {
         private readonly RequestDelegate _next;
-        private readonly IRateLimitStrategy<IClientIdentity> _rateLimitStrategy;
+        private readonly IRateLimitStrategy<TClientIdentity> _rateLimitStrategy;
 
-        public RateLimiterMiddleware(RequestDelegate next, IRateLimitStrategy<IClientIdentity> rateLimitStrategy)
+        public RateLimiterMiddleware(RequestDelegate next, IRateLimitStrategy<TClientIdentity> rateLimitStrategy)
         {
             _next = next;
             _rateLimitStrategy = rateLimitStrategy;
diff --git a/KnowYourLimits.AspNetCore/RateLimiterMiddlewareExtensions.cs b/KnowYourLimits.AspNetCore/RateLimiterMiddlewareExtensions.cs
index 086e754..0beea16 100644
--- a/KnowYourLimits.AspNetCore/RateLimiterMiddlewareExtensions.cs
+++ b/KnowYourLimits.AspNetCore/RateLimiterMiddlewareExtensions.cs
@@ -1,23 +1,40 @@
 using KnowYourLimits.Identity;
 using KnowYourLimits.Strategies;
+using KnowYourLimits.Strategies.LeakyBucket;
 using Microsoft.AspNetCore.Builder;
+using Microsoft.Extensions.DependencyInjection;
+
+// These are all used outside of this lib, and so are 'unused'
+// ReSharper disable UnusedMember.Global
 
 namespace KnowYourLimits.AspNetCore
 {
-    // ReSharper disable once UnusedMember.Global
     public static class RateLimiterMiddlewareExtensions
     {
-        // ReSharper disable once UnusedMember.Global
-        public static void UseRateLimiting<TClientIdentity>(this IApplicationBuilder applicationBuilder,
-            IRateLimitStrategy<TClientIdentity> rateLimitStrategy)
+        public static void UseRateLimiting<TClientIdentity>(this IApplicationBuilder applicationBuilder)
+            where TClientIdentity : IClientIdentity, new()
+        {
+            applicationBuilder.UseMiddleware<RateLimiterMiddleware<TClientIdentity>>();
+        }
+
+        public static void AddRateLimiting<TClientIdentity>(
+            this IServiceCollection serviceCollection,
+            IRateLimitStrategy<TClientIdentity> strategy)
             where TClientIdentity : IClientIdentity, new()
         {
-            if (rateLimitStrategy.IdentityProvider == null)
+            if (strategy.IdentityProvider == null)
             {
-                rateLimitStrategy.IdentityProvider = new IpClientIdentityProvider<TClientIdentity>();
+                strategy.IdentityProvider = new IpClientIdentityProvider<TClientIdentity>();
             }
 
-            applicationBuilder.UseMiddleware<RateLimiterMiddleware<TClientIdentity>>(rateLimitStrategy);
+            serviceCollection.AddSingleton(strategy);
+        }
+
+        public static void AddLeakyBucketRateLimiting(
+            this IServiceCollection serviceCollection,
+            LeakyBucketConfiguration config)
+        {
+            serviceCollection.AddRateLimiting(new LeakyBucketRateLimitStrategy(config));
         }
     }
 }
\ No newline at end of file
diff --git a/KnowYourLimits.sln b/KnowYourLimits.sln
index fb7524a..9f54b1e 100644
--- a/KnowYourLimits.sln
+++ b/KnowYourLimits.sln
@@ -9,6 +9,13 @@ Project("{FAE04EC0-301F-11D3-BF4B-00C04F79EFBC}") = "KnowYourLimits.AspNetCore",
 EndProject
 Project("{FAE04EC0-301F-11D3-BF4B-00C04F79EFBC}") = "KnowYourLimits.UnitTests", "KnowYourLimits.UnitTests\KnowYourLimits.UnitTests.csproj", "{D8B3BB5F-AB4F-4CF2-84C5-2439BB8546CB}"
 EndProject
+Project("{2150E333-8FDC-42A3-9474-1A3956D46DE8}") = "Documentation", "Documentation", "{AEB63280-2D5C-4A2C-BC2D-A2D7AC6F3871}"
+ProjectSection(SolutionItems) = preProject
+	README.md = README.md
+	CODE_OF_CONDUCT.md = CODE_OF_CONDUCT.md
+	LICENSE = LICENSE
+EndProjectSection
+EndProject
 Global
 	GlobalSection(SolutionConfigurationPlatforms) = preSolution
 		Debug|Any CPU = Debug|Any CPU
diff --git a/README.md b/README.md
index 77580a5..b6af61a 100644
--- a/README.md
+++ b/README.md
@@ -17,17 +17,29 @@ The middleware should be attached to the application as high in the order as pos
 For example, to configure the rate limiter to allow 4 requests per second, with a maximum of 100 requests, see below:
 
 ```cs
-var rateLimitingConfiguration = new LeakyBucketConfiguration
+public void ConfigureServices(IServiceCollection services)
 {
-  MaxRequests = 100,
-  LeakRate = TimeSpan.FromSeconds(1),
-  LeakAmount = 4,
-  IdentityProvider = new CustomIdentityProvider(), // If not set, defaults to using the remote address
-  EnableHeaders = true, // If true, a set of headers, documented below, will be returned on all responses describing the rate limits
-  HeaderPrefix = "X-MYORG-" // This will be prepended to all generated headers.
-};
-
-app.UseRateLimiting(new LeakyBucketRateLimitStrategy(rateLimitingConfiguration));
+    ...
+    var rateLimitingConfiguration = new LeakyBucketConfiguration
+    {
+        MaxRequests = 100,
+        LeakRate = TimeSpan.FromSeconds(1),
+        LeakAmount = 4,
+        IdentityProvider = new CustomIdentityProvider(), // If not set, defaults to using the remote address
+        EnableHeaders = true, // If true, a set of headers, documented below, will be returned on all responses describing the rate limits
+        HeaderPrefix = "X-MYORG-" // This will be prepended to all generated headers.
+    };
+    services.AddLeakyBucketRateLimiting(rateLimitingConfiguration);
+    ...
+}
+
+
+public void Configure(IApplicationBuilder app, IWebHostEnvironment env, IApiVersionDescriptionProvider provider)
+{
+    ...
+    app.UseRateLimiting<LeakyBucketClientIdentity>();
+    ...
+}
 ```
 
 By default client requests will be identified using the remote address of the request. To implement a custom identity provider, implement the `IClientIdentityProvider` interface and pass it in to the configuration.