msw 2.5.1 → 2.6.0

package/src/core/ws/WebSocketClientManager.test.ts

@@ -0,0 +1,164 @@
+// @vitest-environment node-websocket
+import { setMaxListeners } from 'node:events'
+import {
+  WebSocketClientConnection,
+  WebSocketData,
+  WebSocketTransport,
+} from '@mswjs/interceptors/WebSocket'
+import {
+  WebSocketClientManager,
+  WebSocketBroadcastChannelMessage,
+} from './WebSocketClientManager'
+
+const channel = new BroadcastChannel('test:channel')
+
+/**
+ * @note Increase the number of maximum event listeners
+ * because the same channel is shared between different
+ * manager instances in different tests.
+ */
+setMaxListeners(Number.MAX_SAFE_INTEGER, channel)
+
+vi.spyOn(channel, 'postMessage')
+
+const socket = new WebSocket('ws://localhost')
+
+class TestWebSocketTransport extends EventTarget implements WebSocketTransport {
+  send(_data: WebSocketData): void {}
+  close(_code?: number | undefined, _reason?: string | undefined): void {}
+}
+
+afterEach(() => {
+  vi.resetAllMocks()
+})
+
+it('adds a client from this runtime to the list of clients', async () => {
+  const manager = new WebSocketClientManager(channel)
+  const connection = new WebSocketClientConnection(
+    socket,
+    new TestWebSocketTransport(),
+  )
+
+  await manager.addConnection(connection)
+
+  // Must add the client to the list of clients.
+  expect(Array.from(manager.clients.values())).toEqual([connection])
+})
+
+it('adds multiple clients from this runtime to the list of clients', async () => {
+  const manager = new WebSocketClientManager(channel)
+  const connectionOne = new WebSocketClientConnection(
+    socket,
+    new TestWebSocketTransport(),
+  )
+  await manager.addConnection(connectionOne)
+
+  // Must add the client to the list of clients.
+  expect(Array.from(manager.clients.values())).toEqual([connectionOne])
+
+  const connectionTwo = new WebSocketClientConnection(
+    socket,
+    new TestWebSocketTransport(),
+  )
+  await manager.addConnection(connectionTwo)
+
+  // Must add the new cilent to the list as well.
+  expect(Array.from(manager.clients.values())).toEqual([
+    connectionOne,
+    connectionTwo,
+  ])
+})
+
+it('replays a "send" event coming from another runtime', async () => {
+  const manager = new WebSocketClientManager(channel)
+  const connection = new WebSocketClientConnection(
+    socket,
+    new TestWebSocketTransport(),
+  )
+  await manager.addConnection(connection)
+  vi.spyOn(connection, 'send')
+
+  // Emulate another runtime signaling this connection to receive data.
+  channel.dispatchEvent(
+    new MessageEvent<WebSocketBroadcastChannelMessage>('message', {
+      data: {
+        type: 'extraneous:send',
+        payload: {
+          clientId: connection.id,
+          data: 'hello',
+        },
+      },
+    }),
+  )
+
+  await vi.waitFor(() => {
+    // Must execute the requested operation on the connection.
+    expect(connection.send).toHaveBeenCalledWith('hello')
+    expect(connection.send).toHaveBeenCalledTimes(1)
+  })
+})
+
+it('replays a "close" event coming from another runtime', async () => {
+  const manager = new WebSocketClientManager(channel)
+  const connection = new WebSocketClientConnection(
+    socket,
+    new TestWebSocketTransport(),
+  )
+  await manager.addConnection(connection)
+  vi.spyOn(connection, 'close')
+
+  // Emulate another runtime signaling this connection to close.
+  channel.dispatchEvent(
+    new MessageEvent<WebSocketBroadcastChannelMessage>('message', {
+      data: {
+        type: 'extraneous:close',
+        payload: {
+          clientId: connection.id,
+          code: 1000,
+          reason: 'Normal closure',
+        },
+      },
+    }),
+  )
+
+  await vi.waitFor(() => {
+    // Must execute the requested operation on the connection.
+    expect(connection.close).toHaveBeenCalledWith(1000, 'Normal closure')
+    expect(connection.close).toHaveBeenCalledTimes(1)
+  })
+})
+
+it('removes the extraneous message listener when the connection closes', async () => {
+  const manager = new WebSocketClientManager(channel)
+  const transport = new TestWebSocketTransport()
+  const connection = new WebSocketClientConnection(socket, transport)
+  vi.spyOn(connection, 'close').mockImplementationOnce(() => {
+    /**
+     * @note This is a nasty hack so we don't have to uncouple
+     * the connection from transport. Creating a mock transport
+     * is difficult because it relies on the `WebSocketOverride` class.
+     * All we care here is that closing the connection triggers
+     * the transport closure, which it always does.
+     */
+    transport.dispatchEvent(new Event('close'))
+  })
+  vi.spyOn(connection, 'send')
+
+  await manager.addConnection(connection)
+  connection.close()
+
+  // Signals from other runtimes have no effect on the closed connection.
+  channel.dispatchEvent(
+    new MessageEvent<WebSocketBroadcastChannelMessage>('message', {
+      data: {
+        type: 'extraneous:send',
+        payload: {
+          clientId: connection.id,
+          data: 'hello',
+        },
+      },
+    }),
+  )
+
+  expect(connection.send).not.toHaveBeenCalled()
+})

