@vellumai/vellum-gateway 0.7.0 → 0.7.2

package/src/velay/websocket-bridge.test.ts

@@ -0,0 +1,367 @@
+import { Buffer } from "node:buffer";
+import { afterAll, beforeEach, describe, expect, mock, test } from "bun:test";
+
+import {
+  VELAY_FRAME_TYPES,
+  VELAY_WEBSOCKET_MESSAGE_TYPES,
+  type VelayFrame,
+  type VelayWebSocketOpenFrame,
+} from "./protocol.js";
+import { FakeWebSocket } from "./test-fake-websocket.js";
+import { VelayWebSocketBridge } from "./websocket-bridge.js";
+
+const WS_OPEN = WebSocket.OPEN;
+
+const OriginalWebSocket = globalThis.WebSocket;
+let fakeSocket: FakeWebSocket;
+let sentFrames: VelayFrame[];
+let bridge: VelayWebSocketBridge;
+let WebSocketMock: ReturnType<typeof mock>;
+
+beforeEach(() => {
+  fakeSocket = new FakeWebSocket("", undefined, { validateReason: true });
+  sentFrames = [];
+  WebSocketMock = mock(() => fakeSocket);
+  Object.assign(WebSocketMock, FakeWebSocket);
+  globalThis.WebSocket = WebSocketMock as unknown as typeof WebSocket;
+  bridge = new VelayWebSocketBridge("http://127.0.0.1:7830", (frame) => {
+    sentFrames.push(frame);
+  });
+});
+
+afterAll(() => {
+  globalThis.WebSocket = OriginalWebSocket;
+});
+
+function makeOpenFrame(
+  overrides: Partial<VelayWebSocketOpenFrame> = {},
+): VelayWebSocketOpenFrame {
+  return {
+    type: VELAY_FRAME_TYPES.websocketOpen,
+    connection_id: "conn-123",
+    path: "/webhooks/twilio/relay",
+    raw_query: "callSessionId=session-123&token=edge-token",
+    headers: {},
+    ...overrides,
+  };
+}
+
+function base64(text: string | Uint8Array): string {
+  return Buffer.from(text).toString("base64");
+}
+
+async function flushPromises(): Promise<void> {
+  await Promise.resolve();
+  await Promise.resolve();
+}
+
+describe("VelayWebSocketBridge", () => {
+  test("opens a loopback WebSocket and sends websocket_opened after local open", () => {
+    bridge.handleFrame(
+      makeOpenFrame({
+        headers: {
+          authorization: ["Bearer edge-token"],
+          connection: ["upgrade"],
+          host: ["public.example.com"],
+          "sec-websocket-key": ["client-key"],
+          "x-twilio-signature": ["sig-123"],
+        },
+        subprotocol: "twilio-relay",
+      }),
+    );
+
+    expect(sentFrames).toEqual([]);
+    expect(WebSocketMock).toHaveBeenCalledTimes(1);
+
+    const [url, options] = WebSocketMock.mock.calls[0] as [
+      string,
+      { headers: Record<string, string>; protocols?: string[] },
+    ];
+    expect(url).toBe(
+      "ws://127.0.0.1:7830/webhooks/twilio/relay?callSessionId=session-123&token=edge-token",
+    );
+    expect(options.protocols).toEqual(["twilio-relay"]);
+    expect(options.headers.authorization).toBe("Bearer edge-token");
+    expect(options.headers.host).toBe("public.example.com");
+    expect(options.headers["x-twilio-signature"]).toBe("sig-123");
+    expect(options.headers.connection).toBeUndefined();
+    expect(options.headers["sec-websocket-key"]).toBeUndefined();
+
+    fakeSocket.readyState = WS_OPEN;
+    fakeSocket.emit("open");
+
+    expect(sentFrames).toEqual([
+      {
+        type: VELAY_FRAME_TYPES.websocketOpened,
+        connection_id: "conn-123",
+      },
+    ]);
+  });
+
+  test("sends websocket_open_error when the local upgrade fails before open", () => {
+    bridge.open(makeOpenFrame());
+
+    fakeSocket.emit("error");
+
+    expect(sentFrames).toEqual([
+      {
+        type: VELAY_FRAME_TYPES.websocketOpenError,
+        connection_id: "conn-123",
+        reason: "WebSocket connection failed",
+      },
+    ]);
+    expect(bridge.getConnectionCount()).toBe(0);
+    expect(fakeSocket.closes).toEqual([{ code: undefined, reason: undefined }]);
+  });
+
+  test("forwards text frames from Velay to the local gateway WebSocket", () => {
+    bridge.open(makeOpenFrame());
+    fakeSocket.readyState = WS_OPEN;
+    fakeSocket.emit("open");
+
+    bridge.message({
+      type: VELAY_FRAME_TYPES.websocketMessage,
+      connection_id: "conn-123",
+      message_type: VELAY_WEBSOCKET_MESSAGE_TYPES.text,
+      body_base64: base64('{"event":"start"}'),
+    });
+
+    expect(fakeSocket.sent).toEqual(['{"event":"start"}']);
+  });
+
+  test("preserves leading UTF-8 BOM bytes in Velay text frames", () => {
+    bridge.open(makeOpenFrame());
+    fakeSocket.readyState = WS_OPEN;
+    fakeSocket.emit("open");
+
+    bridge.message({
+      type: VELAY_FRAME_TYPES.websocketMessage,
+      connection_id: "conn-123",
+      message_type: VELAY_WEBSOCKET_MESSAGE_TYPES.text,
+      body_base64: base64(new Uint8Array([0xef, 0xbb, 0xbf, 0x68, 0x69])),
+    });
+
+    expect(fakeSocket.sent).toEqual(["\ufeffhi"]);
+  });
+
+  test("forwards local text frames back to Velay as websocket_message frames", async () => {
+    bridge.open(makeOpenFrame());
+    fakeSocket.readyState = WS_OPEN;
+    fakeSocket.emit("open");
+
+    fakeSocket.emit("message", { data: "hello from gateway" });
+    await flushPromises();
+
+    expect(sentFrames.at(-1)).toEqual({
+      type: VELAY_FRAME_TYPES.websocketMessage,
+      connection_id: "conn-123",
+      message_type: VELAY_WEBSOCKET_MESSAGE_TYPES.text,
+      body_base64: base64("hello from gateway"),
+    });
+  });
+
+  test("preserves binary frames in both directions", async () => {
+    bridge.open(makeOpenFrame());
+    fakeSocket.readyState = WS_OPEN;
+    fakeSocket.emit("open");
+
+    bridge.message({
+      type: VELAY_FRAME_TYPES.websocketMessage,
+      connection_id: "conn-123",
+      message_type: VELAY_WEBSOCKET_MESSAGE_TYPES.binary,
+      body_base64: base64(new Uint8Array([1, 2, 3])),
+    });
+
+    expect(fakeSocket.sent[0]).toEqual(new Uint8Array([1, 2, 3]));
+
+    fakeSocket.emit("message", { data: new Uint8Array([4, 5, 6]) });
+    await flushPromises();
+
+    expect(sentFrames.at(-1)).toEqual({
+      type: VELAY_FRAME_TYPES.websocketMessage,
+      connection_id: "conn-123",
+      message_type: VELAY_WEBSOCKET_MESSAGE_TYPES.binary,
+      body_base64: base64(new Uint8Array([4, 5, 6])),
+    });
+  });
+
+  test("forwards close frames and cleans up the connection map", () => {
+    bridge.open(makeOpenFrame());
+    fakeSocket.readyState = WS_OPEN;
+    fakeSocket.emit("open");
+
+    bridge.close({
+      type: VELAY_FRAME_TYPES.websocketClose,
+      connection_id: "conn-123",
+      code: 1000,
+      reason: "done",
+    });
+
+    expect(fakeSocket.closes).toEqual([{ code: 1000, reason: "done" }]);
+    expect(bridge.getConnectionCount()).toBe(0);
+
+    fakeSocket.emit("close", { code: 1000, reason: "done" });
+    expect(
+      sentFrames.filter(
+        (frame) => frame.type === VELAY_FRAME_TYPES.websocketClose,
+      ),
+    ).toEqual([]);
+  });
+
+  test("remaps protocol close codes from Velay before closing locally", () => {
+    bridge.open(makeOpenFrame());
+    fakeSocket.readyState = WS_OPEN;
+    fakeSocket.emit("open");
+
+    expect(() => {
+      bridge.close({
+        type: VELAY_FRAME_TYPES.websocketClose,
+        connection_id: "conn-123",
+        code: 1008,
+        reason: "Policy violation",
+      });
+    }).not.toThrow();
+
+    expect(fakeSocket.closes).toEqual([
+      { code: 4008, reason: "Policy violation" },
+    ]);
+    expect(bridge.getConnectionCount()).toBe(0);
+  });
+
+  test("sanitizes invalid reserved close codes from Velay before closing locally", () => {
+    bridge.open(makeOpenFrame());
+    fakeSocket.readyState = WS_OPEN;
+    fakeSocket.emit("open");
+
+    expect(() => {
+      bridge.close({
+        type: VELAY_FRAME_TYPES.websocketClose,
+        connection_id: "conn-123",
+        code: 1006,
+        reason: "invalid reserved code",
+      });
+    }).not.toThrow();
+
+    expect(fakeSocket.closes).toEqual([{ code: undefined, reason: undefined }]);
+    expect(bridge.getConnectionCount()).toBe(0);
+  });
+
+  test("truncates overlong close reasons from Velay without splitting UTF-8 characters", () => {
+    bridge.open(makeOpenFrame());
+    fakeSocket.readyState = WS_OPEN;
+    fakeSocket.emit("open");
+
+    bridge.close({
+      type: VELAY_FRAME_TYPES.websocketClose,
+      connection_id: "conn-123",
+      code: 1000,
+      reason: "🙂".repeat(40),
+    });
+
+    expect(fakeSocket.closes).toEqual([
+      { code: 1000, reason: "🙂".repeat(30) },
+    ]);
+    expect(
+      new TextEncoder().encode(fakeSocket.closes[0].reason).byteLength,
+    ).toBe(120);
+    expect(bridge.getConnectionCount()).toBe(0);
+  });
+
+  test("keeps empty text frames distinct from invalid payloads", () => {
+    bridge.open(makeOpenFrame());
+    fakeSocket.readyState = WS_OPEN;
+    fakeSocket.emit("open");
+
+    bridge.message({
+      type: VELAY_FRAME_TYPES.websocketMessage,
+      connection_id: "conn-123",
+      message_type: VELAY_WEBSOCKET_MESSAGE_TYPES.text,
+      body_base64: "",
+    });
+
+    expect(fakeSocket.sent).toEqual([""]);
+    expect(fakeSocket.closes).toEqual([]);
+  });
+
+  test("rejects invalid base64 payloads and cleans up", () => {
+    bridge.open(makeOpenFrame());
+    fakeSocket.readyState = WS_OPEN;
+    fakeSocket.emit("open");
+
+    bridge.message({
+      type: VELAY_FRAME_TYPES.websocketMessage,
+      connection_id: "conn-123",
+      message_type: VELAY_WEBSOCKET_MESSAGE_TYPES.text,
+      body_base64: "not base64",
+    });
+
+    expect(fakeSocket.closes).toEqual([
+      { code: 4003, reason: "Invalid message" },
+    ]);
+    expect(bridge.getConnectionCount()).toBe(0);
+  });
+
+  test("buffers Velay messages until local open", () => {
+    bridge.open(makeOpenFrame());
+
+    bridge.message({
+      type: VELAY_FRAME_TYPES.websocketMessage,
+      connection_id: "conn-123",
+      message_type: VELAY_WEBSOCKET_MESSAGE_TYPES.text,
+      body_base64: base64("early"),
+    });
+
+    expect(fakeSocket.sent).toEqual([]);
+
+    fakeSocket.readyState = WS_OPEN;
+    fakeSocket.emit("open");
+
+    expect(fakeSocket.sent).toEqual(["early"]);
+  });
+
+  test("sends websocket_open_error without connecting for unsafe paths", () => {
+    bridge.open(
+      makeOpenFrame({
+        path: "https://example.com/webhooks/twilio/relay",
+      }),
+    );
+
+    expect(WebSocketMock).not.toHaveBeenCalled();
+    expect(sentFrames).toEqual([
+      {
+        type: VELAY_FRAME_TYPES.websocketOpenError,
+        connection_id: "conn-123",
+        reason: "Invalid WebSocket path",
+      },
+    ]);
+  });
+
+  test("forwards local close frames back to Velay and cleans up", () => {
+    bridge.open(makeOpenFrame());
+    fakeSocket.readyState = WS_OPEN;
+    fakeSocket.emit("open");
+
+    fakeSocket.emit("close", { code: 1001, reason: "going away" });
+
+    expect(bridge.getConnectionCount()).toBe(0);
+    expect(sentFrames.at(-1)).toEqual({
+      type: VELAY_FRAME_TYPES.websocketClose,
+      connection_id: "conn-123",
+      code: 1001,
+      reason: "going away",
+    });
+  });
+
+  test("remaps closeAll going-away semantics to an application close code", () => {
+    bridge.open(makeOpenFrame());
+    fakeSocket.readyState = WS_OPEN;
+    fakeSocket.emit("open");
+
+    bridge.closeAll();
+
+    expect(fakeSocket.closes).toEqual([
+      { code: 4001, reason: "Tunnel closed" },
+    ]);
+    expect(bridge.getConnectionCount()).toBe(0);
+  });
+});

