trade-message-center/apps/server/test/websocket.test.ts

// 验证 OneTalk WebSocket 握手和授权边界

import assert from "node:assert/strict";
import test from "node:test";

import {
    ONETALK_ERROR_CODES,
    ONETALK_PROTOCOL_VERSION,
    createMockAuthorizationReader,
    type MockAuthorizationRecord,
    type OneTalkAuthorizationRequest,
    type OneTalkMindScope,
    type OneTalkPluginScope,
    type OneTalkScope,
} from "@trade-message-center/onetalk-contract";
import Fastify from "fastify";

import { createApp } from "../src/app.ts";
import type { DatabaseConnection } from "../src/database/index.ts";
import type { OneTalkProfileService, OneTalkService } from "../src/onetalk/index.ts";
import { createMindAuthorizationReader } from "../src/mind-authorization.ts";
import { installWebsocket } from "../src/websocket/index.ts";
import { createOneTalkCutoverPolicy } from "../src/cutover-policy.ts";

const testConfig = {
    host: "127.0.0.1",
    port: 3000,
    databaseUrl: "postgres://test:test@localhost:5432/test",
    environment: "non_development" as const,
    mindAuthorization: {
        baseUrl: "https://mind.example.com",
        mindPageOrigin: "http://mind.localhost",
        pluginOrigins: ["http://plugin.localhost"],
        cutoverPolicy: createOneTalkCutoverPolicy(),
        timeoutMs: 100,
    },
};

const bindingScope: OneTalkPluginScope = {
    channelAccountId: "onetalk-account-1",
    deviceId: "device-1",
};

const developmentConfig = {
    ...testConfig,
    environment: "development" as const,
};

const mindScope: OneTalkMindScope = {
    mindUserId: "mind-user-1",
    workspaceId: "workspace-1",
    channelAccountId: bindingScope.channelAccountId,
};

const authorizationRecord: MockAuthorizationRecord = {
    scope: bindingScope,
    mindScope,
    binding: "binding-1",
    permissions: ["read", "send"],
    authorizationVersion: "version-1",
    active: true,
};

const createDatabaseStub = (): DatabaseConnection => {
    return {
        db: {} as DatabaseConnection["db"],
        close: async () => {},
    };
};

const createServiceStub = (): OneTalkService => {
    return {
        discoverConversation: async (_context, conversationId, _commitGuard, conversationKind) => ({
            channelAccountId: bindingScope.channelAccountId,
            conversationId,
            conversationKind,
            syncPhase: "initial",
            syncResult: "incomplete",
            latestMessageId: null,
            historyComplete: false,
            messageCount: 0,
        }),
        listAnchors: async () => [],
        listConversations: async () => [],
        readConversation: async () => null,
        readHistory: async () => null,
        observeMessage: async () => ({
            status: "anomaly" as const,
            anomalyCode: "test_only",
        }),
        completeSync: async () => ({
            status: "rejected" as const,
            reason: "conversation_not_discovered" as const,
        }),
    };
};

const createProfileServiceStub = (): OneTalkProfileService => ({
    ingestProfiles: async (input) => ({
        status: "accepted",
        inputProfileCount: input.profiles.length,
        profileCount: input.profiles.length,
        writtenProfileCount: input.profiles.length,
        staleProfileCount: 0,
    }),
});

const helloFrame = (
    scope: OneTalkScope = bindingScope,
    binding = "binding-1",
): Record<string, unknown> => {
    return {
        protocolVersion: ONETALK_PROTOCOL_VERSION,
        connectionType: "plugin",
        type: "ws.hello",
        requestId: "hello-1",
        scope,
        payload: {
            binding,
            requestedPermissions: ["read", "send"],
        },
    };
};

