66 lines
5.3 KiB
TypeScript
66 lines
5.3 KiB
TypeScript
import test from "node:test";
|
|
import assert from "node:assert/strict";
|
|
import { randomUUID } from "node:crypto";
|
|
import { WebSocket } from "ws";
|
|
import { createHttpServer } from "../src/app/http.ts";
|
|
import { attachTaskWebSocket } from "../src/app/ws.ts";
|
|
import { createStore } from "../src/store.ts";
|
|
import { hashRefreshToken, issueAccessToken } from "../src/shared/auth.ts";
|
|
|
|
test("task websocket authenticates and replays a subscribed task cursor", async () => {
|
|
const secret = "test";
|
|
const store = createStore();
|
|
const server = createHttpServer(store, { host: "127.0.0.1", port: 0, accessTokenSecret: secret, accessTokenTtlSeconds: 900, refreshCookieName: "refresh", adminEmail: "admin@example.com", adminPassword: "change-me-now", adminMfaRequired: true });
|
|
attachTaskWebSocket(server, store, secret, "refresh");
|
|
await new Promise<void>((resolve) => server.listen(0, "127.0.0.1", resolve));
|
|
const address = server.address();
|
|
assert.ok(address && typeof address !== "string");
|
|
const taskId = "ws-task";
|
|
store.users.set("ws-user", { id: "ws-user", version: 1, password: "", verified: true, status: "active", mfaRequired: false, mfaEnabled: false, roles: ["user"], failedLoginCount: 0, createdAt: new Date().toISOString() });
|
|
store.tasks.set(taskId, { id: taskId, ownerId: "ws-user", taskType: "image", modelProductId: "basic-image", status: "queued", estimatedCost: 1, reservedCost: 0, createdAt: new Date().toISOString(), updatedAt: new Date().toISOString(), eventSequence: 0, routeSnapshot: { groupId: "image-default", groupVersion: 1, publicModelId: "basic-image-v1", routeVersion: 1, orderedCandidates: [], strategy: "strict", totalRetryBudget: 0 }, modelSnapshot: { publicModelId: "basic-image-v1", modelVersion: 1, name: "基础模型", tier: "basic", capabilities: ["image"] }, pricingSnapshot: { ruleId: "price-basic-image", ruleVersion: 1, publicModelId: "basic-image-v1", unitPrice: 1, quantity: 1, billingUnit: "output", multiplier: 1, balanceUnitVersion: 1 }, planSnapshot: { queuePriority: 0, maxConcurrent: 1 }, attempts: [], outputs: [] });
|
|
const token = issueAccessToken("ws-user", "user", secret, 900);
|
|
const socket = new WebSocket(`ws://127.0.0.1:${address.port}/api/v1/ws/tasks?access_token=${encodeURIComponent(token)}`);
|
|
const messages: Array<{ type: string; event?: { taskId: string } }> = [];
|
|
socket.on("message", (raw) => messages.push(JSON.parse(raw.toString())));
|
|
await new Promise<void>((resolve) => socket.once("open", () => resolve()));
|
|
socket.send(JSON.stringify({ type: "subscribe", taskId, cursor: 0 }));
|
|
await new Promise((resolve) => setTimeout(resolve, 25));
|
|
assert.equal(messages[0]?.type, "ready");
|
|
store.events.set(taskId, []);
|
|
const { appendEvent } = await import("../src/store.ts");
|
|
appendEvent(store, taskId, { taskId, type: "task.progress", payload: { progress: 20 } });
|
|
await new Promise((resolve) => setTimeout(resolve, 25));
|
|
assert.ok(messages.some((message) => message.type === "event" && message.event?.taskId === taskId));
|
|
socket.close();
|
|
await new Promise<void>((resolve, reject) => server.close((error) => (error ? reject(error) : resolve())));
|
|
});
|
|
|
|
test("task websocket rejects a revoked refresh session family", async () => {
|
|
const secret = "test";
|
|
const store = createStore();
|
|
const server = createHttpServer(store, { host: "127.0.0.1", port: 0, accessTokenSecret: secret, accessTokenTtlSeconds: 900, refreshCookieName: "refresh", adminEmail: "admin@example.com", adminPassword: "change-me-now", adminMfaRequired: true });
|
|
attachTaskWebSocket(server, store, secret, "refresh");
|
|
await new Promise<void>((resolve) => server.listen(0, "127.0.0.1", resolve));
|
|
const address = server.address(); assert.ok(address && typeof address !== "string");
|
|
const userId = "ws-cookie-user"; const familyId = randomUUID(); const refresh = randomUUID();
|
|
store.users.set(userId, { id: userId, version: 1, password: "", verified: true, status: "active", mfaRequired: false, mfaEnabled: false, roles: ["user"], failedLoginCount: 0, createdAt: new Date().toISOString() });
|
|
store.sessionFamilies.set(familyId, { userId, scope: "user" });
|
|
store.sessions.set(hashRefreshToken(refresh), { userId, scope: "user", expiresAt: Date.now() + 60_000, absoluteExpiresAt: Date.now() + 60_000, createdAt: Date.now(), rotationId: randomUUID(), familyId });
|
|
const first = new WebSocket(`ws://127.0.0.1:${address.port}/api/v1/ws/tasks`, { headers: { Cookie: `refresh=${refresh}` } });
|
|
await new Promise<void>((resolve, reject) => { first.once("open", () => resolve()); first.once("error", reject); });
|
|
store.sessionFamilies.get(familyId)!.revokedAt = Date.now();
|
|
first.send(JSON.stringify({ type: "ping" }));
|
|
await Promise.race([
|
|
new Promise<void>((resolve) => first.once("close", () => resolve())),
|
|
new Promise<void>((resolve) => setTimeout(resolve, 1_000)),
|
|
]);
|
|
assert.equal(first.readyState, WebSocket.CLOSED);
|
|
const second = new WebSocket(`ws://127.0.0.1:${address.port}/api/v1/ws/tasks`, { headers: { Cookie: `refresh=${refresh}` } });
|
|
await Promise.race([
|
|
new Promise<void>((resolve) => { second.once("close", () => resolve()); second.once("error", () => resolve()); }),
|
|
new Promise<void>((resolve) => setTimeout(resolve, 1_000)),
|
|
]);
|
|
second.terminate();
|
|
await new Promise<void>((resolve, reject) => server.close((error) => error ? reject(error) : resolve()));
|
|
});
|