Add an option to select status codes based on exceptions on the ExceptionHandlerOptions

Author: latonz

src/Middleware/Diagnostics/src/ExceptionHandler/ExceptionHandlerMiddlewareImpl.cs

@@ -165,7 +165,7 @@ private async Task HandleException(HttpContext context, ExceptionDispatchInfo ed
 
             context.Features.Set<IExceptionHandlerFeature>(exceptionHandlerFeature);
             context.Features.Set<IExceptionHandlerPathFeature>(exceptionHandlerFeature);
-            context.Response.StatusCode = DefaultStatusCode;
+            context.Response.StatusCode = _options.StatusCodeSelector?.Invoke(edi.SourceException) ?? DefaultStatusCode;
             context.Response.OnStarting(_clearCacheHeadersDelegate, context.Response);
 
             string? handler = null;
@@ -192,7 +192,7 @@ private async Task HandleException(HttpContext context, ExceptionDispatchInfo ed
                     {
                         HttpContext = context,
                         AdditionalMetadata = exceptionHandlerFeature.Endpoint?.Metadata,
-                        ProblemDetails = { Status = DefaultStatusCode },
+                        ProblemDetails = { Status = context.Response.StatusCode },
                         Exception = edi.SourceException,
                     });
                     if (handled)
@@ -202,7 +202,7 @@ private async Task HandleException(HttpContext context, ExceptionDispatchInfo ed
                 }
             }
             // If the response has already started, assume exception handler was successful.
-            if (context.Response.HasStarted || handled || context.Response.StatusCode != StatusCodes.Status404NotFound || _options.AllowStatusCode404Response)
+            if (context.Response.HasStarted || handled || _options.StatusCodeSelector != null || context.Response.StatusCode != StatusCodes.Status404NotFound || _options.AllowStatusCode404Response)
             {
                 const string eventName = "Microsoft.AspNetCore.Diagnostics.HandledException";
                 if (_diagnosticListener.IsEnabled() && _diagnosticListener.IsEnabled(eventName))

src/Middleware/Diagnostics/src/ExceptionHandler/ExceptionHandlerOptions.cs

@@ -38,4 +38,12 @@ public class ExceptionHandlerOptions
     /// the original exception.
     /// </summary>
     public bool AllowStatusCode404Response { get; set; }
+
+    /// <summary>
+    /// Gets or sets a delegate used to map an exception to a http status code.
+    /// </summary>
+    /// <remarks>
+    /// If <see cref="StatusCodeSelector"/> is <c>null</c>, the default exception status code 500 is used.
+    /// </remarks>
+    public Func<Exception, int>? StatusCodeSelector { get; set; }
 }

src/Middleware/Diagnostics/src/PublicAPI.Unshipped.txt

@@ -1 +1,3 @@
 #nullable enable
+Microsoft.AspNetCore.Builder.ExceptionHandlerOptions.StatusCodeSelector.get -> System.Func<System.Exception!, int>?
+Microsoft.AspNetCore.Builder.ExceptionHandlerOptions.StatusCodeSelector.set -> void

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

@@ -656,6 +656,80 @@ public async Task ExceptionHandler_CanReturn404Responses_WhenAllowed()
             && w.Message == "No exception handler was found, rethrowing original exception.");
     }
 
+    [Fact]
+    public async Task ExceptionHandler_SelectsStatusCode()
+    {
+        using var host = new HostBuilder()
+            .ConfigureWebHost(webHostBuilder =>
+            {
+                webHostBuilder
+                    .UseTestServer()
+                    .ConfigureServices(services => services.AddProblemDetails())
+                    .Configure(app =>
+                    {
+                        app.UseExceptionHandler(new ExceptionHandlerOptions
+                        {
+                            StatusCodeSelector = ex => ex is ApplicationException
+                                ? StatusCodes.Status409Conflict
+                                : StatusCodes.Status500InternalServerError,
+                        });
+
+                        app.Map("/throw", innerAppBuilder =>
+                        {
+                            innerAppBuilder.Run(_ => throw new ApplicationException("Something bad happened."));
+                        });
+                    });
+            }).Build();
+
+        await host.StartAsync();
+
+        using (var server = host.GetTestServer())
+        {
+            var client = server.CreateClient();
+            var response = await client.GetAsync("throw");
+            Assert.Equal(HttpStatusCode.Conflict, response.StatusCode);
+        }
+    }
+
+    [Fact]
+    public async Task ExceptionHandler_SelectsStatusCode404_When404ThrowedAndNotNotAllowed()
+    {
+        using var host = new HostBuilder()
+            .ConfigureWebHost(webHostBuilder =>
+            {
+                webHostBuilder
+                    .UseTestServer()
+                    .ConfigureServices(services => services.AddProblemDetails())
+                    .Configure(app =>
+                    {
+                        app.UseExceptionHandler(new ExceptionHandlerOptions
+                        {
+                            // 404 is not allowed,
+                            // but as the exception is explicitly mapped to 404 by the StatusCodeSelector,
+                            // it should be set anyway.
+                            AllowStatusCode404Response = false,
+                            StatusCodeSelector = ex => ex is ApplicationException
+                                ? StatusCodes.Status404NotFound
+                                : StatusCodes.Status500InternalServerError,
+                        });
+
+                        app.Map("/throw", innerAppBuilder =>
+                        {
+                            innerAppBuilder.Run(_ => throw new ApplicationException("Something bad happened."));
+                        });
+                    });
+            }).Build();
+
+        await host.StartAsync();
+
+        using (var server = host.GetTestServer())
+        {
+            var client = server.CreateClient();
+            var response = await client.GetAsync("throw");
+            Assert.Equal(HttpStatusCode.NotFound, response.StatusCode);
+        }
+    }
+
     [Fact]
     public async Task ExceptionHandlerWithOwnBuilder()
     {