From 1d8507a13b5a1a3170eb8f0488369d3c370b38ef Mon Sep 17 00:00:00 2001
From: Khafra <maitken033380023@gmail.com>
Date: Wed, 25 Sep 2024 00:22:37 -0400
Subject: [PATCH] fixup

---
 lib/web/websocket/connection.js             |  9 +++++----
 lib/web/websocket/receiver.js               |  2 +-
 lib/web/websocket/stream/websocketstream.js | 10 ++++++----
 3 files changed, 12 insertions(+), 9 deletions(-)

diff --git a/lib/web/websocket/connection.js b/lib/web/websocket/connection.js
index 80554163b29..223e8c445cb 100644
--- a/lib/web/websocket/connection.js
+++ b/lib/web/websocket/connection.js
@@ -242,9 +242,10 @@ function closeWebSocketConnection (object, code, reason, validate = false) {
     // Fail the WebSocket connection and set object’s ready state to CLOSING (2). [WSP]
     failWebsocketConnection(object)
     object.readyState = states.CLOSING
-  } else if (object.closeState === sentCloseFrameState.NOT_SENT) {
-    // Start the WebSocket closing handshake and set object’s ready state to CLOSING (2). [WSP]
-    object.closeState = sentCloseFrameState.PROCESSING
+  } else if (!object.closeState.has(sentCloseFrameState.SENT) && !object.closeState.has(sentCloseFrameState.RECEIVED)) {
+    // Upon either sending or receiving a Close control frame, it is said
+    // that _The WebSocket Closing Handshake is Started_ and that the
+    // WebSocket connection is in the CLOSING state.
 
     const frame = new WebsocketFrameSend()
 
@@ -280,7 +281,7 @@ function closeWebSocketConnection (object, code, reason, validate = false) {
 
     object.socket.write(frame.createFrame(opcodes.CLOSE))
 
-    object.closeState = sentCloseFrameState.SENT
+    object.closeState.add(sentCloseFrameState.SENT)
 
     // Upon either sending or receiving a Close control frame, it is said
     // that _The WebSocket Closing Handshake is Started_ and that the
diff --git a/lib/web/websocket/receiver.js b/lib/web/websocket/receiver.js
index bf3eb547b72..2638abe3bfd 100644
--- a/lib/web/websocket/receiver.js
+++ b/lib/web/websocket/receiver.js
@@ -356,7 +356,7 @@ class ByteParser extends Writable {
 
       // Upon receiving such a frame, the other peer sends a
       // Close frame in response, if it hasn't already sent one.
-      if (!this.#handler.closeState.has(sentCloseFrameState.SENT)) {
+      if (!this.#handler.closeState.has(sentCloseFrameState.SENT) && !this.#handler.closeState.has(sentCloseFrameState.RECEIVED)) {
         // If an endpoint receives a Close frame and did not previously send a
         // Close frame, the endpoint MUST send a Close frame in response.  (When
         // sending a Close frame in response, the endpoint typically echos the
diff --git a/lib/web/websocket/stream/websocketstream.js b/lib/web/websocket/stream/websocketstream.js
index 3902c28e685..9f0dea25f01 100644
--- a/lib/web/websocket/stream/websocketstream.js
+++ b/lib/web/websocket/stream/websocketstream.js
@@ -69,7 +69,7 @@ class WebSocketStream {
 
     readyState: states.CONNECTING,
     socket: null,
-    closeState: sentCloseFrameState.NOT_SENT,
+    closeState: new Set(),
     receivedClose: false
   }
 
@@ -242,7 +242,7 @@ class WebSocketStream {
     // 6.1. Wait until there is sufficient buffer space in stream to send the message.
 
     // 6.2. If the closing handshake has not yet started , Send a WebSocket Message to stream comprised of data using opcode .
-    if (this.#handler.closeState === sentCloseFrameState.NOT_SENT) {
+    if (!this.#handler.closeState.has(sentCloseFrameState.SENT) && !this.#handler.closeState.has(sentCloseFrameState.RECEIVED)) {
       const frame = new WebsocketFrameSend(data)
 
       this.#handler.socket.write(frame.createFrame(opcode), () => {
@@ -353,7 +353,9 @@ class WebSocketStream {
 
   /** @type {import('../websocket').Handler['onSocketClose']} */
   #onSocketClose () {
-    const wasClean = this.#handler.closeState === sentCloseFrameState.SENT && this.#handler.receivedClose
+    const wasClean =
+      this.#handler.closeState.has(sentCloseFrameState.SENT) &&
+      this.#handler.closeState.has(sentCloseFrameState.RECEIVED)
 
     // 1. Change the ready state to CLOSED (3).
     this.#handler.readyState = states.CLOSED
@@ -380,7 +382,7 @@ class WebSocketStream {
     // 1006.
     let code = result?.code ?? 1005
 
-    if (this.#handler.closeState === sentCloseFrameState.NOT_SENT) {
+    if (!this.#handler.closeState.has(sentCloseFrameState.SENT) && !this.#handler.closeState.has(sentCloseFrameState.RECEIVED)) {
       code = 1006
     }