Make RateLimitingMiddleware endpoint aware

eng/Versions.props

@@ -50,6 +50,8 @@
     <UsingToolNetFrameworkReferenceAssemblies Condition="'$(OS)' != 'Windows_NT'">true</UsingToolNetFrameworkReferenceAssemblies>
     <!-- Disable XLIFF tasks -->
     <UsingToolXliff>false</UsingToolXliff>
+    <!-- Use custom version of Roslyn Compiler -->
+    <UsingToolMicrosoftNetCompilers>true</UsingToolMicrosoftNetCompilers>
   </PropertyGroup>
   <!--
 
@@ -194,6 +196,10 @@
       framework in current .NET SDKs.
     -->
     <MicrosoftNETTestSdkVersion>17.1.0-preview-20211109-03</MicrosoftNETTestSdkVersion>
+    <!--
+      Use a compiler new enough to use required keyword
+    -->
+    <MicrosoftNetCompilersToolsetVersion>4.4.0-1.22315.13</MicrosoftNetCompilersToolsetVersion>
     <!--
       Versions of Microsoft.CodeAnalysis packages referenced by analyzers shipped in the SDK.
       This need to be pinned since they're used in 3.1 apps and need to be loadable in VS 2019.

src/Middleware/RateLimiting/src/AspNetKey.T.cs

@@ -0,0 +1,18 @@
+// 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;
+internal sealed class AspNetKey<TKey>: AspNetKey
+{
+    private readonly TKey _key;
+
+    public AspNetKey(TKey key)
+    {
+        _key = key;
+    }
+
+    public override object? GetKey()
+    {
+        return _key;
+    }
+}

src/Middleware/RateLimiting/src/AspNetKey.cs

@@ -0,0 +1,8 @@
+// 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;
+internal abstract class AspNetKey
+{
+    public abstract object? GetKey();
+}

src/Middleware/RateLimiting/src/AspNetKeyEqualityComparer.cs

@@ -0,0 +1,44 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+using System.Diagnostics.CodeAnalysis;
+
+namespace Microsoft.AspNetCore.RateLimiting;
+internal class AspNetKeyEqualityComparer : IEqualityComparer<AspNetKey>
+{
+    public bool Equals(AspNetKey? x, AspNetKey? y)
+    {
+        if (x == null && y == null)
+        {
+            return true;
+        }
+        else if (x == null || y == null)
+        {
+            return false;
+        }
+
+        var xKey = x.GetKey();
+        var yKey = y.GetKey();
+        if (xKey == null && yKey == null)
+        {
+            return true;
+        }
+        else if (xKey == null || yKey == null)
+        {
+            return false;
+        }
+
+        return xKey.Equals(yKey);
+    }
+
+    public int GetHashCode([DisallowNull] AspNetKey obj)
+    {
+        var key = obj.GetKey();
+        if (key is not null)
+        {
+            return key.GetHashCode();
+        }
+        // REVIEW - is this reasonable?
+        return default;
+    }
+}

src/Middleware/RateLimiting/src/AspNetLease.cs

@@ -0,0 +1,100 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+using System.Threading.RateLimiting;
+
+namespace Microsoft.AspNetCore.RateLimiting;
+internal sealed class AspNetLease : RateLimitLease
+{
+    private readonly RateLimitLease? _globalLease;
+    private readonly RateLimitLease _endpointLease;
+    private HashSet<string>? _metadataNames;
+
+    public AspNetLease(RateLimitLease? globalLease, RateLimitLease endpointLease)
+    {
+        _globalLease = globalLease;
+        _endpointLease = endpointLease;
+    }
+
+    public override bool IsAcquired => true;
+
+    public override IEnumerable<string> MetadataNames
+    {
+        get
+        {
+            if (_metadataNames is null)
+            {
+                _metadataNames = new HashSet<string>();
+                if (_globalLease is not null)
+                {
+                    foreach (string metadataName in _globalLease.MetadataNames)
+                    {
+                        _metadataNames.Add(metadataName);
+                    }
+                }
+                foreach (string metadataName in _endpointLease.MetadataNames)
+                {
+                    _metadataNames.Add(metadataName);
+                }
+            }
+            return _metadataNames;
+        }
+    }
+
+    public override bool TryGetMetadata(string metadataName, out object? metadata)
+    {
+        if (_globalLease is not null)
+        {
+            // Use the first metadata item of a given name, ignore duplicates, we can't reliably merge arbitrary metadata
+            // Creating an object[] if there are multiple of the same metadataName could work, but makes consumption of metadata messy
+            // and makes MetadataName.Create<T>(...) uses no longer work
+            if (_globalLease.TryGetMetadata(metadataName, out metadata))
+            {
+                return true;
+            }
+        }
+        if (_endpointLease.TryGetMetadata(metadataName, out metadata))
+        {
+            return true;
+        }
+
+        metadata = null;
+        return false;
+    }
+
+    protected override void Dispose(bool disposing)
+    {
+        List<Exception>? exceptions = null;
+        // Dispose endpoint lease first, then global lease (reverse order of when they were acquired)
+        // Avoids issues where dispose might unblock a queued acquire and then the acquire fails when acquiring the next limiter.
+        // When disposing in reverse order there wont be any issues of unblocking an acquire that affects acquires on limiters in the chain after it
+
+        try
+        {
+            _endpointLease.Dispose();
+        }
+        catch (Exception ex)
+        {
+            exceptions ??= new List<Exception>();
+            exceptions.Add(ex);
+        }
+
+        if (_globalLease is not null)
+        {
+            try
+            {
+                _globalLease.Dispose();
+            }
+            catch (Exception ex)
+            {
+                exceptions ??= new List<Exception>();
+                exceptions.Add(ex);
+            }
+        }
+
+        if (exceptions is not null)
+        {
+            throw new AggregateException(exceptions);
+        }
+    }
+}

