From 7c542cfc2db0d6cd1155abbe29fd3459be8d3a66 Mon Sep 17 00:00:00 2001 From: Artem Zakharchenko Date: Sun, 28 Jan 2024 21:14:26 +0100 Subject: [PATCH] feat: add WebSocket class interceptor --- README.md | 80 +++++ src/interceptors/WebSocket/WebSocketClient.ts | 65 ++++ src/interceptors/WebSocket/WebSocketServer.ts | 35 +++ .../WebSocket/WebSocketTransport.ts | 31 ++ .../WebSocketClass/WebSocketClassClient.ts | 24 ++ .../WebSocketClassInterceptor.ts | 288 ++++++++++++++++++ .../WebSocketClass/WebSocketClassServer.ts | 76 +++++ .../WebSocketClass/WebSocketClassTransport.ts | 39 +++ src/interceptors/WebSocket/index.ts | 31 ++ src/interceptors/WebSocket/utils/bindEvent.ts | 9 + 10 files changed, 678 insertions(+) create mode 100644 src/interceptors/WebSocket/WebSocketClient.ts create mode 100644 src/interceptors/WebSocket/WebSocketServer.ts create mode 100644 src/interceptors/WebSocket/WebSocketTransport.ts create mode 100644 src/interceptors/WebSocket/implementations/WebSocketClass/WebSocketClassClient.ts create mode 100644 src/interceptors/WebSocket/implementations/WebSocketClass/WebSocketClassInterceptor.ts create mode 100644 src/interceptors/WebSocket/implementations/WebSocketClass/WebSocketClassServer.ts create mode 100644 src/interceptors/WebSocket/implementations/WebSocketClass/WebSocketClassTransport.ts create mode 100644 src/interceptors/WebSocket/index.ts create mode 100644 src/interceptors/WebSocket/utils/bindEvent.ts diff --git a/README.md b/README.md index 280d8cf2..88e4d557 100644 --- a/README.md +++ b/README.md @@ -291,6 +291,86 @@ interceptor.on( > Note that the `isMockedResponse` property will only be set to `true` if you resolved this request in the "request" event listener using the `request.respondWith()` method and providing a mocked `Response` instance. +## WebSocket interception + +```js +import { WebSocketInterceptor } from '@mswjs/interceptors/WebSocket' + +const interceptor = new WebSocketInterceptor() +``` + +The WebSocket interceptor emits a single `connection` event that represents a new client connecting to the server. Within the `connection` event listener, you can use the `client` reference to listen to the outgoing events and send the mock events from the server. + +```js +interceptor.on('connection', ({ client }) => { + console.log(client.url) + + client.on('message', (event) => { + if (event.data === 'hello from client') { + client.send('hello from server') + } + }) +}) +``` + +### Interceptor arguments + +| Name | Type | Description | +| -------- | -------- | -------------------------------------- | +| `client` | `object` | A client WebSocket connection. | +| `server` | `object` | An actual WebSocket server connection. | + +Remember that you write the interceptor from the _server's perspective_. With that in mind, the `client` represents a client connecting to this "server" (your interceptor). The `server`, on the other hand, represents an actual production server connection, if you ever wish to establish one to intercept the incoming events as well. + +### Bypassing events + +By default, the WebSocket interceptor prevents all the outgoing events from reaching the production server. To bypass the event, first establish the actual server connection by calling `await server.connect()`, and then use `server.send()` and provide it the `MessageEvent` to bypass. + +```js +interceptor.on('connection', async ({ client, server }) => { + // First, connect to the actual server. + await server.connect() + + // Forward all outgoing client events to the original server. + client.on('message', (event) => server.send(event)) +}) +``` + +### Incoming server events + +The WebSocket interceptor also intercepts the incoming events sent from the original server. + +```js +interceptor.on('connection', async ({ server }) => { + await server.connect() + + server.on('message', (event) => { + console.log('original server sent:', event.data) + }) +}) +``` + +> Note: If you establish the original server connection, all the incoming server events will be automatically forwarded to the client. + +If you wish to prevent the automatic forwarding of the server events to the client, call `event.preventDefault()` on the incoming event you wish to prevent. + +```js +interceptor.on('connection', async ({ client, server }) => { + await server.connect() + + server.on('message', (event) => { + if (event.data === 'hello from actual server') { + // Never forward this event to the client. + event.preventDefault() + + // Instead, send this mock data. + client.send('hello from mock server') + return + } + }) +}) +``` + ## API ### `Interceptor` diff --git a/src/interceptors/WebSocket/WebSocketClient.ts b/src/interceptors/WebSocket/WebSocketClient.ts new file mode 100644 index 00000000..aec096b8 --- /dev/null +++ b/src/interceptors/WebSocket/WebSocketClient.ts @@ -0,0 +1,65 @@ +import type { + WebSocketSendData, + WebSocketTransport, +} from './WebSocketTransport' + +const kEmitter = Symbol('emitter') + +/** + * The WebSocket client instance represents an incoming + * client connection. The user can control the connection, + * send and receive events. + */ +export class WebSocketClient { + protected [kEmitter]: EventTarget + + constructor( + public readonly url: URL, + protected readonly transport: WebSocketTransport + ) { + this[kEmitter] = new EventTarget() + + /** + * Emit incoming server events so they can be reacted to. + * @note This does NOT forward the events to the client. + * That must be done explicitly via "server.send()". + */ + transport.onIncoming = (event) => { + this[kEmitter].dispatchEvent(event) + } + } + + /** + * Listen for incoming events from the connected client. + */ + public on( + event: string, + listener: (...data: Array) => void + ): void { + this[kEmitter].addEventListener(event, (event) => { + if (event instanceof MessageEvent) { + listener(event.data) + } + }) + } + + /** + * Send data to the connected client. + */ + public send(data: WebSocketSendData): void { + this.transport.send(data) + } + + /** + * Emit the given event to the connected client. + */ + public emit(event: string, data: WebSocketSendData): void { + throw new Error('WebSocketClient#emit is not implemented') + } + + public close(error?: Error): void { + // Don't do any guessing behind the close code's semantics + // and fallback to a generic contrived close code of 3000. + this.transport.close(error ? 3000 : 1000, error?.message) + } +} diff --git a/src/interceptors/WebSocket/WebSocketServer.ts b/src/interceptors/WebSocket/WebSocketServer.ts new file mode 100644 index 00000000..9121a1f3 --- /dev/null +++ b/src/interceptors/WebSocket/WebSocketServer.ts @@ -0,0 +1,35 @@ +import type { WebSocketSendData } from './WebSocketTransport' +import type { WebSocketMessageListener } from './implementations/WebSocketClass/WebSocketClassInterceptor' + +/** + * The WebSocket server instance represents the actual production + * WebSocket server connection. It's idle by default but you can + * establish it by calling `server.connect()`. + */ +export class WebSocketServer { + /** + * Connect to the actual WebSocket server. + */ + public connect(): Promise { + throw new Error('WebSocketServer#connect is not implemented') + } + + /** + * Send the data to the original server. + * The connection to the original server will be opened + * as a part of the first `server.send()` call. + */ + public send(data: WebSocketSendData): void { + throw new Error('WebSocketServer#send is not implemented') + } + + /** + * Listen to the incoming events from the original + * WebSocket server. All the incoming events are automatically + * forwarded to the client connection unless you prevent them + * via `event.preventDefault()`. + */ + public on(event: string, callback: WebSocketMessageListener): void { + throw new Error('WebSocketServer#on is not implemented') + } +} diff --git a/src/interceptors/WebSocket/WebSocketTransport.ts b/src/interceptors/WebSocket/WebSocketTransport.ts new file mode 100644 index 00000000..adbda0b9 --- /dev/null +++ b/src/interceptors/WebSocket/WebSocketTransport.ts @@ -0,0 +1,31 @@ +export type WebSocketSendData = + | string + | ArrayBufferLike + | Blob + | ArrayBufferView + +export type WebSocketTransportOnIncomingCallback = ( + event: MessageEvent +) => void + +export abstract class WebSocketTransport { + /** + * Listener for the incoming server events. + * This is called when the client receives the + * event from the original server connection. + * + * This way, we can trigger the "message" event + * on the mocked connection to let the user know. + */ + abstract onIncoming: WebSocketTransportOnIncomingCallback + + /** + * Send the data from the server to this client. + */ + abstract send(data: WebSocketSendData): void + + /** + * Close the client connection. + */ + abstract close(code: number, reason?: string): void +} diff --git a/src/interceptors/WebSocket/implementations/WebSocketClass/WebSocketClassClient.ts b/src/interceptors/WebSocket/implementations/WebSocketClass/WebSocketClassClient.ts new file mode 100644 index 00000000..2f5fee65 --- /dev/null +++ b/src/interceptors/WebSocket/implementations/WebSocketClass/WebSocketClassClient.ts @@ -0,0 +1,24 @@ +import { invariant } from 'outvariant' +import { WebSocketClient } from '../../WebSocketClient' +import type { WebSocketSendData } from '../../WebSocketTransport' +import { WebSocketClassTransport } from './WebSocketClassTransport' +import type { WebSocketClassOverride } from './WebSocketClassInterceptor' + +export class WebSocketClassClient extends WebSocketClient { + constructor( + readonly ws: WebSocketClassOverride, + readonly transport: WebSocketClassTransport + ) { + super(new URL(ws.url), transport) + } + + public emit(event: string, data: WebSocketSendData): void { + invariant( + event === 'message', + 'Failed to emit unknown WebSocket event "%s": only the "message" event is supported using the standard WebSocket class', + event + ) + + this.transport.send(data) + } +} diff --git a/src/interceptors/WebSocket/implementations/WebSocketClass/WebSocketClassInterceptor.ts b/src/interceptors/WebSocket/implementations/WebSocketClass/WebSocketClassInterceptor.ts new file mode 100644 index 00000000..5bdf56e2 --- /dev/null +++ b/src/interceptors/WebSocket/implementations/WebSocketClass/WebSocketClassInterceptor.ts @@ -0,0 +1,288 @@ +import { invariant } from 'outvariant' +import type { WebSocketEventsMap } from '../../index' +import { Interceptor } from '../../../../Interceptor' +import { WebSocketClassClient } from './WebSocketClassClient' +import { WebSocketClassServer } from './WebSocketClassServer' +import type { + WebSocketSendData, + WebSocketTransportOnIncomingCallback, +} from '../../WebSocketTransport' +import { bindEvent } from '../../utils/bindEvent' +import { WebSocketClassTransport } from './WebSocketClassTransport' + +export class WebSocketClassInterceptor extends Interceptor { + static symbol = Symbol('websocket-class') + + constructor() { + super(WebSocketClassInterceptor.symbol) + } + + protected checkEnvironment(): boolean { + // Enable this interceptor in any environment + // that has a global WebSocket API. + return typeof globalThis.WebSocket !== 'undefined' + } + + protected setup(): void { + const { WebSocket: OriginalWebSocket } = globalThis + + const webSocketProxy = Proxy.revocable(globalThis.WebSocket, { + construct: ( + target, + args: ConstructorParameters, + newTarget + ) => { + const [url, protocols] = args + + const createConnection = (): WebSocket => { + return Reflect.construct(target, args, newTarget) + } + + // All WebSocket instances are mocked and don't forward + // any events to the original server (no connection established). + // To forward the events, the user must use the "server.send()" API. + const mockWs = new WebSocketClassOverride(url, protocols) + + const transport = new WebSocketClassTransport(mockWs) + + // The "globalThis.WebSocket" class stands for + // the client-side connection. Assume it's established + // as soon as the WebSocket instance is constructed. + this.emitter.emit('connection', { + client: new WebSocketClassClient(mockWs, transport), + server: new WebSocketClassServer(mockWs, createConnection), + }) + + return mockWs + }, + }) + + globalThis.WebSocket = webSocketProxy.proxy + + this.subscriptions.push(() => { + webSocketProxy.revoke() + }) + } +} + +const WEBSOCKET_CLOSE_CODE_RANGE_ERROR = + 'InvalidAccessError: close code out of user configurable range' + +const kOnSend = Symbol('kOnSend') +export const kOnReceive = Symbol('kOnReceive') + +type WebSocketEventListener = (this: WebSocket, event: Event) => void +export type WebSocketMessageListener = ( + this: WebSocket, + event: MessageEvent +) => void +type WebSocketCloseListener = (this: WebSocket, event: CloseEvent) => void + +export class WebSocketClassOverride extends EventTarget implements WebSocket { + static readonly CONNECTING = WebSocket.CONNECTING + static readonly OPEN = WebSocket.OPEN + static readonly CLOSING = WebSocket.CLOSING + static readonly CLOSED = WebSocket.CLOSED + readonly CONNECTING = WebSocket.CONNECTING + readonly OPEN = WebSocket.OPEN + readonly CLOSING = WebSocket.CLOSING + readonly CLOSED = WebSocket.CLOSED + + public url: string + public protocol: string + public extensions: string + public binaryType: BinaryType + public readyState: number + + private _onopen: WebSocketEventListener | null = null + private _onmessage: WebSocketMessageListener | null = null + private _onerror: WebSocketEventListener | null = null + private _onclose: WebSocketCloseListener | null = null + + private buffer: Array + private [kOnSend]?: (data: WebSocketSendData) => void + private [kOnReceive]?: WebSocketTransportOnIncomingCallback + + constructor(url: string | URL, protocols?: string | Array) { + super() + this.url = url.toString() + this.protocol = protocols ? protocols[0] : 'ws' + this.extensions = '' + this.binaryType = 'arraybuffer' + this.readyState = this.CONNECTING + + this.buffer = [] + + this.addEventListener( + 'open', + () => { + // As soon as the connection opens, send any buffered data. + if (this.buffer.length > 0) { + this.buffer.map(this.send) + this.buffer = [] + } + }, + { + once: true, + } + ) + } + + get bufferedAmount(): number { + return this.buffer.reduce((totalSize, data) => { + if (typeof data === 'string') { + return totalSize + data.length + } + + if (data instanceof Blob) { + return totalSize + data.size + } + + return totalSize + data.byteLength + }, 0) + } + + set onopen(listener: WebSocketEventListener) { + this.removeEventListener('open', this._onopen) + this._onopen = listener + if (listener !== null) { + this.addEventListener('open', listener) + } + } + get onopen(): WebSocketEventListener | null { + return this._onopen + } + + set onmessage(listener: WebSocketMessageListener) { + this.removeEventListener( + 'message', + this._onmessage as WebSocketEventListener + ) + this.onmessage = listener + if (listener !== null) { + this.addEventListener('message', listener) + } + } + get onmessage(): WebSocketMessageListener | null { + return this._onmessage + } + + set onerror(listener: WebSocketEventListener) { + this.removeEventListener('error', this._onerror) + this._onerror = listener + if (listener !== null) { + this.addEventListener('error', listener) + } + } + get onerror(): WebSocketEventListener | null { + return this._onerror + } + + set onclose(listener: WebSocketCloseListener) { + this.removeEventListener('close', this._onclose as WebSocketEventListener) + this._onclose = listener + if (listener !== null) { + this.addEventListener('close', listener) + } + } + get onclose(): WebSocketCloseListener | null { + return this._onclose + } + + public send(data: WebSocketSendData): void { + if (this.readyState === this.CONNECTING) { + this.close() + throw new Error('InvalidStateError') + } + + if (this.readyState === this.CLOSING || this.readyState === this.CLOSED) { + this.buffer.push(data) + return + } + + /** + * @note Notify the parent about data being sent. + */ + this[kOnSend]?.(data) + } + + public dispatchEvent(event: Event): boolean { + /** + * @note This override class never forwards the incoming + * events to the actual client instance. Instead, it + * forwards the incoming events to the connection + * and lets the "server" API handle the forwarding. + */ + if ( + event.type === 'message' && + // Ignore mocked events sent from the connection. + // This condition is for the original server-sent events only. + !(kOnSend && event.target) + ) { + this[kOnReceive]?.(event as MessageEvent) + return true + } + + // Dispatch the other events (open, close, etc). + return super.dispatchEvent(event) + } + + public close(code?: number, reason?: string): void { + invariant(code, WEBSOCKET_CLOSE_CODE_RANGE_ERROR) + invariant( + code === 1000 || (code >= 3000 && code <= 4999), + WEBSOCKET_CLOSE_CODE_RANGE_ERROR + ) + + if (this.readyState === this.CLOSING || this.readyState === this.CLOSED) { + return + } + + this.dispatchEvent( + bindEvent( + this, + new CloseEvent('close', { + code, + reason, + wasClean: code === 1000, + }) + ) + ) + + // Remove all event listeners once the socket is closed. + this._onopen = null + this._onmessage = null + this._onerror = null + this._onclose = null + } + + public addEventListener( + type: K, + listener: (this: WebSocket, event: WebSocketEventMap[K]) => void, + options?: boolean | AddEventListenerOptions + ): void + public addEventListener( + type: string, + listener: EventListenerOrEventListenerObject, + options?: boolean | AddEventListenerOptions + ): void + public addEventListener( + type: unknown, + listener: unknown, + options?: unknown + ): void { + return super.addEventListener( + type as string, + listener as EventListener, + options as AddEventListenerOptions + ) + } + + removeEventListener( + type: K, + callback: EventListenerOrEventListenerObject | null, + options?: boolean | EventListenerOptions + ): void { + return super.removeEventListener(type, callback, options) + } +} diff --git a/src/interceptors/WebSocket/implementations/WebSocketClass/WebSocketClassServer.ts b/src/interceptors/WebSocket/implementations/WebSocketClass/WebSocketClassServer.ts new file mode 100644 index 00000000..e818775a --- /dev/null +++ b/src/interceptors/WebSocket/implementations/WebSocketClass/WebSocketClassServer.ts @@ -0,0 +1,76 @@ +import { invariant } from 'outvariant' +import { DeferredPromise } from '@open-draft/deferred-promise' +import { WebSocketServer } from '../../WebSocketServer' +import type { + WebSocketClassOverride, + WebSocketMessageListener, +} from './WebSocketClassInterceptor' +import type { WebSocketSendData } from '../../WebSocketTransport' +import { bindEvent } from '../../utils/bindEvent' + +export class WebSocketClassServer extends WebSocketServer { + private prodWs?: WebSocket + + constructor( + private readonly mockWs: WebSocketClassOverride, + private readonly createConnection: () => WebSocket + ) { + super() + } + + public connect(): Promise { + const connectionPromise = new DeferredPromise() + const ws = this.createConnection() + + ws.addEventListener('open', () => connectionPromise.resolve(), { + once: true, + }) + ws.addEventListener('error', () => connectionPromise.reject(), { + once: true, + }) + + return connectionPromise + .then(() => { + this.prodWs = ws + }) + .catch((error) => { + console.error( + 'Failed to connect to the original WebSocket server at "%s"', + this.mockWs.url + ) + console.error(error) + }) + } + + public send(data: WebSocketSendData): void { + invariant( + this.prodWs, + 'Failed to call "server.send()" for "%s": the connection is not open. Did you forget to call "await server.connect()"?', + this.mockWs.url + ) + + // Send the data using the original WebSocket connection. + this.prodWs.send(data) + } + + public on(event: 'message', callback: WebSocketMessageListener): void { + invariant( + this.prodWs, + 'Failed to call "server.on(%s)" for "%s": the connection is not open. Did you forget to call "await server.connect()"?', + this.mockWs.url + ) + + const { prodWs } = this + + prodWs.addEventListener(event, (messageEvent) => { + callback.call(prodWs, messageEvent) + + // Unless the default is prevented, forward the + // messages from the original server to the mock client. + // This is the only way the user can receive them. + if (!messageEvent.defaultPrevented) { + this.mockWs.dispatchEvent(bindEvent(this.mockWs, messageEvent)) + } + }) + } +} diff --git a/src/interceptors/WebSocket/implementations/WebSocketClass/WebSocketClassTransport.ts b/src/interceptors/WebSocket/implementations/WebSocketClass/WebSocketClassTransport.ts new file mode 100644 index 00000000..08494c0d --- /dev/null +++ b/src/interceptors/WebSocket/implementations/WebSocketClass/WebSocketClassTransport.ts @@ -0,0 +1,39 @@ +import { bindEvent } from '../../utils/bindEvent' +import { + WebSocketSendData, + WebSocketTransport, + WebSocketTransportOnIncomingCallback, +} from '../../WebSocketTransport' +import { kOnReceive, WebSocketClassOverride } from './WebSocketClassInterceptor' + +export class WebSocketClassTransport extends WebSocketTransport { + public onIncoming: WebSocketTransportOnIncomingCallback = () => {} + + constructor(protected readonly ws: WebSocketClassOverride) { + super() + this.ws[kOnReceive] = this.onIncoming + } + + public send(data: WebSocketSendData): void { + this.ws.dispatchEvent( + bindEvent( + /** + * @note Setting this event's "target" to the + * WebSocket override instance is important. + * This way it can tell apart original incoming events + * (must be forwarded to the transport) from the + * mocked message events like the one below + * (must be dispatched on the client instance). + */ + this.ws, + new MessageEvent('message', { + data, + }) + ) + ) + } + + public close(code: number, reason?: string): void { + this.ws.close(code, reason) + } +} diff --git a/src/interceptors/WebSocket/index.ts b/src/interceptors/WebSocket/index.ts new file mode 100644 index 00000000..24ba1faa --- /dev/null +++ b/src/interceptors/WebSocket/index.ts @@ -0,0 +1,31 @@ +import { BatchInterceptor } from '../..' +import { WebSocketClient } from './WebSocketClient' +import { WebSocketServer } from './WebSocketServer' +import { WebSocketClassInterceptor } from './implementations/WebSocketClass/WebSocketClassInterceptor' + +export type WebSocketEventsMap = { + connection: [ + args: { + /** + * The connected WebSocket client. + */ + client: WebSocketClient + /** + * The original WebSocket server. + */ + server: WebSocketServer + } + ] +} + +export class WebSocketInterceptor extends BatchInterceptor< + [WebSocketClassInterceptor], + WebSocketEventsMap +> { + constructor() { + super({ + name: 'websocket', + interceptors: [new WebSocketClassInterceptor()], + }) + } +} diff --git a/src/interceptors/WebSocket/utils/bindEvent.ts b/src/interceptors/WebSocket/utils/bindEvent.ts new file mode 100644 index 00000000..05ba1d94 --- /dev/null +++ b/src/interceptors/WebSocket/utils/bindEvent.ts @@ -0,0 +1,9 @@ +type EventWithTarget = E & { target: T } + +export function bindEvent( + target: T, + event: E +): EventWithTarget { + Reflect.set(event, 'target', target) + return event as EventWithTarget +}