Middleware and Manager improvments, Dependencies, Context Design

Author: gedem

src/NetCoreStack.WebSockets.ProxyClient/ClientWebSocketConnector.cs

@@ -38,7 +38,16 @@ public async Task ConnectAsync()
                 var uri = new Uri($"ws://{_options.WebSocketHostAddress}");
                 _webSocket = new ClientWebSocket();
                 await _webSocket.ConnectAsync(uri, CancellationToken.None);
-                await WebSocketReceiver.Receive(_webSocket, _compressor, _invocatorRegistry, (SocketsOptions)_options);
+                var receiverContext = new WebSocketReceiverContext
+                {
+                    Compressor = _compressor,
+                    InvocatorRegistry = _invocatorRegistry,
+                    LoggerFactory = _loggerFactory,
+                    Options = _options,
+                    WebSocket = _webSocket
+                };
+                var receiver = new WebSocketReceiver(receiverContext, Close);
+                await receiver.ReceiveAsync();
             }
             catch (Exception ex)
             {
@@ -68,9 +77,18 @@ public async Task SendBinaryAsync(byte[] bytes)
             await _webSocket.SendAsync(segments, WebSocketMessageType.Binary, true, CancellationToken.None);
         }
 
+        internal void Close(WebSocketReceiverContext context)
+        {
+            context.WebSocket.CloseAsync(WebSocketCloseStatus.NormalClosure, 
+                nameof(WebSocketReceiverContext), 
+                CancellationToken.None);
+        }
+
         internal void Close(string statusDescription)
         {
-            _webSocket.CloseAsync(WebSocketCloseStatus.NormalClosure, statusDescription, CancellationToken.None);
+            _webSocket.CloseAsync(WebSocketCloseStatus.NormalClosure, 
+                statusDescription, 
+                CancellationToken.None);
         }
     }
 }

src/NetCoreStack.WebSockets/ConnectionManager.cs

@@ -1,4 +1,6 @@
-using NetCoreStack.WebSockets.Interfaces;
+using Microsoft.Extensions.Logging;
+using Microsoft.Extensions.Options;
+using NetCoreStack.WebSockets.Interfaces;
 using NetCoreStack.WebSockets.Internal;
 using Newtonsoft.Json;
 using System;