package/src/velay/websocket-bridge.ts

@@ -0,0 +1,324 @@
+import type { OutgoingHttpHeaders } from "node:http";
+
+import {
+  binaryLikeToBytes,
+  buildLoopbackWebSocketUrl,
+  closeWebSocket,
+  decodeBase64Bytes,
+  encodeBase64,
+  websocketHeadersFromVelay,
+} from "./bridge-utils.js";
+import {
+  VELAY_FRAME_TYPES,
+  VELAY_WEBSOCKET_MESSAGE_TYPES,
+  type VelayFrame,
+  type VelayWebSocketCloseFrame,
+  type VelayWebSocketInboundFrame,
+  type VelayWebSocketMessageFrame,
+  type VelayWebSocketOpenErrorFrame,
+  type VelayWebSocketOpenedFrame,
+  type VelayWebSocketOpenFrame,
+} from "./protocol.js";
+
+const MAX_PENDING_MESSAGES = 100;
+const textFrameDecoder = new TextDecoder("utf-8", { ignoreBOM: true });
+
+type PendingMessage = string | Uint8Array;
+
+type WebSocketConstructorWithHeaders = {
+  new (
+    url: string,
+    options?: {
+      headers?: OutgoingHttpHeaders;
+      protocols?: string | string[];
+    },
+  ): WebSocket;
+};
+
+type BridgeConnection = {
+  ws: WebSocket;
+  opened: boolean;
+  openErrorSent: boolean;
+  pendingMessages: PendingMessage[];
+  suppressNextCloseFrame: boolean;
+};
+
+type SendVelayFrame = (frame: VelayFrame) => void;
+
+export class VelayWebSocketBridge {
+  private readonly connections = new Map<string, BridgeConnection>();
+
+  constructor(
+    private readonly gatewayLoopbackBaseUrl: string,
+    private readonly sendFrame: SendVelayFrame,
+  ) {}
+
+  handleFrame(frame: VelayWebSocketInboundFrame): void {
+    switch (frame.type) {
+      case VELAY_FRAME_TYPES.websocketOpen:
+        this.open(frame);
+        return;
+      case VELAY_FRAME_TYPES.websocketMessage:
+        this.message(frame);
+        return;
+      case VELAY_FRAME_TYPES.websocketClose:
+        this.close(frame);
+        return;
+    }
+  }
+
+  open(frame: VelayWebSocketOpenFrame): void {
+    this.closeExisting(frame.connection_id);
+
+    const url = buildLoopbackWebSocketUrl(
+      this.gatewayLoopbackBaseUrl,
+      frame.path,
+      frame.raw_query,
+    );
+    if (!url) {
+      this.sendOpenError(frame.connection_id, "Invalid WebSocket path");
+      return;
+    }
+
+    let ws: WebSocket;
+    try {
+      const WebSocketWithHeaders =
+        WebSocket as unknown as WebSocketConstructorWithHeaders;
+      ws = new WebSocketWithHeaders(url, {
+        headers: websocketHeadersFromVelay(frame.headers),
+        ...(frame.subprotocol ? { protocols: [frame.subprotocol] } : {}),
+      });
+    } catch {
+      this.sendOpenError(frame.connection_id, "WebSocket connection failed");
+      return;
+    }
+
+    const connection: BridgeConnection = {
+      ws,
+      opened: false,
+      openErrorSent: false,
+      pendingMessages: [],
+      suppressNextCloseFrame: false,
+    };
+    this.connections.set(frame.connection_id, connection);
+
+    ws.binaryType = "arraybuffer";
+    ws.addEventListener("open", () => {
+      if (this.connections.get(frame.connection_id) !== connection) return;
+
+      connection.opened = true;
+      this.sendFrame({
+        type: VELAY_FRAME_TYPES.websocketOpened,
+        connection_id: frame.connection_id,
+      } satisfies VelayWebSocketOpenedFrame);
+
+      for (const message of connection.pendingMessages) {
+        ws.send(message);
+      }
+      connection.pendingMessages = [];
+    });
+
+    ws.addEventListener("message", (event) => {
+      void this.forwardLocalMessage(
+        frame.connection_id,
+        connection,
+        event.data,
+      );
+    });
+
+    ws.addEventListener("close", (event) => {
+      this.handleLocalClose(frame.connection_id, connection, event);
+    });
+
+    ws.addEventListener("error", () => {
+      if (connection.opened) return;
+      this.failOpeningConnection(
+        frame.connection_id,
+        connection,
+        "WebSocket connection failed",
+      );
+    });
+  }
+
+  message(frame: VelayWebSocketMessageFrame): void {
+    const connection = this.connections.get(frame.connection_id);
+    if (!connection) return;
+
+    const message = decodeVelayMessage(frame);
+    if (message === undefined) {
+      this.closeConnection(
+        frame.connection_id,
+        connection,
+        1003,
+        "Invalid message",
+      );
+      return;
+    }
+
+    if (connection.opened && connection.ws.readyState === WebSocket.OPEN) {
+      connection.ws.send(message);
+      return;
+    }
+
+    if (connection.pendingMessages.length >= MAX_PENDING_MESSAGES) {
+      this.closeConnection(
+        frame.connection_id,
+        connection,
+        1008,
+        "Buffer overflow",
+      );
+      return;
+    }
+    connection.pendingMessages.push(message);
+  }
+
+  close(frame: VelayWebSocketCloseFrame): void {
+    const connection = this.connections.get(frame.connection_id);
+    if (!connection) return;
+
+    connection.suppressNextCloseFrame = true;
+    this.closeConnection(
+      frame.connection_id,
+      connection,
+      frame.code,
+      frame.reason,
+    );
+  }
+
+  closeAll(code = 1001, reason = "Tunnel closed"): void {
+    for (const [connectionId, connection] of this.connections) {
+      connection.suppressNextCloseFrame = true;
+      this.closeConnection(connectionId, connection, code, reason);
+    }
+  }
+
+  getConnectionCount(): number {
+    return this.connections.size;
+  }
+
+  private async forwardLocalMessage(
+    connectionId: string,
+    connection: BridgeConnection,
+    data: unknown,
+  ): Promise<void> {
+    if (this.connections.get(connectionId) !== connection) return;
+
+    const message = await encodeLocalMessage(connectionId, data);
+    if (this.connections.get(connectionId) !== connection) return;
+    this.sendFrame(message);
+  }
+
+  private handleLocalClose(
+    connectionId: string,
+    connection: BridgeConnection,
+    event: CloseEvent,
+  ): void {
+    if (this.connections.get(connectionId) !== connection) return;
+    this.connections.delete(connectionId);
+    connection.pendingMessages = [];
+
+    if (!connection.opened) {
+      this.sendOpenErrorOnce(
+        connectionId,
+        connection,
+        "WebSocket connection failed",
+      );
+      return;
+    }
+
+    if (connection.suppressNextCloseFrame) return;
+    this.sendFrame({
+      type: VELAY_FRAME_TYPES.websocketClose,
+      connection_id: connectionId,
+      code: event.code,
+      reason: event.reason,
+    } satisfies VelayWebSocketCloseFrame);
+  }
+
+  private failOpeningConnection(
+    connectionId: string,
+    connection: BridgeConnection,
+    reason: string,
+  ): void {
+    if (this.connections.get(connectionId) === connection) {
+      this.connections.delete(connectionId);
+    }
+    connection.pendingMessages = [];
+    this.sendOpenErrorOnce(connectionId, connection, reason);
+    closeWebSocket(connection.ws);
+  }
+
+  private closeExisting(connectionId: string): void {
+    const existing = this.connections.get(connectionId);
+    if (!existing) return;
+    existing.suppressNextCloseFrame = true;
+    this.closeConnection(connectionId, existing, 1000, "Replaced");
+  }
+
+  private closeConnection(
+    connectionId: string,
+    connection: BridgeConnection,
+    code?: number,
+    reason?: string,
+  ): void {
+    if (this.connections.get(connectionId) === connection) {
+      this.connections.delete(connectionId);
+    }
+    connection.pendingMessages = [];
+    closeWebSocket(connection.ws, code, reason);
+  }
+
+  private sendOpenError(connectionId: string, reason: string): void {
+    this.sendFrame({
+      type: VELAY_FRAME_TYPES.websocketOpenError,
+      connection_id: connectionId,
+      reason,
+    } satisfies VelayWebSocketOpenErrorFrame);
+  }
+
+  private sendOpenErrorOnce(
+    connectionId: string,
+    connection: BridgeConnection,
+    reason: string,
+  ): void {
+    if (connection.openErrorSent) return;
+    connection.openErrorSent = true;
+    this.sendOpenError(connectionId, reason);
+  }
+}
+
+function decodeVelayMessage(
+  frame: VelayWebSocketMessageFrame,
+): PendingMessage | undefined {
+  const bytes = decodeBase64Bytes(frame.body_base64 ?? "");
+  if (bytes === undefined) return undefined;
+
+  if (frame.message_type === VELAY_WEBSOCKET_MESSAGE_TYPES.text) {
+    return textFrameDecoder.decode(bytes);
+  }
+  if (frame.message_type === VELAY_WEBSOCKET_MESSAGE_TYPES.binary) {
+    return bytes;
+  }
+  return undefined;
+}
+
+async function encodeLocalMessage(
+  connectionId: string,
+  data: unknown,
+): Promise<VelayWebSocketMessageFrame> {
+  if (typeof data === "string") {
+    return {
+      type: VELAY_FRAME_TYPES.websocketMessage,
+      connection_id: connectionId,
+      message_type: VELAY_WEBSOCKET_MESSAGE_TYPES.text,
+      body_base64: encodeBase64(data),
+    };
+  }
+
+  return {
+    type: VELAY_FRAME_TYPES.websocketMessage,
+    connection_id: connectionId,
+    message_type: VELAY_WEBSOCKET_MESSAGE_TYPES.binary,
+    body_base64: encodeBase64(await binaryLikeToBytes(data)),
+  };
+}