feat(vscodex): add remote Codex collaboration module

2026-10-08 18:37:46 +08:00 · 2026-09-01 20:08:36 +08:00
parent 1aab31a148
commit 5059093d29
102 changed files with 38569 additions and 11 deletions
@@ -0,0 +1,298 @@
+"use strict";
+
+const assert = require("node:assert/strict");
+const fs = require("node:fs");
+const os = require("node:os");
+const path = require("node:path");
+const test = require("node:test");
+const { WebSocket } = require("ws");
+
+const { AetherVscodexCloudServer, RoomManager } = require("../cloud/server.js");
+
+const internalToken = "test-internal-token-with-enough-entropy";
+
+function internalFetch(base, pathname, options = {}) {
+  return fetch(`${base}${pathname}`, {
+    ...options,
+    headers: {
+      Authorization: `Bearer ${internalToken}`,
+      ...(options.body ? { "Content-Type": "application/json" } : {}),
+      ...(options.headers || {}),
+    },
+  });
+}
+
+function websocketClient(base, clientType, token, sessionId) {
+  const socket = new WebSocket(`${base.replace(/^http/, "ws")}/api/vscodex/ws`);
+  const messages = [];
+  const waiters = [];
+  const wait = (predicate, timeout = 5_000, label = "websocket frame") => new Promise((resolve, reject) => {
+    const existing = messages.find(predicate);
+    if (existing) return resolve(existing);
+    const timer = setTimeout(() => {
+      const index = waiters.findIndex((entry) => entry.resolve === resolve);
+      if (index >= 0) waiters.splice(index, 1);
+      reject(new Error(`timed out waiting for ${label}; received: ${JSON.stringify(messages.map((message) => ({ type: message.type, kind: message.kind, commandId: message.commandId })))}`));
+    }, timeout);
+    waiters.push({
+      predicate,
+      resolve: (message) => {
+        clearTimeout(timer);
+        resolve(message);
+      },
+    });
+  });
+  socket.on("message", (data) => {
+    const message = JSON.parse(data.toString("utf8"));
+    messages.push(message);
+    for (let index = waiters.length - 1; index >= 0; index -= 1) {
+      if (!waiters[index].predicate(message)) continue;
+      const waiter = waiters.splice(index, 1)[0];
+      waiter.resolve(message);
+    }
+  });
+  return new Promise((resolve, reject) => {
+    socket.once("open", () => {
+      socket.send(JSON.stringify({ v: 1, kind: "hello", clientType, protocol: 1, ...(sessionId ? { sessionId } : {}) }));
+      socket.send(JSON.stringify(clientType === "host"
+        ? { v: 1, kind: "auth", accessToken: token }
+        : { type: "auth", token }));
+      wait((message) => message.type === "auth.ok").then(() => resolve({ socket, wait, messages }), reject);
+    });
+    socket.once("error", reject);
+  });
+}
+
+async function pairDevice(base, userId, name) {
+  const pairingResponse = await internalFetch(base, `/internal/v1/users/${encodeURIComponent(userId)}/pairings`, {
+    method: "POST",
+    body: JSON.stringify({ name }),
+  });
+  assert.equal(pairingResponse.status, 201);
+  const pairing = await pairingResponse.json();
+  const exchangeResponse = await fetch(`${base}/v1/pairings/exchange`, {
+    method: "POST",
+    headers: { "Content-Type": "application/json" },
+    body: JSON.stringify({ code: pairing.code, name }),
+  });
+  assert.equal(exchangeResponse.status, 201);
+  return exchangeResponse.json();
+}
+
+async function browserTicket(base, userId, deviceId) {
+  const response = await internalFetch(base, `/internal/v1/users/${encodeURIComponent(userId)}/ws-tickets`, {
+    method: "POST",
+    body: JSON.stringify({ device_id: deviceId }),
+  });
+  assert.equal(response.status, 201);
+  return response.json();
+}
+
+function exchangeAttempt(base, headers = {}) {
+  return fetch(`${base}/v1/pairings/exchange`, {
+    method: "POST",
+    headers: { "Content-Type": "application/json", ...headers },
+    body: JSON.stringify({ code: "INVALID-CODE" }),
+  });
+}
+
+test("pairing exchange trusts a gateway client IP only with valid internal authentication", async (t) => {
+  const dataDir = fs.mkdtempSync(path.join(os.tmpdir(), "aether-vscodex-rate-limit-"));
+  const server = new AetherVscodexCloudServer({
+    host: "127.0.0.1",
+    port: 0,
+    internalToken,
+    publicWsUrl: "wss://aether.example/api/vscodex/ws",
+    dataDir,
+  });
+  await server.start();
+  t.after(async () => {
+    await server.stop();
+    fs.rmSync(dataDir, { recursive: true, force: true });
+  });
+  const address = server.address();
+  const base = `http://127.0.0.1:${address.port}`;
+  const trustedHeaders = (clientIp) => ({
+    Authorization: `Bearer ${internalToken}`,
+    "X-Aether-Client-IP": clientIp,
+  });
+
+  for (let attempt = 0; attempt < 10; attempt += 1) {
+    assert.equal((await exchangeAttempt(base, trustedHeaders("198.51.100.10"))).status, 400);
+  }
+  assert.equal((await exchangeAttempt(base, trustedHeaders("198.51.100.10"))).status, 429);
+  assert.equal((await exchangeAttempt(base, trustedHeaders("198.51.100.11"))).status, 400);
+  assert.equal((await exchangeAttempt(base, trustedHeaders("2001:db8::10"))).status, 400);
+
+  server.exchangeAttempts.clear();
+  for (let attempt = 0; attempt < 5; attempt += 1) {
+    assert.equal((await exchangeAttempt(base, { "X-Aether-Client-IP": `198.51.100.${20 + attempt}` })).status, 400);
+  }
+  for (let attempt = 0; attempt < 5; attempt += 1) {
+    assert.equal((await exchangeAttempt(base, {
+      Authorization: "Bearer invalid-internal-token",
+      "X-Aether-Client-IP": `198.51.100.${30 + attempt}`,
+    })).status, 400);
+  }
+  assert.equal((await exchangeAttempt(base, { "X-Aether-Client-IP": "198.51.100.99" })).status, 429);
+
+  server.exchangeAttempts.clear();
+  const invalidForwardedAddresses = ["proxy.internal", "198.51.100.40, 198.51.100.41"];
+  for (let attempt = 0; attempt < 10; attempt += 1) {
+    assert.equal((await exchangeAttempt(base, trustedHeaders(invalidForwardedAddresses[attempt % 2]))).status, 400);
+  }
+  assert.equal((await exchangeAttempt(base, trustedHeaders("198.51.100.42, 198.51.100.43"))).status, 429);
+});
+
+test("cloud sidecar pairs a device and isolates host/browser traffic by Aether user and device", async (t) => {
+  const dataDir = fs.mkdtempSync(path.join(os.tmpdir(), "aether-vscodex-test-"));
+  const server = new AetherVscodexCloudServer({
+    host: "127.0.0.1",
+    port: 0,
+    internalToken,
+    publicWsUrl: "wss://aether.example/api/vscodex/ws",
+    dataDir,
+    pairingTtlMs: 5_000,
+    ticketTtlMs: 5_000,
+  });
+  await server.start();
+  t.after(async () => {
+    await server.stop();
+    fs.rmSync(dataDir, { recursive: true, force: true });
+  });
+  const address = server.address();
+  const base = `http://127.0.0.1:${address.port}`;
+
+  const unauthorized = await fetch(`${base}/internal/v1/users/user-a/devices`);
+  assert.equal(unauthorized.status, 401);
+
+  const paired = await pairDevice(base, "user-a", "MacBook VS Code");
+  assert.match(paired.device_token, /^avx1\./);
+  const devicesResponse = await internalFetch(base, "/internal/v1/users/user-a/devices");
+  assert.equal(devicesResponse.status, 200);
+  const devices = await devicesResponse.json();
+  assert.deepEqual(devices.devices.map((device) => ({ id: device.id, name: device.name, connected: device.connected })), [
+    { id: paired.device_id, name: "MacBook VS Code", connected: false },
+  ]);
+
+  const ticket = await browserTicket(base, "user-a", paired.device_id);
+  assert.equal(ticket.ws_url, "/api/vscodex/ws");
+  const host = await websocketClient(base, "host", paired.device_token, "host-user-a");
+  const browser = await websocketClient(base, "web", ticket.ticket);
+  t.after(() => host.socket.close());
+  t.after(() => browser.socket.close());
+  browser.socket.send(JSON.stringify({ type: "subscribe", fromSeq: 0 }));
+
+  host.socket.send(JSON.stringify({
+    v: 1,
+    kind: "event",
+    type: "connection.opened",
+    id: "connection-a",
+    sessionId: "host-user-a",
+    seq: 1,
+    ts: new Date().toISOString(),
+    payload: {},
+  }));
+  host.socket.send(JSON.stringify({
+    v: 1,
+    kind: "event",
+    type: "session.snapshot",
+    id: "snapshot-a",
+    sessionId: "host-user-a",
+    seq: 2,
+    ts: new Date().toISOString(),
+    payload: { threadId: "thread-a", state: "idle", messages: [{ kind: "assistant", text: "user-a-only" }] },
+  }));
+  const snapshot = await browser.wait((message) => message.kind === "event" && message.type === "session.snapshot", 5_000, "session snapshot");
+  assert.equal(snapshot.payload.threadId, "thread-a");
+  assert.equal(snapshot.payload.messages[0].text, "user-a-only");
+
+  browser.socket.send(JSON.stringify({ type: "command", commandId: "cmd-a", method: "session/list", params: {} }));
+  const command = await host.wait((message) => message.kind === "command" && message.commandId === "cmd-a", 5_000, "browser command");
+  assert.equal(command.type, "session/list");
+
+  const secondUser = await pairDevice(base, "user-b", "Other VS Code");
+  const secondTicket = await browserTicket(base, "user-b", secondUser.device_id);
+  const secondBrowser = await websocketClient(base, "web", secondTicket.ticket);
+  t.after(() => secondBrowser.socket.close());
+  secondBrowser.socket.send(JSON.stringify({ type: "subscribe", fromSeq: 0 }));
+  await new Promise((resolve) => setTimeout(resolve, 50));
+  assert.equal(secondBrowser.messages.some((message) => message.payload?.threadId === "thread-a"), false);
+
+  const reusedTicket = new WebSocket(`${base.replace(/^http/, "ws")}/api/vscodex/ws`);
+  const closed = new Promise((resolve, reject) => {
+    reusedTicket.once("open", () => {
+      reusedTicket.send(JSON.stringify({ v: 1, kind: "hello", clientType: "web", protocol: 1 }));
+      reusedTicket.send(JSON.stringify({ type: "auth", token: ticket.ticket }));
+    });
+    reusedTicket.once("close", (code) => resolve(code));
+    reusedTicket.once("error", reject);
+  });
+  assert.equal(await closed, 1008, "browser tickets are one-time credentials");
+});
+
+test("device revocation closes its room and blocks future host authentication", async (t) => {
+  const dataDir = fs.mkdtempSync(path.join(os.tmpdir(), "aether-vscodex-revoke-"));
+  const server = new AetherVscodexCloudServer({
+    host: "127.0.0.1",
+    port: 0,
+    internalToken,
+    publicWsUrl: "wss://aether.example/api/vscodex/ws",
+    dataDir,
+  });
+  await server.start();
+  t.after(async () => {
+    await server.stop();
+    fs.rmSync(dataDir, { recursive: true, force: true });
+  });
+  const address = server.address();
+  const base = `http://127.0.0.1:${address.port}`;
+  const paired = await pairDevice(base, "user-a", "Revoked device");
+  const host = await websocketClient(base, "host", paired.device_token, "revoked-host");
+
+  const response = await internalFetch(base, `/internal/v1/users/user-a/devices/${paired.device_id}`, { method: "DELETE" });
+  assert.equal(response.status, 204);
+  await new Promise((resolve) => host.socket.once("close", resolve));
+
+  const rejected = new WebSocket(`${base.replace(/^http/, "ws")}/api/vscodex/ws`);
+  const closed = new Promise((resolve, reject) => {
+    rejected.once("open", () => {
+      rejected.send(JSON.stringify({ v: 1, kind: "hello", clientType: "host", protocol: 1, sessionId: "retry" }));
+      rejected.send(JSON.stringify({ v: 1, kind: "auth", accessToken: paired.device_token }));
+    });
+    rejected.once("close", (code) => resolve(code));
+    rejected.once("error", reject);
+  });
+  assert.equal(await closed, 1008);
+});
+
+test("room revocation wins a concurrent room creation", async () => {
+  const rooms = new RoomManager();
+  let releaseCreation;
+  const creationGate = new Promise((resolve) => { releaseCreation = resolve; });
+  let stopped = false;
+  const room = {
+    key: rooms.key("user-a", "device-a"),
+    userId: "user-a",
+    deviceId: "device-a",
+    relay: { stop: async () => { stopped = true; } },
+    connections: 0,
+    lastActiveMs: Date.now(),
+  };
+  rooms.createRoom = async (key) => {
+    await creationGate;
+    rooms.rooms.set(key, room);
+    return room;
+  };
+
+  const pendingGet = rooms.get("user-a", "device-a");
+  await new Promise((resolve) => setImmediate(resolve));
+  const pendingRevoke = rooms.revoke("user-a", "device-a");
+  releaseCreation();
+
+  await assert.rejects(pendingGet, /device revoked/);
+  await pendingRevoke;
+  assert.equal(stopped, true);
+  assert.equal(rooms.rooms.has(room.key), false);
+  await assert.rejects(rooms.get("user-a", "device-a"), /device revoked/);
+});