test: add unit coverage for PostMessageTransport source validation

src/message-transport.test.ts

@@ -0,0 +1,390 @@
+import { describe, it, expect, beforeEach, afterEach, mock } from "bun:test";
+import type { JSONRPCMessage } from "@modelcontextprotocol/sdk/types.js";
+
+import { PostMessageTransport } from "./message-transport";
+
+/**
+ * Minimal `window` stub for bun's DOM-less test environment.
+ *
+ * Captures listeners registered via `addEventListener` so tests can dispatch
+ * fake `MessageEvent`-like objects directly. We deliberately avoid pulling in
+ * jsdom/happy-dom — this file tests the transport's contract, not the browser.
+ */
+type Listener = (event: MessageEvent) => void;
+
+function createFakeWindow() {
+  const listeners = new Map<string, Set<Listener>>();
+  return {
+    addEventListener(type: string, listener: Listener) {
+      if (!listeners.has(type)) listeners.set(type, new Set());
+      listeners.get(type)!.add(listener);
+    },
+    removeEventListener(type: string, listener: Listener) {
+      listeners.get(type)?.delete(listener);
+    },
+    /** Test helper: invoke all listeners for `type` with the given event. */
+    dispatch(type: string, event: unknown) {
+      listeners.get(type)?.forEach((l) => l(event as MessageEvent));
+    },
+    /** Test helper: how many listeners are registered for `type`? */
+    listenerCount(type: string) {
+      return listeners.get(type)?.size ?? 0;
+    },
+  };
+}
+
+/**
+ * A valid JSON-RPC request payload. Used by most happy-path tests.
+ * Kept deliberately minimal — we care that *a* valid message is delivered,
+ * not about exercising every JSON-RPC shape (that's the schema's job).
+ */
+const validRequest: JSONRPCMessage = {
+  jsonrpc: "2.0",
+  id: 1,
+  method: "ping",
+};
+
+describe("PostMessageTransport", () => {
+  let fakeWindow: ReturnType<typeof createFakeWindow>;
+  let trustedSource: object;
+  let untrustedSource: object;
+
+  /** Stand-in for the target window (e.g. window.parent or iframe.contentWindow). */
+  let targetPostMessage: ReturnType<typeof mock>;
+  let eventTarget: { postMessage: typeof targetPostMessage };
+
+  // The transport logs liberally at debug/error level. Silence during tests
+  // so failures are readable; restore afterwards so we don't leak state.
+  let restoreConsole: () => void;
+
+  beforeEach(() => {
+    fakeWindow = createFakeWindow();
+    (globalThis as { window?: unknown }).window = fakeWindow;
+
+    // Distinct object identities — the transport compares with `!==`.
+    trustedSource = { id: "trusted" };
+    untrustedSource = { id: "untrusted" };
+
+    targetPostMessage = mock(() => {});
+    eventTarget = { postMessage: targetPostMessage };
+
+    const origDebug = console.debug;
+    const origError = console.error;
+    console.debug = () => {};
+    console.error = () => {};
+    restoreConsole = () => {
+      console.debug = origDebug;
+      console.error = origError;
+    };
+  });
+
+  afterEach(() => {
+    restoreConsole();
+    delete (globalThis as { window?: unknown }).window;
+  });
+
+  /** Create and start a transport wired to the shared fakes. */
+  async function createStartedTransport() {
+    const transport = new PostMessageTransport(
+      eventTarget as unknown as Window,
+      trustedSource as MessageEventSource,
+    );
+    await transport.start();
+    return transport;
+  }
+
+  // ==========================================================================
+  // Source validation — the security boundary at this layer.
+  // The transport validates `event.source` (window identity), not `event.origin`.
+  // Origin checks live in the sandbox proxy relay between the two endpoints
+  // (see examples/basic-host/src/sandbox.ts). Here, source===contentWindow is
+  // the narrower check; the app side can't know its sandbox's origin anyway.
+  // ==========================================================================
+  describe("source validation", () => {
+    it("delivers messages from the configured eventSource", async () => {
+      const transport = await createStartedTransport();
+      const received: JSONRPCMessage[] = [];
+      transport.onmessage = (msg) => received.push(msg);
+
+      fakeWindow.dispatch("message", {
+        source: trustedSource,
+        data: validRequest,
+      });
+
+      expect(received).toEqual([validRequest]);
+    });
+
+    it("drops messages from a different source", async () => {
+      const transport = await createStartedTransport();
+      const received: JSONRPCMessage[] = [];
+      const errors: Error[] = [];
+      transport.onmessage = (msg) => received.push(msg);
+      transport.onerror = (err) => errors.push(err);
+
+      fakeWindow.dispatch("message", {
+        source: untrustedSource,
+        data: validRequest,
+      });
+
+      // The message is silently dropped — neither delivered nor surfaced as
+      // an error. An attacker flooding us with forged messages should not be
+      // able to DoS the error handler.
+      expect(received).toEqual([]);
+      expect(errors).toEqual([]);
+    });
+
+    it("drops messages with a null source", async () => {
+      // Some browser-injected messages (extensions, devtools) arrive with
+      // source === null. These must not reach the protocol layer.
+      const transport = await createStartedTransport();
+      const received: JSONRPCMessage[] = [];
+      transport.onmessage = (msg) => received.push(msg);
+
+      fakeWindow.dispatch("message", {
+        source: null,
+        data: validRequest,
+      });
+
+      expect(received).toEqual([]);
+    });
+
+    it("continues delivering trusted messages after dropping an untrusted one", async () => {
+      // Regression guard: a rejected message must not break the listener or
+      // close the transport.
+      const transport = await createStartedTransport();
+      const received: JSONRPCMessage[] = [];
+      transport.onmessage = (msg) => received.push(msg);
+
+      fakeWindow.dispatch("message", {
+        source: untrustedSource,
+        data: validRequest,
+      });
+      fakeWindow.dispatch("message", {
+        source: trustedSource,
+        data: validRequest,
+      });
+
+      expect(received).toEqual([validRequest]);
+    });
+  });
+
+  // ==========================================================================
+  // Message format validation.
+  // Three paths: valid → onmessage, non-JSON-RPC → silent, malformed → onerror.
+  // ==========================================================================
+  describe("message format validation", () => {
+    it("delivers valid JSON-RPC requests via onmessage", async () => {
+      const transport = await createStartedTransport();
+      const received: JSONRPCMessage[] = [];
+      transport.onmessage = (msg) => received.push(msg);
+
+      fakeWindow.dispatch("message", {
+        source: trustedSource,
+        data: { jsonrpc: "2.0", id: 42, method: "tools/call" },
+      });
+
+      expect(received).toHaveLength(1);
+      expect(received[0]).toMatchObject({ id: 42, method: "tools/call" });
+    });
+
+    it("delivers valid JSON-RPC notifications via onmessage", async () => {
+      const transport = await createStartedTransport();
+      const received: JSONRPCMessage[] = [];
+      transport.onmessage = (msg) => received.push(msg);
+
+      fakeWindow.dispatch("message", {
+        source: trustedSource,
+        data: { jsonrpc: "2.0", method: "notifications/initialized" },
+      });
+
+      expect(received).toHaveLength(1);
+      expect(received[0]).toMatchObject({
+        method: "notifications/initialized",
+      });
+    });
+
+    it("silently ignores non-JSON-RPC payloads", async () => {
+      // Real iframes receive all sorts of ambient postMessage traffic:
+      // React DevTools, browser extensions, ad frames. These must not crash
+      // the transport or surface as protocol errors.
+      const transport = await createStartedTransport();
+      const received: JSONRPCMessage[] = [];
+      const errors: Error[] = [];
+      transport.onmessage = (msg) => received.push(msg);
+      transport.onerror = (err) => errors.push(err);
+
+      fakeWindow.dispatch("message", {
+        source: trustedSource,
+        data: { type: "react-devtools-hook", payload: {} },
+      });
+      fakeWindow.dispatch("message", { source: trustedSource, data: "hello" });
+      fakeWindow.dispatch("message", { source: trustedSource, data: null });
+      fakeWindow.dispatch("message", { source: trustedSource, data: 42 });
+
+      expect(received).toEqual([]);
+      expect(errors).toEqual([]);
+    });
+
+    it("calls onerror for malformed messages claiming to be JSON-RPC", async () => {
+      // `jsonrpc: "2.0"` but missing required fields — this IS a protocol
+      // violation worth surfacing, unlike ambient noise.
+      const transport = await createStartedTransport();
+      const received: JSONRPCMessage[] = [];
+      const errors: Error[] = [];
+      transport.onmessage = (msg) => received.push(msg);
+      transport.onerror = (err) => errors.push(err);
+
+      fakeWindow.dispatch("message", {
+        source: trustedSource,
+        data: { jsonrpc: "2.0" }, // no id/method/result — invalid
+      });
+
+      expect(received).toEqual([]);
+      expect(errors).toHaveLength(1);
+      expect(errors[0]).toBeInstanceOf(Error);
+    });
+
+    it("does not throw when onerror is unset and a malformed message arrives", async () => {
+      const transport = await createStartedTransport();
+      transport.onmessage = () => {};
+      // onerror deliberately left unset
+
+      expect(() =>
+        fakeWindow.dispatch("message", {
+          source: trustedSource,
+          data: { jsonrpc: "2.0" },
+        }),
+      ).not.toThrow();
+    });
+
+    it("does not throw when onmessage is unset and a valid message arrives", async () => {
+      await createStartedTransport();
+      // onmessage deliberately left unset — may happen briefly during wiring
+
+      expect(() =>
+        fakeWindow.dispatch("message", {
+          source: trustedSource,
+          data: validRequest,
+        }),
+      ).not.toThrow();
+    });
+  });
+
+  // ==========================================================================
+  // send()
+  // ==========================================================================
+  describe("send()", () => {
+    it("posts the message to the configured eventTarget with '*' origin", async () => {
+      const transport = await createStartedTransport();
+
+      const msg: JSONRPCMessage = {
+        jsonrpc: "2.0",
+        id: 7,
+        method: "ui/initialize",
+      };
+      await transport.send(msg);
+
+      expect(targetPostMessage).toHaveBeenCalledTimes(1);
+      expect(targetPostMessage).toHaveBeenCalledWith(msg, "*");
+    });
+
+    it("posts multiple messages in order", async () => {
+      const transport = await createStartedTransport();
+
+      const a: JSONRPCMessage = { jsonrpc: "2.0", id: 1, method: "a" };
+      const b: JSONRPCMessage = { jsonrpc: "2.0", id: 2, method: "b" };
+      await transport.send(a);
+      await transport.send(b);
+
+      expect(targetPostMessage).toHaveBeenCalledTimes(2);
+      expect(targetPostMessage.mock.calls[0][0]).toEqual(a);
+      expect(targetPostMessage.mock.calls[1][0]).toEqual(b);
+    });
+  });
+
+  // ==========================================================================
+  // Lifecycle: start() / close()
+  // ==========================================================================
+  describe("lifecycle", () => {
+    it("start() registers a message listener on window", async () => {
+      const transport = new PostMessageTransport(
+        eventTarget as unknown as Window,
+        trustedSource as MessageEventSource,
+      );
+
+      expect(fakeWindow.listenerCount("message")).toBe(0);
+      await transport.start();
+      expect(fakeWindow.listenerCount("message")).toBe(1);
+    });
+
+    it("close() removes the message listener", async () => {
+      const transport = await createStartedTransport();
+
+      expect(fakeWindow.listenerCount("message")).toBe(1);
+      await transport.close();
+      expect(fakeWindow.listenerCount("message")).toBe(0);
+    });
+
+    it("messages dispatched after close() are not delivered", async () => {
+      const transport = await createStartedTransport();
+      const received: JSONRPCMessage[] = [];
+      transport.onmessage = (msg) => received.push(msg);
+
+      await transport.close();
+      fakeWindow.dispatch("message", {
+        source: trustedSource,
+        data: validRequest,
+      });
+
+      expect(received).toEqual([]);
+    });
+
+    it("close() invokes onclose when set", async () => {
+      const transport = await createStartedTransport();
+      const onclose = mock(() => {});
+      transport.onclose = onclose;
+
+      await transport.close();
+
+      expect(onclose).toHaveBeenCalledTimes(1);
+    });
+
+    it("close() does not throw when onclose is unset", async () => {
+      const transport = await createStartedTransport();
+      // onclose deliberately left unset
+
+      expect(transport.close()).resolves.toBeUndefined();
+    });
+
+    it("two transports listen independently", async () => {
+      // Host scenario: multiple iframes, each with its own transport.
+      // Each transport must only accept messages from its own iframe.
+      const sourceA = { id: "iframe-a" };
+      const sourceB = { id: "iframe-b" };
+
+      const transportA = new PostMessageTransport(
+        eventTarget as unknown as Window,
+        sourceA as unknown as MessageEventSource,
+      );
+      const transportB = new PostMessageTransport(
+        eventTarget as unknown as Window,
+        sourceB as unknown as MessageEventSource,
+      );
+      await transportA.start();
+      await transportB.start();
+
+      const receivedA: JSONRPCMessage[] = [];
+      const receivedB: JSONRPCMessage[] = [];
+      transportA.onmessage = (msg) => receivedA.push(msg);
+      transportB.onmessage = (msg) => receivedB.push(msg);
+
+      fakeWindow.dispatch("message", { source: sourceA, data: validRequest });
+
+      expect(receivedA).toHaveLength(1);
+      expect(receivedB).toHaveLength(0);
+
+      await transportA.close();
+      await transportB.close();
+    });
+  });
+});