import { ConnectedSocket, MessageBody, OnGatewayConnection, OnGatewayDisconnect, SubscribeMessage, WebSocketGateway, } from '@nestjs/websockets'; import { Logger } from '@nestjs/common'; import { IncomingMessage } from 'http'; import { WebSocket } from 'ws'; import { TokenVerificationService } from '../auth/token-verification.service'; import { ChatCaller, ChatService } from './chat.service'; /// Raw `ws` gateway (no socket.io rooms available), so channel membership is /// tracked manually per connected socket. Auth happens once at handshake via /// a `?token=` query param since passport guards don't run for WS upgrades. /// /// The per-socket caller is stored as a *promise*: the token check is async /// and a client can send `chat:join` before it resolves, so handlers await /// the stored promise instead of assuming it's already populated. @WebSocketGateway({ path: '/chat' }) export class ChatGateway implements OnGatewayConnection, OnGatewayDisconnect { private readonly logger = new Logger(ChatGateway.name); private readonly callers = new Map>(); private readonly rooms = new Map>(); constructor( private readonly tokenVerification: TokenVerificationService, private readonly chat: ChatService, ) {} handleConnection(client: WebSocket, request: IncomingMessage) { const token = new URL(request.url ?? '', 'http://localhost').searchParams.get('token'); if (!token) { client.close(4001, 'Missing token'); return; } const pending = this.tokenVerification.verifyEither(token).catch((err) => { this.logger.warn(`WS auth failed: ${(err as Error).message}`); client.close(4001, 'Unauthorized'); throw err; }); this.callers.set(client, pending); } handleDisconnect(client: WebSocket) { this.callers.delete(client); for (const members of this.rooms.values()) { members.delete(client); } } @SubscribeMessage('chat:join') async onJoin( @ConnectedSocket() client: WebSocket, @MessageBody() data: { channelId: string }, ) { const caller = await this.resolveCaller(client); await this.chat.assertCanRead(data.channelId, caller); this.roomFor(data.channelId).add(client); return { event: 'chat:joined', data: { channelId: data.channelId } }; } @SubscribeMessage('chat:send') async onSend( @ConnectedSocket() client: WebSocket, @MessageBody() data: { channelId: string; body: string }, ) { const caller = await this.resolveCaller(client); const message = await this.chat.sendMessage(data.channelId, caller, data.body); this.broadcast(data.channelId, { event: 'chat:message', data: message }); return { event: 'chat:sent', data: { id: message.id } }; } private async resolveCaller(client: WebSocket): Promise { const pending = this.callers.get(client); if (!pending) { client.close(4001, 'Unauthorized'); throw new Error('Unauthorized WS client'); } try { return await pending; } catch { client.close(4001, 'Unauthorized'); throw new Error('Unauthorized WS client'); } } private roomFor(channelId: string): Set { let room = this.rooms.get(channelId); if (!room) { room = new Set(); this.rooms.set(channelId, room); } return room; } private broadcast(channelId: string, payload: unknown) { const room = this.rooms.get(channelId); if (!room) return; const json = JSON.stringify(payload); for (const socket of room) { if (socket.readyState === socket.OPEN) { socket.send(json); } } } /// Called by ChatController after add/removeParticipant so anyone with the /// channel already open (e.g. the creator's participant-management UI) /// gets a live update. Newly added participants join the room themselves /// via `chat:join` once they open the chat. notifyParticipantsChanged(channelId: string) { this.broadcast(channelId, { event: 'chat:participants-changed', data: { channelId } }); } }