package/src/core/ws/WebSocketClientManager.ts

@@ -0,0 +1,211 @@
+import type {
+  WebSocketData,
+  WebSocketClientConnection,
+  WebSocketClientConnectionProtocol,
+} from '@mswjs/interceptors/WebSocket'
+import { WebSocketClientStore } from './WebSocketClientStore'
+import { WebSocketMemoryClientStore } from './WebSocketMemoryClientStore'
+import { WebSocketIndexedDBClientStore } from './WebSocketIndexedDBClientStore'
+
+export type WebSocketBroadcastChannelMessage =
+  | {
+      type: 'extraneous:send'
+      payload: {
+        clientId: string
+        data: WebSocketData
+      }
+    }
+  | {
+      type: 'extraneous:close'
+      payload: {
+        clientId: string
+        code?: number
+        reason?: string
+      }
+    }
+
+/**
+ * A manager responsible for accumulating WebSocket client
+ * connections across different browser runtimes.
+ */
+export class WebSocketClientManager {
+  private store: WebSocketClientStore
+  private runtimeClients: Map<string, WebSocketClientConnectionProtocol>
+  private allClients: Set<WebSocketClientConnectionProtocol>
+
+  constructor(private channel: BroadcastChannel) {
+    // Store the clients in the IndexedDB in the browser,
+    // otherwise, store the clients in memory.
+    this.store =
+      typeof indexedDB !== 'undefined'
+        ? new WebSocketIndexedDBClientStore()
+        : new WebSocketMemoryClientStore()
+
+    this.runtimeClients = new Map()
+    this.allClients = new Set()
+
+    this.channel.addEventListener('message', (message) => {
+      if (message.data?.type === 'db:update') {
+        this.flushDatabaseToMemory()
+      }
+    })
+
+    if (typeof window !== 'undefined') {
+      window.addEventListener('message', async (message) => {
+        if (message.data?.type === 'msw/worker:stop') {
+          await this.removeRuntimeClients()
+        }
+      })
+    }
+  }
+
+  private async flushDatabaseToMemory() {
+    const storedClients = await this.store.getAll()
+
+    this.allClients = new Set(
+      storedClients.map((client) => {
+        const runtimeClient = this.runtimeClients.get(client.id)
+
+        /**
+         * @note For clients originating in this runtime, use their
+         * direct references. No need to wrap them in a remote connection.
+         */
+        if (runtimeClient) {
+          return runtimeClient
+        }
+
+        return new WebSocketRemoteClientConnection(
+          client.id,
+          new URL(client.url),
+          this.channel,
+        )
+      }),
+    )
+  }
+
+  private async removeRuntimeClients(): Promise<void> {
+    await this.store.deleteMany(Array.from(this.runtimeClients.keys()))
+    this.runtimeClients.clear()
+    await this.flushDatabaseToMemory()
+    this.notifyOthersAboutDatabaseUpdate()
+  }
+
+  /**
+   * All active WebSocket client connections.
+   */
+  get clients(): Set<WebSocketClientConnectionProtocol> {
+    return this.allClients
+  }
+
+  /**
+   * Notify other runtimes about the database update
+   * using the shared `BroadcastChannel` instance.
+   */
+  private notifyOthersAboutDatabaseUpdate(): void {
+    this.channel.postMessage({ type: 'db:update' })
+  }
+
+  private async addClient(client: WebSocketClientConnection): Promise<void> {
+    await this.store.add(client)
+    // Sync the in-memory clients in this runtime with the
+    // updated database. This pulls in all the stored clients.
+    await this.flushDatabaseToMemory()
+    this.notifyOthersAboutDatabaseUpdate()
+  }
+
+  /**
+   * Adds the given `WebSocket` client connection to the set
+   * of all connections. The given connection is always the complete
+   * connection object because `addConnection()` is called only
+   * for the opened connections in the same runtime.
+   */
+  public async addConnection(client: WebSocketClientConnection): Promise<void> {
+    // Store this client in the map of clients created in this runtime.
+    // This way, the manager can distinguish between this runtime clients
+    // and extraneous runtime clients when synchronizing clients storage.
+    this.runtimeClients.set(client.id, client)
+
+    // Add the new client to the storage.
+    await this.addClient(client)
+
+    // Handle the incoming BroadcastChannel messages from other runtimes
+    // that attempt to control this runtime (via a remote connection wrapper).
+    // E.g. another runtime calling `client.send()` for the client in this runtime.
+    const handleExtraneousMessage = (
+      message: MessageEvent<WebSocketBroadcastChannelMessage>,
+    ) => {
+      const { type, payload } = message.data
+
+      // Ignore broadcasted messages for other clients.
+      if (
+        typeof payload === 'object' &&
+        'clientId' in payload &&
+        payload.clientId !== client.id
+      ) {
+        return
+      }
+
+      switch (type) {
+        case 'extraneous:send': {
+          client.send(payload.data)
+          break
+        }
+
+        case 'extraneous:close': {
+          client.close(payload.code, payload.reason)
+          break
+        }
+      }
+    }
+
+    const abortController = new AbortController()
+
+    this.channel.addEventListener('message', handleExtraneousMessage, {
+      signal: abortController.signal,
+    })
+
+    // Once closed, this connection cannot be operated on.
+    // This must include the extraneous runtimes as well.
+    client.addEventListener('close', () => abortController.abort(), {
+      once: true,
+    })
+  }
+}
+
+/**
+ * A wrapper class to operate with WebSocket client connections
+ * from other runtimes. This class maintains 1-1 public API
+ * compatibility to the `WebSocketClientConnection` but relies
+ * on the given `BroadcastChannel` to communicate instructions
+ * with the client connections from other runtimes.
+ */
+export class WebSocketRemoteClientConnection
+  implements WebSocketClientConnectionProtocol
+{
+  constructor(
+    public readonly id: string,
+    public readonly url: URL,
+    private channel: BroadcastChannel,
+  ) {}
+
+  send(data: WebSocketData): void {
+    this.channel.postMessage({
+      type: 'extraneous:send',
+      payload: {
+        clientId: this.id,
+        data,
+      },
+    } as WebSocketBroadcastChannelMessage)
+  }
+
+  close(code?: number | undefined, reason?: string | undefined): void {
+    this.channel.postMessage({
+      type: 'extraneous:close',
+      payload: {
+        clientId: this.id,
+        code,
+        reason,
+      },
+    } as WebSocketBroadcastChannelMessage)
+  }
+}

