Create project files

yawaflua.WebSockets/Core/Middleware/WebSocketMiddleware.cs

@@ -0,0 +1,25 @@
+using Microsoft.AspNetCore.Http;
+
+namespace yawaflua.WebSockets.Core.Middleware;
+
+public class WebSocketMiddleware : IMiddleware
+{
+    private readonly WebSocketRouter _router;
+
+    public WebSocketMiddleware(WebSocketRouter router)
+    {
+        _router = router;
+    }
+
+    public async Task InvokeAsync(HttpContext context, RequestDelegate next)
+    {
+        if (context.WebSockets.IsWebSocketRequest)
+        {
+            await _router.HandleRequest(context);
+        }
+        else
+        {
+            await next(context);
+        }
+    }
+}

yawaflua.WebSockets/Core/WebSocket.cs

@@ -0,0 +1,43 @@
+using System.Net.WebSockets;
+using System.Text;
+using yawaflua.WebSockets.Models.Interfaces;
+
+namespace yawaflua.WebSockets.Core;
+
+public class WebSocket : IDisposable
+{
+    private readonly System.Net.WebSockets.WebSocket _webSocket;
+    private readonly WebSocketReceiveResult? _webSocketReceiveResult;
+    private readonly string? _message;
+
+    public WebSocketState State => _webSocket.State;
+    public WebSocketCloseStatus? CloseStatus => _webSocket.CloseStatus;
+    public string? SubProtocol => _webSocket.SubProtocol;
+    public string? CloseStatusDescription => _webSocket.CloseStatusDescription;
+    public string? Message => _message;
+    public WebSocketMessageType? MessageType => _webSocketReceiveResult?.MessageType;
+    public IWebSocketClient Client;
+    internal WebSocket(System.Net.WebSockets.WebSocket webSocket, WebSocketReceiveResult? webSocketReceiveResult, string? message, IWebSocketClient client)
+    {
+        _webSocket = webSocket;
+        _webSocketReceiveResult = webSocketReceiveResult;
+        _message = message;
+        Client = client;
+    }
+
+    public async Task SendAsync(string m, WebSocketMessageType messageType = WebSocketMessageType.Text, CancellationToken cts = default)
+    => await _webSocket.SendAsync(
+            Encoding.UTF8.GetBytes(m),
+            messageType,
+            true,
+            cts);
+
+    public async Task CloseAsync(WebSocketCloseStatus closeStatus = WebSocketCloseStatus.NormalClosure, string? reason = null, CancellationToken cts = default)
+    => await _webSocket.CloseAsync(closeStatus, reason, cts);
+    
+
+    
+    public void Abort() => _webSocket.Abort();
+    public void Dispose() => _webSocket.Dispose();
+    
+}

yawaflua.WebSockets/Core/WebSocketManager.cs

@@ -0,0 +1,34 @@
+using System.Linq.Expressions;
+using System.Net.WebSockets;
+using System.Text;
+using yawaflua.WebSockets.Models.Interfaces;
+
+namespace yawaflua.WebSockets.Core;
+
+internal class WebSocketManager : IWebSocketManager
+{
+    public async Task Broadcast(Func<IWebSocketClient, bool> selector, string message,
+        WebSocketMessageType messageType = WebSocketMessageType.Text, CancellationToken cts = default)
+    {
+        foreach (var client in WebSocketRouter.Clients.Where(selector))
+        {
+            await client.webSocket.SendAsync(Encoding.UTF8.GetBytes(message),
+                messageType,
+                true,
+                cts);
+        }
+    }
+
+    public List<IWebSocketClient> GetAllClients()
+    {
+        return WebSocketRouter.Clients;
+    }
+
+    public async Task SendToUser(Guid id, string message, WebSocketMessageType messageType = WebSocketMessageType.Text, CancellationToken cts = default)
+    {
+        await WebSocketRouter.Clients.First(k => k.Id == id).webSocket.SendAsync(Encoding.UTF8.GetBytes(message),
+            messageType,
+            true,
+            cts);
+    }
+}

yawaflua.WebSockets/Core/WebSocketRouter.cs

