Room-Based Live Notifications over a NestJS WebSocket Gateway

import { Injectable, ForbiddenException } from '@nestjs/common';
import { Socket } from 'socket.io';

export interface AuthUser {
  id: string;
  orgId: string;
  allowedTopics: string[];
}

@Injectable()
export class RoomSubscriptionService {
  userRoom(userId: string): string {
    return `user:${userId}`;
  }

  orgRoom(orgId: string): string {
    return `org:${orgId}`;
  }

  async subscribeToDefaults(socket: Socket, user: AuthUser): Promise<void> {
    await socket.join(this.userRoom(user.id));
    await socket.join(this.orgRoom(user.orgId));
  }

  async join(socket: Socket, user: AuthUser, room: string): Promise<void> {
    if (!this.canAccess(user, room)) {
      throw new ForbiddenException(`cannot subscribe to ${room}`);
    }
    await socket.join(room);
  }

  async leave(socket: Socket, room: string): Promise<void> {
    await socket.leave(room);
  }

  private canAccess(user: AuthUser, room: string): boolean {
    if (room === this.userRoom(user.id)) return true;
    if (room === this.orgRoom(user.orgId)) return true;
    if (room.startsWith('topic:')) {
      const topic = room.slice('topic:'.length);
      return user.allowedTopics.includes(topic);
    }
    return false;
  }
}

import { Injectable } from '@nestjs/common';
import { NotificationsGateway } from './notifications.gateway';
import { RoomSubscriptionService } from './room-subscription.service';

export interface NotificationPayload {
  id: string;
  type: string;
  title: string;
  body: string;
  createdAt: string;
}

@Injectable()
export class NotificationsService {
  constructor(
    private readonly gateway: NotificationsGateway,
    private readonly rooms: RoomSubscriptionService,
  ) {}

  broadcastToUser(userId: string, payload: NotificationPayload): void {
    const room = this.rooms.userRoom(userId);
    this.gateway.server.to(room).emit('notification', payload);
  }

  broadcastToOrg(orgId: string, payload: NotificationPayload): void {
    const room = this.rooms.orgRoom(orgId);
    this.gateway.server.to(room).emit('notification', payload);
  }

  broadcastToTopic(topic: string, payload: NotificationPayload): void {
    this.gateway.server.to(`topic:${topic}`).emit('notification', payload);
  }

  async countListeners(room: string): Promise<number> {
    const sockets = await this.gateway.server.in(room).fetchSockets();
    return sockets.length;
  }
}

import { CanActivate, ExecutionContext, Injectable } from '@nestjs/common';
import { WsException } from '@nestjs/websockets';
import { Socket } from 'socket.io';
import { AuthService } from '../auth/auth.service';

@Injectable()
export class WsAuthGuard implements CanActivate {
  constructor(private readonly auth: AuthService) {}

  async canActivate(context: ExecutionContext): Promise<boolean> {
    const socket = context.switchToWs().getClient<Socket>();
    const token =
      socket.handshake.auth?.token ?? socket.handshake.query?.token;

    const user = await this.auth.verifyToken(String(token ?? ''));
    if (!user) {
      throw new WsException('unauthorized');
    }

    socket.data.user = user;
    return true;
  }
}