Check for HasStarted after running ExceptionHandler (#32198)

Author: John Luo

src/Middleware/Diagnostics/src/ExceptionHandler/ExceptionHandlerMiddleware.cs

@@ -132,7 +132,8 @@ private async Task HandleException(HttpContext context, ExceptionDispatchInfo ed
 
                 await _options.ExceptionHandler!(context);
 
-                if (context.Response.StatusCode != StatusCodes.Status404NotFound || _options.AllowStatusCode404Response)
+                // If the response has already started, assume exception handler was successful.
+                if (context.Response.HasStarted || context.Response.StatusCode != StatusCodes.Status404NotFound || _options.AllowStatusCode404Response)
                 {
                     if (_diagnosticListener.IsEnabled() && _diagnosticListener.IsEnabled("Microsoft.AspNetCore.Diagnostics.HandledException"))
                     {

src/Middleware/Diagnostics/test/UnitTests/ExceptionHandlerTest.cs

@@ -312,6 +312,73 @@ public async Task DoesNotModifyCacheHeaders_WhenNoExceptionIsThrown()
             }
         }
 
+        [Fact]
+        public async Task ExceptionHandlerSucceeded_IfExceptionHandlerResponseHasStarted()
+        {
+            using var host = new HostBuilder()
+                .ConfigureWebHost(webHostBuilder =>
+                {
+                    webHostBuilder
+                    .UseTestServer()
+                    .Configure(app =>
+                    {
+                        app.Use(async (httpContext, next) =>
+                        {
+                            Exception exception = null;
+                            try
+                            {
+                                await next(httpContext);
+                            }
+                            catch (InvalidOperationException ex)
+                            {
+                                exception = ex;
+                            }
+
+                            Assert.Null(exception);
+                        });
+
+                        app.UseExceptionHandler("/handle-errors");
+
+                        app.Map("/handle-errors", (innerAppBuilder) =>
+                        {
+                            innerAppBuilder.Run(async (httpContext) =>
+                            {
+                                httpContext.Response.StatusCode = StatusCodes.Status404NotFound;
+                                await httpContext.Response.WriteAsync("Custom 404");
+                            });
+                        });
+
+                        app.Run(httpContext =>
+                        {
+                            httpContext.Response.Headers.Add("Cache-Control", new[] { "max-age=3600" });
+                            httpContext.Response.Headers.Add("Pragma", new[] { "max-age=3600" });
+                            httpContext.Response.Headers.Add("Expires", new[] { DateTime.UtcNow.AddDays(10).ToString("R") });
+                            httpContext.Response.Headers.Add("ETag", new[] { "abcdef" });
+
+                            throw new InvalidOperationException("Something bad happened");
+                        });
+                    });
+                }).Build();
+
+            await host.StartAsync();
+
+            using (var server = host.GetTestServer())
+            {
+                var client = server.CreateClient();
+                var response = await client.GetAsync(string.Empty);
+                Assert.Equal(HttpStatusCode.NotFound, response.StatusCode);
+                Assert.Equal("Custom 404", await response.Content.ReadAsStringAsync());
+                IEnumerable<string> values;
+                Assert.True(response.Headers.CacheControl.NoCache);
+                Assert.True(response.Headers.CacheControl.NoStore);
+                Assert.True(response.Headers.TryGetValues("Pragma", out values));
+                Assert.Single(values);
+                Assert.Equal("no-cache", values.First());
+                Assert.False(response.Headers.TryGetValues("Expires", out _));
+                Assert.False(response.Headers.TryGetValues("ETag", out _));
+            }
+        }
+
         [Fact]
         public async Task DoesNotClearCacheHeaders_WhenResponseHasAlreadyStarted()
         {