const nextMessages = async (
    socket: {
        once: (event: string, listener: (data: Buffer) => void) => void;
    },
    count: number,
): Promise<Record<string, unknown>[]> => {
    return new Promise((resolve, reject) => {
        const messages: Record<string, unknown>[] = [];
        const listen = (): void => {
            socket.once("message", (data) => {
                try {
                    messages.push(JSON.parse(data.toString()) as Record<string, unknown>);
                    if (messages.length === count) {
                        resolve(messages);
                        return;
                    }
                    listen();
                } catch (error: unknown) {
                    reject(error);
                }
            });
        };
        listen();
        socket.once("error", reject);
    });
};

const nextMessage = async (socket: {
    once: (event: string, listener: (data: Buffer) => void) => void;
}): Promise<Record<string, unknown>> => {
    const [message] = await nextMessages(socket, 1);
    return message;
};

const nextCloseCode = async (socket: {
    once: (event: string, listener: (code: number) => void) => void;
}): Promise<number> => {
    return new Promise((resolve, reject) => {
        socket.once("close", resolve);
        socket.once("error", reject);
    });
};

const openSocket = async (app: ReturnType<typeof createApp>, path = "/ws/plugin") => {
    await app.ready();
    const origin = path === "/ws/mind" ? "http://mind.localhost" : "http://plugin.localhost";
    return app.injectWS(path, { headers: { origin } });
};

const assertAnchorSnapshot = (snapshot: Record<string, unknown>): void => {
    assert.equal(snapshot.type, "anchor.snapshot");
    assert.deepEqual(snapshot.payload, { anchors: [] });
};

const closeApp = async (
    app: ReturnType<typeof createApp>,
    socket: { terminate: () => void },
): Promise<void> => {
    socket.terminate();
    for (const client of app.websocketServer.clients) client.terminate();
    await app.close();
};

test("accepts a mock-authorized plugin handshake and heartbeat", async () => {
    const authorization = createMockAuthorizationReader([authorizationRecord]);
    const app = createApp(testConfig, {
        database: createDatabaseStub(),
        authorization,
        oneTalkService: createServiceStub(),
    });
    const socket = await openSocket(app);

    try {
        const handshake = nextMessages(socket, 2);
        socket.send(JSON.stringify(helloFrame()));
        const [accepted, snapshot] = await handshake;
        assert.deepEqual(accepted, {
            protocolVersion: ONETALK_PROTOCOL_VERSION,
            connectionType: "plugin",
            type: "ws.accepted",
            requestId: "hello-1",
            scope: bindingScope,
            payload: {
                authorizationVersion: "version-1",
                permissions: ["read", "send"],
            },
        });
        assertAnchorSnapshot(snapshot);

        const heartbeat = nextMessage(socket);
        socket.send(
            JSON.stringify({
                protocolVersion: ONETALK_PROTOCOL_VERSION,
                connectionType: "plugin",
                type: "heartbeat",
                requestId: "heartbeat-1",
                scope: bindingScope,
                payload: { sentAtMs: 1_700_000_000_000 },
            }),
        );
        assert.deepEqual(await heartbeat, {
            protocolVersion: ONETALK_PROTOCOL_VERSION,
            connectionType: "plugin",
            type: "heartbeat.ack",
            requestId: "heartbeat-1",
            scope: bindingScope,
            payload: { sentAtMs: 1_700_000_000_000 },
        });
    } finally {
        await closeApp(app, socket);
    }
});

test("accepts a development handshake using an injected authorization reader", async () => {
    const app = createApp(developmentConfig, {
        database: createDatabaseStub(),
        authorization: createMockAuthorizationReader([authorizationRecord]),
        oneTalkService: createServiceStub(),
    });
    const socket = await openSocket(app);

    try {
        const handshake = nextMessages(socket, 2);
        socket.send(JSON.stringify(helloFrame(bindingScope, "binding-1")));
        const [accepted, snapshot] = await handshake;

        assert.deepEqual(accepted.payload, {
            authorizationVersion: "version-1",
            permissions: ["read", "send"],
        });
        assert.equal(JSON.stringify(accepted).includes("mock-session"), false);
        assertAnchorSnapshot(snapshot);

        const heartbeat = nextMessage(socket);
        socket.send(
            JSON.stringify({
                protocolVersion: ONETALK_PROTOCOL_VERSION,
                connectionType: "plugin",
                type: "heartbeat",
                requestId: "development-heartbeat",
                scope: bindingScope,
                payload: { sentAtMs: 1_700_000_000_000 },
            }),
        );
        assert.deepEqual(await heartbeat, {
            protocolVersion: ONETALK_PROTOCOL_VERSION,
            connectionType: "plugin",
            type: "heartbeat.ack",
            requestId: "development-heartbeat",
            scope: bindingScope,
            payload: { sentAtMs: 1_700_000_000_000 },
        });
    } finally {
        await closeApp(app, socket);
    }
});

