Connector name, WebSocket options header set

Author: gedem

src/NetCoreStack.WebSockets.ProxyClient/ClientWebSocketConnector.cs

@@ -19,6 +19,14 @@ internal class ClientWebSocketConnector : IWebSocketConnector
         private readonly ILoggerFactory _loggerFactory;
         private readonly InvocatorRegistry _invocatorRegistry;
 
+        public string ConnectionId
+        {
+            get
+            {
+                return _connectionId;
+            }
+        }
+
         public WebSocketState WebSocketState
         {
             get
@@ -40,26 +48,14 @@ public ClientWebSocketConnector(IOptions<ProxyOptions> options,
             Options = options.Value;
         }
 
-        private WebSocketMessageContext CreateConnectionContext()
-        {
-            var context = new WebSocketMessageContext();
-            context.MessageType = WebSocketMessageType.Text;
-            context.Command = WebSocketCommands.Connect;
-            context.Header = new Dictionary<string, object>
-            {
-                ["Name"] = Options.ConnectorName
-            };
-
-            return context;
-        }
-
         public async Task ConnectAsync()
         {
             try
             {
                 var name = Options.ConnectorName;
                 var uri = new Uri($"ws://{Options.WebSocketHostAddress}");
                 _webSocket = new ClientWebSocket();
+                _webSocket.Options.SetRequestHeader(SocketsConstants.ConnectorName, Options.ConnectorName);
                 await _webSocket.ConnectAsync(uri, CancellationToken.None);
                 var receiverContext = new WebSocketReceiverContext
                 {
@@ -69,7 +65,9 @@ public async Task ConnectAsync()
                     Options = Options,
                     WebSocket = _webSocket
                 };
-                var receiver = new WebSocketReceiver(receiverContext, Close);
+                var receiver = new WebSocketReceiver(receiverContext, Close, (connectionId) => {
+                    _connectionId = connectionId;
+                });
                 await receiver.ReceiveAsync();
             }
             catch (Exception ex)

src/NetCoreStack.WebSockets.ProxyClient/IWebSocketConnector.cs

@@ -5,6 +5,7 @@ namespace NetCoreStack.WebSockets.ProxyClient
 {
     public interface IWebSocketConnector
     {
+        string ConnectionId { get; }
         WebSocketState WebSocketState { get; }
         ProxyOptions Options { get; }
         Task ConnectAsync();

src/NetCoreStack.WebSockets/ConnectionManager.cs

@@ -110,9 +110,9 @@ await transport.WebSocket.SendAsync(segments,
                            token);
         }
 
-        public async Task ConnectAsync(WebSocket webSocket)
+        public async Task ConnectAsync(WebSocket webSocket, string connectorName = "")
         {
-            WebSocketTransport transport = new WebSocketTransport(webSocket);
+            WebSocketTransport transport = new WebSocketTransport(webSocket, connectorName);
             var connectionId = transport.ConnectionId;
             var context = new WebSocketMessageContext();
             context.Command = WebSocketCommands.Handshake;

src/NetCoreStack.WebSockets/Interfaces/IConnectionManager.cs

@@ -10,7 +10,7 @@ public interface IConnectionManager
     {
         ConcurrentDictionary<string, WebSocketTransport> Connections { get; }
 
-        Task ConnectAsync(WebSocket webSocket);
+        Task ConnectAsync(WebSocket webSocket, string connectorName = "");
 
         /// <summary>
         /// Text message broadcaster

src/NetCoreStack.WebSockets/Internal/SocketsConstants.cs

@@ -1,9 +1,10 @@
 namespace NetCoreStack.WebSockets.Internal
 {
-    public class SocketsConstants
+    public static class SocketsConstants
     {
         public static byte[] Splitter = new byte[] { 0x1f };
-        public const string CompressedKey = "Compressed";
         public const int ChunkSize = 1024 * 4;
+        public const string CompressedKey = "Compressed";
+        public const string ConnectorName = "ConnectorName";
     }
 }

src/NetCoreStack.WebSockets/Internal/WebSocketMiddleware.cs

@@ -1,6 +1,7 @@
 using Microsoft.AspNetCore.Http;
 using Microsoft.Extensions.Logging;
 using Microsoft.Extensions.Options;
+using Microsoft.Extensions.Primitives;
 using NetCoreStack.WebSockets.Interfaces;
 using System;
 using System.Net.WebSockets;
@@ -26,8 +27,15 @@ public async Task Invoke(HttpContext httpContext,
         {
             if (httpContext.WebSockets.IsWebSocketRequest)
             {
+                string connectorName = string.Empty;
+                StringValues headerValue = "";
+                if (httpContext.Request.Headers.TryGetValue(SocketsConstants.ConnectorName, out headerValue))
+                {
+                    connectorName = headerValue.ToString();
+                }
+
                 var webSocket = await httpContext.WebSockets.AcceptWebSocketAsync();
-                await manager.ConnectAsync(webSocket);
+                await manager.ConnectAsync(webSocket, connectorName);
             }
             else
             {

src/NetCoreStack.WebSockets/Internal/WebSocketReceiver.cs

@@ -12,11 +12,13 @@ public class WebSocketReceiver
     {
         private readonly WebSocketReceiverContext _context;
         private readonly Action<WebSocketReceiverContext> _closeCallback;
+        private readonly Action<string> _handshakeCallback;
 
-        public WebSocketReceiver(WebSocketReceiverContext context, Action<WebSocketReceiverContext> closeCallback)
+        public WebSocketReceiver(WebSocketReceiverContext context, Action<WebSocketReceiverContext> closeCallback, Action<string> handshakeCallback = null)
         {
             _context = context;
             _closeCallback = closeCallback;
+            _handshakeCallback = handshakeCallback;
         }
 
         private async Task InternalReceiveAsync()
@@ -29,7 +31,10 @@ private async Task InternalReceiveAsync()
                 {
                     var context = result.ToContext(buffer);
                     if (context.Command == WebSocketCommands.Handshake)
+                    {
                         _context.ConnectionId = context.Value?.ToString();
+                        _handshakeCallback?.Invoke(_context.ConnectionId);
+                    }
 
                     var _invocators = _context.InvocatorRegistry.GetInvocators(context, _context.Options);
                     foreach (var invoker in _invocators)

src/NetCoreStack.WebSockets/Internal/WebSocketTransport.cs

@@ -7,11 +7,13 @@ public class WebSocketTransport : IDisposable
     {
         public WebSocket WebSocket { get; }
         public string ConnectionId { get; }
+        public string ConnectorName { get; }
 
-        public WebSocketTransport(WebSocket webSocket)
+        public WebSocketTransport(WebSocket webSocket, string connectorName)
         {
             ConnectionId = Guid.NewGuid().ToString();
             WebSocket = webSocket;
+            ConnectorName = connectorName;
         }
 
         public void Dispose()

test/ServerTestApp/Controllers/DiscoveryController.cs

@@ -120,5 +120,14 @@ public async Task<IActionResult> SendCompressedBinaryAsync([FromBody]Context mod
 
             return Ok();
         }
+
+        [HttpGet(nameof(GetConnections))]
+        public IActionResult GetConnections()
+        {
+            var connections = _connectionManager.Connections
+                .Select(x => new { ConnectionId = x.Value.ConnectionId, ConnectorName = x.Value.ConnectorName });
+
+            return Json(connections);
+        }
     }
 }