CookieBridge/tests/integration.test.ts

import { describe, it, expect, beforeAll, afterAll } from "vitest";
import WebSocket from "ws";
import {
  RelayServer,
  generateKeyPair,
  deviceIdFromKeys,
  sign,
  MESSAGE_TYPES,
} from "../src/index.js";
import { buildEnvelope, openEnvelope } from "../src/sync/envelope.js";
import type { CookieSyncPayload, Envelope } from "../src/protocol/spec.js";

const BASE = (port: number) => `http://127.0.0.1:${port}`;

// Helper: register a device and return its token
async function registerDevice(
  port: number,
  keys: ReturnType<typeof generateKeyPair>,
  name: string,
): Promise<{ token: string; deviceId: string }> {
  const deviceId = deviceIdFromKeys(keys);
  const res = await fetch(`${BASE(port)}/api/devices/register`, {
    method: "POST",
    headers: { "Content-Type": "application/json" },
    body: JSON.stringify({
      deviceId,
      name,
      platform: "test",
      encPub: keys.encPub.toString("hex"),
    }),
  });
  const body = (await res.json()) as { token: string; deviceId: string };
  return { token: body.token, deviceId };
}

// Helper: connect via WebSocket with challenge-response auth
function connectDeviceWs(
  port: number,
  keys: ReturnType<typeof generateKeyPair>,
): Promise<{ ws: WebSocket; deviceId: string }> {
  return new Promise((resolve, reject) => {
    const deviceId = deviceIdFromKeys(keys);
    const ws = new WebSocket(`ws://127.0.0.1:${port}/ws`);
    ws.on("error", reject);
    ws.on("message", (data: Buffer) => {
      const msg = JSON.parse(data.toString());
      if (msg.type === "auth_challenge") {
        const challenge = Buffer.from(msg.challenge, "hex");
        const sig = sign(challenge, keys.signSec);
        ws.send(JSON.stringify({ type: "auth_response", deviceId, sig: sig.toString("hex") }));
      } else if (msg.type === "auth_ok") {
        resolve({ ws, deviceId });
      } else if (msg.type === "error") {
        reject(new Error(msg.error));
      }
    });
  });
}

// Helper: connect via WebSocket with token auth
function connectDeviceWsToken(
  port: number,
  token: string,
): Promise<{ ws: WebSocket; deviceId: string }> {
  return new Promise((resolve, reject) => {
    const ws = new WebSocket(`ws://127.0.0.1:${port}/ws`);
    ws.on("error", reject);
    ws.on("message", (data: Buffer) => {
      const msg = JSON.parse(data.toString());
      if (msg.type === "auth_challenge") {
        // Use token auth instead
        ws.send(JSON.stringify({ type: "auth_token", token }));
      } else if (msg.type === "auth_ok") {
        resolve({ ws, deviceId: msg.deviceId as string });
      }
    });
  });
}

function waitForMessage(ws: WebSocket, type: string, timeoutMs = 5000): Promise<Record<string, unknown>> {
  return new Promise((resolve, reject) => {
    const timer = setTimeout(() => reject(new Error(`Timeout waiting for ${type}`)), timeoutMs);
    const handler = (data: Buffer) => {
      const msg = JSON.parse(data.toString());
      if (msg.type === type) {
        clearTimeout(timer);
        ws.off("message", handler);
        resolve(msg);
      }
    };
    ws.on("message", handler);
  });
}

