#1305 RateLimiting headers

src/Ocelot/RateLimiting/Middleware/RateLimitingMiddleware.cs

@@ -1,4 +1,6 @@
 using Microsoft.AspNetCore.Http;
+using Microsoft.Extensions.DependencyInjection;
+using Microsoft.Net.Http.Headers;
 using Ocelot.Configuration;
 using Ocelot.Logging;
 using Ocelot.Middleware;
@@ -10,15 +12,18 @@ public class RateLimitingMiddleware : OcelotMiddleware
     {
         private readonly RequestDelegate _next;
         private readonly IRateLimiting _limiter;
+        private readonly IServiceProvider _container;
 
         public RateLimitingMiddleware(
             RequestDelegate next,
             IOcelotLoggerFactory factory,
-            IRateLimiting limiter)
+            IRateLimiting limiter,
+            IServiceProvider container)
             : base(factory.CreateLogger<RateLimitingMiddleware>())
         {
             _next = next;
             _limiter = limiter;
+            _container = container;
         }
 
         public async Task Invoke(HttpContext httpContext)
@@ -71,8 +76,13 @@ public async Task Invoke(HttpContext httpContext)
             //set X-Rate-Limit headers for the longest period
             if (!options.DisableRateLimitHeaders)
             {
-                var headers = _limiter.GetHeaders(httpContext, identity, options);
-                httpContext.Response.OnStarting(SetRateLimitHeaders, state: headers);
+                var httpContextAccessor = _container.GetService<IHttpContextAccessor>();
+                if (httpContextAccessor != null)
+                {
+                    var originalHttpContext = httpContextAccessor.HttpContext;
+                    var headers = _limiter.GetHeaders(originalHttpContext, identity, options);
+                    originalHttpContext.Response.OnStarting(SetRateLimitHeaders, state: headers);
+                }
             }
 
             await _next.Invoke(httpContext);
@@ -119,7 +129,8 @@ public virtual DownstreamResponse ReturnQuotaExceededResponse(HttpContext httpCo
 
             if (!option.DisableRateLimitHeaders)
             {
-                http.Headers.TryAddWithoutValidation("Retry-After", retryAfter); // in seconds, not date string
+                http.Headers.TryAddWithoutValidation(HeaderNames.RetryAfter, retryAfter);        // in seconds, not date string
+                httpContext.Response.Headers[HeaderNames.RetryAfter] = retryAfter;
             }
 
             return new DownstreamResponse(http);

test/Ocelot.AcceptanceTests/RateLimiting/ClientRateLimitingTests.cs

@@ -1,4 +1,5 @@
 using Microsoft.AspNetCore.Http;
+using Microsoft.Net.Http.Headers;
 using Ocelot.Configuration.File;
 
 namespace Ocelot.AcceptanceTests.RateLimiting;
@@ -130,6 +131,92 @@ public void StatusShouldNotBeEqualTo429_PeriodTimespanValueIsGreaterThanPeriod()
             .BDDfy();
     }
 