package/src/core/ws/WebSocketClientStore.ts

@@ -0,0 +1,14 @@
+import type { WebSocketClientConnectionProtocol } from '@mswjs/interceptors/WebSocket'
+
+export interface SerializedWebSocketClient {
+  id: string
+  url: string
+}
+
+export abstract class WebSocketClientStore {
+  public abstract add(client: WebSocketClientConnectionProtocol): Promise<void>
+
+  public abstract getAll(): Promise<Array<SerializedWebSocketClient>>
+
+  public abstract deleteMany(clientIds: Array<string>): Promise<void>
+}

package/src/core/ws/WebSocketIndexedDBClientStore.ts

@@ -0,0 +1,145 @@
+import { DeferredPromise } from '@open-draft/deferred-promise'
+import { WebSocketClientConnectionProtocol } from '@mswjs/interceptors/lib/browser/interceptors/WebSocket'
+import {
+  type SerializedWebSocketClient,
+  WebSocketClientStore,
+} from './WebSocketClientStore'
+
+const DB_NAME = 'msw-websocket-clients'
+const DB_STORE_NAME = 'clients'
+
+export class WebSocketIndexedDBClientStore implements WebSocketClientStore {
+  private db: Promise<IDBDatabase>
+
+  constructor() {
+    this.db = this.createDatabase()
+  }
+
+  public async add(client: WebSocketClientConnectionProtocol): Promise<void> {
+    const promise = new DeferredPromise<void>()
+    const store = await this.getStore()
+
+    /**
+     * @note Use `.put()` instead of `.add()` to allow setting clients
+     * that already exist in the database. This can happen if a single page
+     * has multiple event handlers. Each handler will receive the "connection"
+     * event in parallel, and try to set that WebSocket client in the database.
+     */
+    const request = store.put({
+      id: client.id,
+      url: client.url.href,
+    } satisfies SerializedWebSocketClient)
+
+    request.onsuccess = () => {
+      promise.resolve()
+    }
+    request.onerror = () => {
+      // eslint-disable-next-line no-console
+      console.error(request.error)
+      promise.reject(
+        new Error(
+          `Failed to add WebSocket client "${client.id}". There is likely an additional output above.`,
+        ),
+      )
+    }
+
+    return promise
+  }
+
+  public async getAll(): Promise<Array<SerializedWebSocketClient>> {
+    const promise = new DeferredPromise<Array<SerializedWebSocketClient>>()
+    const store = await this.getStore()
+    const request = store.getAll() as IDBRequest<
+      Array<SerializedWebSocketClient>
+    >
+
+    request.onsuccess = () => {
+      promise.resolve(request.result)
+    }
+    request.onerror = () => {
+      // eslint-disable-next-line no-console
+      console.log(request.error)
+      promise.reject(
+        new Error(
+          `Failed to get all WebSocket clients. There is likely an additional output above.`,
+        ),
+      )
+    }
+
+    return promise
+  }
+
+  public async deleteMany(clientIds: Array<string>): Promise<void> {
+    const promise = new DeferredPromise<void>()
+    const store = await this.getStore()
+
+    for (const clientId of clientIds) {
+      store.delete(clientId)
+    }
+
+    store.transaction.oncomplete = () => {
+      promise.resolve()
+    }
+    store.transaction.onerror = () => {
+      // eslint-disable-next-line no-console
+      console.error(store.transaction.error)
+      promise.reject(
+        new Error(
+          `Failed to delete WebSocket clients [${clientIds.join(', ')}]. There is likely an additional output above.`,
+        ),
+      )
+    }
+
+    return promise
+  }
+
+  private async createDatabase(): Promise<IDBDatabase> {
+    const promise = new DeferredPromise<IDBDatabase>()
+    const request = indexedDB.open(DB_NAME, 1)
+
+    request.onsuccess = ({ currentTarget }) => {
+      const db = Reflect.get(currentTarget!, 'result') as IDBDatabase
+
+      if (db.objectStoreNames.contains(DB_STORE_NAME)) {
+        return promise.resolve(db)
+      }
+    }
+
+    request.onupgradeneeded = async ({ currentTarget }) => {
+      const db = Reflect.get(currentTarget!, 'result') as IDBDatabase
+      if (db.objectStoreNames.contains(DB_STORE_NAME)) {
+        return
+      }
+
+      const store = db.createObjectStore(DB_STORE_NAME, { keyPath: 'id' })
+      store.transaction.oncomplete = () => {
+        promise.resolve(db)
+      }
+      store.transaction.onerror = () => {
+        // eslint-disable-next-line no-console
+        console.error(store.transaction.error)
+        promise.reject(
+          new Error(
+            'Failed to create WebSocket client store. There is likely an additional output above.',
+          ),
+        )
+      }
+    }
+    request.onerror = () => {
+      // eslint-disable-next-line no-console
+      console.error(request.error)
+      promise.reject(
+        new Error(
+          'Failed to open an IndexedDB database. There is likely an additional output above.',
+        ),
+      )
+    }
+
+    return promise
+  }
+
+  private async getStore(): Promise<IDBObjectStore> {
+    const db = await this.db
+    return db.transaction(DB_STORE_NAME, 'readwrite').objectStore(DB_STORE_NAME)
+  }
+}

