From eaf71827bee24cdf497c3ccd7d19d630af5a8560 Mon Sep 17 00:00:00 2001 From: Artem Zakharchenko Date: Mon, 18 Nov 2024 13:09:42 +0100 Subject: [PATCH] fix(WebSocket): buffer sending the data until connection is open (#678) --- .../WebSocket/WebSocketClassTransport.ts | 53 ++++++--- .../compliance/websocket.client.send.test.ts | 104 ++++++++++++++++++ .../compliance/websocket.connection.test.ts | 2 +- .../intercept/websocket.dispose.test.ts | 4 +- .../intercept/websocket.send.test.ts | 27 +++-- .../intercept/websocket.server.events.test.ts | 4 +- 6 files changed, 164 insertions(+), 30 deletions(-) create mode 100644 test/modules/WebSocket/compliance/websocket.client.send.test.ts diff --git a/src/interceptors/WebSocket/WebSocketClassTransport.ts b/src/interceptors/WebSocket/WebSocketClassTransport.ts index a0b3b196..70f21665 100644 --- a/src/interceptors/WebSocket/WebSocketClassTransport.ts +++ b/src/interceptors/WebSocket/WebSocketClassTransport.ts @@ -64,23 +64,44 @@ export class WebSocketClassTransport public send(data: WebSocketData): void { queueMicrotask(() => { - this.socket.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.socket, - new MessageEvent('message', { - data, - origin: this.socket.url, - }) + if ( + this.socket.readyState === this.socket.CLOSING || + this.socket.readyState === this.socket.CLOSED + ) { + return + } + + const dispatchEvent = () => { + this.socket.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.socket, + new MessageEvent('message', { + data, + origin: this.socket.url, + }) + ) ) - ) + } + + if (this.socket.readyState === this.socket.CONNECTING) { + this.socket.addEventListener( + 'open', + () => { + dispatchEvent() + }, + { once: true } + ) + } else { + dispatchEvent() + } }) } diff --git a/test/modules/WebSocket/compliance/websocket.client.send.test.ts b/test/modules/WebSocket/compliance/websocket.client.send.test.ts new file mode 100644 index 00000000..d26606f1 --- /dev/null +++ b/test/modules/WebSocket/compliance/websocket.client.send.test.ts @@ -0,0 +1,104 @@ +// @vitest-environment node-with-websocket +import { beforeAll, afterEach, afterAll, vi, it, expect } from 'vitest' +import { WebSocketInterceptor } from '../../../../src/interceptors/WebSocket' + +const interceptor = new WebSocketInterceptor() + +beforeAll(() => { + interceptor.apply() +}) + +afterEach(() => { + interceptor.removeAllListeners() +}) + +afterAll(() => { + interceptor.dispose() +}) + +it('buffers client sends until the connection is open', async () => { + const events: Array = [] + + interceptor.on('connection', ({ client }) => { + client.socket.addEventListener('open', () => { + events.push('open') + }) + client.socket.addEventListener('message', () => { + events.push('send') + }) + + client.send('hello world') + }) + + const socket = new WebSocket('ws://localhost') + const messageListener = vi.fn() + socket.addEventListener('message', messageListener) + + await vi.waitFor(() => { + expect(messageListener).toHaveBeenCalledWith( + expect.objectContaining({ + data: 'hello world', + }) + ) + }) + + expect(events).toEqual(['open', 'send']) +}) + +it('does not send data if the client connection is closing', async () => { + const events: Array = [] + + interceptor.on('connection', ({ client }) => { + client.socket.addEventListener('close', () => { + events.push('close') + }) + client.socket.addEventListener('message', () => { + events.push('send') + }) + + client.close() + client.send('hello world') + }) + + const socket = new WebSocket('ws://localhost') + const messageListener = vi.fn() + const closeListener = vi.fn() + socket.addEventListener('message', messageListener) + socket.addEventListener('close', closeListener) + + await vi.waitFor(() => { + expect(closeListener).toHaveBeenCalledOnce() + }) + + expect(messageListener).not.toHaveBeenCalled() + expect(events).toEqual(['close']) +}) + +it('does not send data if the client connection is closed', async () => { + const events: Array = [] + + interceptor.on('connection', ({ client }) => { + client.socket.addEventListener('close', () => { + events.push('close') + }) + client.socket.addEventListener('message', () => { + events.push('send') + }) + + client.close() + queueMicrotask(() => client.send('hello world')) + }) + + const socket = new WebSocket('ws://localhost') + const messageListener = vi.fn() + const closeListener = vi.fn() + socket.addEventListener('message', messageListener) + socket.addEventListener('close', closeListener) + + await vi.waitFor(() => { + expect(closeListener).toHaveBeenCalledOnce() + }) + + expect(messageListener).not.toHaveBeenCalled() + expect(events).toEqual(['close']) +}) diff --git a/test/modules/WebSocket/compliance/websocket.connection.test.ts b/test/modules/WebSocket/compliance/websocket.connection.test.ts index 4b704593..bed8b2ef 100644 --- a/test/modules/WebSocket/compliance/websocket.connection.test.ts +++ b/test/modules/WebSocket/compliance/websocket.connection.test.ts @@ -15,7 +15,7 @@ const wsServer = new WebSocketServer({ port: 0, }) -beforeAll(async () => { +beforeAll(() => { interceptor.apply() }) diff --git a/test/modules/WebSocket/intercept/websocket.dispose.test.ts b/test/modules/WebSocket/intercept/websocket.dispose.test.ts index 2d11ddaa..7f0dc673 100644 --- a/test/modules/WebSocket/intercept/websocket.dispose.test.ts +++ b/test/modules/WebSocket/intercept/websocket.dispose.test.ts @@ -1,6 +1,4 @@ -/** - * @vitest-environment node-with-websocket - */ +// @vitest-environment node-with-websocket import { vi, it, expect, beforeAll, afterAll } from 'vitest' import { WebSocketServer } from 'ws' import { WebSocketInterceptor } from '../../../../src/interceptors/WebSocket/index' diff --git a/test/modules/WebSocket/intercept/websocket.send.test.ts b/test/modules/WebSocket/intercept/websocket.send.test.ts index 1bd08983..f86a5a8a 100644 --- a/test/modules/WebSocket/intercept/websocket.send.test.ts +++ b/test/modules/WebSocket/intercept/websocket.send.test.ts @@ -1,6 +1,4 @@ -/** - * @vitest-environment node-with-websocket - */ +// @vitest-environment node-with-websocket import { it, expect, beforeAll, afterAll, afterEach } from 'vitest' import { DeferredPromise } from '@open-draft/deferred-promise' import { WebSocketInterceptor } from '../../../../src/interceptors/WebSocket' @@ -30,7 +28,11 @@ it('intercepts text sent over websocket', async () => { interceptor.once('connection', ({ client }) => { client.addEventListener('message', (event) => { - messageReceivedPromise.resolve(event.data) + if (typeof event.data === 'string') { + messageReceivedPromise.resolve(event.data) + } else { + messageReceivedPromise.reject(new Error('Expected string data')) + } }) }) @@ -45,7 +47,11 @@ it('intercepts Blob sent over websocket', async () => { interceptor.once('connection', ({ client }) => { client.addEventListener('message', (event) => { - messageReceivedPromise.resolve(event.data) + if (event.data instanceof Blob) { + messageReceivedPromise.resolve(event.data) + } else { + messageReceivedPromise.reject(new Error('Expected Blob data')) + } }) }) @@ -57,11 +63,18 @@ it('intercepts Blob sent over websocket', async () => { }) it('intercepts ArrayBuffer sent over websocket', async () => { - const messageReceivedPromise = new DeferredPromise() + const messageReceivedPromise = new DeferredPromise() interceptor.once('connection', ({ client }) => { client.addEventListener('message', (event) => { - messageReceivedPromise.resolve(event.data) + /** + * @note ArrayBuffer data is represented as Buffer in Node.js. + */ + if (event.data instanceof Uint8Array) { + messageReceivedPromise.resolve(event.data) + } else { + messageReceivedPromise.reject(new Error('Expected ArrayBuffer data')) + } }) }) diff --git a/test/modules/WebSocket/intercept/websocket.server.events.test.ts b/test/modules/WebSocket/intercept/websocket.server.events.test.ts index 1ab5872b..bb8dd5ff 100644 --- a/test/modules/WebSocket/intercept/websocket.server.events.test.ts +++ b/test/modules/WebSocket/intercept/websocket.server.events.test.ts @@ -1,6 +1,4 @@ -/** - * @vitest-environment node-with-websocket - */ +// @vitest-environment node-with-websocket import { vi, it, expect, beforeAll, afterAll } from 'vitest' import { WebSocketServer } from 'ws' import { WebSocketInterceptor } from '../../../../src/interceptors/WebSocket'