mancala/apps/backend/tests/rtmt/rtmt_websocket.test.ts

import * as http from "http";
import { AddressInfo } from "net";
import WebSocket from "ws";
import { AuthUser, InvalidTokenError, TokenVerifier } from "../../src/auth";
import { channel_auth, channel_on_auth, channel_on_auth_error, CLOSE_AUTH_FAILED } from "../../src/auth/channels";
import { RTMTWS } from "../../src/rtmt/rtmt_websocket";

const LEGACY_KEY = "0b9f6d1e-7c1a-4f7e-9a51-3f0c2b6d8e11";

// Accepts tokens of the form "valid:<id>".
const fakeVerifier: TokenVerifier = {
    async verify(token: string): Promise<AuthUser> {
        if (!token.startsWith("valid:")) throw new InvalidTokenError("bad token");
        const id = token.slice("valid:".length);
        return { id, name: `name-of-${id}`, isGuest: false };
    },
};

type Received = { channel: string, message: any };

class TestClient {
    readonly messages: Received[] = [];
    readonly closed: Promise<number>;
    constructor(readonly ws: WebSocket) {
        ws.on("message", (data) => this.messages.push(JSON.parse(data.toString())));
        this.closed = new Promise((resolve) => ws.on("close", (code) => resolve(code)));
    }
    send(channel: string, message: any) {
        this.ws.send(JSON.stringify({ channel, message }));
    }
    async waitFor(channel: string): Promise<any> {
        for (let i = 0; i < 100; i++) {
            const found = this.messages.find((m) => m.channel === channel);
            if (found) return found.message;
            await sleep(20);
        }
        throw new Error(`no message on ${channel}, got ${JSON.stringify(this.messages)}`);
    }
}

const sleep = (ms: number) => new Promise((resolve) => setTimeout(resolve, ms));

