✨ Handler.

Signed-off-by: 秋雨落 <i@rain.cx>

CrystalConnector/Connector/WebSocketConnectionManager.cs

@@ -0,0 +1,17 @@
+using System.Net.WebSockets;
+using CrystalConnector.Utilities;
+
+namespace CrystalConnector.Connector;
+
+public class WebSocketConnectionManager
+{
+    internal static Dictionary<WebSocket, WebSocketConnectionInfo> Info { get; } = new();
+
+    public static void DisconnectAll()
+    {
+        foreach (var (webSocket, info) in Info)
+        {
+            webSocket.Disconnect();
+        }
+    }
+}

CrystalConnector/CrystalConnector.csproj

@@ -31,4 +31,5 @@
         </None>
     </ItemGroup>
 
+
 </Project>

CrystalConnector/CrystalService.cs

@@ -1,4 +1,5 @@
-using CrystalConnector.Utilities;
+using CrystalConnector.Connector;
+using CrystalConnector.Utilities;
 using Microsoft.Extensions.Configuration;
 using Microsoft.Extensions.Hosting;
 using Microsoft.Extensions.Logging;
@@ -29,7 +30,7 @@ public class CrystalService : IHostedService
     private void OnStopping()
     {
         Logger.LogInformation("Stopping!");
-        WebSocketExtension.DisconnectAll();
+        WebSocketConnectionManager.DisconnectAll();
     }
     
     private void OnStopped()

CrystalConnector/Handlers/WebSocketHandler.cs

@@ -1,42 +1,36 @@
-using System.Net;
+using System.Formats.Cbor;
 using System.Net.WebSockets;
 using CrystalConnector.Utilities;
-using Microsoft.AspNetCore.Http;
+using Microsoft.Extensions.Configuration;
 using Microsoft.Extensions.Logging;
 
 namespace CrystalConnector.Handlers;
 
-public class WebSocketHandler : IWebSocketHandler
+public class WebSocketHandler
 {
+    private IConfiguration Config { get; }
+
     private ILogger<WebSocketHandler> Logger { get; }
     
-    public WebSocketHandler(ILogger<WebSocketHandler> logger)
+    public WebSocketHandler(IConfiguration config, ILogger<WebSocketHandler> logger)
     {
+        Config = config;
         Logger = logger;
     }
     
-    public async void StartHandle(HttpContext context, WebSocket webSocket, TaskCompletionSource<WebSocketHandleResult> taskCompletionSource)
+    public void Handle(WebSocket webSocket, CborReader reader)
     {
-        var remoteEndPoint = new IPEndPoint(context.Connection.RemoteIpAddress!, context.Connection.RemotePort);    // qyl27: we believe it is a tcp connection.
-        var receiveCancellationTokenSource = new CancellationTokenSource();
-        webSocket.AddData(context.Connection.Id, remoteEndPoint, taskCompletionSource, receiveCancellationTokenSource);
-        await Handle(webSocket, receiveCancellationTokenSource.Token);
-    }
-
-    private async Task Handle(WebSocket webSocket, CancellationToken cancellationToken)
-    {
-        var buffer = new byte[1024 * 4];
-        
-        WebSocketReceiveResult receiveResult;
-        
-        do
+        var type = reader.ReadTextString();
+        if (type == "Authenticate")
         {
-            receiveResult = await webSocket.ReceiveAsync(new ArraySegment<byte>(buffer), cancellationToken);
-        } while (!cancellationToken.IsCancellationRequested && !receiveResult.CloseStatus.HasValue);
-    }
-
-    public void Disconnect(WebSocket webSocket)
-    {
-        webSocket.RemoveData();
+            if (!webSocket.GetConnectionInfo().Authorized)
+            {
+                var secret = reader.ReadTextString();
+                if (secret == Config.GetValue<string>("Connector:Auth:Secret"))
+                {
+                    webSocket.GetConnectionInfo().Authorized = true;
+                }
+            }
+        }
     }
 }

CrystalConnector/Program.cs

