Files
openclaw-studio/tests/unit/gatewayProxy.test.ts
T

713 lines
25 KiB
TypeScript

// @vitest-environment node
import { EventEmitter } from "node:events";
import { afterEach, describe, expect, it, vi } from "vitest";
import { WebSocket, WebSocketServer } from "ws";
const waitForEvent = <T = unknown>(
target: { once: (event: string, cb: (...args: unknown[]) => void) => void },
event: string
) =>
new Promise<T>((resolve) => {
target.once(event, (...args: unknown[]) => resolve(args as unknown as T));
});
const closeHttpServer = (server: import("node:http").Server) =>
new Promise<void>((resolve) => server.close(() => resolve()));
const closeWebSocketServer = (server: WebSocketServer) =>
new Promise<void>((resolve) => server.close(() => resolve()));
const closeWebSocket = (ws: WebSocket) =>
new Promise<void>((resolve) => {
if (ws.readyState === WebSocket.CLOSED) {
resolve();
return;
}
ws.once("close", () => resolve());
ws.close();
});
describe("createGatewayProxy", () => {
afterEach(() => {
vi.restoreAllMocks();
});
it("injects gateway token into connect request", async () => {
const upstream = new WebSocketServer({ port: 0 });
const address = upstream.address();
if (!address || typeof address === "string") {
throw new Error("expected upstream server to have a port");
}
const upstreamUrl = `ws://127.0.0.1:${address.port}`;
let seenToken: string | null = null;
let seenOrigin: string | undefined;
upstream.on("connection", (ws, req) => {
seenOrigin = req.headers.origin;
ws.on("message", (raw) => {
const parsed = JSON.parse(String(raw));
if (parsed?.method === "connect") {
seenToken = parsed?.params?.auth?.token ?? null;
ws.send(
JSON.stringify({
type: "res",
id: parsed.id,
ok: true,
payload: { type: "hello-ok", protocol: 3, auth: {} },
})
);
}
});
});
const { createGatewayProxy } = await import("../../server/gateway-proxy");
const proxyHttp = await import("node:http").then((m) => m.createServer());
const proxy = createGatewayProxy({
loadUpstreamSettings: async () => ({ url: upstreamUrl, token: "token-123" }),
allowWs: (req: { url?: string }) => req.url === "/api/gateway/ws",
logError: () => {},
});
proxyHttp.on("upgrade", (req, socket, head) => proxy.handleUpgrade(req, socket, head));
await new Promise<void>((resolve) => proxyHttp.listen(0, "127.0.0.1", resolve));
const proxyAddr = proxyHttp.address();
if (!proxyAddr || typeof proxyAddr === "string") {
throw new Error("expected proxy server to have a port");
}
const browser = new WebSocket(`ws://127.0.0.1:${proxyAddr.port}/api/gateway/ws`);
try {
await waitForEvent(browser, "open");
browser.send(
JSON.stringify({
type: "req",
id: "connect-1",
method: "connect",
params: { auth: {} },
})
);
await waitForEvent(browser, "message");
expect(seenToken).toBe("token-123");
expect(seenOrigin).toBe(`http://localhost:${address.port}`);
} finally {
for (const client of upstream.clients) {
client.close();
}
await Promise.all([
closeWebSocket(browser),
closeWebSocketServer(upstream),
closeHttpServer(proxyHttp),
]);
}
});
it("allows browser auth token passthrough when host token is missing", async () => {
const upstream = new WebSocketServer({ port: 0 });
const address = upstream.address();
if (!address || typeof address === "string") {
throw new Error("expected upstream server to have a port");
}
const upstreamUrl = `ws://127.0.0.1:${address.port}`;
let seenToken: string | null = null;
upstream.on("connection", (ws) => {
ws.on("message", (raw) => {
const parsed = JSON.parse(String(raw));
if (parsed?.method === "connect") {
seenToken = parsed?.params?.auth?.token ?? null;
ws.send(
JSON.stringify({
type: "res",
id: parsed.id,
ok: true,
payload: { type: "hello-ok", protocol: 3, auth: {} },
})
);
}
});
});
const { createGatewayProxy } = await import("../../server/gateway-proxy");
const proxyHttp = await import("node:http").then((m) => m.createServer());
const proxy = createGatewayProxy({
loadUpstreamSettings: async () => ({ url: upstreamUrl, token: "" }),
allowWs: (req: { url?: string }) => req.url === "/api/gateway/ws",
logError: () => {},
});
proxyHttp.on("upgrade", (req, socket, head) => proxy.handleUpgrade(req, socket, head));
await new Promise<void>((resolve) => proxyHttp.listen(0, "127.0.0.1", resolve));
const proxyAddr = proxyHttp.address();
if (!proxyAddr || typeof proxyAddr === "string") {
throw new Error("expected proxy server to have a port");
}
const browser = new WebSocket(`ws://127.0.0.1:${proxyAddr.port}/api/gateway/ws`);
try {
await waitForEvent(browser, "open");
browser.send(
JSON.stringify({
type: "req",
id: "connect-pass-token",
method: "connect",
params: { auth: { token: "browser-token-123" } },
})
);
const [rawMessage] = await waitForEvent<[WebSocket.RawData]>(browser, "message");
const response = JSON.parse(String(rawMessage ?? ""));
expect(response).toMatchObject({ type: "res", id: "connect-pass-token", ok: true });
expect(seenToken).toBe("browser-token-123");
} finally {
for (const client of upstream.clients) {
client.close();
}
await Promise.all([
closeWebSocket(browser),
closeWebSocketServer(upstream),
closeHttpServer(proxyHttp),
]);
}
});
it("preserves browser auth token when both browser and host tokens are present", async () => {
const upstream = new WebSocketServer({ port: 0 });
const address = upstream.address();
if (!address || typeof address === "string") {
throw new Error("expected upstream server to have a port");
}
const upstreamUrl = `ws://127.0.0.1:${address.port}`;
let seenToken: string | null = null;
upstream.on("connection", (ws) => {
ws.on("message", (raw) => {
const parsed = JSON.parse(String(raw));
if (parsed?.method === "connect") {
seenToken = parsed?.params?.auth?.token ?? null;
ws.send(
JSON.stringify({
type: "res",
id: parsed.id,
ok: true,
payload: { type: "hello-ok", protocol: 3, auth: {} },
})
);
}
});
});
const { createGatewayProxy } = await import("../../server/gateway-proxy");
const proxyHttp = await import("node:http").then((m) => m.createServer());
const proxy = createGatewayProxy({
loadUpstreamSettings: async () => ({ url: upstreamUrl, token: "host-token-456" }),
allowWs: (req: { url?: string }) => req.url === "/api/gateway/ws",
logError: () => {},
});
proxyHttp.on("upgrade", (req, socket, head) => proxy.handleUpgrade(req, socket, head));
await new Promise<void>((resolve) => proxyHttp.listen(0, "127.0.0.1", resolve));
const proxyAddr = proxyHttp.address();
if (!proxyAddr || typeof proxyAddr === "string") {
throw new Error("expected proxy server to have a port");
}
const browser = new WebSocket(`ws://127.0.0.1:${proxyAddr.port}/api/gateway/ws`);
try {
await waitForEvent(browser, "open");
browser.send(
JSON.stringify({
type: "req",
id: "connect-browser-precedence",
method: "connect",
params: { auth: { token: "browser-token-789" } },
})
);
const [rawMessage] = await waitForEvent<[WebSocket.RawData]>(browser, "message");
const response = JSON.parse(String(rawMessage ?? ""));
expect(response).toMatchObject({ type: "res", id: "connect-browser-precedence", ok: true });
expect(seenToken).toBe("browser-token-789");
} finally {
for (const client of upstream.clients) {
client.close();
}
await Promise.all([
closeWebSocket(browser),
closeWebSocketServer(upstream),
closeHttpServer(proxyHttp),
]);
}
});
it("allows browser device signature passthrough when host token is missing", async () => {
const upstream = new WebSocketServer({ port: 0 });
const address = upstream.address();
if (!address || typeof address === "string") {
throw new Error("expected upstream server to have a port");
}
const upstreamUrl = `ws://127.0.0.1:${address.port}`;
let seenToken: string | null = null;
let seenDeviceSignature: string | null = null;
let seenDeviceId: string | null = null;
let seenDevicePublicKey: string | null = null;
let seenDeviceNonce: string | null = null;
let seenDeviceSignedAt: number | null = null;
upstream.on("connection", (ws) => {
ws.on("message", (raw) => {
const parsed = JSON.parse(String(raw));
if (parsed?.method === "connect") {
seenToken = parsed?.params?.auth?.token ?? null;
seenDeviceSignature = parsed?.params?.device?.signature ?? null;
seenDeviceId = parsed?.params?.device?.id ?? null;
seenDevicePublicKey = parsed?.params?.device?.publicKey ?? null;
seenDeviceNonce = parsed?.params?.device?.nonce ?? null;
seenDeviceSignedAt = parsed?.params?.device?.signedAt ?? null;
ws.send(
JSON.stringify({
type: "res",
id: parsed.id,
ok: true,
payload: { type: "hello-ok", protocol: 3, auth: {} },
})
);
}
});
});
const { createGatewayProxy } = await import("../../server/gateway-proxy");
const proxyHttp = await import("node:http").then((m) => m.createServer());
const proxy = createGatewayProxy({
loadUpstreamSettings: async () => ({ url: upstreamUrl, token: "" }),
allowWs: (req: { url?: string }) => req.url === "/api/gateway/ws",
logError: () => {},
});
proxyHttp.on("upgrade", (req, socket, head) => proxy.handleUpgrade(req, socket, head));
await new Promise<void>((resolve) => proxyHttp.listen(0, "127.0.0.1", resolve));
const proxyAddr = proxyHttp.address();
if (!proxyAddr || typeof proxyAddr === "string") {
throw new Error("expected proxy server to have a port");
}
const browser = new WebSocket(`ws://127.0.0.1:${proxyAddr.port}/api/gateway/ws`);
try {
await waitForEvent(browser, "open");
browser.send(
JSON.stringify({
type: "req",
id: "connect-pass-device",
method: "connect",
params: {
device: {
id: "device-id-123",
publicKey: "device-public-key-123",
signature: "device-signature-123",
signedAt: Date.now(),
nonce: "device-nonce-123",
},
},
})
);
const [rawMessage] = await waitForEvent<[WebSocket.RawData]>(browser, "message");
const response = JSON.parse(String(rawMessage ?? ""));
expect(response).toMatchObject({ type: "res", id: "connect-pass-device", ok: true });
expect(seenDeviceSignature).toBe("device-signature-123");
expect(seenDeviceId).toBe("device-id-123");
expect(seenDevicePublicKey).toBe("device-public-key-123");
expect(seenDeviceNonce).toBe("device-nonce-123");
expect(typeof seenDeviceSignedAt).toBe("number");
expect(seenToken).toBeNull();
} finally {
for (const client of upstream.clients) {
client.close();
}
await Promise.all([
closeWebSocket(browser),
closeWebSocketServer(upstream),
closeHttpServer(proxyHttp),
]);
}
});
it("allows browser password passthrough when host token is missing", async () => {
const upstream = new WebSocketServer({ port: 0 });
const address = upstream.address();
if (!address || typeof address === "string") {
throw new Error("expected upstream server to have a port");
}
const upstreamUrl = `ws://127.0.0.1:${address.port}`;
let seenPassword: string | null = null;
let seenToken: string | null = null;
upstream.on("connection", (ws) => {
ws.on("message", (raw) => {
const parsed = JSON.parse(String(raw));
if (parsed?.method === "connect") {
seenPassword = parsed?.params?.auth?.password ?? null;
seenToken = parsed?.params?.auth?.token ?? null;
ws.send(
JSON.stringify({
type: "res",
id: parsed.id,
ok: true,
payload: { type: "hello-ok", protocol: 3, auth: {} },
})
);
}
});
});
const { createGatewayProxy } = await import("../../server/gateway-proxy");
const proxyHttp = await import("node:http").then((m) => m.createServer());
const proxy = createGatewayProxy({
loadUpstreamSettings: async () => ({ url: upstreamUrl, token: "" }),
allowWs: (req: { url?: string }) => req.url === "/api/gateway/ws",
logError: () => {},
});
proxyHttp.on("upgrade", (req, socket, head) => proxy.handleUpgrade(req, socket, head));
await new Promise<void>((resolve) => proxyHttp.listen(0, "127.0.0.1", resolve));
const proxyAddr = proxyHttp.address();
if (!proxyAddr || typeof proxyAddr === "string") {
throw new Error("expected proxy server to have a port");
}
const browser = new WebSocket(`ws://127.0.0.1:${proxyAddr.port}/api/gateway/ws`);
try {
await waitForEvent(browser, "open");
browser.send(
JSON.stringify({
type: "req",
id: "connect-pass-password",
method: "connect",
params: { auth: { password: "browser-password-123" } },
})
);
const [rawMessage] = await waitForEvent<[WebSocket.RawData]>(browser, "message");
const response = JSON.parse(String(rawMessage ?? ""));
expect(response).toMatchObject({ type: "res", id: "connect-pass-password", ok: true });
expect(seenPassword).toBe("browser-password-123");
expect(seenToken).toBeNull();
} finally {
for (const client of upstream.clients) {
client.close();
}
await Promise.all([
closeWebSocket(browser),
closeWebSocketServer(upstream),
closeHttpServer(proxyHttp),
]);
}
});
it("allows browser deviceToken passthrough when host token is missing", async () => {
const upstream = new WebSocketServer({ port: 0 });
const address = upstream.address();
if (!address || typeof address === "string") {
throw new Error("expected upstream server to have a port");
}
const upstreamUrl = `ws://127.0.0.1:${address.port}`;
let seenDeviceToken: string | null = null;
let seenToken: string | null = null;
upstream.on("connection", (ws) => {
ws.on("message", (raw) => {
const parsed = JSON.parse(String(raw));
if (parsed?.method === "connect") {
seenDeviceToken = parsed?.params?.auth?.deviceToken ?? null;
seenToken = parsed?.params?.auth?.token ?? null;
ws.send(
JSON.stringify({
type: "res",
id: parsed.id,
ok: true,
payload: { type: "hello-ok", protocol: 3, auth: {} },
})
);
}
});
});
const { createGatewayProxy } = await import("../../server/gateway-proxy");
const proxyHttp = await import("node:http").then((m) => m.createServer());
const proxy = createGatewayProxy({
loadUpstreamSettings: async () => ({ url: upstreamUrl, token: "" }),
allowWs: (req: { url?: string }) => req.url === "/api/gateway/ws",
logError: () => {},
});
proxyHttp.on("upgrade", (req, socket, head) => proxy.handleUpgrade(req, socket, head));
await new Promise<void>((resolve) => proxyHttp.listen(0, "127.0.0.1", resolve));
const proxyAddr = proxyHttp.address();
if (!proxyAddr || typeof proxyAddr === "string") {
throw new Error("expected proxy server to have a port");
}
const browser = new WebSocket(`ws://127.0.0.1:${proxyAddr.port}/api/gateway/ws`);
try {
await waitForEvent(browser, "open");
browser.send(
JSON.stringify({
type: "req",
id: "connect-pass-device-token",
method: "connect",
params: { auth: { deviceToken: "browser-device-token-123" } },
})
);
const [rawMessage] = await waitForEvent<[WebSocket.RawData]>(browser, "message");
const response = JSON.parse(String(rawMessage ?? ""));
expect(response).toMatchObject({ type: "res", id: "connect-pass-device-token", ok: true });
expect(seenDeviceToken).toBe("browser-device-token-123");
expect(seenToken).toBeNull();
} finally {
for (const client of upstream.clients) {
client.close();
}
await Promise.all([
closeWebSocket(browser),
closeWebSocketServer(upstream),
closeHttpServer(proxyHttp),
]);
}
});
it("returns studio.gateway_token_missing when browser auth and host token are both missing", async () => {
const upstream = new WebSocketServer({ port: 0 });
const address = upstream.address();
if (!address || typeof address === "string") {
throw new Error("expected upstream server to have a port");
}
const upstreamUrl = `ws://127.0.0.1:${address.port}`;
let upstreamConnectionCount = 0;
upstream.on("connection", () => {
upstreamConnectionCount += 1;
});
const { createGatewayProxy } = await import("../../server/gateway-proxy");
const proxyHttp = await import("node:http").then((m) => m.createServer());
const proxy = createGatewayProxy({
loadUpstreamSettings: async () => ({ url: upstreamUrl, token: "" }),
allowWs: (req: { url?: string }) => req.url === "/api/gateway/ws",
logError: () => {},
});
proxyHttp.on("upgrade", (req, socket, head) => proxy.handleUpgrade(req, socket, head));
await new Promise<void>((resolve) => proxyHttp.listen(0, "127.0.0.1", resolve));
const proxyAddr = proxyHttp.address();
if (!proxyAddr || typeof proxyAddr === "string") {
throw new Error("expected proxy server to have a port");
}
const browser = new WebSocket(`ws://127.0.0.1:${proxyAddr.port}/api/gateway/ws`);
try {
await waitForEvent(browser, "open");
const closePromise = waitForEvent<[number, Buffer]>(browser, "close");
browser.send(
JSON.stringify({
type: "req",
id: "connect-missing-token",
method: "connect",
params: { auth: {} },
})
);
const [rawMessage] = await waitForEvent<[WebSocket.RawData]>(browser, "message");
const response = JSON.parse(String(rawMessage ?? ""));
expect(response).toMatchObject({
type: "res",
id: "connect-missing-token",
ok: false,
error: { code: "studio.gateway_token_missing" },
});
const [closeCode] = await closePromise;
expect(closeCode).toBe(1011);
expect(upstreamConnectionCount).toBe(0);
} finally {
for (const client of upstream.clients) {
client.close();
}
await Promise.all([
closeWebSocket(browser),
closeWebSocketServer(upstream),
closeHttpServer(proxyHttp),
]);
}
});
it("suppresses expected close-before-open upstream race errors", async () => {
class FakeUpstreamSocket extends EventEmitter {
readyState: number = WebSocket.CONNECTING;
send() {}
close() {
this.readyState = WebSocket.CLOSED;
this.emit("error", new Error("WebSocket was closed before the connection was established"));
this.emit("close", { code: 1000, reason: "closed" });
}
}
const logError = vi.fn();
const log = vi.fn();
let upstreamSocket: FakeUpstreamSocket | null = null;
const { createGatewayProxy } = await import("../../server/gateway-proxy");
const proxyHttp = await import("node:http").then((m) => m.createServer());
const proxy = createGatewayProxy({
loadUpstreamSettings: async () => ({ url: "ws://127.0.0.1:65535", token: "token-123" }),
allowWs: (req: { url?: string }) => req.url === "/api/gateway/ws",
log,
logError,
createUpstreamWebSocket: () => {
upstreamSocket = new FakeUpstreamSocket();
return upstreamSocket;
},
});
proxyHttp.on("upgrade", (req, socket, head) => proxy.handleUpgrade(req, socket, head));
await new Promise<void>((resolve) => proxyHttp.listen(0, "127.0.0.1", resolve));
const proxyAddr = proxyHttp.address();
if (!proxyAddr || typeof proxyAddr === "string") {
throw new Error("expected proxy server to have a port");
}
const browser = new WebSocket(`ws://127.0.0.1:${proxyAddr.port}/api/gateway/ws`);
try {
await waitForEvent(browser, "open");
browser.send(
JSON.stringify({
type: "req",
id: "connect-suppress",
method: "connect",
params: { auth: {} },
})
);
await closeWebSocket(browser);
await new Promise((resolve) => setTimeout(resolve, 20));
expect(upstreamSocket).not.toBeNull();
expect(logError).not.toHaveBeenCalled();
expect(log).toHaveBeenCalledWith("Suppressed upstream close-before-open race.");
} finally {
await closeHttpServer(proxyHttp);
}
});
it("logs and forwards unexpected upstream socket errors", async () => {
class FakeUpstreamSocket extends EventEmitter {
readyState: number = WebSocket.CONNECTING;
send() {}
close() {
this.readyState = WebSocket.CLOSED;
this.emit("close", { code: 1000, reason: "closed" });
}
}
const logError = vi.fn();
const upstreamSocket = new FakeUpstreamSocket();
const { createGatewayProxy } = await import("../../server/gateway-proxy");
const proxyHttp = await import("node:http").then((m) => m.createServer());
const proxy = createGatewayProxy({
loadUpstreamSettings: async () => ({ url: "ws://127.0.0.1:65534", token: "token-123" }),
allowWs: (req: { url?: string }) => req.url === "/api/gateway/ws",
logError,
createUpstreamWebSocket: () => upstreamSocket,
});
proxyHttp.on("upgrade", (req, socket, head) => proxy.handleUpgrade(req, socket, head));
await new Promise<void>((resolve) => proxyHttp.listen(0, "127.0.0.1", resolve));
const proxyAddr = proxyHttp.address();
if (!proxyAddr || typeof proxyAddr === "string") {
throw new Error("expected proxy server to have a port");
}
const browser = new WebSocket(`ws://127.0.0.1:${proxyAddr.port}/api/gateway/ws`);
try {
await waitForEvent(browser, "open");
browser.send(
JSON.stringify({
type: "req",
id: "connect-forward-error",
method: "connect",
params: { auth: {} },
})
);
await new Promise((resolve) => setTimeout(resolve, 0));
upstreamSocket.emit("error", new Error("socket boom"));
const [rawMessage] = await waitForEvent<[WebSocket.RawData]>(browser, "message");
const response = JSON.parse(String(rawMessage ?? ""));
expect(response).toMatchObject({
type: "res",
id: "connect-forward-error",
ok: false,
error: { code: "studio.upstream_error" },
});
expect(logError).toHaveBeenCalledTimes(1);
} finally {
await Promise.all([closeWebSocket(browser), closeHttpServer(proxyHttp)]);
}
});
it("returns structured upstream error when upstream websocket creation throws", async () => {
const logError = vi.fn();
const { createGatewayProxy } = await import("../../server/gateway-proxy");
const proxyHttp = await import("node:http").then((m) => m.createServer());
const proxy = createGatewayProxy({
loadUpstreamSettings: async () => ({ url: "ws://127.0.0.1:65534", token: "token-123" }),
allowWs: (req: { url?: string }) => req.url === "/api/gateway/ws",
logError,
createUpstreamWebSocket: () => {
throw new Error("constructor failed");
},
});
proxyHttp.on("upgrade", (req, socket, head) => proxy.handleUpgrade(req, socket, head));
await new Promise<void>((resolve) => proxyHttp.listen(0, "127.0.0.1", resolve));
const proxyAddr = proxyHttp.address();
if (!proxyAddr || typeof proxyAddr === "string") {
throw new Error("expected proxy server to have a port");
}
const browser = new WebSocket(`ws://127.0.0.1:${proxyAddr.port}/api/gateway/ws`);
try {
await waitForEvent(browser, "open");
browser.send(
JSON.stringify({
type: "req",
id: "connect-creation-throw",
method: "connect",
params: { auth: {} },
})
);
const [rawMessage] = await waitForEvent<[WebSocket.RawData]>(browser, "message");
const response = JSON.parse(String(rawMessage ?? ""));
expect(response).toMatchObject({
type: "res",
id: "connect-creation-throw",
ok: false,
error: { code: "studio.upstream_error" },
});
expect(logError).toHaveBeenCalledTimes(1);
} finally {
await Promise.all([closeWebSocket(browser), closeHttpServer(proxyHttp)]);
}
});
});