psxbase/server/src/ws/signaling.test.ts

import { createServer, type Server } from "node:http";
import { afterEach, beforeEach, describe, expect, it } from "vitest";
import { WebSocket } from "ws";
import type { SignalMessage } from "@psx/shared";
import { PROTOCOL_VERSION } from "@psx/shared";
import { createApp } from "../app.js";
import { InMemoryStore } from "../db/store.js";
import { SignalingServer } from "./signaling.js";

/**
 * End-to-end signalling tests against a real HTTP + WebSocket server.
 *
 * The impersonation test is the important one: the relay stamps `from` off the
 * authenticated socket, so a peer claiming to be someone else must be ignored.
 */

let server: Server;
let signaling: SignalingServer;
let store: InMemoryStore;
let baseUrl: string;
let wsUrl: string;

/** Bridges Hono's fetch handler onto a plain Node server so `ws` can attach. */
function nodeAdapter(app: ReturnType<typeof createApp>["app"]): Server {
  return createServer((req, res) => {
    const chunks: Buffer[] = [];
    req.on("data", (chunk: Buffer) => chunks.push(chunk));
    req.on("end", () => {
      const url = `http://localhost${req.url ?? "/"}`;
      const hasBody = chunks.length > 0;
      const request = new Request(url, {
        method: req.method,
        headers: req.headers as Record<string, string>,
        ...(hasBody ? { body: Buffer.concat(chunks) } : {}),
      });

      // Hono's `fetch` is sync-or-async depending on the route.
      void Promise.resolve(app.fetch(request)).then(async (response: Response) => {
        res.writeHead(response.status, Object.fromEntries(response.headers.entries()));
        res.end(Buffer.from(await response.arrayBuffer()));
      });
    });
  });
}

beforeEach(async () => {
  store = new InMemoryStore();
  const { app } = createApp(store);
  server = nodeAdapter(app);
  signaling = new SignalingServer(server, store);

  await new Promise<void>((resolve) => server.listen(0, resolve));
  const address = server.address();
  const port = typeof address === "object" && address ? address.port : 0;
  baseUrl = `http://localhost:${port}`;
  wsUrl = `ws://localhost:${port}/ws`;
});

afterEach(async () => {
  signaling.close();
  await new Promise<void>((resolve) => server.close(() => resolve()));
});

async function register(username: string): Promise<string> {
  const response = await fetch(`${baseUrl}/auth/register`, {
    method: "POST",
    headers: { "content-type": "application/json" },
    body: JSON.stringify({ username, email: `${username}@example.com`, password: "hunter2hunter2" }),
  });
  expect(response.status).toBe(201);
  return ((await response.json()) as { token: string }).token;
}

/** Opens a socket, sends `hello`, and resolves once `welcome` arrives. */
async function connect(token: string, lobbyCode: string) {
  const socket = new WebSocket(wsUrl);
  const inbox: SignalMessage[] = [];

  await new Promise<void>((resolve, reject) => {
    socket.on("open", resolve);
    socket.on("error", reject);
  });
  socket.on("message", (raw) => inbox.push(JSON.parse(raw.toString()) as SignalMessage));
  socket.send(JSON.stringify({ t: "hello", token, lobbyCode, protocol: PROTOCOL_VERSION }));

  const welcome = await waitFor(inbox, (m) => m.t === "welcome" || m.t === "error");
  return { socket, inbox, welcome };
}

/** Polls `inbox` until a matching message appears, or times out. */
async function waitFor(
  inbox: SignalMessage[],
  predicate: (message: SignalMessage) => boolean,
  timeoutMs = 3000,
): Promise<SignalMessage> {
  const deadline = Date.now() + timeoutMs;
  while (Date.now() < deadline) {
    const found = inbox.find(predicate);
    if (found) return found;
    await new Promise((resolve) => setTimeout(resolve, 10));
  }
  throw new Error(`timed out; inbox was ${JSON.stringify(inbox)}`);
}

