fix(WebSocket): make server "close" event cancelable (#645)

src/interceptors/WebSocket/WebSocketServerConnection.ts

@@ -7,7 +7,11 @@ import {
 import type { WebSocketData } from './WebSocketTransport'
 import type { WebSocketClassTransport } from './WebSocketClassTransport'
 import { bindEvent } from './utils/bindEvent'
-import { CancelableMessageEvent, CloseEvent } from './utils/events'
+import {
+  CancelableMessageEvent,
+  CancelableCloseEvent,
+  CloseEvent,
+} from './utils/events'
 
 const kEmitter = Symbol('kEmitter')
 const kBoundListener = Symbol('kBoundListener')
@@ -265,12 +269,13 @@ export class WebSocketServerConnection {
       this[kEmitter].dispatchEvent(
         bindEvent(
           this.realWebSocket,
-          new CloseEvent('close', {
+          new CancelableCloseEvent('close', {
             /**
              * @note `server.close()` in the interceptor
              * always results in clean closures.
              */
             code: 1000,
+            cancelable: true,
           })
         )
       )
@@ -339,7 +344,12 @@ export class WebSocketServerConnection {
 
     const closeEvent = bindEvent(
       this.realWebSocket,
-      new CloseEvent('close', event)
+      new CancelableCloseEvent('close', {
+        code: event.code,
+        reason: event.reason,
+        wasClean: event.wasClean,
+        cancelable: true,
+      })
     )
 
     this[kEmitter].dispatchEvent(closeEvent)

src/interceptors/WebSocket/utils/events.test.ts

@@ -17,7 +17,7 @@ describe(CancelableMessageEvent, () => {
     expect(event.defaultPrevented).toBe(false)
   })
 
-  it('initiates a canceaable event', () => {
+  it('initiates a cancelable event', () => {
     const event = new CancelableMessageEvent('message', {
       data: 'hello',
       cancelable: true,

src/interceptors/WebSocket/utils/events.ts

@@ -59,3 +59,36 @@ export class CloseEvent extends Event {
     this.wasClean = init.wasClean === undefined ? false : init.wasClean
   }
 }
+
+export class CancelableCloseEvent extends CloseEvent {
+  [kCancelable]: boolean;
+  [kDefaultPrevented]: boolean
+
+  constructor(type: string, init: CloseEventInit = {}) {
+    super(type, init)
+    this[kCancelable] = !!init.cancelable
+    this[kDefaultPrevented] = false
+  }
+
+  get cancelable() {
+    return this[kCancelable]
+  }
+
+  set cancelable(nextCancelable) {
+    this[kCancelable] = nextCancelable
+  }
+
+  get defaultPrevented() {
+    return this[kDefaultPrevented]
+  }
+
+  set defaultPrevented(nextDefaultPrevented) {
+    this[kDefaultPrevented] = nextDefaultPrevented
+  }
+
+  public preventDefault(): void {
+    if (this.cancelable && !this[kDefaultPrevented]) {
+      this[kDefaultPrevented] = true
+    }
+  }
+}

test/modules/WebSocket/compliance/websocket.server.close.test.ts

@@ -170,7 +170,7 @@ it('forwards "close" events from the original server', async () => {
     })
   })
 
-  interceptor.once('connection', ({ client, server }) => {
+  interceptor.once('connection', ({ server }) => {
     server.connect()
     server.addEventListener('close', (event) => {
       interceptorServerCloseListener(event.code, event.reason)

test/modules/WebSocket/compliance/websocket.server.events.test.ts

@@ -3,6 +3,7 @@ import { vi, it, expect, beforeAll, afterEach, afterAll } from 'vitest'
 import { WebSocketServer } from 'ws'
 import { WebSocketInterceptor } from '../../../../src/interceptors/WebSocket/index'
 import { getWsUrl } from '../utils/getWsUrl'
+import { waitForWebSocketEvent } from '../utils/waitForWebSocketEvent'
 
 const interceptor = new WebSocketInterceptor()
 
@@ -155,3 +156,34 @@ it('emits both "error" and "close" events when the server connection errors', as
   // Must emit the correct events on the WebSocket client.
   expect(instanceErrorListener).toHaveBeenCalledTimes(1)
 })
+
+it('prevents "close" event forwarding by calling "event.preventDefault()"', async () => {
+  wsServer.once('connection', (ws) => {
+    ws.close(1003, 'Server reason')
+  })
+
+  interceptor.once('connection', ({ server, client }) => {
+    server.connect()
+    server.addEventListener('close', (event) => {
+      expect(event.defaultPrevented).toBe(false)
+      event.preventDefault()
+      expect(event.defaultPrevented).toBe(true)
+    })
+  })
+
+  const client = new WebSocket(getWsUrl(wsServer))
+  const closeListener = vi.fn()
+  const errorListener = vi.fn()
+  client.addEventListener('close', closeListener)
+  client.addEventListener('error', errorListener)
+
+  await waitForWebSocketEvent('open', client)
+
+  const closePromise = vi.waitFor(() => {
+    expect(closeListener).toHaveBeenCalledTimes(1)
+    return closeListener.mock.calls.length
+  })
+
+  await expect(closePromise).rejects.toThrow()
+  expect(errorListener).not.toHaveBeenCalled()
+})