From 1a7c8762f302754a138f409c538a742639bf56ba Mon Sep 17 00:00:00 2001
From: jvyden <jvyden@jvyden.xyz>
Date: Tue, 30 Apr 2024 17:05:02 -0400
Subject: [PATCH] Always call AfterRequestHandled in MainMiddleware

---
 .../Endpoints/Middlewares/MainMiddleware.cs    | 18 +++++++++++-------
 Bunkum.Core/Services/AuthenticationService.cs  |  2 +-
 Bunkum.Core/Services/DebugService.cs           |  7 ++++---
 Bunkum.Core/Services/Service.cs                |  3 ++-
 4 files changed, 18 insertions(+), 12 deletions(-)

diff --git a/Bunkum.Core/Endpoints/Middlewares/MainMiddleware.cs b/Bunkum.Core/Endpoints/Middlewares/MainMiddleware.cs
index 1a76266..e779912 100644
--- a/Bunkum.Core/Endpoints/Middlewares/MainMiddleware.cs
+++ b/Bunkum.Core/Endpoints/Middlewares/MainMiddleware.cs
@@ -124,16 +124,20 @@ public void HandleRequest(ListenerContext context, Lazy<IDatabaseContext> databa
                     {
                         context.ResponseHeaders.Add("Cache-Control", "max-age=" + cacheSeconds.Value);
                     }
-
-                    object? val = method.Invoke(group, invokeList.ToArray());
-                    Response returnedResponse = GenerateResponseFromEndpoint(val, attribute, method);
                     
-                    foreach (Service service in this._services)
+                    try
                     {
-                        service.AfterRequestHandled(context, returnedResponse, method, database);
+                        object? val = method.Invoke(group, invokeList.ToArray());
+                        Response returnedResponse = GenerateResponseFromEndpoint(val, attribute, method);
+                        return returnedResponse;
+                    }
+                    finally
+                    {
+                        foreach (Service service in this._services)
+                        {
+                            service.AfterRequestHandled(context, null, method, database);
+                        }
                     }
-
-                    return returnedResponse;
                 }
             }
         }
diff --git a/Bunkum.Core/Services/AuthenticationService.cs b/Bunkum.Core/Services/AuthenticationService.cs
index e9dea1b..6a2e7dc 100644
--- a/Bunkum.Core/Services/AuthenticationService.cs
+++ b/Bunkum.Core/Services/AuthenticationService.cs
@@ -73,7 +73,7 @@ public override void Initialize()
     }
 
     /// <inheritdoc />
-    public override void AfterRequestHandled(ListenerContext context, Response response, MethodInfo method, Lazy<IDatabaseContext> database)
+    public override void AfterRequestHandled(ListenerContext context, Response? response, MethodInfo method, Lazy<IDatabaseContext> database)
     {
         this._tokenCache.Value = null;
     }
diff --git a/Bunkum.Core/Services/DebugService.cs b/Bunkum.Core/Services/DebugService.cs
index 8c562ed..97433ad 100644
--- a/Bunkum.Core/Services/DebugService.cs
+++ b/Bunkum.Core/Services/DebugService.cs
@@ -37,11 +37,12 @@ internal DebugService(Logger logger) : base(logger)
     }
     
     /// <inheritdoc/>
-    public override void AfterRequestHandled(ListenerContext context, Response response, MethodInfo method, Lazy<IDatabaseContext> database)
+    public override void AfterRequestHandled(ListenerContext context, Response? response, MethodInfo method, Lazy<IDatabaseContext> database)
     {
         if (!method.HasCustomAttribute<DebugResponseBodyAttribute>()) return;
+        if (response == null) return;
         
-        string body = Encoding.UTF8.GetString(response.Data);
-        this.Logger.LogDebug(BunkumCategory.Request, "Response body for {0} ({1} bytes):\n{2}", context.Uri.AbsolutePath, response.Data.Length, body);
+        string body = Encoding.UTF8.GetString(response.Value.Data);
+        this.Logger.LogDebug(BunkumCategory.Request, "Response body for {0} ({1} bytes):\n{2}", context.Uri.AbsolutePath, response.Value.Data.Length, body);
     }
 }
\ No newline at end of file
diff --git a/Bunkum.Core/Services/Service.cs b/Bunkum.Core/Services/Service.cs
index cf5e25a..7c07639 100644
--- a/Bunkum.Core/Services/Service.cs
+++ b/Bunkum.Core/Services/Service.cs
@@ -33,8 +33,9 @@ public virtual void Initialize()
     
     /// <summary>
     /// Runs after the request has been handled by the main middleware.
+    /// May also be called when a response fails to generate due to an exception. A null response is passed in this case.
     /// </summary>
-    public virtual void AfterRequestHandled(ListenerContext context, Response response, MethodInfo method, Lazy<IDatabaseContext> database)
+    public virtual void AfterRequestHandled(ListenerContext context, Response? response, MethodInfo method, Lazy<IDatabaseContext> database)
     {}
     
     /// <summary>