@@ -2,6 +2,7 @@
 using System.Security.Cryptography.X509Certificates;
 using CrystalConnector;
 using CrystalConnector.Handlers;
+using CrystalConnector.WebSockets;
 using Microsoft.AspNetCore.Hosting;
 using Microsoft.Extensions.Configuration;
 using Microsoft.Extensions.DependencyInjection;
@@ -92,5 +93,6 @@ void ConfigureLogging(ILoggingBuilder logging)
 
 void ConfigureServices(IServiceCollection services)
 {
-    services.AddSingleton<IWebSocketHandler, WebSocketHandler>();
+    services.AddSingleton<IWebSocketManager, WebSocketManager>();
+    services.AddSingleton<WebSocketHandler>();
 }

CrystalConnector/Startup.cs

@@ -1,5 +1,5 @@
-using CrystalConnector.Handlers;
-using CrystalConnector.Utilities;
+using CrystalConnector.Utilities;
+using CrystalConnector.WebSockets;
 using Microsoft.AspNetCore.Builder;
 using Microsoft.Extensions.Configuration;
 using Microsoft.Extensions.DependencyInjection;
@@ -12,7 +12,7 @@ public class Startup
     private IConfiguration Config { get; }
     
     private ILogger<Startup> Logger { get; set; }
-    private IWebSocketHandler WebSocketHandler { get; set; }
+    private IWebSocketManager WebSocketManager { get; set; }
 
     public Startup(IConfiguration config)
     {
@@ -24,7 +24,7 @@ public class Startup
         var provider = services.BuildServiceProvider();
         
         Logger = provider.GetService<ILogger<Startup>>()!;
-        WebSocketHandler = provider.GetService<IWebSocketHandler>()!;
+        WebSocketManager = provider.GetService<IWebSocketManager>()!;
     }
 
     public void Configure(IApplicationBuilder app)
@@ -39,8 +39,7 @@ public class Startup
                 Logger.LogInformation("Got WebSocket connection from {Ip}:{Port}", context.Connection.RemoteIpAddress?.ToString(), context.Connection.RemotePort);
 
                 var taskCompletionSource = new TaskCompletionSource<WebSocketHandleResult>();