package/src/core/ws/WebSocketMemoryClientStore.ts

@@ -0,0 +1,27 @@
+import { WebSocketClientConnectionProtocol } from '@mswjs/interceptors/lib/browser/interceptors/WebSocket'
+import {
+  SerializedWebSocketClient,
+  WebSocketClientStore,
+} from './WebSocketClientStore'
+
+export class WebSocketMemoryClientStore implements WebSocketClientStore {
+  private store: Map<string, SerializedWebSocketClient>
+
+  constructor() {
+    this.store = new Map()
+  }
+
+  public async add(client: WebSocketClientConnectionProtocol): Promise<void> {
+    this.store.set(client.id, { id: client.id, url: client.url.href })
+  }
+
+  public getAll(): Promise<Array<SerializedWebSocketClient>> {
+    return Promise.resolve(Array.from(this.store.values()))
+  }
+
+  public async deleteMany(clientIds: Array<string>): Promise<void> {
+    for (const clientId of clientIds) {
+      this.store.delete(clientId)
+    }
+  }
+}

package/src/core/ws/handleWebSocketEvent.ts

@@ -0,0 +1,82 @@
+import type { WebSocketConnectionData } from '@mswjs/interceptors/lib/browser/interceptors/WebSocket'
+import { RequestHandler } from '../handlers/RequestHandler'
+import { WebSocketHandler, kDispatchEvent } from '../handlers/WebSocketHandler'
+import { webSocketInterceptor } from './webSocketInterceptor'
+import {
+  onUnhandledRequest,
+  UnhandledRequestStrategy,
+} from '../utils/request/onUnhandledRequest'
+
+interface HandleWebSocketEventOptions {
+  getUnhandledRequestStrategy: () => UnhandledRequestStrategy
+  getHandlers: () => Array<RequestHandler | WebSocketHandler>
+  onMockedConnection: (connection: WebSocketConnectionData) => void
+  onPassthroughConnection: (onnection: WebSocketConnectionData) => void
+}
+
+export function handleWebSocketEvent(options: HandleWebSocketEventOptions) {
+  webSocketInterceptor.on('connection', async (connection) => {
+    const handlers = options.getHandlers()
+
+    const connectionEvent = new MessageEvent('connection', {
+      data: connection,
+    })
+
+    // First, filter only those WebSocket handlers that
+    // match the "ws.link()" endpoint predicate. Don't dispatch
+    // anything yet so the logger can be attached to the connection
+    // before it potentially sends events.
+    const matchingHandlers: Array<WebSocketHandler> = []
+
+    for (const handler of handlers) {
+      if (
+        handler instanceof WebSocketHandler &&
+        handler.predicate({
+          event: connectionEvent,
+          parsedResult: handler.parse({
+            event: connectionEvent,
+          }),
+        })
+      ) {
+        matchingHandlers.push(handler)
+      }
+    }
+
+    if (matchingHandlers.length > 0) {
+      options?.onMockedConnection(connection)
+
+      // Iterate over the handlers and forward the connection
+      // event to WebSocket event handlers. This is equivalent
+      // to dispatching that event onto multiple listeners.
+      for (const handler of matchingHandlers) {
+        handler[kDispatchEvent](connectionEvent)
+      }
+    } else {
+      // Construct a request representing this WebSocket connection.
+      const request = new Request(connection.client.url, {
+        headers: {
+          upgrade: 'websocket',
+          connection: 'upgrade',
+        },
+      })
+      await onUnhandledRequest(
+        request,
+        options.getUnhandledRequestStrategy(),
+      ).catch((error) => {
+        const errorEvent = new Event('error')
+        Object.defineProperty(errorEvent, 'cause', {
+          enumerable: true,
+          configurable: false,
+          value: error,
+        })
+        connection.client.socket.dispatchEvent(errorEvent)
+      })
+
+      options?.onPassthroughConnection(connection)
+
+      // If none of the "ws" handlers matched,
+      // establish the WebSocket connection as-is.
+      connection.server.connect()
+    }
+  })
+}