Add an option on the ExceptionHandlerOptions to select status codes based on exceptions (dotnet#56616)

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/FunctionalTests/ProblemDetailsExceptionHandlerSampleTest.cs

@@ -6,6 +6,7 @@
 using Microsoft.AspNetCore.Mvc;
 using System.Net.Http.Json;
 using System.Net.Http.Headers;
+using Microsoft.AspNetCore.Http;
 
 namespace Microsoft.AspNetCore.Diagnostics.FunctionalTests;
 
@@ -34,4 +35,21 @@ public async Task ExceptionHandlerPage_ProducesProblemDetails()
         Assert.NotNull(body);
         Assert.Equal(500, body.Status);
     }
+
+    [Fact]
+    public async Task StatusCodeSelector_ProducesProblemDetailsWithCustomStatusCode()
+    {
+        // Arrange
+        var request = new HttpRequestMessage(HttpMethod.Get, "http://localhost/throw2/conflict");
+        request.Headers.Accept.Add(new MediaTypeWithQualityHeaderValue("application/json"));
+
+        // Act
+        var response = await Client.SendAsync(request);
+
+        // Assert
+        var body = await response.Content.ReadFromJsonAsync<ProblemDetails>();
+        Assert.Equal(HttpStatusCode.Conflict, response.StatusCode);
+        Assert.NotNull(body);
+        Assert.Equal(StatusCodes.Status409Conflict, body.Status);
+    }
 }

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 StatusCodeSelector_CanSelect404()
+    {
+        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()
     {

src/Middleware/Diagnostics/test/testassets/ExceptionHandlerSample/ConflictException.cs

@@ -0,0 +1,6 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+namespace ExceptionHandlerSample;
+
+public class ConflictException(string message) : Exception(message);

src/Middleware/Diagnostics/test/testassets/ExceptionHandlerSample/StartupWithProblemDetails.cs

@@ -16,13 +16,28 @@ public void ConfigureServices(IServiceCollection services)
 
     public void Configure(IApplicationBuilder app)
     {
-        // Configure the error handler to produces a ProblemDetails.
-        app.UseExceptionHandler();
-
         // The broken section of our application.
         app.Map("/throw", throwApp =>
         {
-            throwApp.Run(context => { throw new Exception("Application Exception"); });
+            // Configure the error handler to produce ProblemDetails.
+            throwApp.UseExceptionHandler();
+            throwApp.Run(_ => throw new Exception("Application Exception"));
+        });
+
+        app.Map("/throw2", throwApp =>
+        {
+            // Configure the error handler to produce ProblemDetails with a mapped status code
+            throwApp.UseExceptionHandler(new ExceptionHandlerOptions
+            {
+                StatusCodeSelector = ex => ex is ConflictException
+                    ? StatusCodes.Status409Conflict
+                    : StatusCodes.Status500InternalServerError,
+            });
+
+            throwApp.Map("/conflict", throwConflictApp =>
+            {
+                throwConflictApp.Run(_ => throw new ConflictException("Conflict Exception"));
+            });
         });
 
         app.UseStaticFiles();
@@ -32,7 +47,8 @@ public void Configure(IApplicationBuilder app)
         {
             context.Response.ContentType = "text/html";
             await context.Response.WriteAsync("<html><body>Welcome to the sample<br><br>\r\n");
-            await context.Response.WriteAsync("Click here to throw an exception: <a href=\"/throw\">throw</a>\r\n");
+            await context.Response.WriteAsync("Click here to throw an exception: <a href=\"/throw\">throw</a><br>\r\n");
+            await context.Response.WriteAsync("Click here to throw a conflict exception: <a href=\"/throw2/conflict\">throw conflict</a>\r\n");
             await context.Response.WriteAsync("</body></html>\r\n");
         });
     }