edr-platform/apps/edr-freight-api/src/modules/notification-inbox/notifications.gateway.ts

import {
  NOTIFICATION_WS_EVENTS,
  NOTIFICATION_WS_NAMESPACE,
  NotificationDto,
} from "@edr/types";
import { Logger } from "@nestjs/common";
import {
  OnGatewayConnection,
  WebSocketGateway,
  WebSocketServer,
} from "@nestjs/websockets";
import { Server, Socket } from "socket.io";

import { WsAuthService } from "./ws-auth.service";

/**
 * Server → client push for in-app notifications. Clients only *listen* (no
 * `@SubscribeMessage` handlers), and `@UseGuards(JwtGuard)` on the REST
 * controller does not cover WebSockets; the handshake is authenticated in
 * `handleConnection` and each socket joins a private `user:<id>` room the
 * service targets.
 */
@WebSocketGateway({
  namespace: NOTIFICATION_WS_NAMESPACE,
  cors: { origin: true, credentials: true },
})
export class NotificationsGateway implements OnGatewayConnection {
  private readonly logger = new Logger(NotificationsGateway.name);

  // Not `!`-asserted: Nest only wires this once the WS adapter attaches to a running HTTP
  // listener, which does not happen under `NestFactory.createApplicationContext` — see `skip()`.
  @WebSocketServer()
  private readonly server?: Server;

  constructor(private readonly wsAuth: WsAuthService) {}

  async handleConnection(socket: Socket): Promise<void> {
    const userId = await this.wsAuth.resolveUserId(this.extractToken(socket));
    if (!userId) {
      this.logger.debug(`Rejected notifications handshake ${socket.id}`);
      socket.disconnect(true);
      return;
    }
    socket.data.userId = userId;
    await socket.join(this.room(userId));
  }

  /** Push a freshly-created notification + the new unread count to a user. */
  emitNew(userId: string, notification: NotificationDto, unreadCount: number): void {
    if (!this.server) return this.skip("emitNew");
    const room = this.server.to(this.room(userId));
    room.emit(NOTIFICATION_WS_EVENTS.NEW, notification);
    room.emit(NOTIFICATION_WS_EVENTS.UNREAD_COUNT, unreadCount);
  }

  /** Push only an updated unread count (e.g. after a read on another tab). */
  emitUnreadCount(userId: string, unreadCount: number): void {
    if (!this.server) return this.skip("emitUnreadCount");
    this.server
      .to(this.room(userId))
      .emit(NOTIFICATION_WS_EVENTS.UNREAD_COUNT, unreadCount);
  }

  /**
   * `@WebSocketServer()` only wires `server` once the WS adapter attaches to a running HTTP
   * listener — never under `NestFactory.createApplicationContext` (scripts, one-off jobs), and not
   * for the brief window before `app.listen()` completes in a real boot either. The notification row
   * is already persisted by this point (the caller writes it before pushing), so a missing socket
   * server just means "no live push this time" — skip it rather than throw and lose the caller's
   * own result (e.g. an EIMS registration outcome that already succeeded or failed for real).
   */
  private skip(method: string): void {
    this.logger.debug(`${method}: no WebSocket server attached (non-HTTP context?) — push skipped`);
  }

  private room(userId: string): string {
    return `user:${userId}`;
  }

  private extractToken(socket: Socket): string | undefined {
    const authToken = socket.handshake.auth?.token as string | undefined;
    if (authToken) return authToken;

    const queryToken = socket.handshake.query?.token;
    if (typeof queryToken === "string") return queryToken;

    const header = socket.handshake.headers?.authorization;
    if (header?.startsWith("Bearer ")) return header.slice(7);

    return undefined;
  }
}