Merge pull request #1 from yawaflua/develop

pre-1.0.1 updations

Examples/Program.cs

@@ -1,4 +1,7 @@
-using yawaflua.WebSockets;
+using System.Net.WebSockets;
+using yawaflua.WebSockets;
+using yawaflua.WebSockets.Core;
+using yawaflua.WebSockets.Models.Interfaces;
 
 namespace Examples;
 
@@ -26,6 +29,19 @@ internal class Startup
         services.AddHttpLogging();
         services.AddSingleton<TestWebSocketServer>();
         services.AddSingleton<ChatController>();
+        services.AddSingleton(new WebSocketConfig()
+        {
+            OnOpenHandler = async (socket, context) =>
+            {
+                if (socket.WebSocketManager!.GetAllClients().Count(k =>
+                        Equals(k.ConnectionInfo!.RemoteIpAddress!.MapToIPv4(), 
+                            socket.Client.ConnectionInfo!.RemoteIpAddress!.MapToIPv4())) >= 3)
+                {
+                    await socket.CloseAsync(WebSocketCloseStatus.NormalClosure, "Too many users");
+                }
+                Console.WriteLine($"{socket.Client.Id} has been connected to {socket.Client.Path}");
+            }
+        });
         services.AddScoped<IConfiguration>(k => new ConfigurationBuilder()
             .AddJsonFile("appsettings.json", true)
             .Build());

README.md

@@ -112,6 +112,7 @@ public class ChatController : WebSocketController
 }
 ```
 
+## Run any code on connection to 
 
 ## Lifecycle Management
 1. **Connection** - Automatically handled by middleware

yawaflua.WebSockets/Core/WebSocket.cs

@@ -11,18 +11,24 @@ internal class WebSocket : IWebSocket
     private readonly string? _message;
 
     public WebSocketState State => _webSocket.State;
+    public IWebSocketManager WebSocketManager { get; }
     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 { get; }
-    internal WebSocket(System.Net.WebSockets.WebSocket webSocket, WebSocketReceiveResult? webSocketReceiveResult, string? message, IWebSocketClient client)
+    internal WebSocket(System.Net.WebSockets.WebSocket webSocket, 
+        IWebSocketClient client, 
+        IWebSocketManager webSocketManager,
+        string? message = null, 
+        WebSocketReceiveResult? webSocketReceiveResult = null) 
     {
         _webSocket = webSocket;
         _webSocketReceiveResult = webSocketReceiveResult;
         _message = message;
         Client = client;
+        WebSocketManager = webSocketManager;
     }
 
     public async Task SendAsync(string m, WebSocketMessageType messageType = WebSocketMessageType.Text, CancellationToken cts = default)

yawaflua.WebSockets/Core/WebSocketConfig.cs

@@ -0,0 +1,10 @@
+using Microsoft.AspNetCore.Http;
+using yawaflua.WebSockets.Models.Interfaces;
+
+namespace yawaflua.WebSockets.Core;
+
+public class WebSocketConfig
+{
+    public Func<IWebSocket, HttpContext, Task>? OnOpenHandler { get; set; } = null;
+
+}

yawaflua.WebSockets/Core/WebSocketRouter.cs

@@ -1,4 +1,5 @@
-using System.Net.WebSockets;
+using System.Diagnostics.CodeAnalysis;
+using System.Net.WebSockets;
 using System.Reflection;
 using System.Text;
 using Microsoft.AspNetCore.Http;
@@ -11,17 +12,19 @@ using yawaflua.WebSockets.Models.Interfaces;
 
 namespace yawaflua.WebSockets.Core;
 
+[SuppressMessage("ReSharper", "AsyncVoidLambda")]
 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)
+    private readonly WebSocketConfig? _webSocketConfig;
+    public WebSocketRouter(IServiceProvider serviceProvider, ILogger<WebSocketRouter> logger, WebSocketConfig? webSocketConfig = null)
     {
         _serviceProvider = serviceProvider;
-        this._logger = logger;
+        _logger = logger;
+        _webSocketConfig = webSocketConfig;
         DiscoverHandlers();
         Task.Run(() =>
         {
@@ -62,7 +65,7 @@ public class WebSocketRouter
     
                     var parameters = func.GetParameters();
                     if (parameters.Length != 2 || 
-                        parameters[0].ParameterType != typeof(WebSocket) ||
+                        parameters[0].ParameterType != typeof(IWebSocket) ||
                         parameters[1].ParameterType != typeof(HttpContext) ||
                         func.ReturnType != typeof(Task))
                     {
@@ -144,8 +147,16 @@ public class WebSocketRouter
                 {
                     try
                     {
+                        var webSocketManager = new WebSocketManager();
                         var client = new WebSocketClient(context, webSocket, path);
                         Clients.Add(client);
+                        
+                        await Task.Run(async () =>
+                        {
+                            if (_webSocketConfig?.OnOpenHandler != null)
+                                await _webSocketConfig.OnOpenHandler(new WebSocket(webSocket, client, webSocketManager)!, context);
+                        }, cts);
+                        
                         var buffer = new byte[1024 * 4];
                         while (webSocket.State == WebSocketState.Open)
                         {
@@ -154,13 +165,16 @@ public class WebSocketRouter
                                 await handler(
                                     new WebSocket(
                                         webSocket,
-                                        result,
+                                        client,
+                                        webSocketManager,
                                         Encoding.UTF8.GetString(buffer, 0, result.Count),
-                                        client),
-                                    context);
+                                        result), context);
                             else
                                 Clients.Remove(client);
                         }
+                        
+                        if (Clients.Any(k => k.Id == client.Id))
+                            Clients.Remove(client);
                     }
                     catch (Exception ex)
                     {

yawaflua.WebSockets/Models/Interfaces/IWebSocket.cs

@@ -1,10 +1,12 @@
 using System.Net.WebSockets;
+using yawaflua.WebSockets.Core;
 
 namespace yawaflua.WebSockets.Models.Interfaces;
 
 public interface IWebSocket : IDisposable
 {
     WebSocketState State { get; }
+    IWebSocketManager WebSocketManager { get; }
     WebSocketCloseStatus? CloseStatus { get; }
     string? SubProtocol { get; }
     string? CloseStatusDescription { get; }

yawaflua.WebSockets/ServiceBindings.cs

@@ -8,9 +8,12 @@ namespace yawaflua.WebSockets;
 
 public static class ServiceBindings
 {
-    public static IServiceCollection SettingUpWebSockets(this IServiceCollection isc)
+    public static IServiceCollection SettingUpWebSockets(this IServiceCollection isc, WebSocketConfig? socketOptions = null)
     {
         isc.AddSingleton<WebSocketRouter>();
+        if (socketOptions != null) isc.AddSingleton(socketOptions);
+        if (isc.All(k => k.ServiceType != typeof(WebSocketConfig)))
+            isc.AddSingleton(new WebSocketConfig());
         isc.AddScoped<IWebSocketManager, WebSocketManager>();
         isc.AddSingleton<WebSocketMiddleware>();
         return isc;