KC-APP-Server/src/chat/chat.gateway.ts

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<WebSocket, Promise<ChatCaller>>();
  private readonly rooms = new Map<string, Set<WebSocket>>();

  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<ChatCaller> {
    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<WebSocket> {
    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 } });
  }
}