Apply ConnectionLimitMiddleware to HTTP/3 connections

Author: halter73

src/Servers/Kestrel/Core/src/Internal/ConnectionDispatcherOfT.cs renamed to src/Servers/Kestrel/Core/src/Internal/ConnectionDispatcher.cs

@@ -162,6 +162,9 @@ async Task OnBind(ListenOptions options)
                         options.UseHttp3Server(ServiceContext, application, options.Protocols);
                         var multiplexedConnectionDelegate = ((IMultiplexedConnectionBuilder)options).Build();
 
+                        // Add the connection limit middleware
+                        multiplexedConnectionDelegate = EnforceConnectionLimit(multiplexedConnectionDelegate, Options.Limits.MaxConcurrentConnections, Trace);
+
                         options.EndPoint = await _transportManager.BindAsync(options.EndPoint, multiplexedConnectionDelegate, options.EndpointConfig).ConfigureAwait(false);
                     }
 
@@ -180,10 +183,7 @@ async Task OnBind(ListenOptions options)
                         var connectionDelegate = options.Build();
 
                         // Add the connection limit middleware
-                        if (Options.Limits.MaxConcurrentConnections.HasValue)
-                        {
-                            connectionDelegate = new ConnectionLimitMiddleware(connectionDelegate, Options.Limits.MaxConcurrentConnections.Value, Trace).OnConnectionAsync;
-                        }
+                        connectionDelegate = EnforceConnectionLimit(connectionDelegate, Options.Limits.MaxConcurrentConnections, Trace);
 
                         options.EndPoint = await _transportManager.BindAsync(options.EndPoint, connectionDelegate, options.EndpointConfig).ConfigureAwait(false);
                     }
@@ -359,5 +359,25 @@ private void ValidateOptions()
                     CoreStrings.FormatMaxRequestBufferSmallerThanRequestHeaderBuffer(Options.Limits.MaxRequestBufferSize.Value, Options.Limits.MaxRequestHeadersTotalSize));
             }
         }
+
+        private static ConnectionDelegate EnforceConnectionLimit(ConnectionDelegate innerDelegate, long? connectionLimit, IKestrelTrace trace)
+        {
+            if (!connectionLimit.HasValue)
+            {
+                return innerDelegate;
+            }
+
+            return new ConnectionLimitMiddleware<ConnectionContext>(c => innerDelegate(c), connectionLimit.Value, trace).OnConnectionAsync;
+        }
+
+        private static MultiplexedConnectionDelegate EnforceConnectionLimit(MultiplexedConnectionDelegate innerDelegate, long? connectionLimit, IKestrelTrace trace)
+        {
+            if (!connectionLimit.HasValue)
+            {
+                return innerDelegate;
+            }
+
+            return new ConnectionLimitMiddleware<MultiplexedConnectionContext>(c => innerDelegate(c), connectionLimit.Value, trace).OnConnectionAsync;
+        }
     }
 }

src/Servers/Kestrel/Core/src/KestrelServer.cs

@@ -1,40 +1,40 @@
 // 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;
 using System.Threading.Tasks;
 using Microsoft.AspNetCore.Connections;
 using Microsoft.AspNetCore.Server.Kestrel.Core.Features;
 using Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Infrastructure;
 
 namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal
 {
-    internal class ConnectionLimitMiddleware
+    internal class ConnectionLimitMiddleware<T> where T : BaseConnectionContext
     {
-        private readonly ConnectionDelegate _next;
+        private readonly Func<T, Task> _next;
         private readonly ResourceCounter _concurrentConnectionCounter;
         private readonly IKestrelTrace _trace;
 
-        public ConnectionLimitMiddleware(ConnectionDelegate next, long connectionLimit, IKestrelTrace trace)
+        public ConnectionLimitMiddleware(Func<T, Task> next, long connectionLimit, IKestrelTrace trace)
             : this(next, ResourceCounter.Quota(connectionLimit), trace)
         {
         }
 
         // For Testing
-        internal ConnectionLimitMiddleware(ConnectionDelegate next, ResourceCounter concurrentConnectionCounter, IKestrelTrace trace)
+        internal ConnectionLimitMiddleware(Func<T, Task> next, ResourceCounter concurrentConnectionCounter, IKestrelTrace trace)
         {
             _next = next;
             _concurrentConnectionCounter = concurrentConnectionCounter;
             _trace = trace;
         }
 
-        public async Task OnConnectionAsync(ConnectionContext connection)
+        public async Task OnConnectionAsync(T connection)
         {
             if (!_concurrentConnectionCounter.TryLockOne())
             {
                 KestrelEventSource.Log.ConnectionRejected(connection.ConnectionId);
                 _trace.ConnectionRejected(connection.ConnectionId);
-                connection.Transport.Input.Complete();
-                connection.Transport.Output.Complete();
+                await connection.DisposeAsync();
                 return;
             }
 

src/Servers/Kestrel/Core/src/Middleware/ConnectionLimitMiddleware.cs

@@ -5,6 +5,7 @@
 using System.Net;
 using System.Threading;
 using System.Threading.Tasks;
+using Microsoft.AspNetCore.Connections;
 using Microsoft.AspNetCore.Http;
 using Microsoft.AspNetCore.Http.Features;
 using Microsoft.AspNetCore.Server.Kestrel.Core;
@@ -13,7 +14,6 @@
 using Microsoft.AspNetCore.Server.Kestrel.InMemory.FunctionalTests.TestTransport;
 using Microsoft.AspNetCore.Server.Kestrel.Tests;
 using Microsoft.AspNetCore.Testing;
-using Microsoft.Extensions.Logging.Testing;
 using Xunit;
 
 namespace Microsoft.AspNetCore.Server.Kestrel.InMemory.FunctionalTests
@@ -211,7 +211,7 @@ private TestServer CreateServerWithMaxConnections(RequestDelegate app, ResourceC
             var listenOptions = new ListenOptions(new IPEndPoint(IPAddress.Loopback, 0));
             listenOptions.Use(next =>
             {
-                var middleware = new ConnectionLimitMiddleware(next, concurrentConnectionCounter, serviceContext.Log);
+                var middleware = new ConnectionLimitMiddleware<ConnectionContext>(c => next(c), concurrentConnectionCounter, serviceContext.Log);
                 return middleware.OnConnectionAsync;
             });