describe("signalling", () => {
  it("rejects a socket with an invalid token", async () => {
    const hostToken = await register("host");
    const created = await fetch(`${baseUrl}/lobbies`, {
      method: "POST",
      headers: { authorization: `Bearer ${hostToken}`, "content-type": "application/json" },
      body: "{}",
    });
    const { code } = (await created.json()) as { code: string };

    const { welcome, socket } = await connect("not-a-real-token", code);
    expect(welcome.t).toBe("error");
    socket.close();
  });

  it("refuses a peer that has not joined the lobby over REST", async () => {
    const hostToken = await register("host");
    const strangerToken = await register("stranger");

    const created = await fetch(`${baseUrl}/lobbies`, {
      method: "POST",
      headers: { authorization: `Bearer ${hostToken}`, "content-type": "application/json" },
      body: "{}",
    });
    const { code } = (await created.json()) as { code: string };

    const { welcome, socket } = await connect(strangerToken, code);
    expect(welcome.t).toBe("error");
    socket.close();
  });

  it("welcomes host as peer 0 and announces a guest joining", async () => {
    const hostToken = await register("host");
    const guestToken = await register("guest");

    const created = await fetch(`${baseUrl}/lobbies`, {
      method: "POST",
      headers: { authorization: `Bearer ${hostToken}`, "content-type": "application/json" },
      body: "{}",
    });
    const { code } = (await created.json()) as { code: string };

    await fetch(`${baseUrl}/lobbies/${code}/join`, {
      method: "POST",
      headers: { authorization: `Bearer ${guestToken}` },
    });

    const host = await connect(hostToken, code);
    expect(host.welcome).toMatchObject({ t: "welcome", peerId: 0, isHost: true });

    const guest = await connect(guestToken, code);
    expect(guest.welcome).toMatchObject({ t: "welcome", peerId: 1, isHost: false });

    const announced = await waitFor(host.inbox, (m) => m.t === "peer-joined");
    expect(announced).toMatchObject({ t: "peer-joined", peer: { peerId: 1 } });

    host.socket.close();
    guest.socket.close();
  });

  it("relays SDP between peers and stamps the true sender", async () => {
    const hostToken = await register("host");
    const guestToken = await register("guest");

    const created = await fetch(`${baseUrl}/lobbies`, {
      method: "POST",
      headers: { authorization: `Bearer ${hostToken}`, "content-type": "application/json" },
      body: "{}",
    });
    const { code } = (await created.json()) as { code: string };
    await fetch(`${baseUrl}/lobbies/${code}/join`, {
      method: "POST",
      headers: { authorization: `Bearer ${guestToken}` },
    });

    const host = await connect(hostToken, code);
    const guest = await connect(guestToken, code);

    host.socket.send(
      JSON.stringify({
        t: "signal",
        // Lie about the sender; the relay must overwrite this with peer 0.
        from: 99,
        to: 1,
        payload: { sdp: { type: "offer", sdp: "v=0 fake" } },
      }),
    );

    const relayed = await waitFor(guest.inbox, (m) => m.t === "signal");
    expect(relayed).toMatchObject({
      t: "signal",
      from: 0,
      to: 1,
      payload: { sdp: { type: "offer", sdp: "v=0 fake" } },
    });

    host.socket.close();
    guest.socket.close();
  });

  it("rejects a protocol version mismatch", async () => {
    const hostToken = await register("host");
    const created = await fetch(`${baseUrl}/lobbies`, {
      method: "POST",
      headers: { authorization: `Bearer ${hostToken}`, "content-type": "application/json" },
      body: "{}",
    });
    const { code } = (await created.json()) as { code: string };

    const socket = new WebSocket(wsUrl);
    const inbox: SignalMessage[] = [];
    await new Promise<void>((resolve) => socket.on("open", () => resolve()));
    socket.on("message", (raw) => inbox.push(JSON.parse(raw.toString()) as SignalMessage));
    socket.send(JSON.stringify({ t: "hello", token: hostToken, lobbyCode: code, protocol: 999 }));

    const error = await waitFor(inbox, (m) => m.t === "error");
    expect(error).toMatchObject({ t: "error" });
    expect((error as { message: string }).message).toContain("protocol mismatch");
    socket.close();
  });

  it("tells remaining peers when the host disconnects", async () => {
    const hostToken = await register("host");
    const guestToken = await register("guest");

    const created = await fetch(`${baseUrl}/lobbies`, {
      method: "POST",
      headers: { authorization: `Bearer ${hostToken}`, "content-type": "application/json" },
      body: "{}",
    });
    const { code } = (await created.json()) as { code: string };
    await fetch(`${baseUrl}/lobbies/${code}/join`, {
      method: "POST",
      headers: { authorization: `Bearer ${guestToken}` },
    });

    const host = await connect(hostToken, code);
    const guest = await connect(guestToken, code);

    host.socket.close();

    const left = await waitFor(guest.inbox, (m) => m.t === "peer-left");
    expect(left).toMatchObject({ t: "peer-left", peerId: 0 });

    guest.socket.close();
  });
});