psxbase/server/src/ws/signaling.ts

import type { Server } from "node:http";
import { WebSocketServer, type WebSocket } from "ws";
import type { PeerInfo, SignalMessage, SignalPayload } from "@psx/shared";
import { validateSignalMessage } from "@psx/shared";
import { verifyToken } from "../auth.js";
import type { LobbyRecord, Store } from "../db/store.js";

/**
 * SDP/ICE exchange — GDD §8.
 *
 * The server is a relay and nothing more. It never sees gameplay: once peers
 * have exchanged descriptors they talk over data channels directly, and this
 * socket carries only membership changes for the rest of the session.
 *
 * Two rules make the relay safe to expose:
 *   1. `from` on a relayed signal is stamped from the authenticated socket,
 *      never read from the payload, so peers cannot impersonate one another.
 *   2. A socket may only address peers inside the lobby it authenticated into.
 */

interface Connection {
  socket: WebSocket;
  peerId: number;
  userId: string;
  lobbyCode: string;
  isHost: boolean;
  alive: boolean;
}

const HEARTBEAT_INTERVAL_MS = 30_000;

export class SignalingServer {
  private readonly wss: WebSocketServer;
  /** lobby code → peer id → connection */
  private readonly rooms = new Map<string, Map<number, Connection>>();
  private readonly heartbeat: NodeJS.Timeout;

  constructor(
    server: Server,
    private readonly store: Store,
    path = "/ws",
  ) {
    this.wss = new WebSocketServer({ server, path });
    this.wss.on("connection", (socket) => this.onConnection(socket));

    // Half-open TCP connections would otherwise linger as phantom lobby members.
    this.heartbeat = setInterval(() => {
      for (const room of this.rooms.values()) {
        for (const connection of room.values()) {
          if (!connection.alive) {
            connection.socket.terminate();
            continue;
          }
          connection.alive = false;
          connection.socket.ping();
        }
      }
    }, HEARTBEAT_INTERVAL_MS);
  }

  private onConnection(socket: WebSocket): void {
    let connection: Connection | null = null;

    socket.on("pong", () => {
      if (connection) connection.alive = true;
    });

    socket.on("message", async (raw) => {
      let parsed: unknown;
      try {
        parsed = JSON.parse(raw.toString());
      } catch {
        return send(socket, { t: "error", message: "malformed JSON" });
      }

      const validation = validateSignalMessage(parsed);
      if (!validation.ok) {
        return send(socket, { t: "error", message: validation.errors.join("; ") });
      }
      const message = validation.value;

      // Everything except `hello` requires an authenticated socket.
      if (!connection && message.t !== "hello") {
        return send(socket, { t: "error", message: "expected hello first" });
      }

      switch (message.t) {
        case "hello": {
          if (connection) return send(socket, { t: "error", message: "already authenticated" });
          connection = await this.authenticate(socket, message.token, message.lobbyCode);
          break;
        }
        case "signal":
          this.relaySignal(connection!, message.to, message.payload);
          break;
        case "ready":
          await this.setReady(connection!, message.isReady);
          break;
        default:
          send(socket, { t: "error", message: `unsupported message: ${message.t}` });
      }
    });

    socket.on("close", () => {
      if (connection) void this.disconnect(connection);
    });
  }

  private async authenticate(
    socket: WebSocket,
    token: string,
    lobbyCode: string,
  ): Promise<Connection | null> {
    const verified = verifyToken(token);
    if (!verified) {
      send(socket, { t: "error", message: "invalid or expired token" });
      socket.close();
      return null;
    }

    const lobby = await this.store.findLobbyByCode(lobbyCode);
    if (!lobby || lobby.status === "closed") {
      send(socket, { t: "error", message: "lobby not found or closed" });
      socket.close();
      return null;
    }

    // Membership is established over REST (`POST /lobbies/:code/join`) so that
    // capacity and duplicate-join rules live in exactly one place.
    const member = lobby.players.find((player) => player.userId === verified.userId);
    if (!member) {
      send(socket, { t: "error", message: "join the lobby before connecting" });
      socket.close();
      return null;
    }

    const room = this.rooms.get(lobby.code) ?? new Map<number, Connection>();
    this.rooms.set(lobby.code, room);

    // A reconnect from the same peer supersedes the stale socket.
    room.get(member.peerId)?.socket.terminate();

    const connection: Connection = {
      socket,
      peerId: member.peerId,
      userId: verified.userId,
      lobbyCode: lobby.code,
      isHost: verified.userId === lobby.hostId,
      alive: true,
    };
    room.set(member.peerId, connection);

    send(socket, {
      t: "welcome",
      peerId: connection.peerId,
      isHost: connection.isHost,
      peers: this.peerList(lobby, room),
    });

    this.broadcast(
      lobby.code,
      { t: "peer-joined", peer: this.peerInfo(lobby, member.peerId) },
      connection.peerId,
    );

    return connection;
  }

  private relaySignal(from: Connection, to: number, payload: SignalPayload): void {
    const target = this.rooms.get(from.lobbyCode)?.get(to);
    if (!target) return;

    send(target.socket, {
      t: "signal",
      // Stamped from the authenticated socket — never trusted from the payload.
      from: from.peerId,
      to,
      payload,
    });
  }

  private async setReady(connection: Connection, isReady: boolean): Promise<void> {
    const lobby = await this.store.setReady(connection.lobbyCode, connection.userId, isReady);
    if (!lobby) return;
    this.broadcast(connection.lobbyCode, {
      t: "peer-updated",
      peer: this.peerInfo(lobby, connection.peerId),
    });
  }

  private async disconnect(connection: Connection): Promise<void> {
    const room = this.rooms.get(connection.lobbyCode);
    // Guard against a superseded socket removing its replacement.
    if (room?.get(connection.peerId)?.socket === connection.socket) {
      room.delete(connection.peerId);
    }
    if (room && room.size === 0) this.rooms.delete(connection.lobbyCode);

    await this.store.leaveLobby(connection.lobbyCode, connection.userId);
    this.broadcast(connection.lobbyCode, { t: "peer-left", peerId: connection.peerId });

    // The authoritative simulation lived in the host's browser, so its departure
    // ends the session for everyone rather than silently freezing the world.
    if (connection.isHost) {
      this.broadcast(connection.lobbyCode, { t: "error", message: "the host left the session" });
      for (const remaining of room?.values() ?? []) remaining.socket.close();
      this.rooms.delete(connection.lobbyCode);
    }
  }

  private peerInfo(lobby: LobbyRecord, peerId: number): PeerInfo {
    const player = lobby.players.find((entry) => entry.peerId === peerId);
    return {
      peerId,
      displayName: player?.displayName ?? `Player ${peerId}`,
      isHost: player ? player.userId === lobby.hostId : false,
      isReady: player?.isReady ?? false,
    };
  }

  private peerList(lobby: LobbyRecord, room: Map<number, Connection>): PeerInfo[] {
    return [...room.keys()].map((peerId) => this.peerInfo(lobby, peerId));
  }

  private broadcast(lobbyCode: string, message: SignalMessage, exceptPeerId?: number): void {
    for (const [peerId, connection] of this.rooms.get(lobbyCode) ?? []) {
      if (peerId === exceptPeerId) continue;
      send(connection.socket, message);
    }
  }

  close(): void {
    clearInterval(this.heartbeat);
    this.wss.close();
  }
}

function send(socket: WebSocket, message: SignalMessage): void {
  if (socket.readyState === socket.OPEN) socket.send(JSON.stringify(message));
}