Skip to content

Commit

Permalink
fix(WebSocket): add "server.close()" method (#528)
Browse files Browse the repository at this point in the history
  • Loading branch information
@kettanaito
kettanaito authored Mar 24, 2024
1 parent 1337546 commit 4cd11bb
Show file tree
Hide file tree
Showing 5 changed files with 379 additions and 96 deletions.
148 changes: 74 additions & 74 deletions README.md
View file
Original file line number Diff line number Diff line change
Expand Up @@ -292,8 +292,8 @@ interceptor.on(

You can intercept a WebSocket communication using the `WebSocketInterceptor` class.

> [!WARNING]
> In practice, WebSocket connections can use different mechanisms to work (called "transports"). At this moment, the WebSocket interceptor only supports connections established using the `globalThis.WebSocket` class. Supporting third-party transports is challenging because they are non-standard and specific to individual WebSocket client libraries.
> [!IMPORTANT]
> This library only supports intercepting WebSocket connections created using the global WHATWG `WebSocket` class. Third-party transports, such as HTTP/XHR polling, are not supported by design due to their contrived nature.
```js
import { WebSocketInterceptor } from '@mswjs/interceptors/WebSocket'
Expand All @@ -303,6 +303,13 @@ const interceptor = new WebSocketInterceptor()

Unlike the HTTP-based interceptors that share the same `request`/`response` events, the WebSocket interceptor only emits the `connection` event and let's you handle the incoming/outgoing events in its listener.

### Important defaults

1. Intercepted WebSocket connections are _not_ opened. To open the actual WebSocket connection, call [`server.connect()`](#connect) in the interceptor.
1. Once connected to the actual server, the outgoing client events are _not_ forwarded to that server. You must add a message listener for the `client` and forward the events manually using [`server.send()`](#senddata-1).
1. Once connected to the actual server, the incoming server events _are_ forwarded to the client by default. You can prevent that forwarding by calling `event.preventDefault()` in the message event listener for the `server`.
1. Once connected to the actual server, the `close` event received from that server is forwarded to the intercepted client by default. You can prevent that forwarding by calling `event.preventDefault()` in the close event listener for the `server`.

### WebSocket connection

Whenever a WebSocket instance is constructed, the `connection` event is emitted on the WebSocket interceptor.
Expand All @@ -315,121 +322,114 @@ intereceptor.on('connection', ({ client }) => {

The `connection` event exposes the following arguments:

| Name | Type | Description |
| -------- | -------- | ---------------------------------------------------------------- |
| `client` | `object` | An object representing a connected WebSocket client instance. |
| `server` | `object` | An object representing the original WebSocket server connection. |
| Name | Type | Description |
| -------- | --------------------------------------------------------- | ---------------------------------------------------------------- |
| `client` | [`WebSocketClientConnection`](#websocketclientconnection) | An object representing a connected WebSocket client instance. |
| `server` | [`WebSocketServerConnection`](#websocketserverconnection) | An object representing the original WebSocket server connection. |

### Intercepting outgoing client events
### `WebSocketClientConnection`

To intercept an event sent by the WebSocket client, add a `message` listener to the `client` object.
#### `.addEventListener(type, listener)`

```js
interceptor.on('connection', ({ client }) => {
// Intercept all outgoing events from the client.
client.addEventListener('message', (event) => {
console.log('received:', event.data)
})
})
```
- `type`, `"message"`
- `listener`, `(event: MessageEvent<WebSocketData>) => void`

### Mocking incoming server events
Adds an event listener to the given event type of the WebSocket client.

To mock an event sent from the server to the client, call `client.send()` with the event payload.
| Event name | Description |
| ---------- | ----------------------------------------------------------------- |
| `message` | Dispatched when data is sent by the intercepted WebSocket client. |

```js
interceptor.on('connection', ({ client }) => {
// Send a mocked "MessageEvent" to the client
// with the given "hello" string as the data.
client.send('hello')
client.addEventListener('message', (event) => {
console.log('outgoing:', event.data)
})
```

> The WebSocket interceptor respects the [WebSocket WHATWG standard](https://websockets.spec.whatwg.org/) and supports sending all supported data types (string, Blob, ArrayBuffer, etc).
#### `.removeEventListener(type, listener)`

### Bypassing events
- `type`, `"message"`
- `listener`, `(event: MessageEvent<WebSocketData>) => void`

By default, the WebSocket interceptor **prevents all the outgoing client events from hitting the production server**. This is a sensible default to support mocking a WebSocket communication when a WebSocket server doesn't exist.
Removes the listener for the given event type.

To bypass an event, first establish the actual server connection by calling `server.connect()`, and then call `server.send()` with the data you wish to forward to the original server.
#### `.send(data)`

```js
interceptor.on('connection', ({ client, server }) => {
// First, connect to the original server.
server.connect()
- `data`, `string | Blob | ArrayBuffer`

// Forward all outgoing client events to the original server.
client.addEventListener('message', (event) => server.send(event.data))
})
Sends the data to the intercepted WebSocket client.

```js
client.send('text')
client.send(new Blob(['blob']))
client.send(new TextEncoder().encode('array buffer'))
```

### Intercepting incoming server events
#### `.close(code, reason)`

The WebSocket communication is duplex and the WebSocket interceptor allows you to intercept both outgoing (client) events and incoming (original server) events.
- `code`, close [status code](https://www.rfc-editor.org/rfc/rfc6455#section-7.4.1).
- `reason`, [close reason](https://www.rfc-editor.org/rfc/rfc6455#section-7.1.6).

To intercept an incoming event from the original server, first establish the original connection by calling `server.connect()`, and then add a `message` listener to the `server` object.
Closes the client connection. Unlike the regular `WebSocket.prototype.close()`, the `client.close()` method can accept a non-configurable status codes, such as 1001, 1003, etc.

```js
interceptor.on('connection', ({ server }) => {
server.connect()
// Gracefully close the connection with the
// intercepted WebSocket client.
client.close()
```

server.addEventListener('message', (event) => {
console.log('original server sent:', event.data)
})
})
```js
// Terminate the connection by emulating
// the server unable to process the received data.
client.close(1003)
```

Unlike the outgoing client events, incoming server events **are automatically forwarded to the client as soon as you call `server.connect()`**. This keeps the original server connection authentic if you ever decide to open one.
### `WebSocketServerConnection`

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. This can be handy for observing as well as modifying incoming events.
#### `.connect()`

```js
interceptor.on('connection', ({ client, server }) => {
server.connect()
Establishes the connection to the original WebSocket server. Connection cannot be awaited. Any data sent via `server.send()` while connecting is buffered and flushed once the connection is open.

server.addEventListener('message', (event) => {
if (event.data === 'hello from server') {
// Never forward this event to the client.
event.preventDefault()
#### `.addEventListener(type, listener)`

// Instead, send this mock data.
client.send('greetings, client')
return
}
})
})
```
- `type`, `"message"`
- `listener`, `(event: MessageEvent<WebSocketData>) => void`

### Closing the connection
Adds an event listener to the given event type of the WebSocket server.

You can terminate the open WebSocket client connection by calling `client.close()`.
| Event name | Description |
| ---------- | -------------------------------------------------------------------- |
| `message` | Dispatched when data is received from the original WebSocket server. |

```js
interceptor.on('connection', ({ client }) => {
client.close()
server.addEventListener('message', (event) => {
console.log('incoming:', event.data)
})
```

By default, this will close the connection with the `1000` code, meaning a graceful disconnect.
#### `.removeEventListener(type, listener)`

You can provide a custom close [Status code](https://www.rfc-editor.org/rfc/rfc6455#section-7.4.1) and [Close reason](https://www.rfc-editor.org/rfc/rfc6455#section-7.1.6) to the `client.close()` to mock different connection close scenarios.
- `type`, `"message"`
- `listener`, `(event: MessageEvent<WebSocketData>) => void`

```js
interceptor.on('connection', ({ client }) => {
client.close(3000, 'Close reason')
})
```
Removes the listener for the given event type.

#### `.send(data)`

You can also close the connection with the termination status code (1001 - 1015), which are not configurable by the user otherwise.
- `data`, `string | Blob | ArrayBuffer`

Sends the data to the original WebSocket server. Useful in a combination with the client-sent events forwarding:

```js
interceptor.on('connection', ({ client }) => {
// Terminate the connection because the "server"
// cannot accept the data sent from the client.
client.close(1003)
client.addEventListener('message', (event) => {
server.send(event.data)
})
```

#### `.close()`

Closes the connection with the original WebSocket server. Unlike `client.close()`, closing the server connection does not accept any arguments and always asumes a graceful closure. Sending data via `server.send()` after the connection has been closed will have no effect.

## API

### `Interceptor`
Expand Down
108 changes: 87 additions & 21 deletions src/interceptors/WebSocket/WebSocketServerConnection.ts
View file
Original file line number Diff line number Diff line change
@@ -1,9 +1,9 @@
import { invariant } from 'outvariant'
import type { WebSocketOverride } from './WebSocketOverride'
import { kClose, WebSocketOverride } from './WebSocketOverride'
import type { WebSocketData } from './WebSocketTransport'
import type { WebSocketClassTransport } from './WebSocketClassTransport'
import { bindEvent } from './utils/bindEvent'
import { CancelableMessageEvent } from './utils/events'
import { CancelableMessageEvent, CloseEvent } from './utils/events'

const kEmitter = Symbol('kEmitter')

Expand All @@ -17,6 +17,7 @@ export class WebSocketServerConnection {
* A WebSocket instance connected to the original server.
*/
private realWebSocket?: WebSocket
private mockCloseController: AbortController
private [kEmitter]: EventTarget

constructor(
Expand All @@ -25,6 +26,7 @@ export class WebSocketServerConnection {
private readonly createConnection: () => WebSocket
) {
this[kEmitter] = new EventTarget()
this.mockCloseController = new AbortController()

// Handle incoming events from the actual server.
// The (mock) WebSocket instance will call this
Expand Down Expand Up @@ -97,36 +99,36 @@ export class WebSocketServerConnection {
*/
public connect(): void {
invariant(
this.readyState === -1,
!this.realWebSocket || this.realWebSocket.readyState !== WebSocket.OPEN,
'Failed to call "connect()" on the original WebSocket instance: the connection already open'
)

const ws = this.createConnection()
const realWebSocket = this.createConnection()

// Inherit the binary type from the mock WebSocket client.
ws.binaryType = this.socket.binaryType

// Close the original connection when the (mock)
// client closes, regardless of the reason.
this.socket.addEventListener(
'close',
(event) => {
ws.close(event.code, event.reason)
},
{ once: true }
)
realWebSocket.binaryType = this.socket.binaryType

ws.addEventListener('message', (event) => {
realWebSocket.addEventListener('message', (event) => {
this.transport.onIncoming(event)
})

// Close the original connection when the mock client closes.
// E.g. "client.close()" was called.
this.socket.addEventListener('close', this.handleMockClose.bind(this), {
signal: this.mockCloseController.signal,
})

// Forward the "close" event to let the interceptor handle
// closures initiated by the original server.
realWebSocket.addEventListener('close', this.handleRealClose.bind(this))

// Forward server errors to the WebSocket client as-is.
// We may consider exposing them to the interceptor in the future.
ws.addEventListener('error', () => {
realWebSocket.addEventListener('error', () => {
this.socket.dispatchEvent(bindEvent(this.socket, new Event('error')))
})

this.realWebSocket = ws
this.realWebSocket = realWebSocket
}

/**
Expand All @@ -145,7 +147,7 @@ export class WebSocketServerConnection {
}

/**
* Removes the listener for the given event.
* Remove the listener for the given event.
*/
public removeEventListener<K extends keyof WebSocketEventMap>(
event: K,
Expand All @@ -168,16 +170,25 @@ export class WebSocketServerConnection {
*/
public send(data: WebSocketData): void {
const { realWebSocket } = this

invariant(
realWebSocket,
'Failed to call "server.send()" for "%s": the connection is not open. Did you forget to call "await server.connect()"?',
'Failed to call "server.send()" for "%s": the connection is not open. Did you forget to call "server.connect()"?',
this.socket.url
)

// Silently ignore writes on the closed original WebSocket.
if (
realWebSocket.readyState === WebSocket.CLOSING ||
realWebSocket.readyState === WebSocket.CLOSED
) {
return
}

// Delegate the send to when the original connection is open.
// Unlike the mock, connecting to the original server may take time
// so we cannot call this on the next tick.
if (realWebSocket.readyState === realWebSocket.CONNECTING) {
if (realWebSocket.readyState === WebSocket.CONNECTING) {
realWebSocket.addEventListener(
'open',
() => {
Expand All @@ -191,4 +202,59 @@ export class WebSocketServerConnection {
// Send the data to the original WebSocket server.
realWebSocket.send(data)
}

/**
* Close the actual server connection.
*/
public close(): void {
const { realWebSocket } = this

invariant(
realWebSocket,
'Failed to close server connection for "%s": the connection is not open. Did you forget to call "server.connect()"?',
this.socket.url
)

realWebSocket.removeEventListener('close', this.handleRealClose)

if (
realWebSocket.readyState === WebSocket.CLOSING ||
realWebSocket.readyState === WebSocket.CLOSED
) {
return
}

realWebSocket.close()
}

private handleMockClose(_event: Event): void {
// Close the original connection if the mock client closes.
if (this.realWebSocket) {
this.realWebSocket.close()
}
}

private handleRealClose(event: CloseEvent): void {
// For closures originating from the original server,
// remove the "close" listener from the mock client.
// original close -> (?) client[kClose]() --X-> "close" (again).
this.mockCloseController.abort()

const closeEvent = bindEvent(
this.realWebSocket,
new CloseEvent('close', event)
)

this[kEmitter].dispatchEvent(closeEvent)

// If the close event from the server hasn't been prevented,
// forward the closure to the mock client.
if (!closeEvent.defaultPrevented) {
// Close the intercepted client forcefully to
// allow non-configurable status codes from the server.
// If the socket has been closed by now, no harm calling
// this again—it will have no effect.
this.socket[kClose](event.code, event.reason)
}
}
}
Loading

0 comments on commit 4cd11bb

Please sign in to comment.