Lesson 14: WebSockets and Real-time Communication

NestJS: From Basics to Advanced

Part 4: Advanced Features

xdev.asia

1. Install WebSocket

npm install @nestjs/websockets @nestjs/platform-socket.io
npm install -D @types/socket.io

2. WebSocket Gateway

// chat/chat.gateway.ts
import {
  WebSocketGateway,
  WebSocketServer,
  SubscribeMessage,
  MessageBody,
  ConnectedSocket,
  OnGatewayInit,
  OnGatewayConnection,
  OnGatewayDisconnect,
} from '@nestjs/websockets';
import { Server, Socket } from 'socket.io';
import { Logger } from '@nestjs/common';

@WebSocketGateway({
  cors: { origin: process.env.FRONTEND_URL },
  namespace: '/chat',
})
export class ChatGateway
  implements OnGatewayInit, OnGatewayConnection, OnGatewayDisconnect
{
  @WebSocketServer()
  server: Server;

  private logger = new Logger('ChatGateway');

  afterInit(server: Server) {
    this.logger.log('WebSocket Gateway initialized');
  }

  handleConnection(client: Socket) {
    this.logger.log(`Client connected: ${client.id}`);
  }

  handleDisconnect(client: Socket) {
    this.logger.log(`Client disconnected: ${client.id}`);
  }

  @SubscribeMessage('sendMessage')
  handleMessage(
    @ConnectedSocket() client: Socket,
    @MessageBody() payload: { room: string; message: string },
  ) {
    this.server.to(payload.room).emit('newMessage', {
      userId: client.data.userId,
      message: payload.message,
      timestamp: new Date(),
    });
  }
}

3. Room Management

@SubscribeMessage('joinRoom')
async handleJoinRoom(
  @ConnectedSocket() client: Socket,
  @MessageBody() roomId: string,
) {
  await client.join(roomId);
  
  // Thông báo cho room
  client.to(roomId).emit('userJoined', {
    userId: client.data.userId,
    roomId,
  });

  return { event: 'joinedRoom', data: roomId };
}

@SubscribeMessage('leaveRoom')
async handleLeaveRoom(
  @ConnectedSocket() client: Socket,
  @MessageBody() roomId: string,
) {
  await client.leave(roomId);
  client.to(roomId).emit('userLeft', {
    userId: client.data.userId,
    roomId,
  });
}

// Typing indicator
@SubscribeMessage('typing')
handleTyping(
  @ConnectedSocket() client: Socket,
  @MessageBody() data: { room: string; isTyping: boolean },
) {
  client.to(data.room).emit('userTyping', {
    userId: client.data.userId,
    isTyping: data.isTyping,
  });
}

4. WebSocket Authentication

// chat/ws-auth.guard.ts
import { CanActivate, ExecutionContext, Injectable } from '@nestjs/common';
import { JwtService } from '@nestjs/jwt';
import { WsException } from '@nestjs/websockets';

@Injectable()
export class WsAuthGuard implements CanActivate {
  constructor(private jwtService: JwtService) {}

  canActivate(context: ExecutionContext): boolean {
    const client = context.switchToWs().getClient();
    const token =
      client.handshake.auth?.token ||
      client.handshake.headers?.authorization?.split(' ')[1];

    if (!token) throw new WsException('Unauthorized');

    try {
      const payload = this.jwtService.verify(token);
      client.data.userId = payload.sub;
      client.data.email = payload.email;
      return true;
    } catch {
      throw new WsException('Invalid token');
    }
  }
}

// Middleware cấp Gateway — verify trước khi connect
@WebSocketGateway()
export class ChatGateway implements OnGatewayInit {
  afterInit(server: Server) {
    server.use((socket, next) => {
      const token = socket.handshake.auth?.token;
      if (!token) return next(new Error('Authentication required'));

      try {
        const payload = this.jwtService.verify(token);
        socket.data.userId = payload.sub;
        next();
      } catch {
        next(new Error('Invalid token'));
      }
    });
  }
}

5. Notification System

// notifications/notification.gateway.ts
@WebSocketGateway({ namespace: '/notifications' })
export class NotificationGateway implements OnGatewayConnection {
  @WebSocketServer()
  server: Server;

  private userSockets = new Map<string, Set<string>>();

  handleConnection(client: Socket) {
    const userId = client.data.userId;
    if (!this.userSockets.has(userId)) {
      this.userSockets.set(userId, new Set());
    }
    this.userSockets.get(userId).add(client.id);
  }

  handleDisconnect(client: Socket) {
    const userId = client.data.userId;
    this.userSockets.get(userId)?.delete(client.id);
  }

  // Gửi notification tới user cụ thể
  sendToUser(userId: string, notification: any) {
    const sockets = this.userSockets.get(userId);
    if (sockets) {
      sockets.forEach((socketId) => {
        this.server.to(socketId).emit('notification', notification);
      });
    }
  }

  // Broadcast cho tất cả
  broadcast(event: string, data: any) {
    this.server.emit(event, data);
  }
}

// Sử dụng từ service khác
@Injectable()
export class OrderService {
  constructor(private notificationGateway: NotificationGateway) {}

  async createOrder(dto: CreateOrderDto) {
    const order = await this.orderRepo.save(dto);
    
    // Gửi real-time notification
    this.notificationGateway.sendToUser(order.userId, {
      type: 'ORDER_CREATED',
      message: `Đơn hàng #${order.id} đã được tạo`,
      data: order,
    });
    
    return order;
  }
}

6. Frontend Client

// React/Next.js client
import { io, Socket } from 'socket.io-client';

const socket = io('http://localhost:3000/chat', {
  auth: { token: localStorage.getItem('accessToken') },
  autoConnect: false,
});

socket.connect();

socket.on('connect', () => console.log('Connected'));
socket.on('newMessage', (msg) => console.log('New message:', msg));
socket.on('notification', (n) => console.log('Notification:', n));

// Gửi message
socket.emit('sendMessage', { room: 'general', message: 'Hello!' });

7. Summary

• @WebSocketGateway: Decorator creates WebSocket endpoint
• @SubscribeMessage: Listen to events from the client
• Rooms: Group clients, broadcast by group
• WsAuthGuard: Authenticate WebSocket connections
• Namespaces: Separate logic (chat, notifications)

The next article will explore GraphQL with NestJS.