From 21accfd767ca2e773be3f805c071b5596eb9efe3 Mon Sep 17 00:00:00 2001
From: Gabriel Scatolin <17441745+CypherPotato@users.noreply.github.com>
Date: Mon, 25 May 2026 14:19:10 -0300
Subject: [PATCH] fix: bound websocket receive size and cleanup ws collection

---
 src/Http/Streams/HttpWebSocket.cs                     | 8 ++++++++
 src/Http/Streams/HttpWebSocketConnectionCollection.cs | 6 +++++-
 2 files changed, 13 insertions(+), 1 deletion(-)

diff --git a/src/Http/Streams/HttpWebSocket.cs b/src/Http/Streams/HttpWebSocket.cs
index d99b739..27438fa 100644
--- a/src/Http/Streams/HttpWebSocket.cs
+++ b/src/Http/Streams/HttpWebSocket.cs
@@ -118,6 +118,9 @@ internal HttpWebSocket ( HttpServerEngineWebSocket ctx, HttpRequest req, string?
         private async ValueTask<WebSocketMessage?> ReceiveInternalAsync ( CancellationToken cancellation ) {
             ArraySegment<byte> buffer = new ArraySegment<byte> ( receiveBuffer );
             ValueWebSocketReceiveResult result;
+            long maxMessageLength = request.baseServer.ServerConfiguration.MaximumContentLength <= 0
+                ? Int32.MaxValue
+                : request.baseServer.ServerConfiguration.MaximumContentLength;
 
             if (IsClosed)
                 return null;
@@ -137,6 +140,11 @@ internal HttpWebSocket ( HttpServerEngineWebSocket ctx, HttpRequest req, string?
 
                         do {
                             result = await ctx.ReceiveAsync ( buffer, cancellation );
+                            if ((ms.Length + result.Count) > maxMessageLength) {
+                                await ctx.CloseOutputAsync ( WebSocketCloseStatus.MessageTooBig, null, cancellation );
+                                await CloseAsync ( cancellation );
+                                return null;
+                            }
                             ms.Write ( buffer.Array!, buffer.Offset, result.Count );
                         } while (!result.EndOfMessage);
 
diff --git a/src/Http/Streams/HttpWebSocketConnectionCollection.cs b/src/Http/Streams/HttpWebSocketConnectionCollection.cs
index c168e92..fc3acce 100644
--- a/src/Http/Streams/HttpWebSocketConnectionCollection.cs
+++ b/src/Http/Streams/HttpWebSocketConnectionCollection.cs
@@ -37,6 +37,7 @@ internal void RegisterWebSocket ( HttpWebSocket src ) {
                     // close another websockets with same identifier
                     HttpWebSocket [] wsId = Find ( s => s == src._identifier );
                     foreach (HttpWebSocket ws in wsId) {
+                        _ws.Remove ( ws );
                         ws.Dispose ();
                     }
                     _ws.Add ( src );
@@ -96,7 +97,10 @@ public HttpWebSocket [] All () {
         /// </summary>
         public void DropAll () {
             lock (_ws) {
-                foreach (HttpWebSocket es in _ws)
+                HttpWebSocket [] allWs = _ws.ToArray ();
+                _ws.Clear ();
+
+                foreach (HttpWebSocket es in allWs)
                     es.Dispose ();
             }
         }