RateLimitingMiddleware updates (#43053)

* RateLimitingMiddleware updates

* Feedback

* Feedback, API

* Fix sample

* Feedback

* Add tests

* Fixup

* Fixup2

* disambiguate

* Feedback

Author: wtgodbe

src/Middleware/RateLimiting/samples/RateLimitingSample/Program.cs

@@ -16,35 +16,38 @@
 // Inject an ILogger<SampleRateLimiterPolicy>
 builder.Services.AddLogging();
 
-var app = builder.Build();
-
 var todoName = "todoPolicy";
 var completeName = "completePolicy";
 var helloName = "helloPolicy";
 
-// Define endpoint limiters and a global limiter.
-var options = new RateLimiterOptions()
-        .AddTokenBucketLimiter(todoName, new TokenBucketRateLimiterOptions
-        {
-            TokenLimit = 1,
-            QueueProcessingOrder = QueueProcessingOrder.OldestFirst,
-            QueueLimit = 1,
-            ReplenishmentPeriod = TimeSpan.FromSeconds(10),
-            TokensPerPeriod = 1
-        })
-        .AddPolicy<string>(completeName, new SampleRateLimiterPolicy(NullLogger<SampleRateLimiterPolicy>.Instance))
-        .AddPolicy<string, SampleRateLimiterPolicy>(helloName);
-// The global limiter will be a concurrency limiter with a max permit count of 10 and a queue depth of 5.
-options.GlobalLimiter = PartitionedRateLimiter.Create<HttpContext, string>(context =>
+builder.Services.AddRateLimiter(options =>
+{
+    // Define endpoint limiters and a global limiter.
+    options.AddTokenBucketLimiter(todoName, options =>
+    {
+        options.TokenLimit = 1;
+        options.QueueProcessingOrder = QueueProcessingOrder.OldestFirst;
+        options.QueueLimit = 1;
+        options.ReplenishmentPeriod = TimeSpan.FromSeconds(10);
+        options.TokensPerPeriod = 1;
+    })
+    .AddPolicy<string>(completeName, new SampleRateLimiterPolicy(NullLogger<SampleRateLimiterPolicy>.Instance))
+    .AddPolicy<string, SampleRateLimiterPolicy>(helloName);
+    // The global limiter will be a concurrency limiter with a max permit count of 10 and a queue depth of 5.
+    options.GlobalLimiter = PartitionedRateLimiter.Create<HttpContext, string>(context =>
+    {
+        return RateLimitPartition.GetConcurrencyLimiter<string>("globalLimiter", key => new ConcurrencyLimiterOptions
         {
-            return RateLimitPartition.GetConcurrencyLimiter<string>("globalLimiter", key => new ConcurrencyLimiterOptions
-            {
-                PermitLimit = 10,
-                QueueProcessingOrder = QueueProcessingOrder.NewestFirst,
-                QueueLimit = 5
-            });
+            PermitLimit = 10,
+            QueueProcessingOrder = QueueProcessingOrder.NewestFirst,
+            QueueLimit = 5
         });
-app.UseRateLimiter(options);
+    });
+});
+
+var app = builder.Build();
+
+app.UseRateLimiter();
 
 // The limiter on this endpoint allows 1 request every 5 seconds
 app.MapGet("/", () => "Hello World!").RequireRateLimiting(helloName);

src/Middleware/RateLimiting/src/DefaultKeyType.cs

@@ -5,14 +5,14 @@ namespace Microsoft.AspNetCore.RateLimiting;
 
 internal struct DefaultKeyType
 {
-    public DefaultKeyType(string policyName, object? key, object? factory = null)
+    public DefaultKeyType(string? policyName, object? key, object? factory = null)
     {
         PolicyName = policyName;
         Key = key;
         Factory = factory;
     }
 
-    public string PolicyName { get; }
+    public string? PolicyName { get; }
 
     public object? Key { get; }
 

src/Middleware/RateLimiting/src/DisableRateLimitingAttribute.cs

@@ -0,0 +1,16 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+namespace Microsoft.AspNetCore.RateLimiting;
+
+/// <summary>
+/// Metadata that disables request rate limiting on an endpoint.
+/// </summary>
+/// <remarks>
+/// Completely disables the rate limiting middleware from applying to this endpoint.
+/// </remarks>
+[AttributeUsage(AttributeTargets.Class | AttributeTargets.Method, AllowMultiple = false, Inherited = true)]
+public sealed class DisableRateLimitingAttribute : Attribute
+{
+    internal static DisableRateLimitingAttribute Instance { get; } = new DisableRateLimitingAttribute();
+}

src/Middleware/RateLimiting/src/EnableRateLimitingAttribute.cs

@@ -0,0 +1,41 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+namespace Microsoft.AspNetCore.RateLimiting;
+
+/// <summary>
+/// Metadata that provides endpoint-specific request rate limiting.
+/// </summary>
+/// <remarks>
+/// Replaces any policies currently applied to the endpoint.
+/// The global limiter will still run on endpoints with this attribute applied.
+/// </remarks>
+[AttributeUsage(AttributeTargets.Class | AttributeTargets.Method, AllowMultiple = false, Inherited = true)]
+public sealed class EnableRateLimitingAttribute : Attribute
+{
+    /// <summary>
+    /// Creates a new instance of <see cref="EnableRateLimitingAttribute"/> using the specified policy.
+    /// </summary>
+    /// <param name="policyName">The name of the policy which needs to be applied.</param>
+    public EnableRateLimitingAttribute(string policyName)
+    {
+        ArgumentNullException.ThrowIfNull(policyName);
+
+        PolicyName = policyName;
+    }
+
+    internal EnableRateLimitingAttribute(DefaultRateLimiterPolicy policy)
+    {
+        Policy = policy;
+    }
+
+    /// <summary>
+    /// The name of the policy which needs to be applied.
+    /// </summary>
+    public string? PolicyName { get; }
+
+    /// <summary>
+    /// The policy which needs to be applied, if present.
+    /// </summary>
+    internal DefaultRateLimiterPolicy? Policy { get; }
+}

src/Middleware/RateLimiting/src/IRateLimiterMetadata.cs

@@ -18,4 +18,7 @@
     <Compile Include="$(SharedSourceRoot)ValueStopwatch\*.cs" />
   </ItemGroup>
 
+  <ItemGroup>
+    <InternalsVisibleTo Include="Microsoft.AspNetCore.RateLimiting.Tests" />
+  </ItemGroup>
 </Project>

src/Middleware/RateLimiting/src/Microsoft.AspNetCore.RateLimiting.csproj

@@ -1,5 +1,11 @@
 Microsoft.AspNetCore.Builder.RateLimiterApplicationBuilderExtensions
 Microsoft.AspNetCore.Builder.RateLimiterEndpointConventionBuilderExtensions
+Microsoft.AspNetCore.Builder.RateLimiterServiceCollectionExtensions
+Microsoft.AspNetCore.RateLimiting.DisableRateLimitingAttribute
+Microsoft.AspNetCore.RateLimiting.DisableRateLimitingAttribute.DisableRateLimitingAttribute() -> void
+Microsoft.AspNetCore.RateLimiting.EnableRateLimitingAttribute
+Microsoft.AspNetCore.RateLimiting.EnableRateLimitingAttribute.EnableRateLimitingAttribute(string! policyName) -> void
+Microsoft.AspNetCore.RateLimiting.EnableRateLimitingAttribute.PolicyName.get -> string?
 Microsoft.AspNetCore.RateLimiting.IRateLimiterPolicy<TPartitionKey>
 Microsoft.AspNetCore.RateLimiting.IRateLimiterPolicy<TPartitionKey>.GetPartition(Microsoft.AspNetCore.Http.HttpContext! httpContext) -> System.Threading.RateLimiting.RateLimitPartition<TPartitionKey>
 Microsoft.AspNetCore.RateLimiting.IRateLimiterPolicy<TPartitionKey>.OnRejected.get -> System.Func<Microsoft.AspNetCore.RateLimiting.OnRejectedContext!, System.Threading.CancellationToken, System.Threading.Tasks.ValueTask>?
@@ -23,9 +29,11 @@ Microsoft.AspNetCore.RateLimiting.RateLimiterOptions.RejectionStatusCode.set ->
 Microsoft.AspNetCore.RateLimiting.RateLimiterOptionsExtensions
 static Microsoft.AspNetCore.Builder.RateLimiterApplicationBuilderExtensions.UseRateLimiter(this Microsoft.AspNetCore.Builder.IApplicationBuilder! app) -> Microsoft.AspNetCore.Builder.IApplicationBuilder!
 static Microsoft.AspNetCore.Builder.RateLimiterApplicationBuilderExtensions.UseRateLimiter(this Microsoft.AspNetCore.Builder.IApplicationBuilder! app, Microsoft.AspNetCore.RateLimiting.RateLimiterOptions! options) -> Microsoft.AspNetCore.Builder.IApplicationBuilder!
+static Microsoft.AspNetCore.Builder.RateLimiterEndpointConventionBuilderExtensions.DisableRateLimiting<TBuilder>(this TBuilder builder) -> TBuilder
+static Microsoft.AspNetCore.Builder.RateLimiterEndpointConventionBuilderExtensions.RequireRateLimiting<TBuilder, TPartitionKey>(this TBuilder builder, Microsoft.AspNetCore.RateLimiting.IRateLimiterPolicy<TPartitionKey>! policy) -> TBuilder
 static Microsoft.AspNetCore.Builder.RateLimiterEndpointConventionBuilderExtensions.RequireRateLimiting<TBuilder>(this TBuilder builder, string! policyName) -> TBuilder
-static Microsoft.AspNetCore.RateLimiting.RateLimiterOptionsExtensions.AddConcurrencyLimiter(this Microsoft.AspNetCore.RateLimiting.RateLimiterOptions! options, string! policyName, System.Threading.RateLimiting.ConcurrencyLimiterOptions! concurrencyLimiterOptions) -> Microsoft.AspNetCore.RateLimiting.RateLimiterOptions!
-static Microsoft.AspNetCore.RateLimiting.RateLimiterOptionsExtensions.AddFixedWindowLimiter(this Microsoft.AspNetCore.RateLimiting.RateLimiterOptions! options, string! policyName, System.Threading.RateLimiting.FixedWindowRateLimiterOptions! fixedWindowRateLimiterOptions) -> Microsoft.AspNetCore.RateLimiting.RateLimiterOptions!
-static Microsoft.AspNetCore.RateLimiting.RateLimiterOptionsExtensions.AddNoLimiter(this Microsoft.AspNetCore.RateLimiting.RateLimiterOptions! options, string! policyName) -> Microsoft.AspNetCore.RateLimiting.RateLimiterOptions!
-static Microsoft.AspNetCore.RateLimiting.RateLimiterOptionsExtensions.AddSlidingWindowLimiter(this Microsoft.AspNetCore.RateLimiting.RateLimiterOptions! options, string! policyName, System.Threading.RateLimiting.SlidingWindowRateLimiterOptions! slidingWindowRateLimiterOptions) -> Microsoft.AspNetCore.RateLimiting.RateLimiterOptions!
-static Microsoft.AspNetCore.RateLimiting.RateLimiterOptionsExtensions.AddTokenBucketLimiter(this Microsoft.AspNetCore.RateLimiting.RateLimiterOptions! options, string! policyName, System.Threading.RateLimiting.TokenBucketRateLimiterOptions! tokenBucketRateLimiterOptions) -> Microsoft.AspNetCore.RateLimiting.RateLimiterOptions!
+static Microsoft.AspNetCore.Builder.RateLimiterServiceCollectionExtensions.AddRateLimiter(this Microsoft.Extensions.DependencyInjection.IServiceCollection! services, System.Action<Microsoft.AspNetCore.RateLimiting.RateLimiterOptions!>! configureOptions) -> Microsoft.Extensions.DependencyInjection.IServiceCollection!
+static Microsoft.AspNetCore.RateLimiting.RateLimiterOptionsExtensions.AddConcurrencyLimiter(this Microsoft.AspNetCore.RateLimiting.RateLimiterOptions! options, string! policyName, System.Action<System.Threading.RateLimiting.ConcurrencyLimiterOptions!>! configureOptions) -> Microsoft.AspNetCore.RateLimiting.RateLimiterOptions!
+static Microsoft.AspNetCore.RateLimiting.RateLimiterOptionsExtensions.AddFixedWindowLimiter(this Microsoft.AspNetCore.RateLimiting.RateLimiterOptions! options, string! policyName, System.Action<System.Threading.RateLimiting.FixedWindowRateLimiterOptions!>! configureOptions) -> Microsoft.AspNetCore.RateLimiting.RateLimiterOptions!
+static Microsoft.AspNetCore.RateLimiting.RateLimiterOptionsExtensions.AddSlidingWindowLimiter(this Microsoft.AspNetCore.RateLimiting.RateLimiterOptions! options, string! policyName, System.Action<System.Threading.RateLimiting.SlidingWindowRateLimiterOptions!>! configureOptions) -> Microsoft.AspNetCore.RateLimiting.RateLimiterOptions!
+static Microsoft.AspNetCore.RateLimiting.RateLimiterOptionsExtensions.AddTokenBucketLimiter(this Microsoft.AspNetCore.RateLimiting.RateLimiterOptions! options, string! policyName, System.Action<System.Threading.RateLimiting.TokenBucketRateLimiterOptions!>! configureOptions) -> Microsoft.AspNetCore.RateLimiting.RateLimiterOptions!

src/Middleware/RateLimiting/src/PublicAPI.Unshipped.txt

@@ -24,7 +24,44 @@ public static TBuilder RequireRateLimiting<TBuilder>(this TBuilder builder, stri
 
         builder.Add(endpointBuilder =>
         {
-            endpointBuilder.Metadata.Add(new RateLimiterMetadata(policyName));
+            endpointBuilder.Metadata.Add(new EnableRateLimitingAttribute(policyName));
+        });
+
+        return builder;
+    }
+
+    /// <summary>
+    /// Adds the specified rate limiting policy to the endpoint(s).
+    /// </summary>
+    /// <param name="builder">The endpoint convention builder.</param>
+    /// <param name="policy">The rate limiting policy to add to the endpoint.</param>
+    /// <returns>The original convention builder parameter.</returns>
+    public static TBuilder RequireRateLimiting<TBuilder, TPartitionKey>(this TBuilder builder, IRateLimiterPolicy<TPartitionKey> policy) where TBuilder : IEndpointConventionBuilder
+    {
+        ArgumentNullException.ThrowIfNull(builder);
+
+        ArgumentNullException.ThrowIfNull(policy);
+
+        builder.Add(endpointBuilder =>
+        {
+            endpointBuilder.Metadata.Add(new EnableRateLimitingAttribute(new DefaultRateLimiterPolicy(RateLimiterOptions.ConvertPartitioner<TPartitionKey>(null, policy.GetPartition), policy.OnRejected)));
+        });
+        return builder;
+    }
+
+    /// <summary>
+    /// Disables rate limiting on the endpoint(s).
+    /// </summary>
+    /// <param name="builder">The endpoint convention builder.</param>
+    /// <returns>The original convention builder parameter.</returns>
+    /// <remarks>Will skip both the global limiter, and any endpoint-specific limiters that apply to the endpoint(s).</remarks>
+    public static TBuilder DisableRateLimiting<TBuilder>(this TBuilder builder) where TBuilder : IEndpointConventionBuilder
+    {
+        ArgumentNullException.ThrowIfNull(builder);
+
+        builder.Add(endpointBuilder =>
+        {
+            endpointBuilder.Metadata.Add(DisableRateLimitingAttribute.Instance);
         });
 
         return builder;

src/Middleware/RateLimiting/src/RateLimiterEndpointConventionBuilderExtensions.cs

@@ -107,7 +107,7 @@ public RateLimiterOptions AddPolicy<TPartitionKey>(string policyName, IRateLimit
     }
 
     // Converts a Partition<TKey> to a Partition<DefaultKeyType<TKey>> to prevent accidental collisions with the keys we create in the the RateLimiterOptionsExtensions.
-    private static Func<HttpContext, RateLimitPartition<DefaultKeyType>> ConvertPartitioner<TPartitionKey>(string policyName, Func<HttpContext, RateLimitPartition<TPartitionKey>> partitioner)
+    internal static Func<HttpContext, RateLimitPartition<DefaultKeyType>> ConvertPartitioner<TPartitionKey>(string? policyName, Func<HttpContext, RateLimitPartition<TPartitionKey>> partitioner)
     {
         return context =>
         {