From 7c35f1f8ec92a24c79fca51d270666f1b9a474d0 Mon Sep 17 00:00:00 2001 From: Artem Zakharchenko Date: Sun, 7 Apr 2024 18:06:56 +0200 Subject: [PATCH 1/4] fix(ws): keep clients in localStorage --- src/core/ws.ts | 6 +- src/core/ws/WebSocketClientManager.test.ts | 7 +- src/core/ws/WebSocketClientManager.ts | 134 +++++++++----- .../browser/ws-api/ws.clients.browser.test.ts | 166 ++++++++++++++++++ 4 files changed, 263 insertions(+), 50 deletions(-) create mode 100644 test/browser/ws-api/ws.clients.browser.test.ts diff --git a/src/core/ws.ts b/src/core/ws.ts index 718d1ef0d..77fb681ef 100644 --- a/src/core/ws.ts +++ b/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) diff --git a/src/core/ws/WebSocketClientManager.test.ts b/src/core/ws/WebSocketClientManager.test.ts index 9cac29c04..b0b9fe0f1 100644 --- a/src/core/ws/WebSocketClientManager.test.ts +++ b/src/core/ws/WebSocketClientManager.test.ts @@ -17,9 +17,8 @@ vi.spyOn(channel, 'postMessage') const socket = new WebSocket('ws://localhost') const transport = { - onOutgoing: vi.fn(), - onIncoming: vi.fn(), - onClose: vi.fn(), + addEventListener: vi.fn(), + dispatchEvent: vi.fn(), send: vi.fn(), close: vi.fn(), } satisfies WebSocketTransport @@ -135,7 +134,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() + connection['transport'].dispatchEvent(new Event('close')) }) vi.spyOn(connection, 'send') diff --git a/src/core/ws/WebSocketClientManager.ts b/src/core/ws/WebSocketClientManager.ts index 4048878b6..a88b963f1 100644 --- a/src/core/ws/WebSocketClientManager.ts +++ b/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' + +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,104 @@ 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 + + constructor( + private channel: BroadcastChannel, + private url: Path, + ) { + this.inMemoryClients = new Set() + } + /** * All active WebSocket client connections. */ - public clients: Set + get clients(): Set { + // 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) + + 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, + ) + }), + ), + ) + } - constructor(private channel: BroadcastChannel) { - this.clients = new Set() + // In Node.js, the manager acts as a singleton, and all clients + // are kept in-memory. + return this.inMemoryClients + } - this.channel.addEventListener('message', (message) => { - const { type, payload } = message.data as WebSocketBroadcastChannelMessage + private getSerializedClients(): Array { + 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', + ) - 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 - } - } + const clientsJson = localStorage.getItem(MSW_WEBSOCKET_CLIENTS_KEY) + + if (!clientsJson) { + return [] + } + + const allClients = JSON.parse(clientsJson) as Array + 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 +132,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 +175,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), - ) - } } /** diff --git a/test/browser/ws-api/ws.clients.browser.test.ts b/test/browser/ws-api/ws.clients.browser.test.ts new file mode 100644 index 000000000..30d3dd5b1 --- /dev/null +++ b/test/browser/ws-api/ws.clients.browser.test.ts @@ -0,0 +1,166 @@ +import type { WebSocketLink, ws } from 'msw' +import type { setupWorker } from 'msw/browser' +import { test, expect } from '../playwright.extend' + +declare global { + interface Window { + msw: { + ws: typeof ws + setupWorker: typeof setupWorker + } + link: WebSocketLink + ws: WebSocket + messages: string[] + } +} + +test('returns the number of active clients in the same runtime', async ({ + loadExample, + page, +}) => { + await loadExample(require.resolve('./ws.runtime.js'), { + skipActivation: true, + }) + + await page.evaluate(async () => { + const { setupWorker, ws } = window.msw + const api = ws.link('wss://example.com') + const worker = setupWorker(api.on('connection', () => {})) + window.link = api + await worker.start() + }) + + // Must return 0 when no clients are present. + expect( + await page.evaluate(() => { + return window.link.clients.size + }), + ).toBe(0) + + await page.evaluate(async () => { + const ws = new WebSocket('wss://example.com') + await new Promise((done) => (ws.onopen = done)) + }) + + // Must return 1 after a single client joined. + expect( + await page.evaluate(() => { + return window.link.clients.size + }), + ).toBe(1) + + await page.evaluate(async () => { + const ws = new WebSocket('wss://example.com') + await new Promise((done) => (ws.onopen = done)) + }) + + // Must return 2 now that another client has joined. + expect( + await page.evaluate(() => { + return window.link.clients.size + }), + ).toBe(2) +}) + +test('returns the number of active clients across different runtimes', async ({ + loadExample, + context, +}) => { + const { compilation } = await loadExample( + require.resolve('./ws.runtime.js'), + { + skipActivation: true, + }, + ) + + const pageOne = await context.newPage() + const pageTwo = await context.newPage() + + for (const page of [pageOne, pageTwo]) { + await page.goto(compilation.previewUrl) + await page.evaluate(async () => { + const { setupWorker, ws } = window.msw + const api = ws.link('wss://example.com') + const worker = setupWorker(api.on('connection', () => {})) + window.link = api + await worker.start() + }) + } + + await pageOne.bringToFront() + await pageOne.evaluate(async () => { + const ws = new WebSocket('wss://example.com') + await new Promise((done) => (ws.onopen = done)) + }) + + expect(await pageOne.evaluate(() => window.link.clients.size)).toBe(1) + expect(await pageTwo.evaluate(() => window.link.clients.size)).toBe(1) + + await pageTwo.bringToFront() + await pageTwo.evaluate(async () => { + const ws = new WebSocket('wss://example.com') + await new Promise((done) => (ws.onopen = done)) + }) + + expect(await pageTwo.evaluate(() => window.link.clients.size)).toBe(2) + expect(await pageOne.evaluate(() => window.link.clients.size)).toBe(2) +}) + +test('broadcasts messages across runtimes', async ({ + loadExample, + context, +}) => { + const { compilation } = await loadExample( + require.resolve('./ws.runtime.js'), + { + skipActivation: true, + }, + ) + + const pageOne = await context.newPage() + const pageTwo = await context.newPage() + + for (const page of [pageOne, pageTwo]) { + await page.goto(compilation.previewUrl) + await page.evaluate(async () => { + const { setupWorker, ws } = window.msw + const api = ws.link('wss://example.com') + const worker = setupWorker( + api.on('connection', ({ client }) => { + client.addEventListener('message', (event) => { + api.broadcast(event.data) + }) + }), + ) + await worker.start() + }) + + await page.evaluate(() => { + window.messages = [] + const ws = new WebSocket('wss://example.com') + window.ws = ws + ws.onmessage = (event) => { + window.messages.push(event.data) + } + }) + } + + await pageOne.evaluate(() => { + window.ws.send('hi from one') + }) + expect(await pageOne.evaluate(() => window.messages)).toEqual(['hi from one']) + expect(await pageTwo.evaluate(() => window.messages)).toEqual(['hi from one']) + + await pageTwo.evaluate(() => { + window.ws.send('hi from two') + }) + + expect(await pageTwo.evaluate(() => window.messages)).toEqual([ + 'hi from one', + 'hi from two', + ]) + expect(await pageOne.evaluate(() => window.messages)).toEqual([ + 'hi from one', + 'hi from two', + ]) +}) From c253633331b52f15525d3c18cc255de204c035d7 Mon Sep 17 00:00:00 2001 From: Artem Zakharchenko Date: Sun, 7 Apr 2024 18:25:55 +0200 Subject: [PATCH 2/4] test: fix WebSocketClientManager test --- src/core/ws/WebSocketClientManager.test.ts | 73 +++++++--------------- 1 file changed, 24 insertions(+), 49 deletions(-) diff --git a/src/core/ws/WebSocketClientManager.test.ts b/src/core/ws/WebSocketClientManager.test.ts index b0b9fe0f1..7dbd13664 100644 --- a/src/core/ws/WebSocketClientManager.test.ts +++ b/src/core/ws/WebSocketClientManager.test.ts @@ -1,78 +1,49 @@ /** * @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 = { - addEventListener: vi.fn(), - dispatchEvent: 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) - - channel.dispatchEvent( - new MessageEvent('message', { - data: { - type: 'connection:open', - payload: { - clientId, - url: url.href, - }, - }, - }), - ) - - await vi.waitFor(() => { - expect(Array.from(manager.clients.values())).toEqual([ - new WebSocketRemoteClientConnection(clientId, url, channel), - ]) - }) }) 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') @@ -97,8 +68,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') @@ -124,7 +98,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(() => { /** @@ -134,7 +109,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'].dispatchEvent(new Event('close')) + transport.dispatchEvent(new Event('close')) }) vi.spyOn(connection, 'send') From 78c46a074a4f71fc81242698a81fe433f80d7d68 Mon Sep 17 00:00:00 2001 From: Artem Zakharchenko Date: Sun, 7 Apr 2024 18:28:48 +0200 Subject: [PATCH 3/4] test: add multiple clients test --- src/core/ws/WebSocketClientManager.test.ts | 24 ++++++++++++++++++++++ 1 file changed, 24 insertions(+) diff --git a/src/core/ws/WebSocketClientManager.test.ts b/src/core/ws/WebSocketClientManager.test.ts index 7dbd13664..8db322f8e 100644 --- a/src/core/ws/WebSocketClientManager.test.ts +++ b/src/core/ws/WebSocketClientManager.test.ts @@ -38,6 +38,30 @@ it('adds a client from this runtime 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', () => { + const manager = new WebSocketClientManager(channel, '*') + const connectionOne = new WebSocketClientConnection( + socket, + new TestWebSocketTransport(), + ) + 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(), + ) + 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( From ee4ca2e34361165369af92dbe456b849013f5b2c Mon Sep 17 00:00:00 2001 From: Artem Zakharchenko Date: Fri, 12 Apr 2024 12:50:19 -0600 Subject: [PATCH 4/4] fix(WebSocketClientManager): clear inMemoryClients on localStorage removal --- src/browser/setupWorker/stop/createStop.ts | 4 ++ src/core/ws/WebSocketClientManager.ts | 20 ++++++++- .../browser/ws-api/ws.clients.browser.test.ts | 41 ++++++++++++++++++- 3 files changed, 63 insertions(+), 2 deletions(-) diff --git a/src/browser/setupWorker/stop/createStop.ts b/src/browser/setupWorker/stop/createStop.ts index 48c37996d..6c0e7b7d9 100644 --- a/src/browser/setupWorker/stop/createStop.ts +++ b/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 }) } } diff --git a/src/core/ws/WebSocketClientManager.ts b/src/core/ws/WebSocketClientManager.ts index a88b963f1..e9d11d468 100644 --- a/src/core/ws/WebSocketClientManager.ts +++ b/src/core/ws/WebSocketClientManager.ts @@ -6,7 +6,7 @@ import type { } from '@mswjs/interceptors/WebSocket' import { matchRequestUrl, type Path } from '../utils/matching/matchRequestUrl' -const MSW_WEBSOCKET_CLIENTS_KEY = 'msw:ws:clients' +export const MSW_WEBSOCKET_CLIENTS_KEY = 'msw:ws:clients' export type WebSocketBroadcastChannelMessage = | { @@ -42,6 +42,22 @@ export class WebSocketClientManager { 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) + }, + }) + } } /** @@ -53,6 +69,8 @@ export class WebSocketClientManager { if (typeof localStorage !== 'undefined') { const inMemoryClients = Array.from(this.inMemoryClients) + console.log('get clients()', inMemoryClients, this.getSerializedClients()) + return new Set( inMemoryClients.concat( this.getSerializedClients() diff --git a/test/browser/ws-api/ws.clients.browser.test.ts b/test/browser/ws-api/ws.clients.browser.test.ts index 30d3dd5b1..4d0a4a77d 100644 --- a/test/browser/ws-api/ws.clients.browser.test.ts +++ b/test/browser/ws-api/ws.clients.browser.test.ts @@ -1,5 +1,5 @@ import type { WebSocketLink, ws } from 'msw' -import type { setupWorker } from 'msw/browser' +import type { SetupWorker, setupWorker } from 'msw/browser' import { test, expect } from '../playwright.extend' declare global { @@ -8,6 +8,7 @@ declare global { ws: typeof ws setupWorker: typeof setupWorker } + worker: SetupWorker link: WebSocketLink ws: WebSocket messages: string[] @@ -164,3 +165,41 @@ test('broadcasts messages across runtimes', async ({ 'hi from two', ]) }) + +test('clears the list of clients when the worker is stopped', async ({ + loadExample, + page, +}) => { + await loadExample(require.resolve('./ws.runtime.js'), { + skipActivation: true, + }) + + await page.evaluate(async () => { + const { setupWorker, ws } = window.msw + const api = ws.link('wss://example.com') + const worker = setupWorker(api.on('connection', () => {})) + window.link = api + window.worker = worker + await worker.start() + }) + + await page.evaluate(async () => { + const ws = new WebSocket('wss://example.com') + await new Promise((done) => (ws.onopen = done)) + }) + + // Must return 1 after a single client joined. + expect( + await page.evaluate(() => { + return window.link.clients.size + }), + ).toBe(1) + + await page.evaluate(() => { + window.worker.stop() + }) + + // Must return 0. + // The localStorage has been purged, and the in-memory manager clients too. + expect(await page.evaluate(() => window.link.clients.size)).toBe(0) +})