Some feedback

dotnet · May 27, 2021 · b544026 · b544026
1 parent 191575d
commit b544026
Show file tree

Hide file tree

Showing 6 changed files with 74 additions and 30 deletions.
diff --git a/src/Middleware/RequestLimiter/RequestLimiter.slnf b/src/Middleware/RequestLimiter/RequestLimiter.slnf
@@ -5,19 +5,20 @@
       "src\\Hosting\\Abstractions\\src\\Microsoft.AspNetCore.Hosting.Abstractions.csproj",
       "src\\Hosting\\Hosting\\src\\Microsoft.AspNetCore.Hosting.csproj",
       "src\\Hosting\\Server.Abstractions\\src\\Microsoft.AspNetCore.Hosting.Server.Abstractions.csproj",
+      "src\\Http\\Headers\\src\\Microsoft.Net.Http.Headers.csproj",
       "src\\Http\\Http.Abstractions\\src\\Microsoft.AspNetCore.Http.Abstractions.csproj",
       "src\\Http\\Http.Extensions\\src\\Microsoft.AspNetCore.Http.Extensions.csproj",
       "src\\Http\\Http.Features\\src\\Microsoft.AspNetCore.Http.Features.csproj",
+      "src\\Http\\Http\\src\\Microsoft.AspNetCore.Http.csproj",
       "src\\Http\\WebUtilities\\src\\Microsoft.AspNetCore.WebUtilities.csproj",
+      "src\\Middleware\\HttpsPolicy\\src\\Microsoft.AspNetCore.HttpsPolicy.csproj",
+      "src\\Middleware\\RequestLimiter\\sample\\RequestLimiterSample.csproj",
+      "src\\Middleware\\RequestLimiter\\src\\Microsoft.AspNetCore.RequestLimiter.csproj",
+      "src\\ResourceLimits\\src\\System.Threading.ResourceLimits.csproj",
       "src\\Servers\\Connections.Abstractions\\src\\Microsoft.AspNetCore.Connections.Abstractions.csproj",
       "src\\Servers\\Kestrel\\Core\\src\\Microsoft.AspNetCore.Server.Kestrel.Core.csproj",
       "src\\Servers\\Kestrel\\Kestrel\\src\\Microsoft.AspNetCore.Server.Kestrel.csproj",
-      "src\\Servers\\Kestrel\\Transport.Sockets\\src\\Microsoft.AspNetCore.Server.Kestrel.Transport.Sockets.csproj",
-      "src\\http\\Headers\\src\\Microsoft.Net.Http.Headers.csproj",
-      "src\\http\\http\\src\\Microsoft.AspNetCore.Http.csproj",
-      "src\\Middleware\\RequestLimiter\\sample\\RequestLimiterSample.csproj",
-      "src\\Middleware\\RequestLimiter\\src\\Microsoft.AspNetCore.RequestLimiter.csproj",
-      "src\\Middleware\\HttpsPolicy\\src\\Microsoft.AspNetCore.HttpsPolicy.csproj"
+      "src\\Servers\\Kestrel\\Transport.Sockets\\src\\Microsoft.AspNetCore.Server.Kestrel.Transport.Sockets.csproj"
     ]
   }
 }
