joke-app/apps/server/src/ws/gateway.ts

/** WebSocket gateway: auth on upgrade, validation, rate limits, dispatch (docs/03). */
import type { Server, ServerWebSocket } from 'bun';
import { ClientMessage, MAX_CLIENT_MESSAGE_BYTES, PROTOCOL, type PublicUser, type ServerMessage } from '@bg/shared';
import type { RoomManager } from '../rooms/manager';
import type { Conn, Room } from '../rooms/room';
import { TokenBucket } from './rate-limit';

export interface WsData {
  userId: string;
  connId: number;
  bucket: TokenBucket;
  strikes: number;
  room: string | null;
  conn: Conn | null;
}

export interface GatewayDeps {
  rooms: RoomManager;
  me(userId: string): PublicUser | null;
  now?: () => number;
  log?: (level: 'info' | 'warn' | 'error', msg: string, extra?: Record<string, unknown>) => void;
  /** A socket closed (the user may have left the site). */
  onClose?: (userId: string) => void;
}

export const MAX_CONNS_PER_USER = 5;
const MAX_BUFFERED = 1024 * 1024;

export class Gateway {
  private nextId = 1;
  private byUser = new Map<string, Set<ServerWebSocket<WsData>>>();
  /** userId → when their last socket closed (removed again when one opens). */
  private offlineAt = new Map<string, number>();
  private closing = false;

  constructor(private deps: GatewayDeps) {}

  get connectionCount(): number {
    let n = 0;
    for (const s of this.byUser.values()) n += s.size;
    return n;
  }

  /** Called from fetch() after the session was resolved. Returns a Response on refusal. */
  upgrade(req: Request, server: Server<WsData>, userId: string): Response | undefined {
    if (this.closing) return new Response('restarting', { status: 503 });
    if ((this.byUser.get(userId)?.size ?? 0) >= MAX_CONNS_PER_USER) return new Response('too many connections', { status: 429 });
    const data: WsData = { userId, connId: this.nextId++, bucket: new TokenBucket(20, 40), strikes: 0, room: null, conn: null };
    if (server.upgrade(req, { data })) return undefined;
    return new Response('upgrade failed', { status: 400 });
  }

  private send(ws: ServerWebSocket<WsData>, msg: ServerMessage): void {
    if (ws.getBufferedAmount() > MAX_BUFFERED) {
      ws.close(1013, 'slow consumer');
      return;
    }
    ws.send(JSON.stringify(msg));
  }

  private connFor(ws: ServerWebSocket<WsData>): Conn {
    ws.data.conn ??= {
      userId: ws.data.userId,
      send: (m) => this.send(ws, m),
      close: (code, reason) => ws.close(code, reason),
    };
    return ws.data.conn;
  }

  readonly handlers = {
    open: (ws: ServerWebSocket<WsData>) => {
      const set = this.byUser.get(ws.data.userId) ?? new Set();
      set.add(ws);
      this.byUser.set(ws.data.userId, set);
      this.offlineAt.delete(ws.data.userId);
      const me = this.deps.me(ws.data.userId);
      if (!me) {
        ws.close(4401, 'unauthorized');
        return;
      }
      this.send(ws, {
        t: 'welcome',
        me,
        serverTime: (this.deps.now ?? Date.now)(),
        protocol: PROTOCOL,
        activeRoom: this.deps.rooms.activeRoomFor(ws.data.userId),
      });
    },

    message: (ws: ServerWebSocket<WsData>, raw: string | Buffer) => {
      const size = typeof raw === 'string' ? Buffer.byteLength(raw) : raw.byteLength;
      if (size > MAX_CLIENT_MESSAGE_BYTES) {
        ws.close(1009, 'message too big');
        return;
      }
      if (!ws.data.bucket.take()) {
        if (++ws.data.strikes > 50) ws.close(4429, 'rate limited');
        return;
      }
      let parsed: ClientMessage;
      try {
        const json = JSON.parse(typeof raw === 'string' ? raw : raw.toString());
        const r = ClientMessage.safeParse(json);
        if (!r.success) {
          this.send(ws, { t: 'error', code: 'bad-message', message: '잘못된 요청이에요.' });
          return;
        }
        parsed = r.data;
      } catch {
        this.send(ws, { t: 'error', code: 'bad-json', message: '잘못된 요청이에요.' });
        return;
      }
      try {
        this.dispatch(ws, parsed);
      } catch (err) {
        this.deps.log?.('error', 'dispatch failed', { user: ws.data.userId, t: parsed.t, err: String(err), stack: (err as Error).stack });
        this.send(ws, { t: 'error', code: 'internal', message: '알 수 없는 오류가 발생했어요.' });
      }
    },

    close: (ws: ServerWebSocket<WsData>) => {
      this.byUser.get(ws.data.userId)?.delete(ws);
      if (this.byUser.get(ws.data.userId)?.size === 0) {
        this.byUser.delete(ws.data.userId);
        this.offlineAt.set(ws.data.userId, (this.deps.now ?? Date.now)());
      }
      const room = ws.data.room ? this.deps.rooms.get(ws.data.room) : undefined;
      if (room && ws.data.conn) room.disconnect(ws.data.conn);
      this.deps.onClose?.(ws.data.userId);
    },
  };

