websocket-intercept.ts

/**
 * Intercepting a page's own WebSocket feed from a browser extension.
 *
 * A content script in the extension's *isolated world* shares the page's DOM
 * but not its JavaScript heap, so it can't see or wrap the page's own objects.
 * To read the app's socket you inject into the page's *main world* at
 * `document_start` — before the app opens any socket — and wrap the
 * `WebSocket` constructor itself rather than any single instance.
 *
 * The golden rule is: observe, never interfere. The handler is wrapped so a
 * throw can never bubble into the host app's message loop.
 */

export interface InterceptOptions {
  /** Called with every string message the page receives. */
  onMessage: (data: string) => void;
  /**
   * Object that carries the `WebSocket` constructor. Defaults to `globalThis`
   * (i.e. `window`). Injectable so the interceptor can be unit-tested without
   * a real DOM.
   */
  target?: { WebSocket: typeof WebSocket };
}

/**
 * Wrap the target's `WebSocket` constructor so every message on every socket
 * the page opens is forwarded to `onMessage`. Returns an `uninstall` function
 * that restores the native constructor.
 */
export function interceptWebSocket(opts: InterceptOptions): () => void {
  const target = (opts.target ?? (globalThis as unknown as { WebSocket: typeof WebSocket }));
  const Native = target.WebSocket;
  if (!Native) throw new Error('interceptWebSocket: no WebSocket on target');

  function Patched(url: string | URL, protocols?: string | string[]) {
    const ws = new Native(url, protocols);
    ws.addEventListener('message', (event: MessageEvent) => {
      // Never let an interceptor throw inside the host app's message loop.
      try {
        if (typeof event.data === 'string') opts.onMessage(event.data);
      } catch {
        /* swallow — observe, never interfere */
      }
    });
    return ws;
  }

  // Preserve the prototype chain so `x instanceof WebSocket` still holds
  // everywhere in the host app, and copy the static readyState constants.
  Patched.prototype = Native.prototype;
  for (const key of ['CONNECTING', 'OPEN', 'CLOSING', 'CLOSED'] as const) {
    (Patched as unknown as Record<string, unknown>)[key] = (Native as unknown as Record<string, unknown>)[key];
  }

  target.WebSocket = Patched as unknown as typeof WebSocket;
  return () => {
    target.WebSocket = Native;
  };
}

const FRAME = /~m~(\d+)~m~/g;

/**
 * De-frame a length-prefixed, multiplexed text protocol of the shape:
 *
 *     ~m~<byteLength>~m~<payload>   (repeated, heartbeats interleaved)
 *
 * Yields each payload as-is; the caller decides what to do with it (parse
 * JSON, skip heartbeats, etc.). If the input carries no framing at all, the
 * whole string is yielded once so callers can handle unframed feeds uniformly.
 */
export function* deframe(raw: string): Generator<string> {
  FRAME.lastIndex = 0;
  let match: RegExpExecArray | null;
  let matched = false;
  while ((match = FRAME.exec(raw)) !== null) {
    matched = true;
    const len = Number(match[1]);
    const start = match.index + match[0].length;
    yield raw.slice(start, start + len);
  }
  if (!matched && raw.length > 0) yield raw;
}