Set StatusCode before disposing HttpContext (#876)

src/Microsoft.AspNetCore.Server.Kestrel/Internal/Http/Frame.cs

@@ -4,7 +4,6 @@
 using System;
 using System.Collections.Generic;
 using System.Diagnostics;
-using System.Globalization;
 using System.IO;
 using System.Linq;
 using System.Net;

src/Microsoft.AspNetCore.Server.Kestrel/Internal/Http/FrameOfT.cs

@@ -92,55 +92,69 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Internal.Http
                         var context = _application.CreateContext(this);
                         try
                         {
-                            await _application.ProcessRequestAsync(context).ConfigureAwait(false);
+                            try
-                            VerifyResponseContentLength();
+                            {
-                        }
+                                await _application.ProcessRequestAsync(context).ConfigureAwait(false);
-                        catch (Exception ex)
+                                VerifyResponseContentLength();
-                        {
+                            }
-                            ReportApplicationError(ex);
+                            catch (Exception ex)
+                            {
+                                ReportApplicationError(ex);
+                            }
+                            finally
+                            {
+                                // Trigger OnStarting if it hasn't been called yet and the app hasn't
+                                // already failed. If an OnStarting callback throws we can go through
+                                // our normal error handling in ProduceEnd.
+                                // https://github.com/aspnet/KestrelHttpServer/issues/43
+                                if (!HasResponseStarted && _applicationException == null && _onStarting != null)
+                                {
+                                    await FireOnStarting();
+                                }
+
+                                PauseStreams();
+
+                                if (_onCompleted != null)
+                                {
+                                    await FireOnCompleted();
+                                }
+
+                                // If _requestAbort is set, the connection has already been closed.
+                                if (Volatile.Read(ref _requestAborted) == 0)
+                                {
+                                    ResumeStreams();
+
+                                    if (_keepAlive)
+                                    {
+                                        // Finish reading the request body in case the app did not.
+                                        await messageBody.Consume();
+                                    }
+
+                                    // ProduceEnd() must be called before _application.DisposeContext(), to ensure
+                                    // HttpContext.Response.StatusCode is correctly set when
+                                    // IHttpContextFactory.Dispose(HttpContext) is called.
+                                    await ProduceEnd();
+                                }
+                                else if (!HasResponseStarted)
+                                {
+                                    // If the request was aborted and no response was sent, there's no
+                                    // meaningful status code to log.
+                                    StatusCode = 0;
+                                }
+                            }
                         }
                         finally
                         {
-                            // Trigger OnStarting if it hasn't been called yet and the app hasn't
-                            // already failed. If an OnStarting callback throws we can go through
-                            // our normal error handling in ProduceEnd.
-                            // https://github.com/aspnet/KestrelHttpServer/issues/43
-                            if (!HasResponseStarted && _applicationException == null && _onStarting != null)
-                            {
-                                await FireOnStarting();
-                            }
-
-                            PauseStreams();
-
-                            if (_onCompleted != null)
-                            {
-                                await FireOnCompleted();
-                            }
-
                             _application.DisposeContext(context, _applicationException);
                         }
+                    }
 
-                        // If _requestAbort is set, the connection has already been closed.
+                    StopStreams();
-                        if (Volatile.Read(ref _requestAborted) == 0)
-                        {
-                            ResumeStreams();
 
-                            if (_keepAlive)
+                    if (!_keepAlive)
-                            {
+                    {
-                                // Finish reading the request body in case the app did not.
+                        // End the connection for non keep alive as data incoming may have been thrown off
-                                await messageBody.Consume();
+                        return;
-                            }
-
-                            await ProduceEnd();
-                        }
-
-                        StopStreams();
-
-                        if (!_keepAlive)
-                        {
-                            // End the connection for non keep alive as data incoming may have been thrown off
-                            return;
-                        }
                     }
 
                     // Don't reset frame state if we're exiting the loop. This avoids losing request rejection

test/Microsoft.AspNetCore.Server.Kestrel.FunctionalTests/ResponseTests.cs

@@ -6,14 +6,17 @@ using System.IO;
 using System.Linq;
 using System.Net;
 using System.Net.Http;
+using System.Net.Sockets;
 using System.Text;
 using System.Threading.Tasks;
 using Microsoft.AspNetCore.Builder;
 using Microsoft.AspNetCore.Hosting;
 using Microsoft.AspNetCore.Http;
+using Microsoft.AspNetCore.Http.Features;
 using Microsoft.AspNetCore.Server.Kestrel.Internal.Http;
 using Microsoft.AspNetCore.Server.Kestrel.Internal.Infrastructure;
 using Microsoft.AspNetCore.Testing;
