Files
Aether/aether-vscodex/test/cloud-server.test.js
T

299 lines
12 KiB
JavaScript

"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/);
});