test("accepts a development Mind page handshake using an injected authorization reader", async () => {
    const app = createApp(developmentConfig, {
        database: createDatabaseStub(),
        authorization: createMockAuthorizationReader([authorizationRecord]),
        oneTalkService: createServiceStub(),
    });
    const socket = await openSocket(app, "/ws/mind");

    try {
        const accepted = nextMessage(socket);
        socket.send(
            JSON.stringify({
                ...helloFrame(mindScope),
                connectionType: "mind_page",
                payload: { requestedPermissions: ["read"] },
            }),
        );

        const response = await accepted;
        assert.deepEqual(response, {
            protocolVersion: ONETALK_PROTOCOL_VERSION,
            connectionType: "mind_page",
            type: "ws.accepted",
            requestId: "hello-1",
            scope: mindScope,
            payload: {
                authorizationVersion: "version-1",
                permissions: ["read"],
            },
        });
        assert.equal(JSON.stringify(response).includes("binding-1"), false);
        assert.equal(JSON.stringify(response).includes("mock-session"), false);
    } finally {
        await closeApp(app, socket);
    }
});

test("keeps an explicitly injected authorization reader ahead of the development default", async () => {
    const authorization = createMockAuthorizationReader([]);
    const app = createApp(developmentConfig, {
        database: createDatabaseStub(),
        authorization,
        oneTalkService: createServiceStub(),
    });
    const socket = await openSocket(app);

    try {
        const errorMessage = nextMessage(socket);
        socket.send(JSON.stringify(helloFrame()));
        assert.deepEqual(await errorMessage, {
            protocolVersion: ONETALK_PROTOCOL_VERSION,
            connectionType: "plugin",
            type: "ws.error",
            requestId: "hello-1",
            scope: bindingScope,
            payload: { code: ONETALK_ERROR_CODES.scopeMismatch },
        });
    } finally {
        await closeApp(app, socket);
    }
});

test("accepts a Mind page handshake without a plugin device identifier", async () => {
    const authorization = createMockAuthorizationReader([authorizationRecord]);
    const app = createApp(testConfig, {
        database: createDatabaseStub(),
        authorization,
        oneTalkService: createServiceStub(),
    });
    const socket = await openSocket(app, "/ws/mind");

    try {
        const accepted = nextMessage(socket);
        socket.send(
            JSON.stringify({
                ...helloFrame(mindScope),
                connectionType: "mind_page",
                payload: { requestedPermissions: ["read"] },
            }),
        );
        assert.deepEqual(await accepted, {
            protocolVersion: ONETALK_PROTOCOL_VERSION,
            connectionType: "mind_page",
            type: "ws.accepted",
            requestId: "hello-1",
            scope: mindScope,
            payload: {
                authorizationVersion: "version-1",
                permissions: ["read"],
            },
        });
    } finally {
        await closeApp(app, socket);
    }
});