-                WebSocketHandler.StartHandle(context, websocket, taskCompletionSource);
-                await taskCompletionSource.Task;
+                await WebSocketManager.StartHandle(context, websocket, taskCompletionSource);
             }
             else
             {

CrystalConnector/Utilities/WebSocketConnectionInfo.cs

@@ -3,7 +3,7 @@ using CrystalConnector.Connector;
 
 namespace CrystalConnector.Utilities;
 
-public class WebSocketData
+public class WebSocketConnectionInfo
 {
     public required string Id { get; set; }
     
@@ -11,8 +11,6 @@ public class WebSocketData
     
     public required TaskCompletionSource<WebSocketHandleResult> CompletionSource { get; set; }
 
-    public required CancellationTokenSource ReceiveCancellationTokenSource { get; set; }
-    
     public bool Authorized { get; set; } = false;
 
     public List<MessageChannel> RegisteredChannels { get; } = new();

CrystalConnector/Utilities/WebSocketExtension.cs

@@ -1,44 +1,34 @@
 using System.Net;
 using System.Net.WebSockets;
+using CrystalConnector.Connector;
 
 namespace CrystalConnector.Utilities;
 
 public static class WebSocketExtension
 {
-    private static Dictionary<WebSocket, WebSocketData> Data { get; } = new();
-
-    public static void AddData(this WebSocket webSocket, string id, IPEndPoint ip, 
-        TaskCompletionSource<WebSocketHandleResult> taskCompletionSource,
-        CancellationTokenSource receiveCancellationTokenSource)
+    public static void AddConnectionInfo(this WebSocket webSocket, string id, IPEndPoint ip, 
+        TaskCompletionSource<WebSocketHandleResult> taskCompletionSource)
     {
-        Data.Add(webSocket, new WebSocketData
+        WebSocketConnectionManager.Info.Add(webSocket, new WebSocketConnectionInfo
         {
             Id = id, 
             IpEndPoint = ip,
             CompletionSource = taskCompletionSource,
-            ReceiveCancellationTokenSource = receiveCancellationTokenSource
         });
     }
 
-    public static WebSocketData GetData(this WebSocket webSocket)
+    public static WebSocketConnectionInfo GetConnectionInfo(this WebSocket webSocket)
     {
-        return Data[webSocket];
+        return WebSocketConnectionManager.Info[webSocket];
     }
 
-    public static void RemoveData(this WebSocket webSocket)
+    public static void Disconnect(this WebSocket webSocket)
     {
-        var data = GetData(webSocket);
-        data.ReceiveCancellationTokenSource.Cancel();
+        var data = GetConnectionInfo(webSocket);
         data.CompletionSource.SetResult(WebSocketHandleResult.Successful);
 
-        Data.Remove(webSocket);
-    }
+        webSocket.CloseOutputAsync(WebSocketCloseStatus.NormalClosure, "", CancellationToken.None);
         
-    public static void DisconnectAll()
-    {
-        foreach (var (socket, _) in Data)
-        {
-            RemoveData(socket);
-        }
+        WebSocketConnectionManager.Info.Remove(webSocket);
     }
 }

CrystalConnector/WebSockets/IWebSocketManager.cs

@@ -2,11 +2,11 @@
 using CrystalConnector.Utilities;
 using Microsoft.AspNetCore.Http;
 
-namespace CrystalConnector.Handlers;
+namespace CrystalConnector.WebSockets;
 
-public interface IWebSocketHandler
+public interface IWebSocketManager
 {
-    public void StartHandle(HttpContext context, WebSocket socket, TaskCompletionSource<WebSocketHandleResult> taskCompletionSource);
+    public Task StartHandle(HttpContext context, WebSocket socket, TaskCompletionSource<WebSocketHandleResult> taskCompletionSource);
 
     public void Disconnect(WebSocket socket);
 }

CrystalConnector/WebSockets/WebSocketManager.cs

@@ -0,0 +1,49 @@
+using System.Formats.Cbor;
+using System.Net;
+using System.Net.WebSockets;
+using CrystalConnector.Handlers;
+using CrystalConnector.Utilities;
+using Microsoft.AspNetCore.Http;
+using Microsoft.Extensions.Logging;
+
+namespace CrystalConnector.WebSockets;
+
+public class WebSocketManager : IWebSocketManager
+{
+    private ILogger<WebSocketManager> Logger { get; }
+
+    private WebSocketHandler Handler { get; }
+
+    public WebSocketManager(ILogger<WebSocketManager> logger, WebSocketHandler handler)
+    {
+        Logger = logger;
+        Handler = handler;
+    }
+
+    public async Task StartHandle(HttpContext context, WebSocket webSocket, TaskCompletionSource<WebSocketHandleResult> taskCompletionSource)
+    {
+        var remoteEndPoint = new IPEndPoint(context.Connection.RemoteIpAddress!, context.Connection.RemotePort);    // qyl27: we believe it is a tcp connection.
+        webSocket.AddConnectionInfo(context.Connection.Id, remoteEndPoint, taskCompletionSource);
+        await Receive(webSocket);
+    }
+
+    private async Task Receive(WebSocket webSocket)
+    {
+        var buffer = new byte[1024 * 4];
+        var receiveResult = await webSocket.ReceiveAsync(new ArraySegment<byte>(buffer), CancellationToken.None);
+
+        while (!receiveResult.CloseStatus.HasValue)
+        {
+            var reader = new CborReader(new ReadOnlyMemory<byte>(buffer, 0, receiveResult.Count));
+
+            Handler.Handle(webSocket, reader);
+            
+            receiveResult = await webSocket.ReceiveAsync(new ArraySegment<byte>(buffer), CancellationToken.None);
+        }
+    }
+
+    public void Disconnect(WebSocket webSocket)
+    {
+        webSocket.Disconnect();
+    }
+}