From 5d5181da5b1a3108c993b6dcd1797e69002d42b8 Mon Sep 17 00:00:00 2001
From: Alex Macocian <amacocian@yahoo.com>
Date: Sun, 25 May 2025 02:54:56 +0200
Subject: [PATCH] Use DefferedWebSocketResult to defer Task execution for
 NativeAOT context

---
 .../Net.Sdk.Web.Extensions.csproj             |  2 +-
 .../WebApplicationExtensions.cs               | 16 +++++++------
 .../Websockets/DeferredWebSocketResult.cs     | 23 +++++++++++++++++++
 3 files changed, 33 insertions(+), 8 deletions(-)
 create mode 100644 Net.Sdk.Web.Extensions/Websockets/DeferredWebSocketResult.cs

diff --git a/Net.Sdk.Web.Extensions/Net.Sdk.Web.Extensions.csproj b/Net.Sdk.Web.Extensions/Net.Sdk.Web.Extensions.csproj
index 0e50418..74d0ba5 100644
--- a/Net.Sdk.Web.Extensions/Net.Sdk.Web.Extensions.csproj
+++ b/Net.Sdk.Web.Extensions/Net.Sdk.Web.Extensions.csproj
@@ -6,7 +6,7 @@
     <Nullable>enable</Nullable>
     <OutputType>Library</OutputType>
     <IsPackable>true</IsPackable>
-    <Version>0.8.9</Version>
+    <Version>0.8.10</Version>
     <Authors>Alexandru Macocian</Authors>
     <PackageLicenseFile>LICENSE</PackageLicenseFile>
     <PackageRequireLicenseAcceptance>true</PackageRequireLicenseAcceptance>
diff --git a/Net.Sdk.Web.Extensions/WebApplicationExtensions.cs b/Net.Sdk.Web.Extensions/WebApplicationExtensions.cs
index c7fcbb5..0ea806e 100644
--- a/Net.Sdk.Web.Extensions/WebApplicationExtensions.cs
+++ b/Net.Sdk.Web.Extensions/WebApplicationExtensions.cs
@@ -34,16 +34,22 @@ public static class WebApplicationExtensions
         where TWebSocketRoute : WebSocketRouteBase
     {
         app.ThrowIfNull();
-        return app.MapGet(route, HandleWebSocketRoute<TWebSocketRoute>);
+        return app.MapGet(route, HandleDeferredWebSocketRoute<TWebSocketRoute>);
     }
 
-    private static async Task<IResult> HandleWebSocketRoute<TWebSocketRoute>(HttpContext context, TWebSocketRoute route, ILogger<TWebSocketRoute> logger)
+    private static DeferredWebSocketResult HandleDeferredWebSocketRoute<TWebSocketRoute>(HttpContext context, TWebSocketRoute route, ILogger<TWebSocketRoute> logger)
+        where TWebSocketRoute : WebSocketRouteBase
+    {
+        return new DeferredWebSocketResult(route, HandleWebSocketRoute(context, route, logger));
+    }
+
+    private static async Task HandleWebSocketRoute<TWebSocketRoute>(HttpContext context, TWebSocketRoute route, ILogger<TWebSocketRoute> logger)
         where TWebSocketRoute : WebSocketRouteBase
     {
         if (!context.WebSockets.IsWebSocketRequest)
         {
             logger.LogError("WebSocket request expected");
-            return Results.BadRequest("WebSocket request expected");
+            return;
         }
 
         var routeFilters = GetRouteFilters<TWebSocketRoute>(context).ToList();
@@ -64,22 +70,18 @@ public static class WebApplicationExtensions
         {
             var processingTask = new Func<Task>(() => ProcessWebSocketRequest(route, context));
             await BeginProcessingPipeline(actionExecutingContext, actionExecutedContext, processingTask);
-            return Results.Empty;
         }
         catch (WebSocketException ex) when (ex.WebSocketErrorCode == WebSocketError.ConnectionClosedPrematurely)
         {
             logger.LogInformation("Websocket closed prematurely. Marking as closed");
-            return Results.Empty;
         }
         catch (OperationCanceledException)
         {
             logger.LogInformation("Websocket closed prematurely. Marking as closed");
-            return Results.Empty;
         }
         catch (Exception ex)
         {
             logger.LogError(ex, "Encountered exception while handling websocket. Closing");
-            return Results.Empty;
         }
         finally
         {
diff --git a/Net.Sdk.Web.Extensions/Websockets/DeferredWebSocketResult.cs b/Net.Sdk.Web.Extensions/Websockets/DeferredWebSocketResult.cs
new file mode 100644
index 0000000..fe05a1c
--- /dev/null
+++ b/Net.Sdk.Web.Extensions/Websockets/DeferredWebSocketResult.cs
@@ -0,0 +1,23 @@
+using System.Core.Extensions;
+
+namespace Net.Sdk.Web.Websockets;
+
+public sealed class DeferredWebSocketResult(
+    WebSocketRouteBase route,
+    Task inner) : IResult
+{
+    private readonly WebSocketRouteBase route = route.ThrowIfNull();
+    private readonly Task inner = inner.ThrowIfNull();
+
+    public async Task ExecuteAsync(HttpContext httpContext)
+    {
+        try
+        {
+            await this.inner;
+        }
+        finally
+        {
+            await this.route.SocketClosed();
+        }
+    }
+}