Merge branch 'develop'

Author: LucHeart

Common/Swagger/AttributeFilter.cs

@@ -1,6 +1,4 @@
 using Microsoft.AspNetCore.Authorization;
-using Microsoft.Net.Http.Headers;
-using Microsoft.OpenApi.Any;
 using Microsoft.OpenApi.Models;
 using OpenShock.Common.Authentication;
 using OpenShock.Common.DataAnnotations.Interfaces;
@@ -36,15 +34,6 @@ public void Apply(OpenApiSchema schema, SchemaFilterContext context)
 
     public void Apply(OpenApiOperation operation, OperationFilterContext context)
     {
-        operation.Parameters.Add(new OpenApiParameter
-        {
-            Required = false,
-            In = ParameterLocation.Header,
-            Name = HeaderNames.UserAgent,
-            AllowEmptyValue = false,
-            Example = new OpenApiString("MyApiClient/1.0")
-        });
-        
         // Apply OpenShock Parameter Attributes
         foreach (var attribute in context.MethodInfo?.GetCustomAttributes(true).OfType<IOperationAttribute>() ?? [])
         {

Common/Websocket/WebsockBaseController.cs

@@ -34,7 +34,7 @@ public abstract class WebsocketBaseController<T> : OpenShockControllerBase, IAsy
     private CancellationTokenSource? _linkedSource;
 
     protected CancellationToken LinkedToken;
-
+    
     /// <summary>
     /// Channel for multithreading thread safety of the websocket, MessageLoop is the only reader for this channel
     /// </summary>
@@ -53,7 +53,6 @@ protected WebsocketBaseController(ILogger<WebsocketBaseController<T>> logger)
         Logger = logger;
     }
 
-
     /// <inheritdoc />
     [NonAction]
     public ValueTask QueueMessage(T data)
@@ -210,14 +209,18 @@ protected virtual Task SendWebSocketMessage(T message, WebSocket websocket, Canc
 
     #endregion
 
+    private CancellationTokenSource _receiveCancellationTokenSource = new();
+    
     /// <summary>
     /// Main receiver logic for the websocket
     /// </summary>
     /// <returns></returns>
     [NonAction]
     private async Task Logic()
     {
-        while (!LinkedToken.IsCancellationRequested)
+        using var linkedReceiverToken = CancellationTokenSource.CreateLinkedTokenSource(LinkedToken, _receiveCancellationTokenSource.Token);
+        
+        while (!linkedReceiverToken.IsCancellationRequested)
         {
             try
             {
@@ -241,7 +244,7 @@ private async Task Logic()
                     return;
                 }
 
-                if (!await HandleReceive())
+                if (!await HandleReceive(linkedReceiverToken.Token))
                 {
                     // HandleReceive returned false, we will close the connection after this
                     Logger.LogDebug("HandleReceive returned false, closing connection");
@@ -273,7 +276,18 @@ private async Task Logic()
     /// </summary>
     /// <returns>True if you want to continue the receiver loop, false if you want to terminate</returns>
     [NonAction]
-    protected abstract Task<bool> HandleReceive();
+    protected abstract Task<bool> HandleReceive(CancellationToken cancellationToken);
+    
+    [NonAction]
+    protected async Task ForceClose(WebSocketCloseStatus closeStatus, string? statusDescription)
+    {
+        await _receiveCancellationTokenSource.CancelAsync();
+
+        if (WebSocket is { State: WebSocketState.CloseReceived or WebSocketState.Open })
+        {
+            await WebSocket.CloseOutputAsync(closeStatus, statusDescription, LinkedToken);
+        }
+    }
 
     /// <summary>
     /// Send initial data to the client

LiveControlGateway/Controllers/HubControllerBase.cs

@@ -102,13 +102,7 @@ ILogger<FlatbuffersWebsocketBaseController<TIn, TOut>> logger
             try
             {
                 Logger.LogInformation("Keep alive timeout reached, closing websocket connection");
-
-                if (WebSocket is { State: WebSocketState.Open })
-                {
-                    await WebSocket!.CloseOutputAsync(WebSocketCloseStatus.ProtocolError, "Keep alive timeout reached",
-                        LinkedToken);
-                }
-                
+                await ForceClose(WebSocketCloseStatus.ProtocolError, "Keep alive timeout reached");
                 WebSocket?.Abort();
             }
             catch (Exception ex)
@@ -187,14 +181,7 @@ public async Task DisconnectOld()
         if (WebSocket == null)
             return;
 
-        if (WebSocket is { State: WebSocketState.Open })
-        {
-            await WebSocket.CloseOutputAsync(WebSocketCloseStatus.NormalClosure,
-                "Hub is connecting from a different location",
-                LinkedToken);
-        }
-        
-        WebSocket?.Abort();
+        await ForceClose(WebSocketCloseStatus.NormalClosure, "Hub is connecting from a different location");
     }
 
     private static DateTimeOffset? GetBootedAtFromUptimeMs(ulong uptimeMs)

LiveControlGateway/Controllers/LiveControlController.cs

@@ -255,7 +255,7 @@ await QueueMessage(new LiveControlResponse<LiveResponseType>
     }
 
     /// <inheritdoc />
-    protected override async Task<bool> HandleReceive()
+    protected override async Task<bool> HandleReceive(CancellationToken cancellationToken)
     {
         var message =
             await JsonWebSocketUtils.ReceiveFullMessageAsyncNonAlloc<BaseRequest<LiveRequestType>>(WebSocket!,
@@ -266,8 +266,7 @@ await JsonWebSocketUtils.ReceiveFullMessageAsyncNonAlloc<BaseRequest<LiveRequest
                 if (request?.Data == null)
                 {
                     Logger.LogWarning("Received null data from client");
-                    await WebSocket!.CloseOutputAsync(WebSocketCloseStatus.InvalidPayloadData,
-                        "Invalid json message received", LinkedToken);
+                    await ForceClose(WebSocketCloseStatus.InvalidPayloadData, "Invalid json message received");
                     return false;
                 }
 
@@ -278,8 +277,7 @@ await JsonWebSocketUtils.ReceiveFullMessageAsyncNonAlloc<BaseRequest<LiveRequest
             async failed =>
             {
                 Logger.LogWarning(failed.Exception, "Deserialization failed for websocket message");
-                await WebSocket!.CloseOutputAsync(WebSocketCloseStatus.InvalidPayloadData,
-                    "Invalid json message received", LinkedToken);
+                await ForceClose(WebSocketCloseStatus.InvalidPayloadData, "Invalid json message received");
                 return false;
             }, closure =>
             {
@@ -550,8 +548,8 @@ await SendWebSocketMessage(new LiveControlResponse<LiveResponseType>
                 ResponseType = LiveResponseType.DeviceNotConnected,
             }, WebSocket!, LinkedToken);
 
-            await WebSocket!.CloseOutputAsync(WebSocketCloseStatus.NormalClosure,
-                "Hub is disconnected", LinkedToken);
+            
+            await ForceClose(WebSocketCloseStatus.NormalClosure, "Hub is disconnected");
         }
         catch (Exception e)
         {

LiveControlGateway/Websocket/FlatbuffersWebsocketBaseController.cs

@@ -48,18 +48,17 @@ protected override Task
     protected abstract Task<bool> Handle(TIn data);
 
     /// <inheritdoc />
-    protected override async Task<bool> HandleReceive()
+    protected override async Task<bool> HandleReceive(CancellationToken cancellationToken)
     {
         var message =
             await FlatbufferWebSocketUtils.ReceiveFullMessageAsyncNonAlloc(WebSocket!,
-                _incomingSerializer, LinkedToken);
+                _incomingSerializer, cancellationToken);
 
         var continueLoop = await message.Match(
             Handle,
             async _ =>
             {
-                await WebSocket!.CloseAsync(WebSocketCloseStatus.InvalidPayloadData, "Invalid flatbuffers message",
-                    LinkedToken);
+                await ForceClose(WebSocketCloseStatus.InvalidPayloadData, "Invalid flatbuffers message");
                 return false;
             },
             _ =>