More request validation out of the AsyncAcceptContext (#28452)

* More request validation out of the AsyncAcceptContext
- Moved request validation logic to the accept loop. This simplifies the logic inside of the accept context.
- This also moves feature initializaton to after the threadpool dispatch.
- This does make make the async state machine yield when requests exceeed the header limit or if auth fails.

* PR feedback

Author: davidfowl

src/Servers/HttpSys/src/AsyncAcceptContext.cs

@@ -5,7 +5,6 @@
 using System.Threading;
 using System.Threading.Tasks;
 using System.Threading.Tasks.Sources;
-using Microsoft.AspNetCore.Http;
 using Microsoft.AspNetCore.HttpSys.Internal;
 
 namespace Microsoft.AspNetCore.Server.HttpSys
@@ -14,6 +13,8 @@ internal unsafe class AsyncAcceptContext : IValueTaskSource<RequestContext>, IDi
     {
         private static readonly IOCompletionCallback IOCallback = IOWaitCallback;
         private readonly PreAllocatedOverlapped _preallocatedOverlapped;
+        private readonly IRequestContextFactory _requestContextFactory;
+
         private NativeOverlapped* _overlapped;
 
         // mutable struct; do not make this readonly
@@ -24,7 +25,6 @@ internal unsafe class AsyncAcceptContext : IValueTaskSource<RequestContext>, IDi
         };
 
         private RequestContext _requestContext;
-        private readonly IRequestContextFactory _requestContextFactory;
 
         internal AsyncAcceptContext(HttpSysListener server, IRequestContextFactory requestContextFactory)
         {
@@ -55,7 +55,6 @@ internal ValueTask<RequestContext> AcceptAsync()
 
         private static void IOCompleted(AsyncAcceptContext asyncContext, uint errorCode, uint numBytes)
         {
-            var complete = false;
             // This is important to stash a ref to as it's a mutable struct
             ref var mrvts = ref asyncContext._mrvts;
             var requestContext = asyncContext._requestContext;
@@ -73,48 +72,18 @@ private static void IOCompleted(AsyncAcceptContext asyncContext, uint errorCode,
                 HttpSysListener server = asyncContext.Server;
                 if (errorCode == UnsafeNclNativeMethods.ErrorCodes.ERROR_SUCCESS)
                 {
-                    // at this point we have received an unmanaged HTTP_REQUEST and memoryBlob
-                    // points to it we need to hook up our authentication handling code here.
-                    try
-                    {
-                        if (server.ValidateRequest(requestContext) && server.ValidateAuth(requestContext))
-                        {
-                            // It's important that we clear the request context before we set the result
-                            // we want to reuse the acceptContext object for future accepts.
-                            asyncContext._requestContext = null;
+                    // It's important that we clear the request context before we set the result
+                    // we want to reuse the acceptContext object for future accepts.
+                    asyncContext._requestContext = null;
 
-                            // Initialize features here once we're successfully validated the request
-                            // TODO: In the future defer this work to the thread pool so we can get off the IO thread
-                            // as quickly as possible
-                            requestContext.InitializeFeatures();
-
-                            mrvts.SetResult(requestContext);
-
-                            complete = true;
-                        }
-                    }
-                    catch (Exception ex)
-                    {
-                        server.SendError(requestId, StatusCodes.Status400BadRequest);
-                        mrvts.SetException(ex);
-                    }
-                    finally
-                    {
-                        if (!complete)
-                        {
-                            asyncContext.AllocateNativeRequest(size: requestContext.Size);
-                        }
-                    }
+                    mrvts.SetResult(requestContext);
                 }
                 else
                 {
                     //  (uint)backingBuffer.Length - AlignmentPadding
                     asyncContext.AllocateNativeRequest(numBytes, requestId);
-                }
 
-                // We need to issue a new request, either because auth failed, or because our buffer was too small the first time.
-                if (!complete)
-                {
+                    // We need to issue a new request, either because auth failed, or because our buffer was too small the first time.
                     uint statusCode = asyncContext.QueueBeginGetContext();
 
                     if (statusCode != UnsafeNclNativeMethods.ErrorCodes.ERROR_SUCCESS &&

src/Servers/HttpSys/src/HttpSysListener.cs

@@ -293,17 +293,14 @@ internal unsafe bool ValidateRequest(NativeRequestContext requestMemory)
                 SendError(requestMemory.RequestId, StatusCodes.Status400BadRequest, authChallenges: null);
                 return false;
             }
-            return true;
-        }
 
-        internal unsafe bool ValidateAuth(NativeRequestContext requestMemory)
-        {
             if (!Options.Authentication.AllowAnonymous && !requestMemory.CheckAuthenticated())
             {
                 SendError(requestMemory.RequestId, StatusCodes.Status401Unauthorized,
                     AuthenticationManager.GenerateChallenges(Options.Authentication.Schemes));
                 return false;
             }
+
             return true;
         }
 

src/Servers/HttpSys/src/MessagePump.cs

@@ -189,6 +189,12 @@ private async Task ProcessRequestsWorker()
                 try
                 {
                     requestContext = await Listener.AcceptAsync(acceptContext);
+
+                    if (!Listener.ValidateRequest(requestContext))
+                    {
+                        // If either of these is false then a response has already been sent to the client, so we can accept the next request
+                        continue;
+                    }
                 }
                 catch (Exception exception)
                 {

src/Servers/HttpSys/src/RequestProcessing/RequestContext.FeatureCollection.cs

@@ -93,7 +93,7 @@ private enum Fields
             TraceIdentifier = 0x200,
         }
 
-        public void InitializeFeatures()
+        protected internal void InitializeFeatures()
         {
             _initialized = true;
 

src/Servers/HttpSys/src/RequestProcessing/RequestContextOfT.cs

@@ -25,6 +25,8 @@ protected override async Task ExecuteAsync()
 
             try
             {
+                InitializeFeatures();
+
                 if (messagePump.Stopping)
                 {
                     SetFatalResponse(503);

src/Servers/HttpSys/test/FunctionalTests/Listener/Utilities.cs

@@ -114,8 +114,23 @@ internal static HttpSysListener CreateServerOnExistingQueue(AuthenticationScheme
         internal static async Task<RequestContext> AcceptAsync(this HttpSysListener server, TimeSpan timeout)
         {
             var factory = new TestRequestContextFactory(server);
-            var acceptContext = new AsyncAcceptContext(server, factory);
-            var acceptTask = server.AcceptAsync(acceptContext).AsTask();
+            using var acceptContext = new AsyncAcceptContext(server, factory);
+            
+            async Task<RequestContext> AcceptAsync()
+            {
+                while (true)
+                {
+                    var requestContext = await server.AcceptAsync(acceptContext);
+
+                    if (server.ValidateRequest(requestContext))
+                    {
+                        requestContext.InitializeFeatures();
+                        return requestContext;
+                    }
+                }
+            }
+
+            var acceptTask = AcceptAsync();
             var completedTask = await Task.WhenAny(acceptTask, Task.Delay(timeout));
 
             if (completedTask == acceptTask)