@@ -14,13 +16,25 @@ namespace NetCoreStack.WebSockets
 {
     public class ConnectionManager : IConnectionManager
     {
-        protected IStreamCompressor Compressor { get; }
-        protected ConcurrentDictionary<string, WebSocketTransport> Connections { get; }
-
-        public ConnectionManager(IStreamCompressor compressor)
+        private readonly InvocatorRegistry _invocatorRegistry;
+        private readonly ServerSocketsOptions _options;
+        private readonly IHandshakeStateTransport _initState;
+        private readonly ILoggerFactory _loggerFactory;
+        private readonly IStreamCompressor _compressor;
+        private readonly ConcurrentDictionary<string, WebSocketTransport> _connections;
+
+        public ConnectionManager(IStreamCompressor compressor,
+            InvocatorRegistry invocatorRegistry,
+            IOptions<ServerSocketsOptions> options,
+            IHandshakeStateTransport initState,
+            ILoggerFactory loggerFactory)
         {
-            Compressor = compressor;
-            Connections = new ConcurrentDictionary<string, WebSocketTransport>();
+            _invocatorRegistry = invocatorRegistry;
+            _options = options.Value;
+            _initState = initState;
+            _loggerFactory = loggerFactory;
+            _compressor = compressor;
+            _connections = new ConcurrentDictionary<string, WebSocketTransport>();
         }
 
         private async Task<byte[]> PrepareBytesAsync(byte[] input, JsonObject properties)
@@ -36,17 +50,7 @@ private async Task<byte[]> PrepareBytesAsync(byte[] input, JsonObject properties
             var bytesCount = input.Length;
             input = propsBytes.Concat(input).ToArray();
 
-            return await Compressor.CompressAsync(input);
-
-            //if (input.Length > SocketsConstants.CompressorThreshold)
-            //{
-            //    using (MemoryStream ms = new MemoryStream(input))
-            //    {
-            //        return await Compressor.CompressAsync(ms);
-            //    }
-            //}
-            //else
-            //    return input;           
+            return await _compressor.CompressAsync(input);         
         }
 
         private async Task SendAsync(WebSocketTransport transport, WebSocketMessageDescriptor descriptor)
@@ -82,6 +86,31 @@ await transport.WebSocket.SendAsync(segments,
                            token);
         }
 
+        public async Task ConnectAsync(WebSocket webSocket)
+        {
+            WebSocketTransport transport = new WebSocketTransport(webSocket);
+            var connectionId = transport.ConnectionId;
+            var context = new WebSocketMessageContext();
+            context.Command = WebSocketCommands.Handshake;
+            context.Value = connectionId;
+            context.State = await _initState.GetStateAsync();
+            _connections.TryAdd(connectionId, transport);
+
+            await SendAsync(connectionId, context);
+
+            var receiverContext = new WebSocketReceiverContext
+            {
+                Compressor = _compressor,
+                ConnectionId = connectionId,
+                InvocatorRegistry = _invocatorRegistry,
+                LoggerFactory = _loggerFactory,
+                Options = _options,
+                WebSocket = webSocket
+            };
+            var receiver = new WebSocketReceiver(receiverContext, CloseConnection);
+            await receiver.ReceiveAsync();
+        }
+
         public async Task BroadcastAsync(WebSocketMessageContext context)
         {
             if (context == null)
@@ -94,7 +123,7 @@ public async Task BroadcastAsync(WebSocketMessageContext context)
                 throw new ArgumentNullException(nameof(context.Value));
             }
 
-            if (!Connections.Any())
+            if (!_connections.Any())
             {
                 return;
             }
@@ -107,15 +136,15 @@ public async Task BroadcastAsync(WebSocketMessageContext context)
                 MessageType = WebSocketMessageType.Text
             };
 