test("keeps Mind page authorization input separate from plugin binding", async () => {
    const credentials: Array<string | undefined> = [];
    const baseAuthorization = createMockAuthorizationReader([authorizationRecord]);
    const authorization = {
        authorize: async (request: OneTalkAuthorizationRequest) => {
            if (request.connectionType === "mind_page") credentials.push(request.cookie);
            return baseAuthorization.authorize(request);
        },
        readAuthorizationVersion: baseAuthorization.readAuthorizationVersion,
    };
    const app = createApp(testConfig, {
        database: createDatabaseStub(),
        authorization,
        oneTalkService: createServiceStub(),
    });
    const socket = await openSocket(app, "/ws/mind");

    try {
        const accepted = nextMessage(socket);
        socket.send(
            JSON.stringify({
                ...helloFrame(mindScope),
                connectionType: "mind_page",
                payload: { requestedPermissions: ["read"] },
            }),
        );
        assert.equal((await accepted).type, "ws.accepted");

        const heartbeat = nextMessage(socket);
        socket.send(
            JSON.stringify({
                protocolVersion: ONETALK_PROTOCOL_VERSION,
                connectionType: "mind_page",
                type: "heartbeat",
                requestId: "mind-heartbeat",
                scope: mindScope,
                payload: { sentAtMs: 1_700_000_000_000 },
            }),
        );
        assert.equal((await heartbeat).type, "heartbeat.ack");
        assert.deepEqual(credentials, [undefined, undefined]);
    } finally {
        await closeApp(app, socket);
    }
});

test("keeps direct WebSocket installation fail-closed by default", async () => {
    const app = Fastify({ logger: false });
    installWebsocket(app, {
        service: createServiceStub(),
        profileService: createProfileServiceStub(),
        mindPageOrigin: "http://mind.localhost",
        pluginOrigins: ["http://plugin.localhost"],
    });
    const socket = await openSocket(app);

    try {
        const errorMessage = nextMessage(socket);
        socket.send(JSON.stringify(helloFrame()));

        assert.deepEqual(await errorMessage, {
            protocolVersion: ONETALK_PROTOCOL_VERSION,
            connectionType: "plugin",
            type: "ws.error",
            requestId: "hello-1",
            scope: bindingScope,
            payload: { code: ONETALK_ERROR_CODES.authorizationUnavailable },
        });
    } finally {
        await closeApp(app, socket);
    }
});

test("does not treat an arbitrary Bearer header as a Mind WebSocket identity", async () => {
    const bearerToken = "retired-internal-read-token";
    const authorization = createMindAuthorizationReader({
        baseUrl: "https://mind.example.com",
        timeoutMs: 100,
        fetch: async () => {
            throw new Error("a missing Cookie must reject before upstream access");
        },
    });
    const app = createApp(testConfig, {
        database: createDatabaseStub(),
        authorization,
        oneTalkService: createServiceStub(),
    });
    await app.ready();
    const socket = await app.injectWS("/ws/mind", {
        headers: {
            origin: "http://mind.localhost",
            authorization: `Bearer ${bearerToken}`,
        },
    });

    try {
        const errorMessage = nextMessage(socket);
        socket.send(
            JSON.stringify({
                protocolVersion: ONETALK_PROTOCOL_VERSION,
                connectionType: "mind_page",
                type: "ws.hello",
                requestId: "bearer-token-hello",
                scope: mindScope,
                payload: { requestedPermissions: ["read", "send"] },
            }),
        );
        assert.deepEqual(await errorMessage, {
            protocolVersion: ONETALK_PROTOCOL_VERSION,
            connectionType: "mind_page",
            type: "ws.error",
            requestId: "bearer-token-hello",
            scope: mindScope,
            payload: { code: ONETALK_ERROR_CODES.authRequired },
        });
    } finally {
        await closeApp(app, socket);
    }
});

test("fails closed when the authorization dependency is unavailable", async () => {
    const app = createApp(testConfig, {
        database: createDatabaseStub(),
        oneTalkService: createServiceStub(),
    });
    const socket = await openSocket(app);

    try {
        const errorMessage = nextMessage(socket);
        socket.send(
            JSON.stringify({
                ...helloFrame(),
                payload: { binding: "secret-binding", requestedPermissions: ["read", "send"] },
            }),
        );
        const error = await errorMessage;
        assert.deepEqual(error, {
            protocolVersion: ONETALK_PROTOCOL_VERSION,
            connectionType: "plugin",
            type: "ws.error",
            requestId: "hello-1",
            scope: bindingScope,
            payload: { code: ONETALK_ERROR_CODES.authorizationUnavailable },
        });
        assert.equal(JSON.stringify(error).includes("secret-binding"), false);
    } finally {
        await closeApp(app, socket);
    }
});

