diff --git a/CrystalConnector/Connector/WebSocketConnectionManager.cs b/CrystalConnector/Connector/WebSocketConnectionManager.cs new file mode 100644 index 0000000..c4269fa --- /dev/null +++ b/CrystalConnector/Connector/WebSocketConnectionManager.cs @@ -0,0 +1,17 @@ +using System.Net.WebSockets; +using CrystalConnector.Utilities; + +namespace CrystalConnector.Connector; + +public class WebSocketConnectionManager +{ + internal static Dictionary Info { get; } = new(); + + public static void DisconnectAll() + { + foreach (var (webSocket, info) in Info) + { + webSocket.Disconnect(); + } + } +} \ No newline at end of file diff --git a/CrystalConnector/CrystalConnector.csproj b/CrystalConnector/CrystalConnector.csproj index 1608de9..3508238 100644 --- a/CrystalConnector/CrystalConnector.csproj +++ b/CrystalConnector/CrystalConnector.csproj @@ -31,4 +31,5 @@ + diff --git a/CrystalConnector/CrystalService.cs b/CrystalConnector/CrystalService.cs index dce6bad..8ceebcb 100644 --- a/CrystalConnector/CrystalService.cs +++ b/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() diff --git a/CrystalConnector/Handlers/WebSocketHandler.cs b/CrystalConnector/Handlers/WebSocketHandler.cs index 40cd485..09034fe 100644 --- a/CrystalConnector/Handlers/WebSocketHandler.cs +++ b/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 ILogger Logger { get; } + private IConfiguration Config { get; } - public WebSocketHandler(ILogger logger) + private ILogger Logger { get; } + + public WebSocketHandler(IConfiguration config, ILogger logger) { + Config = config; Logger = logger; } - - public async void StartHandle(HttpContext context, WebSocket webSocket, TaskCompletionSource 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(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("Connector:Auth:Secret")) + { + webSocket.GetConnectionInfo().Authorized = true; + } + } + } } } \ No newline at end of file diff --git a/CrystalConnector/Program.cs b/CrystalConnector/Program.cs index 89ee558..be6366b 100644 --- a/CrystalConnector/Program.cs +++ b/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(); + services.AddSingleton(); + services.AddSingleton(); } diff --git a/CrystalConnector/Startup.cs b/CrystalConnector/Startup.cs index 01aa160..41b74a9 100644 --- a/CrystalConnector/Startup.cs +++ b/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 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>()!; - WebSocketHandler = provider.GetService()!; + WebSocketManager = provider.GetService()!; } 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(); - WebSocketHandler.StartHandle(context, websocket, taskCompletionSource); - await taskCompletionSource.Task; + await WebSocketManager.StartHandle(context, websocket, taskCompletionSource); } else { diff --git a/CrystalConnector/Utilities/WebSocketData.cs b/CrystalConnector/Utilities/WebSocketConnectionInfo.cs similarity index 77% rename from CrystalConnector/Utilities/WebSocketData.cs rename to CrystalConnector/Utilities/WebSocketConnectionInfo.cs index f4b2d1a..24895ca 100644 --- a/CrystalConnector/Utilities/WebSocketData.cs +++ b/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 CompletionSource { get; set; } - public required CancellationTokenSource ReceiveCancellationTokenSource { get; set; } - public bool Authorized { get; set; } = false; public List RegisteredChannels { get; } = new(); diff --git a/CrystalConnector/Utilities/WebSocketExtension.cs b/CrystalConnector/Utilities/WebSocketExtension.cs index ea9cf1d..5cef381 100644 --- a/CrystalConnector/Utilities/WebSocketExtension.cs +++ b/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 Data { get; } = new(); - - public static void AddData(this WebSocket webSocket, string id, IPEndPoint ip, - TaskCompletionSource taskCompletionSource, - CancellationTokenSource receiveCancellationTokenSource) + public static void AddConnectionInfo(this WebSocket webSocket, string id, IPEndPoint ip, + TaskCompletionSource 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); - } - public static void DisconnectAll() - { - foreach (var (socket, _) in Data) - { - RemoveData(socket); - } + webSocket.CloseOutputAsync(WebSocketCloseStatus.NormalClosure, "", CancellationToken.None); + + WebSocketConnectionManager.Info.Remove(webSocket); } } \ No newline at end of file diff --git a/CrystalConnector/Handlers/IWebSocketHandler.cs b/CrystalConnector/WebSockets/IWebSocketManager.cs similarity index 61% rename from CrystalConnector/Handlers/IWebSocketHandler.cs rename to CrystalConnector/WebSockets/IWebSocketManager.cs index 8344e47..8319773 100644 --- a/CrystalConnector/Handlers/IWebSocketHandler.cs +++ b/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 taskCompletionSource); + public Task StartHandle(HttpContext context, WebSocket socket, TaskCompletionSource taskCompletionSource); public void Disconnect(WebSocket socket); } \ No newline at end of file diff --git a/CrystalConnector/WebSockets/WebSocketManager.cs b/CrystalConnector/WebSockets/WebSocketManager.cs new file mode 100644 index 0000000..e163bbb --- /dev/null +++ b/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 Logger { get; } + + private WebSocketHandler Handler { get; } + + public WebSocketManager(ILogger logger, WebSocketHandler handler) + { + Logger = logger; + Handler = handler; + } + + public async Task StartHandle(HttpContext context, WebSocket webSocket, TaskCompletionSource 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(buffer), CancellationToken.None); + + while (!receiveResult.CloseStatus.HasValue) + { + var reader = new CborReader(new ReadOnlyMemory(buffer, 0, receiveResult.Count)); + + Handler.Handle(webSocket, reader); + + receiveResult = await webSocket.ReceiveAsync(new ArraySegment(buffer), CancellationToken.None); + } + } + + public void Disconnect(WebSocket webSocket) + { + webSocket.Disconnect(); + } +} \ No newline at end of file