describe("Integration: CookieBridge server v2", () => {
  let server: RelayServer;
  let port: number;

  beforeAll(async () => {
    server = new RelayServer({ port: 0 });
    await server.start();
    port = server.port;
  });

  afterAll(async () => {
    await server.stop();
  });

  it("health check", async () => {
    const res = await fetch(`${BASE(port)}/health`);
    const body = await res.json();
    expect(body.status).toBe("ok");
  });

  // --- Device Registration ---

  it("registers a device and returns a token", async () => {
    const keys = generateKeyPair();
    const { token, deviceId } = await registerDevice(port, keys, "My Laptop");
    expect(token).toMatch(/^cb_/);
    expect(deviceId).toBe(deviceIdFromKeys(keys));
  });

  it("returns existing device on re-register", async () => {
    const keys = generateKeyPair();
    const first = await registerDevice(port, keys, "Phone");
    const second = await registerDevice(port, keys, "Phone");
    expect(first.token).toBe(second.token);
  });

  // --- Pairing ---

  it("pairing flow: create and accept", async () => {
    const alice = generateKeyPair();
    const bob = generateKeyPair();

    const createRes = await fetch(`${BASE(port)}/api/pair`, {
      method: "POST",
      headers: { "Content-Type": "application/json" },
      body: JSON.stringify({
        deviceId: deviceIdFromKeys(alice),
        x25519PubKey: alice.encPub.toString("hex"),
      }),
    });
    expect(createRes.status).toBe(201);
    const { pairingCode } = (await createRes.json()) as { pairingCode: string };

    const acceptRes = await fetch(`${BASE(port)}/api/pair/accept`, {
      method: "POST",
      headers: { "Content-Type": "application/json" },
      body: JSON.stringify({
        deviceId: deviceIdFromKeys(bob),
        x25519PubKey: bob.encPub.toString("hex"),
        pairingCode,
      }),
    });
    expect(acceptRes.status).toBe(200);
    const result = (await acceptRes.json()) as {
      initiator: { deviceId: string };
      acceptor: { deviceId: string };
    };
    expect(result.initiator.deviceId).toBe(deviceIdFromKeys(alice));
    expect(result.acceptor.deviceId).toBe(deviceIdFromKeys(bob));
  });

  // --- Cookie Storage (HTTP) ---

  it("pushes and pulls encrypted cookies via HTTP", async () => {
    const keys = generateKeyPair();
    const { token } = await registerDevice(port, keys, "Test Device");

    // Push
    const pushRes = await fetch(`${BASE(port)}/api/cookies`, {
      method: "POST",
      headers: {
        "Content-Type": "application/json",
        Authorization: `Bearer ${token}`,
      },
      body: JSON.stringify({
        cookies: [
          {
            domain: "example.com",
            cookieName: "session",
            path: "/",
            nonce: "aabbcc",
            ciphertext: "ZW5jcnlwdGVk", // base64 of "encrypted"
            lamportTs: 1,
          },
        ],
      }),
    });
    expect(pushRes.status).toBe(200);

    // Pull
    const pullRes = await fetch(`${BASE(port)}/api/cookies?domain=example.com`, {
      headers: { Authorization: `Bearer ${token}` },
    });
    expect(pullRes.status).toBe(200);
    const pullBody = (await pullRes.json()) as { cookies: Array<{ domain: string; cookieName: string }> };
    expect(pullBody.cookies).toHaveLength(1);
    expect(pullBody.cookies[0].domain).toBe("example.com");
    expect(pullBody.cookies[0].cookieName).toBe("session");
  });

  it("paired devices can see each other's cookies", async () => {
    const alice = generateKeyPair();
    const bob = generateKeyPair();
    const aliceReg = await registerDevice(port, alice, "Alice Laptop");
    const bobReg = await registerDevice(port, bob, "Bob Phone");

    // Pair them
    const createRes = await fetch(`${BASE(port)}/api/pair`, {
      method: "POST",
      headers: { "Content-Type": "application/json" },
      body: JSON.stringify({
        deviceId: aliceReg.deviceId,
        x25519PubKey: alice.encPub.toString("hex"),
      }),
    });
    const { pairingCode } = (await createRes.json()) as { pairingCode: string };
    await fetch(`${BASE(port)}/api/pair/accept`, {
      method: "POST",
      headers: { "Content-Type": "application/json" },
      body: JSON.stringify({
        deviceId: bobReg.deviceId,
        x25519PubKey: bob.encPub.toString("hex"),
        pairingCode,
      }),
    });

    // Alice pushes a cookie
    await fetch(`${BASE(port)}/api/cookies`, {
      method: "POST",
      headers: {
        "Content-Type": "application/json",
        Authorization: `Bearer ${aliceReg.token}`,
      },
      body: JSON.stringify({
        cookies: [
          { domain: "shared.com", cookieName: "tok", path: "/", nonce: "112233", ciphertext: "c2hhcmVk", lamportTs: 1 },
        ],
      }),
    });

    // Bob can pull it
    const pullRes = await fetch(`${BASE(port)}/api/cookies?domain=shared.com`, {
      headers: { Authorization: `Bearer ${bobReg.token}` },
    });
    const body = (await pullRes.json()) as { cookies: Array<{ cookieName: string }> };
    expect(body.cookies).toHaveLength(1);
    expect(body.cookies[0].cookieName).toBe("tok");
  });

  it("polls for cookie updates since timestamp", async () => {
    const keys = generateKeyPair();
    const { token } = await registerDevice(port, keys, "Poller");
    const since = new Date(Date.now() - 1000).toISOString(); // 1s in the past

    // Push after timestamp
    await fetch(`${BASE(port)}/api/cookies`, {
      method: "POST",
      headers: {
        "Content-Type": "application/json",
        Authorization: `Bearer ${token}`,
      },
      body: JSON.stringify({
        cookies: [
          { domain: "poll.com", cookieName: "s", path: "/", nonce: "ff", ciphertext: "cA==", lamportTs: 1 },
        ],
      }),
    });

    const pollRes = await fetch(`${BASE(port)}/api/cookies/updates?since=${since}`, {
      headers: { Authorization: `Bearer ${token}` },
    });
    expect(pollRes.status).toBe(200);
    const body = (await pollRes.json()) as { cookies: unknown[]; serverTime: string };
    expect(body.cookies.length).toBeGreaterThanOrEqual(1);
    expect(body.serverTime).toBeTruthy();
  });

  // --- Agent Skill API ---

  it("creates agent token and retrieves cookies", async () => {
    const keys = generateKeyPair();
    const agentKeys = generateKeyPair();
    const { token: deviceToken } = await registerDevice(port, keys, "Agent Host");

    // Push a cookie first
    await fetch(`${BASE(port)}/api/cookies`, {
      method: "POST",
      headers: {
        "Content-Type": "application/json",
        Authorization: `Bearer ${deviceToken}`,
      },
      body: JSON.stringify({
        cookies: [
          { domain: "agent-test.com", cookieName: "auth", path: "/", nonce: "aa", ciphertext: "dGVzdA==", lamportTs: 1 },
        ],
      }),
    });

    // Create agent token
    const agentRes = await fetch(`${BASE(port)}/api/agent/tokens`, {
      method: "POST",
      headers: {
        "Content-Type": "application/json",
        Authorization: `Bearer ${deviceToken}`,
      },
      body: JSON.stringify({
        name: "My AI Assistant",
        encPub: agentKeys.encPub.toString("hex"),
        allowedDomains: ["agent-test.com"],
      }),
    });
    expect(agentRes.status).toBe(201);
    const agentBody = (await agentRes.json()) as { id: string; token: string };
    expect(agentBody.token).toMatch(/^cb_/);

    // Agent retrieves cookies
    const cookieRes = await fetch(`${BASE(port)}/api/agent/cookies?domain=agent-test.com`, {
      headers: { Authorization: `Bearer ${agentBody.token}` },
    });
    expect(cookieRes.status).toBe(200);
    const cookieBody = (await cookieRes.json()) as { cookies: Array<{ domain: string }> };
    expect(cookieBody.cookies).toHaveLength(1);
    expect(cookieBody.cookies[0].domain).toBe("agent-test.com");
  });

  it("agent domain restriction works", async () => {
    const keys = generateKeyPair();
    const agentKeys = generateKeyPair();
    const { token: deviceToken } = await registerDevice(port, keys, "Restricted Host");

    const agentRes = await fetch(`${BASE(port)}/api/agent/tokens`, {
      method: "POST",
      headers: {
        "Content-Type": "application/json",
        Authorization: `Bearer ${deviceToken}`,
      },
      body: JSON.stringify({
        name: "Restricted Agent",
        encPub: agentKeys.encPub.toString("hex"),
        allowedDomains: ["allowed.com"],
      }),
    });
    const agentBody = (await agentRes.json()) as { token: string };

    const res = await fetch(`${BASE(port)}/api/agent/cookies?domain=forbidden.com`, {
      headers: { Authorization: `Bearer ${agentBody.token}` },
    });
    expect(res.status).toBe(403);
  });

  // --- WebSocket ---

  it("authenticates via challenge-response", async () => {
    const keys = generateKeyPair();
    const { ws, deviceId } = await connectDeviceWs(port, keys);
    expect(deviceId).toBe(deviceIdFromKeys(keys));
    ws.close();
  });

  it("authenticates via token", async () => {
    const keys = generateKeyPair();
    const { token } = await registerDevice(port, keys, "WS Token");
    const { ws, deviceId } = await connectDeviceWsToken(port, token);
    expect(deviceId).toBe(deviceIdFromKeys(keys));
    ws.close();
  });

  it("rejects bad auth", async () => {
    const alice = generateKeyPair();
    const eve = generateKeyPair();

    await expect(
      new Promise<void>((resolve, reject) => {
        const ws = new WebSocket(`ws://127.0.0.1:${port}/ws`);
        ws.on("message", (data: Buffer) => {
          const msg = JSON.parse(data.toString());
          if (msg.type === "auth_challenge") {
            const challenge = Buffer.from(msg.challenge, "hex");
            const sig = sign(challenge, eve.signSec);
            ws.send(JSON.stringify({ type: "auth_response", deviceId: deviceIdFromKeys(alice), sig: sig.toString("hex") }));
          }
        });
        ws.on("close", (code: number) => {
          if (code === 4003) reject(new Error("Auth failed as expected"));
          else resolve();
        });
      }),
    ).rejects.toThrow("Auth failed");
  });

  it("relays encrypted cookie sync between two devices over WebSocket", async () => {
    const alice = generateKeyPair();
    const bob = generateKeyPair();

    const aliceConn = await connectDeviceWs(port, alice);
    const bobConn = await connectDeviceWs(port, bob);
    const bobDeviceId = deviceIdFromKeys(bob);

    const payload: CookieSyncPayload = {
      action: "set",
      cookies: [
        { domain: "ws-test.com", name: "session", value: "tok_ws", path: "/", secure: true, httpOnly: true, sameSite: "lax", expiresAt: null },
      ],
      lamportTs: 1,
    };

    const envelope = buildEnvelope(MESSAGE_TYPES.COOKIE_SYNC, payload, alice, bob.encPub, bobDeviceId);
    const bobMsg = waitForMessage(bobConn.ws, MESSAGE_TYPES.COOKIE_SYNC);

    aliceConn.ws.send(JSON.stringify(envelope));

    const ack = await waitForMessage(aliceConn.ws, MESSAGE_TYPES.ACK);
    expect(ack.delivered).toBe(true);

    const received = (await bobMsg) as unknown as Envelope;
    const decrypted = openEnvelope(received, bob, alice.encPub);
    expect(decrypted.cookies[0].value).toBe("tok_ws");

    aliceConn.ws.close();
    bobConn.ws.close();
  });

  it("queues messages for offline devices", async () => {
    const alice = generateKeyPair();
    const bob = generateKeyPair();
    const bobDeviceId = deviceIdFromKeys(bob);

    const aliceConn = await connectDeviceWs(port, alice);

    const payload: CookieSyncPayload = {
      action: "set",
      cookies: [
        { domain: "offline.com", name: "q", value: "queued", path: "/", secure: true, httpOnly: false, sameSite: "none", expiresAt: null },
      ],
      lamportTs: 1,
    };

    const envelope = buildEnvelope(MESSAGE_TYPES.COOKIE_SYNC, payload, alice, bob.encPub, bobDeviceId);
    aliceConn.ws.send(JSON.stringify(envelope));

    const ack = await waitForMessage(aliceConn.ws, MESSAGE_TYPES.ACK);
    expect(ack.delivered).toBe(false);

    // Bob comes online
    const bobMsg = new Promise<Envelope>((resolve) => {
      const ws = new WebSocket(`ws://127.0.0.1:${port}/ws`);
      ws.on("message", (data: Buffer) => {
        const msg = JSON.parse(data.toString());
        if (msg.type === "auth_challenge") {
          const challenge = Buffer.from(msg.challenge, "hex");
          const sig = sign(challenge, bob.signSec);
          ws.send(JSON.stringify({ type: "auth_response", deviceId: bobDeviceId, sig: sig.toString("hex") }));
        } else if (msg.type === MESSAGE_TYPES.COOKIE_SYNC) {
          resolve(msg as unknown as Envelope);
          ws.close();
        }
      });
    });

    const received = await bobMsg;
    const decrypted = openEnvelope(received, bob, alice.encPub);
    expect(decrypted.cookies[0].value).toBe("queued");

    aliceConn.ws.close();
  });
});