-            foreach (var connection in Connections)
+            foreach (var connection in _connections)
             {
                 await SendAsync(connection.Value, descriptor);
             }
         }
 
         public async Task BroadcastBinaryAsync(byte[] inputs, JsonObject properties)
         {
-            if (!Connections.Any())
+            if (!_connections.Any())
             {
                 return;
             }
@@ -136,7 +165,7 @@ public async Task BroadcastBinaryAsync(byte[] inputs, JsonObject properties)
                         if (chunkedBytes.Length < SocketsConstants.ChunkSize)
                             endOfMessage = true;
 
-                        foreach (var connection in Connections)
+                        foreach (var connection in _connections)
                         {
                             await SendBinaryAsync(connection.Value, chunkedBytes, endOfMessage, CancellationToken.None);
                         }
@@ -151,13 +180,13 @@ public async Task BroadcastBinaryAsync(byte[] inputs, JsonObject properties)
 
         public async Task SendAsync(string connectionId, WebSocketMessageContext context)
         {
-            if (!Connections.Any())
+            if (!_connections.Any())
             {
                 return;
             }
 
             WebSocketTransport transport = null;
-            if (!Connections.TryGetValue(connectionId, out transport))
+            if (!_connections.TryGetValue(connectionId, out transport))
             {
                 throw new ArgumentOutOfRangeException(nameof(transport));
             }
@@ -175,13 +204,13 @@ public async Task SendAsync(string connectionId, WebSocketMessageContext context
 
         public async Task SendBinaryAsync(string connectionId, byte[] input, JsonObject properties)
         {
-            if (!Connections.Any())
+            if (!_connections.Any())
             {
                 return;
             }
 
             WebSocketTransport transport = null;
-            if (!Connections.TryGetValue(connectionId, out transport))
+            if (!_connections.TryGetValue(connectionId, out transport))
             {
                 throw new ArgumentOutOfRangeException(nameof(transport));
             }
@@ -216,39 +245,18 @@ await transport.WebSocket.SendAsync(segments,
             }
         }
 
-        public async Task SendAsync(string connectionId, WebSocketMessageContext context, WebSocket webSocket)
+        public void CloseConnection(string connectionId)
         {
-            if (string.IsNullOrEmpty(connectionId))
-            {
-                throw new ArgumentNullException(nameof(connectionId));
-            }
-
-            if (context == null)
-            {
-                throw new ArgumentNullException(nameof(context));
-            }
-
-            if (webSocket == null)
-            {
-                throw new ArgumentNullException(nameof(webSocket));
-            }
-
             WebSocketTransport transport = null;
-            if (!Connections.TryGetValue(connectionId, out transport))
+            if (_connections.TryRemove(connectionId, out transport))
             {
-                Connections.TryAdd(connectionId, new WebSocketTransport(webSocket));
+                transport.Dispose();
             }
-
-            await SendAsync(connectionId, context);
         }
 
-        public void CloseConnection(string connectionId)
+        public void CloseConnection(WebSocketReceiverContext context)
         {
-            WebSocketTransport transport = null;
-            if (Connections.TryRemove(connectionId, out transport))
-            {
-                transport.Dispose();
-            }
+            CloseConnection(context.ConnectionId);
         }
     }
 }

src/NetCoreStack.WebSockets/Interfaces/IConnectionManager.cs

@@ -5,6 +5,8 @@ namespace NetCoreStack.WebSockets
 {
     public interface IConnectionManager
     {
+        Task ConnectAsync(WebSocket webSocket);
+
         Task BroadcastAsync(WebSocketMessageContext context);
 
         Task BroadcastBinaryAsync(byte[] input, JsonObject properties);
@@ -13,8 +15,6 @@ public interface IConnectionManager
 
         Task SendBinaryAsync(string connectionId, byte[] input, JsonObject properties);
 
-        Task SendAsync(string connectionId, WebSocketMessageContext context, WebSocket webSocket);
-
         void CloseConnection(string connectionId);
     }
 }

src/NetCoreStack.WebSockets/Interfaces/WebSocketReceiverContext.cs

@@ -0,0 +1,16 @@
+using Microsoft.Extensions.Logging;
+using NetCoreStack.WebSockets.Internal;
+using System.Net.WebSockets;
+
+namespace NetCoreStack.WebSockets.Interfaces
+{
+    public class WebSocketReceiverContext
+    {
+        public string ConnectionId { get; set; }
+        public WebSocket WebSocket { get; set; }
+        public InvocatorRegistry InvocatorRegistry { get; set; }
+        public SocketsOptions Options { get; set; }
+        public ILoggerFactory LoggerFactory { get; set; }
+        public IStreamCompressor Compressor { get; set; }
+    }
+}

src/NetCoreStack.WebSockets/Internal/ConnectionManagerExtensions.cs

@@ -1,6 +1,9 @@
 using Microsoft.AspNetCore.Http;
+using Microsoft.Extensions.Logging;
 using Microsoft.Extensions.Options;
 using NetCoreStack.WebSockets.Interfaces;
+using System;
+using System.Net.WebSockets;
 using System.Threading.Tasks;
 
 namespace NetCoreStack.WebSockets.Internal
@@ -19,12 +22,12 @@ public async Task Invoke(HttpContext httpContext,
             IStreamCompressor compressor,
             InvocatorRegistry invocatorRegistry,
             IOptions<ServerSocketsOptions> options,
-            IHandshakeStateTransport initState)
+            ILoggerFactory loggerFactory)
         {
             if (httpContext.WebSockets.IsWebSocketRequest)
             {
                 var webSocket = await httpContext.WebSockets.AcceptWebSocketAsync();
-                await manager.Handshake(webSocket, compressor, invocatorRegistry, options.Value, initState);
+                await manager.ConnectAsync(webSocket);
             }
             else
             {