websocket: improve .close() (#2865)

Author: Uzlopak

lib/web/websocket/connection.js

@@ -1,6 +1,6 @@
 'use strict'
 
-const { uid, states } = require('./constants')
+const { uid, states, sentCloseFrameState } = require('./constants')
 const {
   kReadyState,
   kSentClose,
@@ -230,7 +230,7 @@ function onSocketClose () {
   // If the TCP connection was closed after the
   // WebSocket closing handshake was completed, the WebSocket connection
   // is said to have been closed _cleanly_.
-  const wasClean = ws[kSentClose] && ws[kReceivedClose]
+  const wasClean = ws[kSentClose] === sentCloseFrameState.SENT && ws[kReceivedClose]
 
   let code = 1005
   let reason = ''
@@ -240,7 +240,7 @@ function onSocketClose () {
   if (result) {
     code = result.code ?? 1005
     reason = result.reason
-  } else if (!ws[kSentClose]) {
+  } else if (ws[kSentClose] !== sentCloseFrameState.SENT) {
     // If _The WebSocket
     // Connection is Closed_ and no Close control frame was received by the
     // endpoint (such as could occur if the underlying transport connection

lib/web/websocket/constants.js

@@ -20,6 +20,12 @@ const states = {
   CLOSED: 3
 }
 
+const sentCloseFrameState = {
+  NOT_SENT: 0,
+  PROCESSING: 1,
+  SENT: 2
+}
+
 const opcodes = {
   CONTINUATION: 0x0,
   TEXT: 0x1,
@@ -42,6 +48,7 @@ const emptyBuffer = Buffer.allocUnsafe(0)
 
 module.exports = {
   uid,
+  sentCloseFrameState,
   staticPropertyDescriptors,
   states,
   opcodes,

lib/web/websocket/receiver.js

@@ -1,7 +1,7 @@
 'use strict'
 
 const { Writable } = require('node:stream')
-const { parserStates, opcodes, states, emptyBuffer } = require('./constants')
+const { parserStates, opcodes, states, emptyBuffer, sentCloseFrameState } = require('./constants')
 const { kReadyState, kSentClose, kResponse, kReceivedClose } = require('./symbols')
 const { channels } = require('../../core/diagnostics')
 const { isValidStatusCode, failWebsocketConnection, websocketMessageReceived } = require('./util')
@@ -104,7 +104,7 @@ class ByteParser extends Writable {
 
           this.#info.closeInfo = this.parseCloseBody(body)
 
-          if (!this.ws[kSentClose]) {
+          if (this.ws[kSentClose] !== sentCloseFrameState.SENT) {
             // 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
@@ -120,7 +120,7 @@ class ByteParser extends Writable {
               closeFrame.createFrame(opcodes.CLOSE),
               (err) => {
                 if (!err) {
-                  this.ws[kSentClose] = true
+                  this.ws[kSentClose] = sentCloseFrameState.SENT
                 }
               }
             )

lib/web/websocket/util.js

@@ -8,6 +8,17 @@ const { MessageEvent, ErrorEvent } = require('./events')
 
 /**
  * @param {import('./websocket').WebSocket} ws
+ * @returns {boolean}
+ */
+function isConnecting (ws) {
+  // If the WebSocket connection is not yet established, and the connection
+  // is not yet closed, then the WebSocket connection is in the CONNECTING state.
+  return ws[kReadyState] === states.CONNECTING
+}
+
+/**
+ * @param {import('./websocket').WebSocket} ws
+ * @returns {boolean}
  */
 function isEstablished (ws) {
   // If the server's response is validated as provided for above, it is
@@ -18,6 +29,7 @@ function isEstablished (ws) {
 
 /**
  * @param {import('./websocket').WebSocket} ws
+ * @returns {boolean}
  */
 function isClosing (ws) {
   // Upon either sending or receiving a Close control frame, it is said
@@ -28,6 +40,7 @@ function isClosing (ws) {
 
 /**
  * @param {import('./websocket').WebSocket} ws
+ * @returns {boolean}
  */
 function isClosed (ws) {
   return ws[kReadyState] === states.CLOSED
@@ -190,6 +203,7 @@ function failWebsocketConnection (ws, reason) {
 }
 
 module.exports = {
+  isConnecting,
   isEstablished,
   isClosing,
   isClosed,

lib/web/websocket/websocket.js

@@ -3,7 +3,7 @@
 const { webidl } = require('../fetch/webidl')
 const { URLSerializer } = require('../fetch/data-url')
 const { getGlobalOrigin } = require('../fetch/global')
-const { staticPropertyDescriptors, states, opcodes, emptyBuffer } = require('./constants')
+const { staticPropertyDescriptors, states, sentCloseFrameState, opcodes, emptyBuffer } = require('./constants')
 const {
   kWebSocketURL,
   kReadyState,
@@ -13,7 +13,15 @@ const {
   kSentClose,
   kByteParser
 } = require('./symbols')
-const { isEstablished, isClosing, isValidSubprotocol, failWebsocketConnection, fireEvent } = require('./util')
+const {
+  isConnecting,
+  isEstablished,
+  isClosed,
+  isClosing,
+  isValidSubprotocol,
+  failWebsocketConnection,
+  fireEvent
+} = require('./util')
 const { establishWebSocketConnection } = require('./connection')
 const { WebsocketFrameSend } = require('./frame')
 const { ByteParser } = require('./receiver')
@@ -132,6 +140,8 @@ class WebSocket extends EventTarget {
     // be CONNECTING (0).
     this[kReadyState] = WebSocket.CONNECTING
 
+    this[kSentClose] = sentCloseFrameState.NOT_SENT
+
     // The extensions attribute must initially return the empty string.
 
     // The protocol attribute must initially return the empty string.
@@ -184,7 +194,7 @@ class WebSocket extends EventTarget {
     }
 
     // 3. Run the first matching steps from the following list:
-    if (this[kReadyState] === WebSocket.CLOSING || this[kReadyState] === WebSocket.CLOSED) {
+    if (isClosing(this) || isClosed(this)) {
       // If this's ready state is CLOSING (2) or CLOSED (3)
       // Do nothing.
     } else if (!isEstablished(this)) {
@@ -193,7 +203,7 @@ class WebSocket extends EventTarget {
       // to CLOSING (2).
       failWebsocketConnection(this, 'Connection was closed before it was established.')
       this[kReadyState] = WebSocket.CLOSING
-    } else if (!isClosing(this)) {
+    } else if (this[kSentClose] === sentCloseFrameState.NOT_SENT) {
       // If the WebSocket closing handshake has not yet been started
       // Start the WebSocket closing handshake and set this's ready
       // state to CLOSING (2).
@@ -204,6 +214,8 @@ class WebSocket extends EventTarget {
       // - If reason is also present, then reasonBytes must be
       //   provided in the Close message after the status code.
 
+      this[kSentClose] = sentCloseFrameState.PROCESSING
+
       const frame = new WebsocketFrameSend()
 
       // If neither code nor reason is present, the WebSocket Close
@@ -230,7 +242,7 @@ class WebSocket extends EventTarget {
 
       socket.write(frame.createFrame(opcodes.CLOSE), (err) => {
         if (!err) {
-          this[kSentClose] = true
+          this[kSentClose] = sentCloseFrameState.SENT
         }
       })
 
@@ -258,7 +270,7 @@ class WebSocket extends EventTarget {
 
     // 1. If this's ready state is CONNECTING, then throw an
     //    "InvalidStateError" DOMException.
-    if (this[kReadyState] === WebSocket.CONNECTING) {
+    if (isConnecting(this)) {
       throw new DOMException('Sent before connected.', 'InvalidStateError')
     }
 

test/websocket/close.js

@@ -1,6 +1,7 @@
 'use strict'
 
-const { describe, test } = require('node:test')
+const { tspl } = require('@matteo.collina/tspl')
+const { describe, test, after } = require('node:test')
 const assert = require('node:assert')
 const { WebSocketServer } = require('ws')
 const { WebSocket } = require('../..')
@@ -128,4 +129,26 @@ describe('Close', () => {
       ws.addEventListener('open', () => ws.close(3000))
     })
   })
+
+  test('calling close twice will only trigger the close event once', async (t) => {
+    t = tspl(t, { plan: 1 })
+
+    const server = new WebSocketServer({ port: 0 })
+
+    after(() => server.close())
+
+    server.on('connection', (ws) => {
+      ws.on('close', (code) => {
+        t.strictEqual(code, 1000)
+      })
+    })
+
+    const ws = new WebSocket(`ws://localhost:${server.address().port}`)
+    ws.addEventListener('open', () => {
+      ws.close(1000)
+      ws.close(1000)
+    })
+
+    await t.completed
+  })
 })