test("rejects the legacy plugin credential field and mismatched binding", async () => {
    const authorization = createMockAuthorizationReader([authorizationRecord]);
    const app = createApp(testConfig, {
        database: createDatabaseStub(),
        authorization,
        oneTalkService: createServiceStub(),
    });
    const socket = await openSocket(app);

    try {
        const legacyError = nextMessage(socket);
        socket.send(
            JSON.stringify({
                ...helloFrame(),
                payload: { credential: "legacy-plugin-secret", requestedPermissions: ["read"] },
            }),
        );
        const error = await legacyError;
        assert.deepEqual(error.payload, { code: ONETALK_ERROR_CODES.invalidMessage });
        assert.equal(JSON.stringify(error).includes("legacy-plugin-secret"), false);
    } finally {
        await closeApp(app, socket);
    }

    const mismatchApp = createApp(testConfig, {
        database: createDatabaseStub(),
        authorization: createMockAuthorizationReader([authorizationRecord]),
        oneTalkService: createServiceStub(),
    });
    const mismatchSocket = await openSocket(mismatchApp);
    try {
        const errorMessage = nextMessage(mismatchSocket);
        mismatchSocket.send(
            JSON.stringify({
                ...helloFrame(),
                payload: { binding: "wrong-binding", requestedPermissions: ["read"] },
            }),
        );
        const error = await errorMessage;
        assert.deepEqual(error.payload, { code: ONETALK_ERROR_CODES.bindingRevoked });
        assert.equal(JSON.stringify(error).includes("wrong-binding"), false);
    } finally {
        await closeApp(mismatchApp, mismatchSocket);
    }
});

test("rejects a send operation when the authorized binding lacks send permission", async () => {
    const authorization = createMockAuthorizationReader([
        { ...authorizationRecord, permissions: ["read"] },
    ]);
    const app = createApp(testConfig, {
        database: createDatabaseStub(),
        authorization,
        oneTalkService: createServiceStub(),
    });
    const socket = await openSocket(app, "/ws/mind");

    try {
        const handshake = nextMessage(socket);
        socket.send(
            JSON.stringify({
                ...helloFrame(mindScope),
                connectionType: "mind_page",
                payload: { requestedPermissions: ["read"] },
            }),
        );
        assert.equal((await handshake).type, "ws.accepted");

        const errorMessage = nextMessage(socket);
        socket.send(
            JSON.stringify({
                protocolVersion: ONETALK_PROTOCOL_VERSION,
                connectionType: "mind_page",
                type: "send.request",
                requestId: "send-1",
                sendRequestId: "send-request-1",
                scope: mindScope,
                payload: { conversationId: "conversation-1", content: "hello" },
            }),
        );
        assert.deepEqual(await errorMessage, {
            protocolVersion: ONETALK_PROTOCOL_VERSION,
            connectionType: "mind_page",
            type: "send.result",
            requestId: "send-1",
            sendRequestId: "send-request-1",
            scope: mindScope,
            payload: { status: "rejected_before_send", reason: "waiting_for_page" },
        });
    } finally {
        await closeApp(app, socket);
    }
});

test("allows heartbeat for a send-only connection", async () => {
    const authorization = createMockAuthorizationReader([
        { ...authorizationRecord, permissions: ["send"] },
    ]);
    const app = createApp(testConfig, {
        database: createDatabaseStub(),
        authorization,
        oneTalkService: createServiceStub(),
    });
    const socket = await openSocket(app);

    try {
        const accepted = nextMessage(socket);
        socket.send(JSON.stringify(helloFrame()));
        assert.deepEqual((await accepted).payload, {
            authorizationVersion: "version-1",
            permissions: ["send"],
        });

        const heartbeat = nextMessage(socket);
        socket.send(
            JSON.stringify({
                protocolVersion: ONETALK_PROTOCOL_VERSION,
                connectionType: "plugin",
                type: "heartbeat",
                requestId: "heartbeat-send-only",
                scope: bindingScope,
                payload: { sentAtMs: 1_700_000_000_000 },
            }),
        );
        assert.equal((await heartbeat).type, "heartbeat.ack");
    } finally {
        await closeApp(app, socket);
    }
});