describe("RTMTWS", () => {
    let server: http.Server;
    let rtmt: RTMTWS;
    let clients: TestClient[];
    let connectionChanges: [string, boolean][];

    const start = async (legacyUserKeyAuth: boolean) => {
        server = http.createServer();
        rtmt = new RTMTWS({ tokenVerifier: fakeVerifier, legacyUserKeyAuth, authTimeoutMs: 300 });
        const changes: [string, boolean][] = [];
        connectionChanges = changes;
        rtmt.listenOnClientConnectionChange((id, isOnline) => changes.push([id, isOnline]));
        rtmt.initWebSocket(server);
        await new Promise<void>((resolve) => server.listen(0, resolve));
    };

    const connect = async (query = ""): Promise<TestClient> => {
        const port = (server.address() as AddressInfo).port;
        const ws = new WebSocket(`ws://localhost:${port}/${query}`);
        const client = new TestClient(ws);
        clients.push(client);
        await new Promise((resolve, reject) => { ws.on("open", resolve); ws.on("error", reject); });
        return client;
    };

    const connectAs = async (id: string): Promise<TestClient> => {
        const client = await connect();
        client.send(channel_auth, { token: `valid:${id}` });
        await client.waitFor(channel_on_auth);
        return client;
    };

    beforeEach(() => { clients = []; });

    afterEach(async () => {
        clients.forEach((c) => c.ws.terminate());
        rtmt.dispose();
        await new Promise((resolve) => server.close(resolve));
    });

    describe("token auth", () => {
        beforeEach(() => start(false));

        it("authenticates a socket and routes its messages as that user", async () => {
            const received: [string, any][] = [];
            rtmt.listenMessage("hello", (id, message) => received.push([id, message]));

            const client = await connect();
            client.send(channel_auth, { token: "valid:alice" });
            expect(await client.waitFor(channel_on_auth)).toEqual({ user: { id: "alice", name: "name-of-alice", isGuest: false } });

            client.send("hello", { n: 1 });
            await sleep(50);
            expect(received).toEqual([["alice", { n: 1 }]]);
            expect(rtmt.isClientOnline("alice")).toBe(true);
            expect(rtmt.getUser("alice")?.name).toBe("name-of-alice");
        });

        it("ignores messages before authentication", async () => {
            const received: string[] = [];
            rtmt.listenMessage("hello", (id) => received.push(id));
            const client = await connect();
            client.send("hello", {});
            await sleep(50);
            expect(received).toEqual([]);
        });

        it("closes the socket on an invalid token", async () => {
            const client = await connect();
            client.send(channel_auth, { token: "forged" });
            expect(await client.waitFor(channel_on_auth_error)).toEqual({ reason: "invalid token" });
            expect(await client.closed).toBe(CLOSE_AUTH_FAILED);
        });

        it("closes the socket when no token arrives in time", async () => {
            const client = await connect();
            expect(await client.closed).toBe(CLOSE_AUTH_FAILED);
        });

        it("accepts a refreshed token for the same user but not for another user", async () => {
            const client = await connectAs("alice");
            client.messages.length = 0;
            client.send(channel_auth, { token: "valid:alice" });
            await client.waitFor(channel_on_auth);

            client.send(channel_auth, { token: "valid:mallory" });
            expect(await client.waitFor(channel_on_auth_error)).toEqual({ reason: "token belongs to another user" });
            expect(await client.closed).toBe(CLOSE_AUTH_FAILED);
        });

        it("rejects userKey connections when legacy auth is off", async () => {
            const client = await connect(`?userKey=${LEGACY_KEY}`);
            expect(await client.closed).toBe(CLOSE_AUTH_FAILED);
        });

        it("survives malformed messages", async () => {
            const client = await connectAs("alice");
            client.ws.send("{not json");
            client.ws.send("null");
            client.messages.length = 0;
            client.send(channel_auth, { token: "valid:alice" });
            await client.waitFor(channel_on_auth);
        });
    });

    describe("several sockets per user", () => {
        beforeEach(() => start(false));

        it("delivers messages to every socket of the user", async () => {
            const phone = await connectAs("alice");
            const laptop = await connectAs("alice");
            rtmt.sendMessage("alice", "news", { n: 1 });
            expect(await phone.waitFor("news")).toEqual({ n: 1 });
            expect(await laptop.waitFor("news")).toEqual({ n: 1 });
        });

        it("keeps the user online until the last socket closes", async () => {
            const phone = await connectAs("alice");
            const laptop = await connectAs("alice");

            phone.ws.close();
            await phone.closed;
            await sleep(50);
            expect(rtmt.isClientOnline("alice")).toBe(true);
            expect(connectionChanges).not.toContainEqual(["alice", false]);

            laptop.ws.close();
            await laptop.closed;
            await sleep(50);
            expect(rtmt.isClientOnline("alice")).toBe(false);
            expect(connectionChanges).toContainEqual(["alice", false]);
        });

        it("does not drop a reconnected socket when the old one closes late", async () => {
            const oldSocket = await connectAs("alice");
            const newSocket = await connectAs("alice");
            oldSocket.ws.terminate();
            await oldSocket.closed;
            await sleep(50);

            rtmt.sendMessage("alice", "news", {});
            expect(await newSocket.waitFor("news")).toEqual({});
        });
    });

    describe("legacy userKey auth", () => {
        beforeEach(() => start(true));

        it("accepts a UUID userKey", async () => {
            await connect(`?userKey=${LEGACY_KEY}`);
            await sleep(50);
            expect(rtmt.isClientOnline(LEGACY_KEY)).toBe(true);
            expect(rtmt.getUser(LEGACY_KEY)).toEqual({ id: LEGACY_KEY, isGuest: true });
        });

        it("cannot claim a Keycloak user's id", async () => {
            const client = await connect("?userKey=kc:0b9f6d1e-7c1a-4f7e-9a51-3f0c2b6d8e11");
            expect(await client.closed).toBe(CLOSE_AUTH_FAILED);
        });

        it("cannot switch a legacy socket to a token identity", async () => {
            const client = await connect(`?userKey=${LEGACY_KEY}`);
            client.send(channel_auth, { token: "valid:kc:alice" });
            expect(await client.closed).toBe(CLOSE_AUTH_FAILED);
        });
    });
});