fix(WebSocketClientManager): use localStorage for clients persistence (…

…#2127)

src/browser/setupWorker/stop/createStop.ts

@@ -1,4 +1,5 @@
 import { devUtils } from '~/core/utils/internal/devUtils'
+import { MSW_WEBSOCKET_CLIENTS_KEY } from '~/core/ws/WebSocketClientManager'
 import { SetupWorkerInternalContext, StopHandler } from '../glossary'
 import { printStopMessage } from './utils/printStopMessage'
 
@@ -24,6 +25,9 @@ export const createStop = (
     context.isMockingEnabled = false
     window.clearInterval(context.keepAliveInterval)
 
+    // Clear the WebSocket clients from the shared storage.
+    localStorage.removeItem(MSW_WEBSOCKET_CLIENTS_KEY)
+
     printStopMessage({ quiet: context.startOptions?.quiet })
   }
 }

src/core/ws.ts

@@ -72,10 +72,12 @@ function createWebSocketLinkHandler(url: Path): WebSocketLink {
     typeof url,
   )
 
-  const clientManager = new WebSocketClientManager(wsBroadcastChannel)
+  const clientManager = new WebSocketClientManager(wsBroadcastChannel, url)
 
   return {
-    clients: clientManager.clients,
+    get clients() {
+      return clientManager.clients
+    },
     on(event, listener) {
       const handler = new WebSocketHandler(url)
 

src/core/ws/WebSocketClientManager.test.ts

@@ -1,79 +1,73 @@
 /**
  * @vitest-environment node-websocket
  */
-import { randomUUID } from 'node:crypto'
 import {
   WebSocketClientConnection,
+  WebSocketData,
   WebSocketTransport,
 } from '@mswjs/interceptors/WebSocket'
 import {
   WebSocketClientManager,
   WebSocketBroadcastChannelMessage,
-  WebSocketRemoteClientConnection,
 } from './WebSocketClientManager'
 
 const channel = new BroadcastChannel('test:channel')
 vi.spyOn(channel, 'postMessage')
 
 const socket = new WebSocket('ws://localhost')
-const transport = {
-  onOutgoing: vi.fn(),
-  onIncoming: vi.fn(),
-  onClose: vi.fn(),
-  send: vi.fn(),
-  close: vi.fn(),
-} satisfies WebSocketTransport
+
+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', () => {
-  const manager = new WebSocketClientManager(channel)
-  const connection = new WebSocketClientConnection(socket, transport)
+  const manager = new WebSocketClientManager(channel, '*')
+  const connection = new WebSocketClientConnection(
+    socket,
+    new TestWebSocketTransport(),
+  )
 
   manager.addConnection(connection)
 
   // Must add the client to the list of clients.
   expect(Array.from(manager.clients.values())).toEqual([connection])
-
-  // Must emit the connection open event to notify other runtimes.
-  expect(channel.postMessage).toHaveBeenCalledWith({
-    type: 'connection:open',
-    payload: {
-      clientId: connection.id,
-      url: socket.url,
-    },
-  } satisfies WebSocketBroadcastChannelMessage)
 })
 