src/Middleware/RateLimiting/src/AspNetPolicy.cs

@@ -0,0 +1,25 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+using System.Threading.RateLimiting;
+using Microsoft.AspNetCore.Http;
+
+namespace Microsoft.AspNetCore.RateLimiting;
+internal sealed class AspNetPolicy : IRateLimiterPolicy<AspNetKey>
+{
+    private readonly Func<HttpContext, RateLimitPartition<AspNetKey>> _partitioner;
+    private readonly Func<OnRejectedContext, CancellationToken, ValueTask>? _onRejected;
+
+    public AspNetPolicy(Func<HttpContext, RateLimitPartition<AspNetKey>> partitioner, Func<OnRejectedContext, CancellationToken, ValueTask>? onRejected)
+    {
+        _partitioner = partitioner;
+        _onRejected = onRejected;
+    }
+
+    public Func<OnRejectedContext, CancellationToken, ValueTask>? OnRejected => _onRejected;
+
+    public RateLimitPartition<AspNetKey> GetPartition(HttpContext httpContext)
+    {
+        return _partitioner(httpContext);
+    }
+}

src/Middleware/RateLimiting/src/IRateLimiterMetadata.cs

@@ -0,0 +1,15 @@
+// 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>
+/// An interface which can be used to identify a type which provides metadata needed for enabling request rate limiting support.
+/// </summary>
+internal interface IRateLimiterMetadata
+{
+    /// <summary>
+    /// The name of the limiter which needs to be applied.
+    /// </summary>
+    string Name { get; }
+}

src/Middleware/RateLimiting/src/IRateLimiterPolicy.cs

@@ -0,0 +1,24 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+using System.Threading.RateLimiting;
+using Microsoft.AspNetCore.Http;
+
+namespace Microsoft.AspNetCore.RateLimiting;
+
+/// <summary>
+/// An interface which is used to represent a RateLimiter policy.
+/// </summary>
+public interface IRateLimiterPolicy<TPartitionKey>
+{
+    /// <summary>
+    /// Gets the <see cref="Func{OnRejectedContext, CancellationToken, ValueTask}"/> that handles requests rejected by this middleware.
+    /// </summary>
+    public Func<OnRejectedContext, CancellationToken, ValueTask>? OnRejected { get; }
+
+    /// <summary>
+    /// Gets the <see cref="RateLimitPartition{TPartitionKey}"/> that applies to the given <see cref="HttpContext"/>.
+    /// </summary>
+    /// <param name="httpContext">The <see cref="HttpContext"/> to get the partition for.</param>
+    public RateLimitPartition<TPartitionKey> GetPartition(HttpContext httpContext);
+}

src/Middleware/RateLimiting/src/LeaseContext.cs

@@ -0,0 +1,17 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+using System.Threading.RateLimiting;
+
+namespace Microsoft.AspNetCore.RateLimiting;
+internal struct LeaseContext : IDisposable
+{
+    public Rejector? Rejector { get; init; }
+
+    public required RateLimitLease Lease { get; init; }
+
+    public void Dispose()
+    {
+        Lease.Dispose();
+    }
+}

src/Middleware/RateLimiting/src/OnRejectedContext.cs

@@ -0,0 +1,23 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+using System.Threading.RateLimiting;
+using Microsoft.AspNetCore.Http;
+
+namespace Microsoft.AspNetCore.RateLimiting;
+
+/// <summary>
+/// Holds state needed for the OnRejected callback in the RateLimitingMiddleware
+/// </summary>
+public sealed class OnRejectedContext
+{
+    /// <summary>
+    /// Gets or sets the <see cref="HttpContext"/> that the OnRejected callback will have access to
+    /// </summary>
+    public required HttpContext HttpContext { get; init; }
+
+    /// <summary>
+    /// Gets or sets the <see cref="RateLimitLease"/> that the OnRejected callback will have access to
+    /// </summary>
+    public required RateLimitLease Lease { get; init; }
+}

src/Middleware/RateLimiting/src/PolicyNameKey.cs

@@ -0,0 +1,22 @@
+// 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;
+internal class PolicyNameKey
+{
+    public required string PolicyName { get; init; }
+
+    public override bool Equals(object? obj)
+    {
+        if (obj is PolicyNameKey key)
+        {
+            return PolicyName.Equals(key.PolicyName);
+        }
+        return false;
+    }
+
+    public override int GetHashCode()
+    {
+        return PolicyName.GetHashCode();
+    }
+}

src/Middleware/RateLimiting/src/PolicyTypeInfo.cs

@@ -0,0 +1,12 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+using System.Diagnostics.CodeAnalysis;
+
+namespace Microsoft.AspNetCore.RateLimiting;
+internal class PolicyTypeInfo
+{
+    [DynamicallyAccessedMembers(DynamicallyAccessedMemberTypes.PublicConstructors)]
+    public required Type PolicyType { get; init; }
+    public required Type PartitionKeyType { get; init; }
+}