From 3637225d236453537929887a3f0b9469fe2c09d6 Mon Sep 17 00:00:00 2001
From: "littlesheep.code" <littlesheep.code@gmail.com>
Date: Tue, 29 Apr 2025 17:07:00 +0000
Subject: [PATCH] :bricks: Vide coded the websocket controller

---
 .../Connection/WebSocketController.cs         | 104 ++++++++++++++----
 1 file changed, 81 insertions(+), 23 deletions(-)

diff --git a/DysonNetwork.Sphere/Connection/WebSocketController.cs b/DysonNetwork.Sphere/Connection/WebSocketController.cs
index d44eb99..02bc99a 100644
--- a/DysonNetwork.Sphere/Connection/WebSocketController.cs
+++ b/DysonNetwork.Sphere/Connection/WebSocketController.cs
@@ -1,4 +1,6 @@
+using System.Collections.Concurrent;
 using System.Net.WebSockets;
+using DysonNetwork.Sphere.Account;
 using Microsoft.AspNetCore.Authorization;
 using Microsoft.AspNetCore.Mvc;
 
@@ -7,44 +9,100 @@ namespace DysonNetwork.Sphere.Connection;
 [ApiController]
 [Route("/ws")]
 public class WebSocketController : ControllerBase
-{
+{    
+    // Concurrent dictionary to store active WebSocket connections.
+    // Key: Tuple (AccountId, DeviceId); Value: WebSocket and CancellationTokenSource
+    private static readonly ConcurrentDictionary<(long AccountId, string DeviceId), (WebSocket Socket, CancellationTokenSource Cts)> ActiveConnections = new ConcurrentDictionary<(long, string), (WebSocket, CancellationTokenSource)>();
+
     [Route("/ws")]
     [Authorize]
-    public async Task TheGateway()
+    public async Task TheGateway([FromQuery] string deviceId)
     {
         if (HttpContext.WebSockets.IsWebSocketRequest)
         {
-            using var webSocket = await HttpContext.WebSockets.AcceptWebSocketAsync();
-            await _ConnectionEventLoop(webSocket);
+            // Get AccountId from HttpContext
+            if (!HttpContext.Items.TryGetValue("CurrentUser", out var currentUserValue) || currentUserValue is not Account.Account currentUser)
+            {
+                HttpContext.Response.StatusCode = StatusCodes.Status401Unauthorized;
+                return;
+            }
+            long accountId = currentUser.Id;
+
+            // Verify deviceId
+            if (string.IsNullOrEmpty(deviceId))
+            {
+                HttpContext.Response.StatusCode = StatusCodes.Status400BadRequest;
+                return;
+            }
+
+            using var webSocket = await HttpContext.WebSockets.AcceptWebSocketAsync();            
+            // Create a CancellationTokenSource for this connection
+            var cts = new CancellationTokenSource();
+            var connectionKey = (accountId, deviceId);
+
+            // Add the connection to the active connections dictionary
+            if (!ActiveConnections.TryAdd(connectionKey, (webSocket, cts)))
+            {
+                // Failed to add
+                await webSocket.CloseAsync(WebSocketCloseStatus.InternalServerError, "Failed to establish connection.", CancellationToken.None);
+                return;
+            }
+            
+            try
+            {
+                await _ConnectionEventLoop(webSocket, connectionKey, cts.Token);
+            }
+            catch (Exception ex)
+            {
+                Console.WriteLine($"WebSocket Error: {ex.Message}");
+            }
+            finally
+            {
+                // Connection is closed, remove it from the active connections dictionary
+                ActiveConnections.TryRemove(connectionKey, out _);
+                cts.Dispose();
+            }
         }
         else
         {
             HttpContext.Response.StatusCode = StatusCodes.Status400BadRequest;
         }
     }
-    
-    private static async Task _ConnectionEventLoop(WebSocket webSocket)
+
+    private static async Task _ConnectionEventLoop(WebSocket webSocket, (long AccountId, string DeviceId) connectionKey, CancellationToken cancellationToken)
     {
-        // For now, it's echo
-        var buffer = new byte[1024 * 4];
-        var receiveResult = await webSocket.ReceiveAsync(
-            new ArraySegment<byte>(buffer), CancellationToken.None);
-
-        while (!receiveResult.CloseStatus.HasValue)
+        // Buffer for receiving messages.
+        var buffer = new byte[1024 * 4]; 
+        try
         {
-            await webSocket.SendAsync(
-                new ArraySegment<byte>(buffer, 0, receiveResult.Count),
-                receiveResult.MessageType,
-                receiveResult.EndOfMessage,
-                CancellationToken.None);
+            // We don't handle receiving data, so we ignore the return.
+            var receiveResult = await webSocket.ReceiveAsync(new ArraySegment<byte>(buffer), cancellationToken);
+            while (!receiveResult.CloseStatus.HasValue)
+            {
+                // Keep connection alive and wait for close requests
+                receiveResult = await webSocket.ReceiveAsync(new ArraySegment<byte>(buffer), cancellationToken);
+            }
 
-            receiveResult = await webSocket.ReceiveAsync(
-                new ArraySegment<byte>(buffer), CancellationToken.None);
+            // Close connection
+            await webSocket.CloseAsync(receiveResult.CloseStatus.Value, receiveResult.CloseStatusDescription, cancellationToken);
         }
+        catch (OperationCanceledException)
+        {
+            // Connection was canceled, close it gracefully
+            if (webSocket.State != WebSocketState.Closed && webSocket.State != WebSocketState.Aborted)
+            {
+                await webSocket.CloseAsync(WebSocketCloseStatus.NormalClosure, "Connection closed by server", CancellationToken.None);
+            }
+        }
+    }
 
-        await webSocket.CloseAsync(
-            receiveResult.CloseStatus.Value,
-            receiveResult.CloseStatusDescription,
-            CancellationToken.None);
+    // This method will be used later to send messages to specific connections
+    public static async Task SendMessageAsync(long accountId, string deviceId, string message)
+    {
+        if (ActiveConnections.TryGetValue((accountId, deviceId), out var connection))
+        {
+            var buffer = System.Text.Encoding.UTF8.GetBytes(message);
+            await connection.Socket.SendAsync(new ArraySegment<byte>(buffer, 0, buffer.Length), WebSocketMessageType.Text, true, connection.Cts.Token);
+        }
     }
 }
\ No newline at end of file