prevent update http header on websocket connection (server side)

alec1o · Dec 10, 2024 · d373eaf · d373eaf
1 parent 3c8a335
commit d373eaf
Show file tree

Hide file tree

Showing 3 changed files with 17 additions and 8 deletions.
diff --git a/src/http/partials/Server/HTTP.ServerRequest.cs b/src/http/partials/Server/HTTP.ServerRequest.cs
@@ -18,7 +18,7 @@ internal ServerRequest(HttpListenerRequest request)
             {
                 {
                     Headers = new Dictionary<string, string>();
-                    foreach (var headerKey in request.Headers.AllKeys)
+                    foreach (var headerKey in request.Headers.AllKeys.Where(x => !string.IsNullOrWhiteSpace(x)))
                         Headers.Add(headerKey, request.Headers[headerKey] ?? string.Empty);
                 }
 

diff --git a/src/http/partials/Server/HTTP.ServerResponse.cs b/src/http/partials/Server/HTTP.ServerResponse.cs
@@ -14,15 +14,17 @@ internal class ServerResponse : IHTTP.ServerResponse
         {
             private readonly HttpListenerResponse _response;
             private readonly List<byte> _bytes = new List<byte>();
+            private readonly bool _isWebsocketRequest;
 
-            public ServerResponse(HttpListenerResponse response)
+            public ServerResponse(HttpListenerResponse response, bool isWebsocketRequest)
             {
                 NativeResponse = response;
                 IsOpened = true;
                 Encoding = Encoding.UTF8;
                 Headers = new Dictionary<string, string>();
                 Cookies = Array.Empty<Cookie>();
                 _response = response;
+                _isWebsocketRequest = isWebsocketRequest;
             }
 
             public HttpListenerResponse NativeResponse { get; }
@@ -92,11 +94,15 @@ public void Close()
 
             private void WriteHeaders()
             {
-                foreach (var header in Headers) _response.AddHeader(header.Key, header.Value);
+                if (Headers.Count <= 0 || _isWebsocketRequest) return;
+
+                foreach (var header in Headers) _response.Headers[header.Key] = header.Value;
             }
 
             private void WriteCookies()
             {
+                if (Cookies.Length <= 0 || _isWebsocketRequest) return;
+
                 foreach (var cookie in Cookies) _response.Cookies.Add(cookie);
             }
 

diff --git a/src/http/partials/Server/HTTP.ServerTo.cs b/src/http/partials/Server/HTTP.ServerTo.cs
@@ -206,7 +206,7 @@ private void HandleConnection(HttpListenerContext context)
             {
                 var request = new ServerRequest(context.Request);
 
-                var response = new ServerResponse(context.Response);
+                var response = new ServerResponse(context.Response, request.IsWebSocket);
 
                 var middlewares = _server.MyMiddleware.Middlewares.Length <= byte.MinValue
                     ? new List<IHTTP.MiddlewareDescriptor>()
@@ -359,16 +359,19 @@ private void MapMiddlewareCallback(HttpListenerContext context, ServerRequest re
 
                     var websocket = new WebSocket(ws.WebSocket, request);
 
-                    lock (_websocketListLock)
+                    websocket.On.Open(() =>
                     {
-                        _websocketList.Add(websocket);
-                    }
+                        lock (_websocketListLock)
+                        {
+                            _websocketList.Add(websocket);
+                        }
+                    });
 
                     websocket.On.Close(() =>
                     {
                         lock (_websocketListLock)
                         {
-                            _websocketList.Add(websocket);
+                            _websocketList.Remove(websocket);
                         }
                     });