-it('adds a client from another runtime to the list of clients', async () => {
-  const clientId = randomUUID()
-  const url = new URL('ws://localhost')
-  const manager = new WebSocketClientManager(channel)
+it('adds multiple clients from this runtime to the list of clients', () => {
+  const manager = new WebSocketClientManager(channel, '*')
+  const connectionOne = new WebSocketClientConnection(
+    socket,
+    new TestWebSocketTransport(),
+  )
+  manager.addConnection(connectionOne)
 
-  channel.dispatchEvent(
-    new MessageEvent<WebSocketBroadcastChannelMessage>('message', {
-      data: {
-        type: 'connection:open',
-        payload: {
-          clientId,
-          url: url.href,
-        },
-      },
-    }),
+  // Must add the client to the list of clients.
+  expect(Array.from(manager.clients.values())).toEqual([connectionOne])
+
+  const connectionTwo = new WebSocketClientConnection(
+    socket,
+    new TestWebSocketTransport(),
   )
+  manager.addConnection(connectionTwo)
 
-  await vi.waitFor(() => {
-    expect(Array.from(manager.clients.values())).toEqual([
-      new WebSocketRemoteClientConnection(clientId, url, channel),
-    ])
-  })
+  // 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, transport)
+  const manager = new WebSocketClientManager(channel, '*')
+  const connection = new WebSocketClientConnection(
+    socket,
+    new TestWebSocketTransport(),
+  )
   manager.addConnection(connection)
   vi.spyOn(connection, 'send')
 
@@ -98,8 +92,11 @@ it('replays a "send" event coming from another runtime', async () => {
 })
 
 it('replays a "close" event coming from another runtime', async () => {
-  const manager = new WebSocketClientManager(channel)
-  const connection = new WebSocketClientConnection(socket, transport)
+  const manager = new WebSocketClientManager(channel, '*')
+  const connection = new WebSocketClientConnection(
+    socket,
+    new TestWebSocketTransport(),
+  )
   manager.addConnection(connection)
   vi.spyOn(connection, 'close')
 
@@ -125,7 +122,8 @@ it('replays a "close" event coming from another runtime', async () => {
 })
 
 it('removes the extraneous message listener when the connection closes', async () => {
-  const manager = new WebSocketClientManager(channel)
+  const manager = new WebSocketClientManager(channel, '*')
+  const transport = new TestWebSocketTransport()
   const connection = new WebSocketClientConnection(socket, transport)
   vi.spyOn(connection, 'close').mockImplementationOnce(() => {
     /**
@@ -135,7 +133,7 @@ it('removes the extraneous message listener when the connection closes', async (
      * All we care here is that closing the connection triggers
      * the transport closure, which it always does.
      */
-    connection['transport'].onClose()
+    transport.dispatchEvent(new Event('close'))
   })
   vi.spyOn(connection, 'send')
 

src/core/ws/WebSocketClientManager.ts

@@ -1,17 +1,14 @@
+import { invariant } from 'outvariant'
 import type {
   WebSocketData,
   WebSocketClientConnection,
   WebSocketClientConnectionProtocol,
 } from '@mswjs/interceptors/WebSocket'
+import { matchRequestUrl, type Path } from '../utils/matching/matchRequestUrl'
+
+export const MSW_WEBSOCKET_CLIENTS_KEY = 'msw:ws:clients'
 
 export type WebSocketBroadcastChannelMessage =
-  | {
-      type: 'connection:open'
-      payload: {
-        clientId: string
-        url: string
-      }
-    }
   | {
       type: 'extraneous:send'
       payload: {
@@ -28,33 +25,122 @@ export type WebSocketBroadcastChannelMessage =
       }
     }
 
-export const kAddByClientId = Symbol('kAddByClientId')
+type SerializedClient = {
+  clientId: string
+  url: string
+}
 
 /**
  * A manager responsible for accumulating WebSocket client
  * connections across different browser runtimes.
  */
 export class WebSocketClientManager {
+  private inMemoryClients: Set<WebSocketClientConnectionProtocol>
+
+  constructor(
+    private channel: BroadcastChannel,
+    private url: Path,
+  ) {
+    this.inMemoryClients = new Set()
+
+    if (typeof localStorage !== 'undefined') {
+      // When the worker clears the local storage key in "worker.stop()",
+      // also clear the in-memory clients map.
+      localStorage.removeItem = new Proxy(localStorage.removeItem, {
+        apply: (target, thisArg, args) => {
+          const [key] = args
+
+          if (key === MSW_WEBSOCKET_CLIENTS_KEY) {
+            this.inMemoryClients.clear()
+          }
+
+          return Reflect.apply(target, thisArg, args)
+        },
+      })
+    }
+  }
+
   /**
    * All active WebSocket client connections.
    */
-  public clients: Set<WebSocketClientConnectionProtocol>
+  get clients(): Set<WebSocketClientConnectionProtocol> {
+    // In the browser, different runtimes use "localStorage"
+    // as the shared source of all the clients.
+    if (typeof localStorage !== 'undefined') {
+      const inMemoryClients = Array.from(this.inMemoryClients)
+
+      console.log('get clients()', inMemoryClients, this.getSerializedClients())
+
+      return new Set(
+        inMemoryClients.concat(
+          this.getSerializedClients()
+            // Filter out the serialized clients that are already present
+            // in this runtime in-memory. This is crucial because a remote client
+            // wrapper CANNOT send a message to the client in THIS runtime
+            // (the "message" event on broadcast channel won't trigger).
+            .filter((serializedClient) => {
+              if (
+                inMemoryClients.every(
+                  (client) => client.id !== serializedClient.clientId,
+                )
+              ) {
+                return serializedClient
+              }
+            })
+            .map((serializedClient) => {
+              return new WebSocketRemoteClientConnection(
+                serializedClient.clientId,
+                new URL(serializedClient.url),
+                this.channel,
+              )
+            }),
+        ),
+      )
+    }
+
+    // In Node.js, the manager acts as a singleton, and all clients
+    // are kept in-memory.
+    return this.inMemoryClients
+  }
 
-  constructor(private channel: BroadcastChannel) {
-    this.clients = new Set()
+  private getSerializedClients(): Array<SerializedClient> {
+    invariant(
+      typeof localStorage !== 'undefined',
+      'Failed to call WebSocketClientManager#getSerializedClients() in a non-browser environment. This is likely a bug in MSW. Please, report it on GitHub: https://github.com/mswjs/msw',
+    )
 
-    this.channel.addEventListener('message', (message) => {
-      const { type, payload } = message.data as WebSocketBroadcastChannelMessage
+    const clientsJson = localStorage.getItem(MSW_WEBSOCKET_CLIENTS_KEY)
 
-      switch (type) {
-        case 'connection:open': {
-          // When another runtime notifies about a new connection,
-          // create a connection wrapper class and add it to the set.
-          this.onRemoteConnection(payload.clientId, new URL(payload.url))
-          break
-        }
-      }
+    if (!clientsJson) {
+      return []
+    }
+
+    const allClients = JSON.parse(clientsJson) as Array<SerializedClient>
+    const matchingClients = allClients.filter((client) => {
+      return matchRequestUrl(new URL(client.url), this.url).matches
     })
+
+    return matchingClients
+  }
+
+  private addClient(client: WebSocketClientConnection): void {
+    this.inMemoryClients.add(client)
+
+    if (typeof localStorage !== 'undefined') {
+      const serializedClients = this.getSerializedClients()
+
+      // Serialize the current client for other runtimes to create
+      // a remote wrapper over it. This has no effect on the current runtime.
+      const nextSerializedClients = serializedClients.concat({
+        clientId: client.id,
+        url: client.url.href,
+      } as SerializedClient)
+
+      localStorage.setItem(
+        MSW_WEBSOCKET_CLIENTS_KEY,
+        JSON.stringify(nextSerializedClients),
+      )
+    }
   }
 
   /**
@@ -64,16 +150,7 @@ export class WebSocketClientManager {
    * for the opened connections in the same runtime.
    */
   public addConnection(client: WebSocketClientConnection): void {
-    this.clients.add(client)
-
-    // Signal to other runtimes about this connection.
-    this.channel.postMessage({
-      type: 'connection:open',
-      payload: {
-        clientId: client.id,
-        url: client.url.toString(),
-      },
-    } as WebSocketBroadcastChannelMessage)
+    this.addClient(client)
 
     // Instruct the current client how to handle events
     // coming from other runtimes (e.g. when calling `.broadcast()`).
@@ -116,19 +193,6 @@ export class WebSocketClientManager {
       once: true,
     })
   }
-
-  /**
-   * Adds a client connection wrapper to operate with
-   * WebSocket client connections in other runtimes.
-   */
-  private onRemoteConnection(id: string, url: URL): void {
-    this.clients.add(
-      // Create a connection-compatible instance that can
-      // operate with this client from a different runtime
-      // using the BroadcastChannel messages.
-      new WebSocketRemoteClientConnection(id, url, this.channel),
-    )
-  }
 }
 
 /**