+using Microsoft.Extensions.DependencyInjection;
 using Microsoft.Extensions.Internal;
 using Microsoft.Extensions.Logging;
 using Microsoft.Extensions.Primitives;
@@ -185,6 +188,123 @@ namespace Microsoft.AspNetCore.Server.Kestrel.FunctionalTests
             }
         }
 
+        [Fact]
+        public Task ResponseStatusCodeSetBeforeHttpContextDisposeAppException()
+        {
+            return ResponseStatusCodeSetBeforeHttpContextDispose(
+                context =>
+                {
+                    throw new Exception();
+                },
+                expectedClientStatusCode: HttpStatusCode.InternalServerError,
+                expectedServerStatusCode: HttpStatusCode.InternalServerError);
+        }
+
+        [Fact]
+        public Task ResponseStatusCodeSetBeforeHttpContextDisposeRequestAborted()
+        {
+            return ResponseStatusCodeSetBeforeHttpContextDispose(
+                context =>
+                {
+                    context.Abort();
+                    return TaskCache.CompletedTask;
+                },
+                expectedClientStatusCode: null,
+                expectedServerStatusCode: 0);
+        }
+
+        [Fact]
+        public Task ResponseStatusCodeSetBeforeHttpContextDisposeRequestAbortedAppException()
+        {
+            return ResponseStatusCodeSetBeforeHttpContextDispose(
+                context =>
+                {
+                    context.Abort();
+                    throw new Exception();
+                },
+                expectedClientStatusCode: null,
+                expectedServerStatusCode: 0);
+        }
+
+        [Fact]
+        public Task ResponseStatusCodeSetBeforeHttpContextDisposedRequestMalformed()
+        {
+            return ResponseStatusCodeSetBeforeHttpContextDispose(
+                context =>
+                {
+                    return TaskCache.CompletedTask;
+                },
+                expectedClientStatusCode: null,
+                expectedServerStatusCode: HttpStatusCode.BadRequest,
+                sendMalformedRequest: true);
+        }
+
+        private static async Task ResponseStatusCodeSetBeforeHttpContextDispose(
+            RequestDelegate handler,
+            HttpStatusCode? expectedClientStatusCode,
+            HttpStatusCode expectedServerStatusCode,
+            bool sendMalformedRequest = false)
+        {
+            var mockHttpContextFactory = new Mock<IHttpContextFactory>();
+            mockHttpContextFactory.Setup(f => f.Create(It.IsAny<IFeatureCollection>()))
+                .Returns<IFeatureCollection>(fc => new DefaultHttpContext(fc));
+
+            var disposedTcs = new TaskCompletionSource<int>();
+            mockHttpContextFactory.Setup(f => f.Dispose(It.IsAny<HttpContext>()))
+                .Callback<HttpContext>(c =>
+                {
+                    disposedTcs.TrySetResult(c.Response.StatusCode);
+                });
+
+            var hostBuilder = new WebHostBuilder()
+                .UseKestrel()
+                .UseUrls("http://127.0.0.1:0")
+                .ConfigureServices(services => services.AddSingleton<IHttpContextFactory>(mockHttpContextFactory.Object))
+                .Configure(app =>
+                {
+                    app.Run(handler);
+                });
+
+            using (var host = hostBuilder.Build())
+            {
+                host.Start();
+
+                if (!sendMalformedRequest)
+                {
+                    using (var client = new HttpClient())
+                    {
+                        try
+                        {
+                            var response = await client.GetAsync($"http://localhost:{host.GetPort()}/");
+                            Assert.Equal(expectedClientStatusCode, response.StatusCode);
+                        }
+                        catch
+                        {
+                            if (expectedClientStatusCode != null)
+                            {
+                                throw;
+                            }
+                        }
+                    }
+                }
+                else
+                {
+                    using (var socket = new Socket(AddressFamily.InterNetwork, SocketType.Stream, ProtocolType.Tcp))
+                    {
+                        socket.Connect(new IPEndPoint(IPAddress.Loopback, host.GetPort()));
+                        socket.Send(Encoding.ASCII.GetBytes(
+                            "POST / HTTP/1.1\r\n" +
+                            "Transfer-Encoding: chunked\r\n" +
+                            "\r\n" +
+                            "wrong"));
+                    }
+                }
+
+                var disposedStatusCode = await disposedTcs.Task.TimeoutAfter(TimeSpan.FromSeconds(10));
+                Assert.Equal(expectedServerStatusCode, (HttpStatusCode)disposedStatusCode);
+            }
+        }
+
         // https://github.com/aspnet/KestrelHttpServer/pull/1111/files#r80584475 explains the reason for this test.
         [Fact]
         public async Task SingleErrorResponseSentWhenAppSwallowsBadRequestException()