@@ -0,0 +1,182 @@
+using System.Net.WebSockets;
+using System.Reflection;
+using System.Text;
+using Microsoft.AspNetCore.Http;
+using Microsoft.Extensions.DependencyInjection;
+using Microsoft.Extensions.Logging;
+using yawaflua.WebSockets.Attributes;
+using yawaflua.WebSockets.Models;
+using yawaflua.WebSockets.Models.Abstracts;
+using yawaflua.WebSockets.Models.Interfaces;
+
+namespace yawaflua.WebSockets.Core;
+
+public class WebSocketRouter
+{
+    internal static readonly Dictionary<string, Func<WebSocket, HttpContext, Task>> Routes = new();
+    internal static readonly List<IWebSocketClient> Clients = new();
+    private readonly IServiceProvider _serviceProvider;
+    private readonly ILogger<WebSocketRouter> _logger;
+
+    public WebSocketRouter(IServiceProvider serviceProvider, ILogger<WebSocketRouter> logger)
+    {
+        _serviceProvider = serviceProvider;
+        this._logger = logger;
+        DiscoverHandlers();
+        Task.Run(() =>
+        {
+            Clients.ForEach(async l =>
+            {
+                await l.webSocket.SendAsync(ArraySegment<byte>.Empty, WebSocketMessageType.Binary,
+                    WebSocketMessageFlags.EndOfMessage, default);
+                await Task.Delay(TimeSpan.FromSeconds(10));
+            });
+        });
+    }
+
+    internal void DiscoverHandlers()
+    {
+        try
+        {
+            var handlerTypes = AppDomain.CurrentDomain.GetAssemblies()
+                .SelectMany(a => a.GetTypes())
+                .Where(t =>
+                    t.IsSubclassOf(typeof(IWebSocketController))
+                    || t.IsSubclassOf(typeof(WebSocketController))
+                    || t.IsInstanceOfType(typeof(WebSocketController))
+                    || t.IsInstanceOfType(typeof(IWebSocketController))
+                );
+            using var scope = _serviceProvider.CreateScope();
+            foreach (var type in handlerTypes.Where(k => k.GetMethods().Length > 0))
+            {
+                var parentAttributeTemplate =
+                    new PathString((type.GetCustomAttribute(typeof(WebSocketAttribute)) as WebSocketAttribute)?.Template ?? "/");
+                var methods = type.GetMethods()
+                    .Where(m => m.GetCustomAttributes(typeof(WebSocketAttribute), false).Length > 0).ToList();
+                
+
+                if (methods.Count == 0 && type.GetMethods().Any(k => k.Name.StartsWith("OnMessage")))
+                {
+                    var func = type.GetMethods()
+                        .First(k => k.Name.StartsWith("OnMessage"));
+    
+                    var parameters = func.GetParameters();
+                    if (parameters.Length != 2 || 
+                        parameters[0].ParameterType != typeof(WebSocket) ||
+                        parameters[1].ParameterType != typeof(HttpContext) ||
+                        func.ReturnType != typeof(Task))
+                    {
+                        throw new InvalidOperationException(
+                            $"Invalid handler signature in {type.Name}.{func.Name}");
+                    }
+
+                    if (func.IsStatic)
+                    {
+                        var delegateFunc = (Func<WebSocket, HttpContext, Task>)Delegate.CreateDelegate(
+                            typeof(Func<WebSocket, HttpContext, Task>), 
+                            func
+                        );
+                        Routes.Add(parentAttributeTemplate, delegateFunc);
+                    }
+                    else
+                    {
+                        Routes.Add(parentAttributeTemplate, async (ws, context) => 
+                        {
+                            var instance = context.RequestServices.GetRequiredService(type);
+                            await (Task)func.Invoke(instance, new object[] { ws, context })!;
+                        });
+                    }
+                }
+                else
+                {
+                    foreach (var method in methods)
+                    {
+                        var attribute =
+                            (WebSocketAttribute)method.GetCustomAttributes(typeof(WebSocketAttribute), false).First();
+                        if (method.IsStatic)
+                        {
+                            var delegateFunc = (Func<WebSocket, HttpContext, Task>)Delegate.CreateDelegate(
+                                typeof(Func<WebSocket, HttpContext, Task>), 
+                                method
+                            );
+                            Routes.Add(parentAttributeTemplate+attribute.Template, delegateFunc);
+                        }
+                        else
+                        {
+                            Routes.Add(parentAttributeTemplate+attribute.Template, async (ws, context) => 
+                            {
+                                var instance = context.RequestServices.GetRequiredService(type);
+                                await (Task)method.Invoke(instance, new object[] { ws, context })!;
+                            });
+                        }
+                    }
+                }
+                var constructors = type.GetConstructors();
+                if (constructors.Length != 0)
+                {
+                    var parameters = constructors[0].GetParameters()
+                        .Select(param => scope.ServiceProvider.GetRequiredService(param.ParameterType))
+                        .ToArray();
+
+                    constructors[0].Invoke(parameters);
+                }
+            }
+        }
+        catch (Exception ex)
+        {
+            _logger.LogCritical(message:"Error when parsing attributes from assemblies: ", exception:ex);
+        }
+    }
+
+    internal async Task HandleRequest(HttpContext context, CancellationToken cts = default)
+    {
+        try
+        {
+            if (!context.WebSockets.IsWebSocketRequest)
+                return;
+            
+            var path = context.Request.Path.Value;
+
+            if (Routes.TryGetValue(path, out var handler))
+            {
+                var webSocket = await context.WebSockets.AcceptWebSocketAsync();
+                await Task.Run(async () =>
+                {
+                    try
+                    {
+                        var client = new WebSocketClient(context, webSocket, path);
+                        Clients.Add(client);
+                        var buffer = new byte[1024 * 4];
+                        while (webSocket.State == WebSocketState.Open)
+                        {
+                            var result = await webSocket.ReceiveAsync(new ArraySegment<byte>(buffer), cts);
+                            if (result.MessageType != WebSocketMessageType.Close)
+                                await handler(
+                                    new WebSocket(
+                                        webSocket,
+                                        result,
+                                        Encoding.UTF8.GetString(buffer, 0, result.Count),
+                                        client),
+                                    context);
+                            else
+                                Clients.Remove(client);
+                        }
+                    }
+                    catch (Exception ex)
+                    {
+                        _logger.LogError(message:"Error with handling request: ",exception: ex);
+                    }
+
+                }, cts);
+            }
+            else
+            { 
+                context.Response.StatusCode = 404;
+            }
+        }
+        catch (Exception ex)
+        {
+            _logger.LogError(ex, $"Error when handle request {context.Connection.Id}: ");
+        }
+    }
+}