diff --git a/src/Middleware/RequestLimiter/sample/Startup.cs b/src/Middleware/RequestLimiter/sample/Startup.cs
@@ -51,37 +51,39 @@ public void Configure(IApplicationBuilder app, IWebHostEnvironment env, ILogger<
                 {
                     await Task.Delay(5000);
                     await context.Response.WriteAsync("Default!");
-                }).EnforceLimit();
+                }).EnforceDefaultRequestLimit();
 
                 endpoints.MapGet("/instance", async context =>
                 {
                     await Task.Delay(5000);
                     await context.Response.WriteAsync("Hello World!");
-                }).EnforceLimit(new TokenBucketRateLimiter(2, 2));
+                }).EnforceRequestLimit(new TokenBucketRateLimiter(2, 2));
 
                 endpoints.MapGet("/concurrent", async context =>
                 {
                     await Task.Delay(5000);
                     await context.Response.WriteAsync("Wrote!");
-                }).EnforceConcurrencyLimit(concurrentRequests: 1);
+                }).EnforceRequestConcurrencyLimit(concurrentRequests: 1);
 
                 endpoints.MapGet("/adhoc", async context =>
                 {
                     await Task.Delay(5000);
                     await context.Response.WriteAsync("Tested!");
-                }).EnforceRateLimit(requestPerSecond: 2);
+                }).EnforceRequestRateLimit(requestPerSecond: 2);
 
                 endpoints.MapGet("/ipFromDI", async context =>
                 {
                     await Task.Delay(5000);
                     await context.Response.WriteAsync("IP limited!");
-                }).EnforceLimit("ipPolicy");
+                }).EnforceRequestLimitPolicy("ipPolicy");
 
                 endpoints.MapGet("/multiple", async context =>
                 {
                     await Task.Delay(5000);
                     await context.Response.WriteAsync("IP limited!");
-                }).EnforceLimit("concurrency").EnforceLimit("rate");
+                })
+                .EnforceRequestLimitPolicy("ipPolicy")
+                .EnforceRequestLimitPolicy("rate");
 
                 endpoints.MapControllerRoute(
                     name: "default",

diff --git a/src/Middleware/RequestLimiter/src/NoRequestLimitAttribute.cs b/src/Middleware/RequestLimiter/src/NoRequestLimitAttribute.cs
@@ -0,0 +1,11 @@
+// Copyright (c) .NET Foundation. All rights reserved.
+// Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information.
+
+using System;
+
+namespace Microsoft.AspNetCore.RequestLimiter
+{
+    // TODO: Double check ordering
+    [AttributeUsage(AttributeTargets.Class | AttributeTargets.Method, AllowMultiple = false, Inherited = true)]
+    public class NoRequestLimitAttribute : Attribute { }
+}
diff --git a/...estLimiter/src/RequestLimiterAttribute.cs → ...questLimiter/src/RequestLimitAttribute.cs b/...estLimiter/src/RequestLimiterAttribute.cs → ...questLimiter/src/RequestLimitAttribute.cs
diff --git a/src/Middleware/RequestLimiter/src/RequestLimiterEndpointExtensions.cs b/src/Middleware/RequestLimiter/src/RequestLimiterEndpointExtensions.cs
@@ -9,7 +9,7 @@ namespace Microsoft.AspNetCore.RequestLimiter
 {
     public static class RequestLimiterEndpointExtensions
     {
-        public static IEndpointConventionBuilder EnforceLimit(this IEndpointConventionBuilder builder)
+        public static IEndpointConventionBuilder EnforceDefaultRequestLimit(this IEndpointConventionBuilder builder)
         {
             builder.Add(endpointBuilder =>
             {
@@ -19,7 +19,7 @@ public static IEndpointConventionBuilder EnforceLimit(this IEndpointConventionBu
             return builder;
         }
 
-        public static IEndpointConventionBuilder EnforceLimit(this IEndpointConventionBuilder builder, string policyName)
+        public static IEndpointConventionBuilder EnforceRequestLimitPolicy(this IEndpointConventionBuilder builder, string policyName)
         {
             builder.Add(endpointBuilder =>
             {
@@ -29,7 +29,7 @@ public static IEndpointConventionBuilder EnforceLimit(this IEndpointConventionBu
             return builder;
         }
 
-        public static IEndpointConventionBuilder EnforceRateLimit(this IEndpointConventionBuilder builder, long requestPerSecond)
+        public static IEndpointConventionBuilder EnforceRequestRateLimit(this IEndpointConventionBuilder builder, long requestPerSecond)
         {
             builder.Add(endpointBuilder =>
             {
@@ -43,7 +43,7 @@ public static IEndpointConventionBuilder EnforceRateLimit(this IEndpointConventi
             return builder;
         }
 
-        public static IEndpointConventionBuilder EnforceConcurrencyLimit(this IEndpointConventionBuilder builder, long concurrentRequests)
+        public static IEndpointConventionBuilder EnforceRequestConcurrencyLimit(this IEndpointConventionBuilder builder, long concurrentRequests)
         {
             builder.Add(endpointBuilder =>
             {
@@ -56,10 +56,10 @@ public static IEndpointConventionBuilder EnforceConcurrencyLimit(this IEndpointC
             return builder;
         }
 
-        public static IEndpointConventionBuilder EnforceLimit(this IEndpointConventionBuilder builder, ResourceLimiter limiter)
-            => builder.EnforceLimit((HttpContextLimiter)limiter);
+        public static IEndpointConventionBuilder EnforceRequestLimit(this IEndpointConventionBuilder builder, ResourceLimiter limiter)
+            => builder.EnforceRequestLimit((HttpContextLimiter)limiter);
 
-        public static IEndpointConventionBuilder EnforceLimit(this IEndpointConventionBuilder builder, AggregatedResourceLimiter<HttpContext> limiter)
+        public static IEndpointConventionBuilder EnforceRequestLimit(this IEndpointConventionBuilder builder, AggregatedResourceLimiter<HttpContext> limiter)
         {
 
             builder.Add(endpointBuilder =>

diff --git a/src/Middleware/RequestLimiter/src/RequestLimiterMiddleware.cs b/src/Middleware/RequestLimiter/src/RequestLimiterMiddleware.cs
@@ -3,6 +3,7 @@
 
 using System;
 using System.Collections.Generic;
+using System.Diagnostics;
 using System.Threading.ResourceLimits;
 using System.Threading.Tasks;
 using Microsoft.AspNetCore.Http;
@@ -24,28 +25,30 @@ public RequestLimiterMiddleware(RequestDelegate next, ILoggerFactory loggerFacto
             _options = options.Value;
         }
 
-        public async Task Invoke(HttpContext context)
+        public Task Invoke(HttpContext context)
         {
             _logger.LogInformation("Resource limiting: " + context.Request.Path);
 
             var endpoint = context.GetEndpoint();
             var attributes = endpoint?.Metadata.GetOrderedMetadata<RequestLimitAttribute>();
+            var noLimitAttributes = endpoint?.Metadata.GetOrderedMetadata<NoRequestLimitAttribute>();
 
-            if (attributes == null)
+            if (attributes == null || noLimitAttributes != null)
             {
-                await _next.Invoke(context);
-                return;
+                return _next.Invoke(context);
             }
 
+            return InvokeAsync(context, attributes);
+        }
+        private async Task InvokeAsync(HttpContext context, IReadOnlyList<RequestLimitAttribute> attributes)
+        {
             var resourceLeases = new Stack<ResourceLease>();
             try
             {
                 foreach (var attribute in attributes)
                 {
-                    if (!string.IsNullOrEmpty(attribute.Policy) && attribute.Limiter != null)
-                    {
-                        throw new InvalidOperationException("Cannot specify both policy and limiter");
-                    }
+                    // At most one of Policy or Limiter can be set.
+                    Debug.Assert(string.IsNullOrEmpty(attribute.Policy) || attribute.Limiter == null);
 
                     if (string.IsNullOrEmpty(attribute.Policy) && attribute.Limiter == null)
                     {
@@ -97,11 +100,38 @@ public async Task Invoke(HttpContext context)
                 }
             };
         }
-
-        private async Task<bool> ApplyLimitAsync(AggregatedResourceLimiter<HttpContext> limiter, HttpContext context, Stack<ResourceLease> obtainedResources)
+        private Task<bool> ApplyLimitAsync(AggregatedResourceLimiter<HttpContext> limiter, HttpContext context, Stack<ResourceLease> obtainedResources)
         {
             _logger.LogInformation("Resource count: " + limiter.EstimatedCount(context));
-            var resourceLease = await limiter.WaitAsync(context);
+            var resourceLeaseTask = limiter.WaitAsync(context);
+
+            if (resourceLeaseTask.IsCompletedSuccessfully)
+            {
+                var resourceLease = resourceLeaseTask.Result;
+                if (!resourceLease.IsAcquired)
+                {
+                    _logger.LogInformation("Resource exhausted");
+                    context.Response.StatusCode = StatusCodes.Status429TooManyRequests;
+                    return OnRejectAsync(context, resourceLease);
+                }
+
+                _logger.LogInformation("Resource obtained");
+                obtainedResources.Push(resourceLease);
+                return Task.FromResult(true);
+            }
+
+            return ApplyLimitAsyncAwaited(resourceLeaseTask, context, obtainedResources);
+        }
+
+        private async Task<bool> OnRejectAsync(HttpContext context, ResourceLease resourceLease)
+        {
+            await _options.OnRejected(context, resourceLease);
+            return false;
+        }
+
+        private async Task<bool> ApplyLimitAsyncAwaited(ValueTask<ResourceLease> resourceLeaseTask, HttpContext context, Stack<ResourceLease> obtainedResources)
+        {
+            var resourceLease = await resourceLeaseTask;
             if (!resourceLease.IsAcquired)
             {
                 _logger.LogInformation("Resource exhausted");