diff --git a/src/chat/chat.gateway.ts b/src/chat/chat.gateway.ts index 27978a9..932d0a1 100644 --- a/src/chat/chat.gateway.ts +++ b/src/chat/chat.gateway.ts @@ -15,10 +15,14 @@ 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 WeakMap(); + private readonly callers = new Map>(); private readonly rooms = new Map>(); constructor( @@ -26,18 +30,18 @@ export class ChatGateway implements OnGatewayConnection, OnGatewayDisconnect { private readonly chat: ChatService, ) {} - async handleConnection(client: WebSocket, request: IncomingMessage) { + handleConnection(client: WebSocket, request: IncomingMessage) { const token = new URL(request.url ?? '', 'http://localhost').searchParams.get('token'); if (!token) { client.close(4001, 'Missing token'); return; } - try { - this.callers.set(client, await this.tokenVerification.verifyEither(token)); - } catch (err) { + 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) { @@ -52,7 +56,7 @@ export class ChatGateway implements OnGatewayConnection, OnGatewayDisconnect { @ConnectedSocket() client: WebSocket, @MessageBody() data: { channelId: string }, ) { - const caller = this.requireCaller(client); + 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 } }; @@ -63,19 +67,24 @@ export class ChatGateway implements OnGatewayConnection, OnGatewayDisconnect { @ConnectedSocket() client: WebSocket, @MessageBody() data: { channelId: string; body: string }, ) { - const caller = this.requireCaller(client); + 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 requireCaller(client: WebSocket): ChatCaller { - const caller = this.callers.get(client); - if (!caller) { + 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'); } - return caller; } private roomFor(channelId: string): Set {