Adding the possibility of defining headers for web socket client.

behnammby · Mar 8, 2022 · cd7f23d · cd7f23d
1 parent 7a2f692
commit cd7f23d
Show file tree

Hide file tree

Showing 3 changed files with 30 additions and 7 deletions.
diff --git a/src/Ninja.WebSocketClient/NinjaWebSocket.cs b/src/Ninja.WebSocketClient/NinjaWebSocket.cs
@@ -1,4 +1,5 @@
-using System.Buffers;
+using Ninja.WebSocketClient.Options;
+using System.Buffers;
 
 namespace Ninja.WebSocketClient
 {
@@ -22,10 +23,10 @@ public class NinjaWebSocket
         public event Func<Exception?, Task>? OnClosed;
         public event Func<ReadOnlySequence<byte>?, Task>? OnReceived;
 
-        public NinjaWebSocket(string url)
+        public NinjaWebSocket(string url, Headers? headers = null)
         {
             _url = url;
-            _webSocketPipe = new WebSocketDuplexPipe();
+            _webSocketPipe = new WebSocketDuplexPipe(headers);
         }
 
         public async Task StartAsync(CancellationToken ct = default)

diff --git a/src/Ninja.WebSocketClient/Options/Headers.cs b/src/Ninja.WebSocketClient/Options/Headers.cs
@@ -0,0 +1,12 @@
+namespace Ninja.WebSocketClient.Options
+{
+    public class Headers : Dictionary<string, string>
+    {
+        public Headers AddHeader(string key, string value)
+        {
+            this[key] = value;
+
+            return this;
+        }
+    }
+}
diff --git a/src/Ninja.WebSocketClient/Transport/WebSocketDuplexPipe.cs b/src/Ninja.WebSocketClient/Transport/WebSocketDuplexPipe.cs
@@ -1,11 +1,13 @@
-using System.IO.Pipelines;
+using Ninja.WebSocketClient.Options;
+using System.IO.Pipelines;
 using System.Net.WebSockets;
 
 namespace Ninja.WebSocketClient
 {
     internal class WebSocketDuplexPipe : IDuplexPipe
     {
         private ClientWebSocket? _webSocket;
+        private readonly Headers? _headers;
         private IDuplexPipe? _transport;
         private IDuplexPipe? _application;
         private volatile bool _aborted;
@@ -18,15 +20,23 @@ internal class WebSocketDuplexPipe : IDuplexPipe
 
         public PipeWriter Output => _transport!.Output;
 
-        public WebSocketDuplexPipe()
+        public WebSocketDuplexPipe(Headers? headers = null)
         {
-
+            _headers = headers;
         }
 
         public async Task StartAsync(string url, CancellationToken ct = default)
         {
             _webSocket = new ClientWebSocket();
 
+            if (_headers != null)
+            {
+                foreach (var header in _headers)
+                {
+                    _webSocket.Options.SetRequestHeader(header.Key, header.Value);
+                }
+            }
+
             try
             {
                 await _webSocket.ConnectAsync(new Uri(url), ct);
@@ -100,7 +110,7 @@ private async Task ProcessSocketAsync(WebSocket socket)
         private async Task StartReceiving(WebSocket socket)
         {
             try
-            { 
+            {
                 while (true)
                 {
                     ValueWebSocketReceiveResult receiveResult;