  private dispatch(ws: ServerWebSocket<WsData>, m: ClientMessage): void {
    const u = ws.data.userId;
    const err = (message: string, code = 'rejected') => this.send(ws, { t: 'error', code, message });
    if (m.t === 'ping') {
      this.send(ws, { t: 'pong', ts: m.ts, serverTime: (this.deps.now ?? Date.now)() });
      return;
    }
    if (m.t === 'join') {
      const room = this.deps.rooms.get(m.code);
      if (!room) return err('방을 찾을 수 없어요. 코드를 다시 확인해 주세요.', 'room-not-found');
      // Leaving a previous room on this socket.
      if (ws.data.room && ws.data.room !== m.code) {
        const prev = this.deps.rooms.get(ws.data.room);
        if (prev && ws.data.conn) prev.disconnect(ws.data.conn);
      }
      if (m.as === 'player') this.deps.rooms.releaseOtherLobbySeats(u, m.code);
      const e = room.join(this.connFor(ws), m.as);
      if (e) return err(e, 'join-failed');
      ws.data.room = m.code;
      return;
    }
    const room: Room | undefined = ws.data.room ? this.deps.rooms.get(ws.data.room) : undefined;
    if (m.t === 'detach') {
      // ws.data.room stays: when this socket closes, the room still learns that the person went offline.
      if (room && ws.data.conn) room.detach(ws.data.conn);
      return;
    }
    if (!room || room.conns.get(u) !== ws.data.conn) return err('먼저 방에 들어가 주세요.', 'not-in-room');
    let e: string | null = null;
    switch (m.t) {
      case 'leave':
        room.leave(u);
        ws.data.room = null;
        return;
      case 'seat':
        e = room.seat(u, m.seat);
        break;
      case 'unseat':
        e = room.unseat(u);
        break;
      case 'ready':
        e = room.ready(u, m.ready);
        break;
      case 'config': {
        const { t: _t, ...patch } = m;
        e = room.configure(u, patch);
        break;
      }
      case 'start':
        e = room.start(u);
        break;
      case 'kick':
        e = room.kick(u, m.userId);
        break;
      case 'host':
        e = room.transferHost(u, m.userId);
        break;
      case 'addBot':
        e = room.addBot(u, m.seat);
        break;
      case 'removeBot':
        e = room.removeBot(u, m.userId);
        break;
      case 'act':
        room.act(u, m.cs, m.a);
        return;
      case 'chat':
        e = room.chat(u, m.text);
        break;
      case 'emote':
        room.emote(u, m.id);
        return;
      case 'rematch':
        e = room.rematch(u);
        break;
      case 'void':
        e = room.voidGame(u);
        break;
      case 'sync':
        room.sendSync(this.connFor(ws));
        return;
    }
    if (e) err(e);
  }

  /** Drop every connection of a user (ban / forced logout). */
  disconnectUser(userId: string): void {
    for (const ws of [...(this.byUser.get(userId) ?? [])]) {
      this.send(ws, { t: 'bye', reason: 'kicked' });
      ws.close(4403, 'banned');
    }
  }

  isOnline(userId: string): boolean {
    return (this.byUser.get(userId)?.size ?? 0) > 0;
  }

  /** When the user's last socket closed; null while one is open, 0 if not seen since the server started. */
  offlineSince(userId: string): number | null {
    return this.isOnline(userId) ? null : (this.offlineAt.get(userId) ?? 0);
  }

  /** The account is gone (guest logged out or expired): close its sockets quietly. */
  closeUser(userId: string): void {
    for (const ws of [...(this.byUser.get(userId) ?? [])]) {
      this.send(ws, { t: 'bye', reason: 'closed' });
      ws.close(4401, 'logged out');
    }
  }

  /** Graceful shutdown: tell everyone to reconnect shortly (docs/03 §10). */
  shutdown(): void {
    this.closing = true;
    for (const set of this.byUser.values()) {
      for (const ws of set) {
        this.send(ws, { t: 'bye', reason: 'restart' });
        ws.close(1012, 'restart');
      }
    }
  }
}