Skip to content

Commit

Permalink
websocket: improve .close() (#2865)
Browse files Browse the repository at this point in the history
  • Loading branch information
Uzlopak authored Mar 3, 2024
1 parent 1216ba0 commit 2505e42
Show file tree
Hide file tree
Showing 6 changed files with 69 additions and 13 deletions.
6 changes: 3 additions & 3 deletions lib/web/websocket/connection.js
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
'use strict'

const { uid, states } = require('./constants')
const { uid, states, sentCloseFrameState } = require('./constants')
const {
kReadyState,
kSentClose,
Expand Down Expand Up @@ -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 = ''
Expand All @@ -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
Expand Down
7 changes: 7 additions & 0 deletions lib/web/websocket/constants.js
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,12 @@ const states = {
CLOSED: 3
}

const sentCloseFrameState = {
NOT_SENT: 0,
PROCESSING: 1,
SENT: 2
}

const opcodes = {
CONTINUATION: 0x0,
TEXT: 0x1,
Expand All @@ -42,6 +48,7 @@ const emptyBuffer = Buffer.allocUnsafe(0)

module.exports = {
uid,
sentCloseFrameState,
staticPropertyDescriptors,
states,
opcodes,
Expand Down
6 changes: 3 additions & 3 deletions lib/web/websocket/receiver.js
Original file line number Diff line number Diff line change
@@ -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')
Expand Down Expand Up @@ -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
Expand All @@ -120,7 +120,7 @@ class ByteParser extends Writable {
closeFrame.createFrame(opcodes.CLOSE),
(err) => {
if (!err) {
this.ws[kSentClose] = true
this.ws[kSentClose] = sentCloseFrameState.SENT
}
}
)
Expand Down
14 changes: 14 additions & 0 deletions lib/web/websocket/util.js
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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
Expand All @@ -28,6 +40,7 @@ function isClosing (ws) {

/**
* @param {import('./websocket').WebSocket} ws
* @returns {boolean}
*/
function isClosed (ws) {
return ws[kReadyState] === states.CLOSED
Expand Down Expand Up @@ -190,6 +203,7 @@ function failWebsocketConnection (ws, reason) {
}

module.exports = {
isConnecting,
isEstablished,
isClosing,
isClosed,
Expand Down
24 changes: 18 additions & 6 deletions lib/web/websocket/websocket.js
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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')
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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)) {
Expand All @@ -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).
Expand All @@ -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
Expand All @@ -230,7 +242,7 @@ class WebSocket extends EventTarget {

socket.write(frame.createFrame(opcodes.CLOSE), (err) => {
if (!err) {
this[kSentClose] = true
this[kSentClose] = sentCloseFrameState.SENT
}
})

Expand Down Expand Up @@ -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')
}

Expand Down
25 changes: 24 additions & 1 deletion test/websocket/close.js
Original file line number Diff line number Diff line change
@@ -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('../..')
Expand Down Expand Up @@ -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
})
})

0 comments on commit 2505e42

Please sign in to comment.