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), so the global HTTP JwtGuard never applies here;
 * 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);

  @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 {
    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 {
    this.server
      .to(this.room(userId))
      .emit(NOTIFICATION_WS_EVENTS.UNREAD_COUNT, unreadCount);
  }

  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;
  }
}