clawdbot/src/gateway/server.auth.test.ts

import { describe, expect, test, vi } from "vitest";
import { WebSocket } from "ws";
import { PROTOCOL_VERSION } from "./protocol/index.js";
import { getHandshakeTimeoutMs } from "./server-constants.js";
import {
  connectReq,
  getFreePort,
  installGatewayTestHooks,
  onceMessage,
  startGatewayServer,
  startServerWithClient,
  testState,
} from "./test-helpers.js";
import { GATEWAY_CLIENT_MODES, GATEWAY_CLIENT_NAMES } from "../utils/message-channel.js";

installGatewayTestHooks({ scope: "suite" });

async function waitForWsClose(ws: WebSocket, timeoutMs: number): Promise<boolean> {
  if (ws.readyState === WebSocket.CLOSED) return true;
  return await new Promise((resolve) => {
    const timer = setTimeout(() => resolve(ws.readyState === WebSocket.CLOSED), timeoutMs);
    ws.once("close", () => {
      clearTimeout(timer);
      resolve(true);
    });
  });
}

describe("gateway server auth/connect", () => {
  test("closes silent handshakes after timeout", { timeout: 60_000 }, async () => {
    vi.useRealTimers();
    const prevHandshakeTimeout = process.env.CLAWDBOT_TEST_HANDSHAKE_TIMEOUT_MS;
    process.env.CLAWDBOT_TEST_HANDSHAKE_TIMEOUT_MS = "50";
    try {
      const { server, ws } = await startServerWithClient();
      const handshakeTimeoutMs = getHandshakeTimeoutMs();
      const closed = await waitForWsClose(ws, handshakeTimeoutMs + 250);
      expect(closed).toBe(true);
      await server.close();
    } finally {
      if (prevHandshakeTimeout === undefined) {
        delete process.env.CLAWDBOT_TEST_HANDSHAKE_TIMEOUT_MS;
      } else {
        process.env.CLAWDBOT_TEST_HANDSHAKE_TIMEOUT_MS = prevHandshakeTimeout;
      }
    }
  });

  test("connect (req) handshake returns hello-ok payload", async () => {
    const { CONFIG_PATH_CLAWDBOT, STATE_DIR_CLAWDBOT } = await import("../config/config.js");
    const port = await getFreePort();
    const server = await startGatewayServer(port);
    const ws = new WebSocket(`ws://127.0.0.1:${port}`);
    await new Promise<void>((resolve) => ws.once("open", resolve));

    const res = await connectReq(ws);
    expect(res.ok).toBe(true);
    const payload = res.payload as
      | {
          type?: unknown;
          snapshot?: { configPath?: string; stateDir?: string };
        }
      | undefined;
    expect(payload?.type).toBe("hello-ok");
    expect(payload?.snapshot?.configPath).toBe(CONFIG_PATH_CLAWDBOT);
    expect(payload?.snapshot?.stateDir).toBe(STATE_DIR_CLAWDBOT);

    ws.close();
    await server.close();
  });

  test("sends connect challenge on open", async () => {
    const port = await getFreePort();
    const server = await startGatewayServer(port);
    const ws = new WebSocket(`ws://127.0.0.1:${port}`);
    const evtPromise = onceMessage<{ payload?: unknown }>(
      ws,
      (o) => o.type === "event" && o.event === "connect.challenge",
    );
    await new Promise<void>((resolve) => ws.once("open", resolve));
    const evt = await evtPromise;
    const nonce = (evt.payload as { nonce?: unknown } | undefined)?.nonce;
    expect(typeof nonce).toBe("string");
    ws.close();
    await server.close();
  });

  test("rejects protocol mismatch", async () => {
    const { server, ws } = await startServerWithClient();
    try {
      const res = await connectReq(ws, {
        minProtocol: PROTOCOL_VERSION + 1,
        maxProtocol: PROTOCOL_VERSION + 2,
      });
      expect(res.ok).toBe(false);
    } catch {
      // If the server closed before we saw the frame, that's acceptable.
    }
    ws.close();
    await server.close();
  });

  test("rejects invalid token", async () => {
    const { server, ws, prevToken } = await startServerWithClient("secret");
    const res = await connectReq(ws, { token: "wrong" });
    expect(res.ok).toBe(false);
    expect(res.error?.message ?? "").toContain("unauthorized");
    ws.close();
    await server.close();
    if (prevToken === undefined) {
      delete process.env.CLAWDBOT_GATEWAY_TOKEN;
    } else {
      process.env.CLAWDBOT_GATEWAY_TOKEN = prevToken;
    }
  });

  test("accepts password auth when configured", async () => {
    testState.gatewayAuth = { mode: "password", password: "secret" };
    const port = await getFreePort();
    const server = await startGatewayServer(port);
    const ws = new WebSocket(`ws://127.0.0.1:${port}`);
    await new Promise<void>((resolve) => ws.once("open", resolve));

    const res = await connectReq(ws, { password: "secret" });
    expect(res.ok).toBe(true);

    ws.close();
    await server.close();
  });

  test("rejects control ui without device identity by default", async () => {
    const { server, ws, prevToken } = await startServerWithClient("secret");
    const res = await connectReq(ws, {
      token: "secret",
      device: null,
      client: {
        id: GATEWAY_CLIENT_NAMES.CONTROL_UI,
        version: "1.0.0",
        platform: "web",
        mode: GATEWAY_CLIENT_MODES.WEBCHAT,
      },
    });
    expect(res.ok).toBe(false);
    expect(res.error?.message ?? "").toContain("secure context");
    ws.close();
    await server.close();
    if (prevToken === undefined) {
      delete process.env.CLAWDBOT_GATEWAY_TOKEN;
    } else {
      process.env.CLAWDBOT_GATEWAY_TOKEN = prevToken;
    }
  });

  test("allows control ui without device identity when insecure auth is enabled", async () => {
    testState.gatewayControlUi = { allowInsecureAuth: true };
    const { server, ws, prevToken } = await startServerWithClient("secret");
    const res = await connectReq(ws, {
      token: "secret",
      device: null,
      client: {
        id: GATEWAY_CLIENT_NAMES.CONTROL_UI,
        version: "1.0.0",
        platform: "web",
        mode: GATEWAY_CLIENT_MODES.WEBCHAT,
      },
    });
    expect(res.ok).toBe(true);
    ws.close();
    await server.close();
    if (prevToken === undefined) {
      delete process.env.CLAWDBOT_GATEWAY_TOKEN;
    } else {
      process.env.CLAWDBOT_GATEWAY_TOKEN = prevToken;
    }
  });

  test("accepts device token auth for paired device", async () => {
    const { loadOrCreateDeviceIdentity } = await import("../infra/device-identity.js");
    const { approveDevicePairing, getPairedDevice, listDevicePairing } =
      await import("../infra/device-pairing.js");
    const { server, ws, port, prevToken } = await startServerWithClient("secret");
    const res = await connectReq(ws, { token: "secret" });
    if (!res.ok) {
      const list = await listDevicePairing();
      const pending = list.pending.at(0);
      expect(pending?.requestId).toBeDefined();
      if (pending?.requestId) {
        await approveDevicePairing(pending.requestId);
      }
    }

    const identity = loadOrCreateDeviceIdentity();
    const paired = await getPairedDevice(identity.deviceId);
    const deviceToken = paired?.tokens?.operator?.token;
    expect(deviceToken).toBeDefined();

    ws.close();

    const ws2 = new WebSocket(`ws://127.0.0.1:${port}`);
    await new Promise<void>((resolve) => ws2.once("open", resolve));
    const res2 = await connectReq(ws2, { token: deviceToken });
    expect(res2.ok).toBe(true);

    ws2.close();
    await server.close();
    if (prevToken === undefined) {
      delete process.env.CLAWDBOT_GATEWAY_TOKEN;
    } else {
      process.env.CLAWDBOT_GATEWAY_TOKEN = prevToken;
    }
  });

  test("requires pairing for scope upgrades", async () => {
    const { mkdtemp } = await import("node:fs/promises");
    const { tmpdir } = await import("node:os");
    const { join } = await import("node:path");
    const { buildDeviceAuthPayload } = await import("./device-auth.js");
    const { loadOrCreateDeviceIdentity, publicKeyRawBase64UrlFromPem, signDevicePayload } =
      await import("../infra/device-identity.js");
    const { approveDevicePairing, getPairedDevice, listDevicePairing } =
      await import("../infra/device-pairing.js");
    const { GATEWAY_CLIENT_MODES, GATEWAY_CLIENT_NAMES } =
      await import("../utils/message-channel.js");
    const { server, ws, port, prevToken } = await startServerWithClient("secret");
    const identityDir = await mkdtemp(join(tmpdir(), "clawdbot-device-scope-"));
    const identity = loadOrCreateDeviceIdentity(join(identityDir, "device.json"));
    const client = {
      id: GATEWAY_CLIENT_NAMES.TEST,
      version: "1.0.0",
      platform: "test",
      mode: GATEWAY_CLIENT_MODES.TEST,
    };
    const buildDevice = (scopes: string[]) => {
      const signedAtMs = Date.now();
      const payload = buildDeviceAuthPayload({
        deviceId: identity.deviceId,
        clientId: client.id,
        clientMode: client.mode,
        role: "operator",
        scopes,
        signedAtMs,
        token: "secret",
      });
      return {
        id: identity.deviceId,
        publicKey: publicKeyRawBase64UrlFromPem(identity.publicKeyPem),
        signature: signDevicePayload(identity.privateKeyPem, payload),
        signedAt: signedAtMs,
      };
    };
    const initial = await connectReq(ws, {
      token: "secret",
      scopes: ["operator.read"],
      client,
      device: buildDevice(["operator.read"]),
    });
    if (!initial.ok) {
      const list = await listDevicePairing();
      const pending = list.pending.at(0);
      expect(pending?.requestId).toBeDefined();
      if (pending?.requestId) {
        await approveDevicePairing(pending.requestId);
      }
    }

    let paired = await getPairedDevice(identity.deviceId);
    expect(paired?.scopes).toContain("operator.read");

    ws.close();

    const ws2 = new WebSocket(`ws://127.0.0.1:${port}`);
    await new Promise<void>((resolve) => ws2.once("open", resolve));
    const res = await connectReq(ws2, {
      token: "secret",
      scopes: ["operator.admin"],
      client,
      device: buildDevice(["operator.admin"]),
    });
    expect(res.ok).toBe(true);
    paired = await getPairedDevice(identity.deviceId);
    expect(paired?.scopes).toContain("operator.admin");

    ws2.close();
    await server.close();
    if (prevToken === undefined) {
      delete process.env.CLAWDBOT_GATEWAY_TOKEN;
    } else {
      process.env.CLAWDBOT_GATEWAY_TOKEN = prevToken;
    }
  });

  test("rejects revoked device token", async () => {
    const { loadOrCreateDeviceIdentity } = await import("../infra/device-identity.js");
    const { approveDevicePairing, getPairedDevice, listDevicePairing, revokeDeviceToken } =
      await import("../infra/device-pairing.js");
    const { server, ws, port, prevToken } = await startServerWithClient("secret");
    const res = await connectReq(ws, { token: "secret" });
    if (!res.ok) {
      const list = await listDevicePairing();
      const pending = list.pending.at(0);
      expect(pending?.requestId).toBeDefined();
      if (pending?.requestId) {
        await approveDevicePairing(pending.requestId);
      }
    }

    const identity = loadOrCreateDeviceIdentity();
    const paired = await getPairedDevice(identity.deviceId);
    const deviceToken = paired?.tokens?.operator?.token;
    expect(deviceToken).toBeDefined();

    await revokeDeviceToken({ deviceId: identity.deviceId, role: "operator" });

    ws.close();

    const ws2 = new WebSocket(`ws://127.0.0.1:${port}`);
    await new Promise<void>((resolve) => ws2.once("open", resolve));
    const res2 = await connectReq(ws2, { token: deviceToken });
    expect(res2.ok).toBe(false);

    ws2.close();
    await server.close();
    if (prevToken === undefined) {
      delete process.env.CLAWDBOT_GATEWAY_TOKEN;
    } else {
      process.env.CLAWDBOT_GATEWAY_TOKEN = prevToken;
    }
  });

  test("rejects invalid password", async () => {
    testState.gatewayAuth = { mode: "password", password: "secret" };
    const port = await getFreePort();
    const server = await startGatewayServer(port);
    const ws = new WebSocket(`ws://127.0.0.1:${port}`);
    await new Promise<void>((resolve) => ws.once("open", resolve));

    const res = await connectReq(ws, { password: "wrong" });
    expect(res.ok).toBe(false);
    expect(res.error?.message ?? "").toContain("unauthorized");

    ws.close();
    await server.close();
  });

  test("rejects non-connect first request", async () => {
    const { server, ws } = await startServerWithClient();
    ws.send(JSON.stringify({ type: "req", id: "h1", method: "health" }));
    const res = await onceMessage<{ ok: boolean; error?: unknown }>(
      ws,
      (o) => o.type === "res" && o.id === "h1",
    );
    expect(res.ok).toBe(false);
    await new Promise<void>((resolve) => ws.once("close", () => resolve()));
    await server.close();
  });

  test(
    "invalid connect params surface in response and close reason",
    { timeout: 60_000 },
    async () => {
      const { server, ws } = await startServerWithClient();
      const closeInfoPromise = new Promise<{ code: number; reason: string }>((resolve) => {
        ws.once("close", (code, reason) => resolve({ code, reason: reason.toString() }));
      });

      ws.send(
        JSON.stringify({
          type: "req",
          id: "h-bad",
          method: "connect",
          params: {
            minProtocol: PROTOCOL_VERSION,
            maxProtocol: PROTOCOL_VERSION,
            client: {
              id: "bad-client",
              version: "dev",
              platform: "web",
              mode: "webchat",
            },
            device: {
              id: 123,
              publicKey: "bad",
              signature: "bad",
              signedAt: "bad",
            },
          },
        }),
      );

      const res = await onceMessage<{
        ok: boolean;
        error?: { message?: string };
      }>(
        ws,
        (o) => (o as { type?: string }).type === "res" && (o as { id?: string }).id === "h-bad",
      );
      expect(res.ok).toBe(false);
      expect(String(res.error?.message ?? "")).toContain("invalid connect params");

      const closeInfo = await closeInfoPromise;
      expect(closeInfo.code).toBe(1008);
      expect(closeInfo.reason).toContain("invalid connect params");

      await server.close();
    },
  );
});