feat(ws): support handleMessage option on ws.link

src/core/ws.ts

@@ -1,6 +1,7 @@
 import { invariant } from 'outvariant'
 import type {
   WebSocketClientConnectionProtocol,
+  WebSocketConnectionData,
   WebSocketData,
 } from '@mswjs/interceptors/WebSocket'
 import {
@@ -81,6 +82,13 @@ export type WebSocketLink = {
   ): void
 }
 
+export interface WebSocketLinkOptions {
+  handleMessage?: (
+    type: 'incoming' | 'outgoing',
+    data: WebSocketData,
+  ) => WebSocketData | Promise<WebSocketData>
+}
+
 /**
  * Intercepts outgoing WebSocket connections to the given URL.
  *
@@ -90,7 +98,10 @@ export type WebSocketLink = {
  *   client.send('hello from server!')
  * })
  */
-function createWebSocketLinkHandler(url: Path): WebSocketLink {
+function createWebSocketLinkHandler(
+  url: Path,
+  options: WebSocketLinkOptions = {},
+): WebSocketLink {
   invariant(url, 'Expected a WebSocket server URL but got undefined')
 
   invariant(
@@ -99,7 +110,80 @@ function createWebSocketLinkHandler(url: Path): WebSocketLink {
     typeof url,
   )
 
-  const clientManager = new WebSocketClientManager(wsBroadcastChannel, url)
+  const clientManager = new WebSocketClientManager({
+    url,
+    channel: wsBroadcastChannel,
+  })
+
+  const withCustomMessageHandler = (
+    connection: WebSocketConnectionData,
+  ): void => {
+    const { handleMessage } = options
+    const { client, server } = connection
+
+    if (!handleMessage) {
+      return
+    }
+
+    // Handle data sent to the client from the interceptor.
+    client.send = new Proxy(client.send, {
+      apply: (target, thisArg, args) => {
+        Promise.resolve(handleMessage('outgoing', args[0])).then((nextData) => {
+          Reflect.apply(target, thisArg, [nextData].concat(args.slice(1)))
+        })
+      },
+    })
+
+    // Handle data sent to the original server from the interceptor.
+    server.send = new Proxy(server.send, {
+      apply: (target, thisArg, args) => {
+        Promise.resolve(handleMessage('outgoing', args[0])).then((nextData) => {
+          Reflect.apply(target, thisArg, [nextData].concat(args.slice(1)))
+        })
+      },
+    })
+
+    client['transport'].dispatchEvent = new Proxy(
+      client['transport'].dispatchEvent,
+      {
+        apply(target, thisArg, args: [Event]) {
+          const [event] = args
+          const dispatchEvent = (event: Event) => {
+            return Reflect.apply(target, thisArg, [event])
+          }
+
+          if (
+            event instanceof MessageEvent &&
+            (event.type === 'outgoing' || event.type === 'incoming')
+          ) {
+            /**
+             * @note Treat both "incoming" and "outgoing" transport events
+             * as "incoming" events in `handleMessage`. These two transport events
+             * are different to distinguish between the client sending data
+             * ("outgoing") and the original server sending data ("incoming").
+             * But for the context of `handleMessage`, both those sends are incoming
+             * because both of them MUST be decoded, not encoded.
+             */
+            Promise.resolve(handleMessage('incoming', event.data)).then(
+              (nextData) => {
+                // Modify the original message event's `data` before dispatching.
+                // Can't do this in the event listener because those are already
+                // added internally by Interceptors.
+                Object.defineProperty(event, 'data', {
+                  value: nextData,
+                  enumerable: true,
+                })
+                dispatchEvent(event)
+              },
+            )
+            return
+          }
+
+          return dispatchEvent(event)
+        },
+      },
+    )
+  }
 
   return {
     get clients() {
@@ -112,8 +196,9 @@ function createWebSocketLinkHandler(url: Path): WebSocketLink {
       // handler matches and emits a connection event.
       // When that happens, store that connection in the
       // set of all connections for reference.
-      handler[kEmitter].on('connection', ({ client }) => {
-        clientManager.addConnection(client)
+      handler[kEmitter].on('connection', (connection) => {
+        withCustomMessageHandler(connection)
+        clientManager.addConnection(connection.client)
       })
 
       // The "handleWebSocketEvent" function will invoke

src/core/ws/WebSocketClientManager.test.ts

@@ -26,7 +26,7 @@ afterEach(() => {
 })
 
 it('adds a client from this runtime to the list of clients', () => {
-  const manager = new WebSocketClientManager(channel, '*')
+  const manager = new WebSocketClientManager({ url: '*', channel })
   const connection = new WebSocketClientConnection(
     socket,
     new TestWebSocketTransport(),
@@ -39,7 +39,7 @@ it('adds a client from this runtime to the list of clients', () => {
 })
 
 it('adds multiple clients from this runtime to the list of clients', () => {
-  const manager = new WebSocketClientManager(channel, '*')
+  const manager = new WebSocketClientManager({ url: '*', channel })
   const connectionOne = new WebSocketClientConnection(
     socket,
     new TestWebSocketTransport(),
@@ -63,7 +63,7 @@ it('adds multiple clients from this runtime to the list of clients', () => {
 })
 
 it('replays a "send" event coming from another runtime', async () => {
-  const manager = new WebSocketClientManager(channel, '*')
+  const manager = new WebSocketClientManager({ url: '*', channel })
   const connection = new WebSocketClientConnection(
     socket,
     new TestWebSocketTransport(),
@@ -92,7 +92,7 @@ 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 manager = new WebSocketClientManager({ url: '*', channel })
   const connection = new WebSocketClientConnection(
     socket,
     new TestWebSocketTransport(),
@@ -122,7 +122,7 @@ 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({ url: '*', channel })
   const transport = new TestWebSocketTransport()
   const connection = new WebSocketClientConnection(socket, transport)
   vi.spyOn(connection, 'close').mockImplementationOnce(() => {

src/core/ws/WebSocketClientManager.ts

@@ -38,8 +38,10 @@ export class WebSocketClientManager {
   private inMemoryClients: Set<WebSocketClientConnectionProtocol>
 
   constructor(
-    private channel: BroadcastChannel,
-    private url: Path,
+    private options: {
+      url: Path
+      channel: BroadcastChannel
+    },
   ) {
     this.inMemoryClients = new Set()
 
@@ -88,7 +90,7 @@ export class WebSocketClientManager {
               return new WebSocketRemoteClientConnection(
                 serializedClient.clientId,
                 new URL(serializedClient.url),
-                this.channel,
+                this.options.channel,
               )
             }),
         ),
@@ -114,7 +116,7 @@ export class WebSocketClientManager {
 
     const allClients = JSON.parse(clientsJson) as Array<SerializedClient>
     const matchingClients = allClients.filter((client) => {
-      return matchRequestUrl(new URL(client.url), this.url).matches
+      return matchRequestUrl(new URL(client.url), this.options.url).matches
     })
 
     return matchingClients
@@ -156,7 +158,7 @@ export class WebSocketClientManager {
     ) => {
       const { type, payload } = message.data
 
-      // Ignore broadcasted messages for other clients.
+      // Ignore messages broadcasted to other clients.
       if (
         typeof payload === 'object' &&
         'clientId' in payload &&
@@ -180,7 +182,7 @@ export class WebSocketClientManager {
 
     const abortController = new AbortController()
 
-    this.channel.addEventListener('message', handleExtraneousMessage, {
+    this.options.channel.addEventListener('message', handleExtraneousMessage, {
       signal: abortController.signal,
     })