/** 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) => 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>>(); /** userId → when their last socket closed (removed again when one opens). */ private offlineAt = new Map(); 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, 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, msg: ServerMessage): void { if (ws.getBufferedAmount() > MAX_BUFFERED) { ws.close(1013, 'slow consumer'); return; } ws.send(JSON.stringify(msg)); } private connFor(ws: ServerWebSocket): 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) => { 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, 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) => { 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, 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'); } } } }