From a757332d5c8b6d6fe20a877f591692d1cf10d313 Mon Sep 17 00:00:00 2001
From: Jay Malhotra <5047192+SapiensAnatis@users.noreply.github.com>
Date: Sun, 27 Oct 2024 19:21:49 +0000
Subject: [PATCH] Refactor exception handler middleware (#1129)

After adding metrics, our 'unhandled exception' counter never increases,
probably because we handle all exceptions in ExceptionHandlerMiddleware

Put the exception handler middleware as a dedicated exception handler
rather than just a generic middleware with try/catch to ensure ASP.NET
increments the unhandled exception counter
---
 .../Controllers/DragaliaControllerBase.cs     |  1 -
 .../SerializeExceptionAttribute.cs            |  4 --
 .../Middleware/ExceptionHandlerMiddleware.cs  | 43 ++++++++++---------
 DragaliaAPI/DragaliaAPI/Program.cs            |  4 +-
 4 files changed, 25 insertions(+), 27 deletions(-)
 delete mode 100644 DragaliaAPI/DragaliaAPI/Controllers/SerializeExceptionAttribute.cs

diff --git a/DragaliaAPI/DragaliaAPI/Controllers/DragaliaControllerBase.cs b/DragaliaAPI/DragaliaAPI/Controllers/DragaliaControllerBase.cs
index 10e971447..8ffc94e9d 100644
--- a/DragaliaAPI/DragaliaAPI/Controllers/DragaliaControllerBase.cs
+++ b/DragaliaAPI/DragaliaAPI/Controllers/DragaliaControllerBase.cs
@@ -16,7 +16,6 @@ namespace DragaliaAPI.Controllers;
 /// should be used.
 /// </remarks>
 [ApiController]
-[SerializeException]
 [Authorize(AuthenticationSchemes = SchemeName.Session)]
 [Consumes("application/octet-stream")]
 [Produces("application/x-msgpack")]
diff --git a/DragaliaAPI/DragaliaAPI/Controllers/SerializeExceptionAttribute.cs b/DragaliaAPI/DragaliaAPI/Controllers/SerializeExceptionAttribute.cs
deleted file mode 100644
index dc699d191..000000000
--- a/DragaliaAPI/DragaliaAPI/Controllers/SerializeExceptionAttribute.cs
+++ /dev/null
@@ -1,4 +0,0 @@
-namespace DragaliaAPI.Controllers;
-
-[AttributeUsage(AttributeTargets.Class)]
-public class SerializeExceptionAttribute : Attribute { }
diff --git a/DragaliaAPI/DragaliaAPI/Infrastructure/Middleware/ExceptionHandlerMiddleware.cs b/DragaliaAPI/DragaliaAPI/Infrastructure/Middleware/ExceptionHandlerMiddleware.cs
index 953dd2d81..a9e0c2106 100644
--- a/DragaliaAPI/DragaliaAPI/Infrastructure/Middleware/ExceptionHandlerMiddleware.cs
+++ b/DragaliaAPI/DragaliaAPI/Infrastructure/Middleware/ExceptionHandlerMiddleware.cs
@@ -3,36 +3,38 @@
 using DragaliaAPI.Models;
 using DragaliaAPI.Services.Exceptions;
 using MessagePack;
+using Microsoft.AspNetCore.Diagnostics;
 using Microsoft.Extensions.Caching.Distributed;
 using Microsoft.Extensions.Primitives;
 using Microsoft.IdentityModel.Tokens;
 
 namespace DragaliaAPI.Infrastructure.Middleware;
 
-public class ExceptionHandlerMiddleware(RequestDelegate next)
+public static class ExceptionHandlerMiddleware
 {
     private static readonly DistributedCacheEntryOptions CacheOptions =
         new() { AbsoluteExpirationRelativeToNow = TimeSpan.FromMinutes(5) };
 
-    public async Task InvokeAsync(
-        HttpContext context,
-        ILogger<ExceptionHandlerMiddleware> logger,
-        IDistributedCache cache
-    )
+    public static async Task HandleAsync(HttpContext context)
     {
-        Endpoint? endpoint = context.GetEndpoint();
-        bool serializeException =
-            endpoint?.Metadata.GetMetadata<SerializeExceptionAttribute>() is not null;
-
         string? deviceId = null;
         if (context.Request.Headers.TryGetValue("DeviceId", out StringValues values))
-            deviceId = values.FirstOrDefault();
-
-        try
         {
-            await next(context);
+            deviceId = values.FirstOrDefault();
         }
-        catch (SecurityTokenExpiredException ex) when (serializeException && deviceId is not null)
+
+        IExceptionHandlerPathFeature? exceptionHandlerPathFeature =
+            context.Features.Get<IExceptionHandlerPathFeature>();
+
+        Exception? exception = exceptionHandlerPathFeature?.Error;
+
+        ILogger logger = context
+            .RequestServices.GetRequiredService<ILoggerFactory>()
+            .CreateLogger(typeof(ExceptionHandlerMiddleware));
+
+        IDistributedCache cache = context.RequestServices.GetRequiredService<IDistributedCache>();
+
+        if (exception is SecurityTokenExpiredException ex && deviceId is not null)
         {
             logger.LogDebug("ID token was expired. Expiry: {expiry}", ex.Expires);
 
@@ -55,20 +57,19 @@ IDistributedCache cache
             context.Response.StatusCode = StatusCodes.Status400BadRequest;
             context.Response.Headers.Append("Is-Required-Refresh-Id-Token", "true");
         }
-        catch (Exception ex)
-            when (serializeException && context.RequestAborted.IsCancellationRequested)
+        else if (context.RequestAborted.IsCancellationRequested)
         {
-            logger.LogWarning(ex, "Client cancelled request.");
+            logger.LogWarning(exception, "Client cancelled request.");
             context.Response.StatusCode = StatusCodes.Status400BadRequest;
         }
-        catch (Exception ex) when (serializeException)
+        else
         {
-            ResultCode code = ex is DragaliaException dragaliaException
+            ResultCode code = exception is DragaliaException dragaliaException
                 ? dragaliaException.Code
                 : ResultCode.CommonServerError;
 
             logger.LogError(
-                ex,
+                exception,
                 "Encountered unhandled exception. Returning result_code {code}",
                 code
             );
diff --git a/DragaliaAPI/DragaliaAPI/Program.cs b/DragaliaAPI/DragaliaAPI/Program.cs
index 6b42061ab..b52e3bf5e 100644
--- a/DragaliaAPI/DragaliaAPI/Program.cs
+++ b/DragaliaAPI/DragaliaAPI/Program.cs
@@ -160,7 +160,9 @@
         applicationBuilder.UseMiddleware<ResultCodeLoggingMiddleware>();
         applicationBuilder.UseOutputCache();
         applicationBuilder.UseMiddleware<NotFoundHandlerMiddleware>();
-        applicationBuilder.UseMiddleware<ExceptionHandlerMiddleware>();
+        applicationBuilder.UseExceptionHandler(cfg =>
+            cfg.Run(ExceptionHandlerMiddleware.HandleAsync)
+        );
         applicationBuilder.UseMiddleware<DailyResetMiddleware>();
         applicationBuilder.UseEndpoints(endpoints =>
         {