test("returns explicit protocol-upgrade rejection for an old frame", async () => {
    const authorization = createMockAuthorizationReader([authorizationRecord]);
    const app = createApp(testConfig, {
        database: createDatabaseStub(),
        authorization,
        oneTalkService: createServiceStub(),
    });
    const socket = await openSocket(app);

    try {
        const errorMessage = nextMessage(socket);
        const closeCode = nextCloseCode(socket);
        socket.send(JSON.stringify({ ...helloFrame(), protocolVersion: 0 }));
        assert.deepEqual(await errorMessage, {
            protocolVersion: ONETALK_PROTOCOL_VERSION,
            connectionType: "plugin",
            type: "ws.error",
            requestId: "hello-1",
            scope: bindingScope,
            payload: { code: ONETALK_ERROR_CODES.protocolUpgradeRequired },
        });
        assert.equal(await closeCode, 1003);
    } finally {
        await closeApp(app, socket);
    }
});

test("rejects a scope change after authentication", async () => {
    const authorization = createMockAuthorizationReader([authorizationRecord]);
    const app = createApp(testConfig, {
        database: createDatabaseStub(),
        authorization,
        oneTalkService: createServiceStub(),
    });
    const socket = await openSocket(app);

    try {
        const handshake = nextMessages(socket, 2);
        socket.send(JSON.stringify(helloFrame()));
        const [, snapshot] = await handshake;
        assertAnchorSnapshot(snapshot);

        const errorMessage = nextMessage(socket);
        const changedScope = { ...bindingScope, deviceId: "other-device" };
        socket.send(
            JSON.stringify({
                protocolVersion: ONETALK_PROTOCOL_VERSION,
                connectionType: "plugin",
                type: "heartbeat",
                requestId: "heartbeat-2",
                scope: changedScope,
                payload: { sentAtMs: 1_700_000_000_000 },
            }),
        );
        assert.deepEqual(await errorMessage, {
            protocolVersion: ONETALK_PROTOCOL_VERSION,
            connectionType: "plugin",
            type: "ws.error",
            requestId: "heartbeat-2",
            scope: changedScope,
            payload: { code: ONETALK_ERROR_CODES.scopeMismatch },
        });
    } finally {
        await closeApp(app, socket);
    }
});

test("closes a connection when the authorization version changes", async () => {
    const authorization = createMockAuthorizationReader([authorizationRecord]);
    const app = createApp(testConfig, {
        database: createDatabaseStub(),
        authorization,
        oneTalkService: createServiceStub(),
    });
    const socket = await openSocket(app);

    try {
        const handshake = nextMessages(socket, 2);
        socket.send(JSON.stringify(helloFrame()));
        const [, snapshot] = await handshake;
        assertAnchorSnapshot(snapshot);
        authorization.upsert({ ...authorizationRecord, authorizationVersion: "version-2" });

        const errorMessage = nextMessage(socket);
        socket.send(
            JSON.stringify({
                protocolVersion: ONETALK_PROTOCOL_VERSION,
                connectionType: "plugin",
                type: "heartbeat",
                requestId: "heartbeat-3",
                scope: bindingScope,
                payload: { sentAtMs: 1_700_000_000_000 },
            }),
        );
        assert.deepEqual(await errorMessage, {
            protocolVersion: ONETALK_PROTOCOL_VERSION,
            connectionType: "plugin",
            type: "ws.error",
            requestId: "heartbeat-3",
            scope: bindingScope,
            payload: { code: ONETALK_ERROR_CODES.authorizationVersionChanged },
        });
    } finally {
        await closeApp(app, socket);
    }
});