+    [Fact]
+    public void should_set_ratelimiting_headers_on_response_when_DisableRateLimitHeaders_set_to_false()
+    {
+        int port = PortFinder.GetRandomPort();
+
+        var configuration = CreateConfigurationForCheckingHeaders(port, false);
+
+        this.Given(x => x.GivenThereIsAServiceRunningOn($"http://localhost:{port}", "/api/ClientRateLimit"))
+            .And(x => GivenThereIsAConfiguration(configuration))
+            .And(x => GivenOcelotIsRunning())
+            .When(x => WhenIGetUrlOnTheApiGatewayMultipleTimesForRateLimit("/api/ClientRateLimit", 1))
+            .Then(x => ThenRateLimitingHeadersExistInResponse(true))
+            .And(x => ThenRetryAfterHeaderExistsInResponse(false))
+            .When(x => WhenIGetUrlOnTheApiGatewayMultipleTimesForRateLimit("/api/ClientRateLimit", 2))
+            .Then(x => ThenRateLimitingHeadersExistInResponse(true))
+            .And(x => ThenRetryAfterHeaderExistsInResponse(false))
+            .When(x => WhenIGetUrlOnTheApiGatewayMultipleTimesForRateLimit("/api/ClientRateLimit", 1))
+            .Then(x => ThenRateLimitingHeadersExistInResponse(false))
+            .And(x => ThenRetryAfterHeaderExistsInResponse(true))
+            .BDDfy();
+    }
+
+    [Fact]
+    public void should_not_set_ratelimiting_headers_on_response_when_DisableRateLimitHeaders_set_to_true()
+    {
+        int port = PortFinder.GetRandomPort();
+
+        var configuration = CreateConfigurationForCheckingHeaders(port, true);
+
+        this.Given(x => x.GivenThereIsAServiceRunningOn($"http://localhost:{port}", "/api/ClientRateLimit"))
+            .And(x => GivenThereIsAConfiguration(configuration))
+            .And(x => GivenOcelotIsRunning())
+            .When(x => WhenIGetUrlOnTheApiGatewayMultipleTimesForRateLimit("/api/ClientRateLimit", 1))
+            .Then(x => ThenRateLimitingHeadersExistInResponse(false))
+            .And(x => ThenRetryAfterHeaderExistsInResponse(false))
+            .When(x => WhenIGetUrlOnTheApiGatewayMultipleTimesForRateLimit("/api/ClientRateLimit", 2))
+            .Then(x => ThenRateLimitingHeadersExistInResponse(false))
+            .And(x => ThenRetryAfterHeaderExistsInResponse(false))
+            .When(x => WhenIGetUrlOnTheApiGatewayMultipleTimesForRateLimit("/api/ClientRateLimit", 1))
+            .Then(x => ThenRateLimitingHeadersExistInResponse(false))
+            .And(x => ThenRetryAfterHeaderExistsInResponse(false))
+            .BDDfy();
+    }
+
+    private FileConfiguration CreateConfigurationForCheckingHeaders(int port, bool disableRateLimitHeaders)
+    {
+        return new FileConfiguration
+        {
+            Routes = new List<FileRoute>
+                {
+                    new()
+                    {
+                        DownstreamPathTemplate = "/api/ClientRateLimit",
+                        DownstreamHostAndPorts = new List<FileHostAndPort>
+                        {
+                            new()
+                            {
+                                Host = "localhost",
+                                Port = port,
+                            },
+                        },
+                        DownstreamScheme = "http",
+                        UpstreamPathTemplate = "/api/ClientRateLimit",
+                        UpstreamHttpMethod = new List<string> { "Get" },
+                        RateLimitOptions = new FileRateLimitRule()
+                        {
+                            EnableRateLimiting = true,
+                            ClientWhitelist = new List<string>(),
+                            Limit = 3,
+                            Period = "100s",
+                            PeriodTimespan = 1000,
+                        },
+                    },
+                },
+            GlobalConfiguration = new FileGlobalConfiguration()
+            {
+                RateLimitOptions = new FileRateLimitOptions()
+                {
+                    DisableRateLimitHeaders = disableRateLimitHeaders,
+                    QuotaExceededMessage = "",
+                    HttpStatusCode = 428,
+                },
+            },
+        };
+    }
+
     private void GivenThereIsAServiceRunningOn(string baseUrl, string basePath)
     {
         _serviceHandler.GivenThereIsAServiceRunningOn(baseUrl, basePath, context =>
@@ -179,4 +266,16 @@ private static FileConfiguration GivenConfigurationWithRateLimitOptions(params F
         };
         return config;
     }
+
+    private void ThenRateLimitingHeadersExistInResponse(bool headersExist)
+    {
+        _response.Headers.Contains("X-Rate-Limit-Limit").ShouldBe(headersExist);
+        _response.Headers.Contains("X-Rate-Limit-Remaining").ShouldBe(headersExist);
+        _response.Headers.Contains("X-Rate-Limit-Reset").ShouldBe(headersExist);
+    }
+
+    private void ThenRetryAfterHeaderExistsInResponse(bool headersExist)
+    {
+        _response.Headers.Contains(HeaderNames.RetryAfter).ShouldBe(headersExist);
+    }
 }