feat: add web sockets rate limiting

README.md

@@ -213,25 +213,25 @@ TeachLink Backend follows a **modular microservices architecture** built on Nest
 
 ## 📦 Tech Stack
 
-| Layer             | Technology                     | Purpose                              |
-| ----------------- | ------------------------------ | ------------------------------------ |
-| **Framework**     | NestJS                         | Node.js application framework        |
-| **Language**      | TypeScript                     | Type-safe JavaScript                 |
-| **Database**      | PostgreSQL + TypeORM           | Primary data storage                  |
-| **Caching**       | Redis + IORedis                | Session store, caching, queues       |
-| **Authentication**| JWT + Passport                 | Token-based authentication           |
-| **GraphQL**       | Apollo Server                  | GraphQL API (optional)               |
-| **Real-time**     | Socket.io                      | WebSocket connections                |
-| **File Storage**  | AWS S3 + Cloudinary            | Media file storage and CDN           |
-| **Email**         | SendGrid + Nodemailer          | Email delivery and marketing          |
-| **Payments**      | Stripe                         | Payment processing                   |
-| **Search**        | Elasticsearch                  | Full-text search capabilities        |
-| **Queue**         | BullMQ                         | Background job processing             |
-| **Monitoring**    | OpenTelemetry + Prometheus     | Metrics and observability            |
-| **Testing**       | Jest + Supertest               | Unit and integration tests           |
-| **Documentation** | Swagger                        | API documentation                    |
-| **Validation**    | class-validator + class-transformer | DTO validation                    |
-| **Security**      | Helmet + bcrypt                | Security headers and password hashing |
+| Layer              | Technology                          | Purpose                               |
+| ------------------ | ----------------------------------- | ------------------------------------- |
+| **Framework**      | NestJS                              | Node.js application framework         |
+| **Language**       | TypeScript                          | Type-safe JavaScript                  |
+| **Database**       | PostgreSQL + TypeORM                | Primary data storage                  |
+| **Caching**        | Redis + IORedis                     | Session store, caching, queues        |
+| **Authentication** | JWT + Passport                      | Token-based authentication            |
+| **GraphQL**        | Apollo Server                       | GraphQL API (optional)                |
+| **Real-time**      | Socket.io                           | WebSocket connections                 |
+| **File Storage**   | AWS S3 + Cloudinary                 | Media file storage and CDN            |
+| **Email**          | SendGrid + Nodemailer               | Email delivery and marketing          |
+| **Payments**       | Stripe                              | Payment processing                    |
+| **Search**         | Elasticsearch                       | Full-text search capabilities         |
+| **Queue**          | BullMQ                              | Background job processing             |
+| **Monitoring**     | OpenTelemetry + Prometheus          | Metrics and observability             |
+| **Testing**        | Jest + Supertest                    | Unit and integration tests            |
+| **Documentation**  | Swagger                             | API documentation                     |
+| **Validation**     | class-validator + class-transformer | DTO validation                        |
+| **Security**       | Helmet + bcrypt                     | Security headers and password hashing |
 
 ## �️ Database
 

src/collaboration/gateway/collaboration.gateway.ts

@@ -9,8 +9,9 @@ import {
   ConnectedSocket,
 } from '@nestjs/websockets';
 import { Server, Socket } from 'socket.io';
-import { Logger } from '@nestjs/common';
+import { Logger, UseGuards } from '@nestjs/common';
 import { CollaborationService } from '../collaboration.service';
+import { WsThrottlerGuard } from '../../common/guards/ws-throttler.guard';
 import { SharedDocumentService } from '../documents/shared-document.service';
 import { WhiteboardService } from '../whiteboard/whiteboard.service';
 import { VersionControlService } from '../versioning/version-control.service';
@@ -34,6 +35,7 @@ export interface CollaborativeOperation {
     credentials: true,
   },
 })
+@UseGuards(WsThrottlerGuard)
 export class CollaborationGateway
   implements OnGatewayInit, OnGatewayConnection, OnGatewayDisconnect
 {
@@ -53,6 +55,11 @@ export class CollaborationGateway
   }
 
   async handleConnection(_server: any, @ConnectedSocket() client: Socket) {
+    if (wsManager.getTotalConnections() >= 5000) {
+      client.emit('error', { message: 'Server is at maximum capacity' });
+      client.disconnect(true);
+      return;
+    }
     this.logger.log(`Client connected: ${client.id}`);
 
     // Optionally authenticate the user here based on token
@@ -61,6 +68,7 @@ export class CollaborationGateway
   }
 
   async handleDisconnect(@ConnectedSocket() client: Socket) {
+    wsManager.cleanupSocket(client);
     this.logger.log(`Client disconnected: ${client.id}`);
 
     // Clean up any session associations for this client
@@ -82,6 +90,12 @@ export class CollaborationGateway
         return;
       }
 
+      const registered = wsManager.registerConnection(userId, client);
+      if (!registered) {
+        client.emit('error', { message: 'Connection limit reached' });
+        return;
+      }
+
       // Join the room
       client.join(sessionId);
 

src/common/guards/ws-throttler.guard.ts

@@ -0,0 +1,24 @@
+import { Injectable, ExecutionContext, Logger } from '@nestjs/common';
+import { ThrottlerGuard, ThrottlerLimitDetail } from '@nestjs/throttler';
+import { WsException } from '@nestjs/websockets';
+
+@Injectable()
+export class WsThrottlerGuard extends ThrottlerGuard {
+  private readonly wsLogger = new Logger(WsThrottlerGuard.name);
+
+  protected async getTracker(req: Record<string, any>): Promise<string> {
+    const user = req.user;
+    const ip = req.conn?.remoteAddress || req.request?.connection?.remoteAddress || 'unknown';
+    return user?.sub || user?.id || ip;
+  }
+
+  protected async throwThrottlingException(
+    context: ExecutionContext,
+    _throttlerLimitDetail: ThrottlerLimitDetail,
+  ): Promise<void> {
+    const client = context.switchToWs().getClient();
+    const tracker = await this.getTracker(client);
+    this.wsLogger.warn(`WebSocket rate limit exceeded for ${tracker}`);
+    throw new WsException('Rate limit exceeded');
+  }
+}

src/common/utils/websocket.utils.ts

@@ -11,6 +11,7 @@ class WebSocketManager {
   private meta = new Map<string, ConnectionMeta>(); // socketId -> meta
 
   private MAX_CONNECTIONS_PER_USER = 3;
+  private MAX_GLOBAL_CONNECTIONS = 5000;
   private HEARTBEAT_INTERVAL = 30000; // 30s
   private TIMEOUT = 60000; // 60s
 
@@ -46,12 +47,13 @@ class WebSocketManager {
       this.connections.set(userId, new Set());
     }
 
-    const userConnections = this.connections.get(userId) ?? new Set();
+    const userConnections = this.connections.get(userId) || new Set<Socket>();
 
-    const userConnections = this.connections.get(userId);
-
-    if (!userConnections) {
-      return;
+    // enforce global connection limits
+    if (this.meta.size >= this.MAX_GLOBAL_CONNECTIONS) {
+      socket.emit('error', { message: 'Server is at maximum capacity' });
+      socket.disconnect(true);
+      return false;
     }
 
     // enforce max connections
@@ -68,6 +70,7 @@ class WebSocketManager {
       lastSeen: Date.now(),
       isAlive: true,
     });
+    return true;
   }
 
   cleanupSocket(socket: Socket) {

src/gateways/messaging.gateway.ts

@@ -4,15 +4,47 @@ import {
   MessageBody,
   ConnectedSocket,
   OnGatewayConnection,
+  OnGatewayDisconnect,
 } from '@nestjs/websockets';
-import { UseGuards } from '@nestjs/common';
+import { UseGuards, Logger } from '@nestjs/common';
 import { Socket } from 'socket.io';
 import { WsJwtAuthGuard } from '../auth/guards/ws-jwt-auth.guard';
+import { WsThrottlerGuard } from '../common/guards/ws-throttler.guard';
+import { wsManager } from '../common/utils/websocket.utils';
+import { JwtService } from '@nestjs/jwt';
 
 @WebSocketGateway({ namespace: '/messaging' })
-export class MessagingGateway implements OnGatewayConnection {
-  async handleConnection(@ConnectedSocket() _client: Socket) {
-    // Guard will disconnect unauthorized clients
+@UseGuards(WsThrottlerGuard)
+export class MessagingGateway implements OnGatewayConnection, OnGatewayDisconnect {
+  private readonly logger = new Logger(MessagingGateway.name);
+
+  constructor(private readonly jwtService: JwtService) {}
+
+  async handleConnection(client: Socket) {
+    try {
+      const token =
+        client.handshake.auth?.token || client.handshake.headers?.authorization?.split(' ')[1];
+      if (!token) {
+        throw new Error('No token provided');
+      }
+
+      const payload = await this.jwtService.verifyAsync(token);
+      const userId = payload.sub || payload.id;
+
+      const registered = wsManager.registerConnection(userId, client);
+      if (!registered) {
+        return;
+      }
+      this.logger.log(`Client connected: ${client.id}`);
+    } catch (_error) {
+      client.emit('error', { message: 'Unauthorized' });
+      client.disconnect(true);
+    }
+  }
+
+  handleDisconnect(client: Socket) {
+    wsManager.cleanupSocket(client);
+    this.logger.log(`Client disconnected: ${client.id}`);
   }
 
   @UseGuards(WsJwtAuthGuard)

src/gateways/notifications.gateway.ts

@@ -1,10 +1,52 @@
-import { WebSocketGateway, SubscribeMessage, ConnectedSocket } from '@nestjs/websockets';
-import { UseGuards } from '@nestjs/common';
+import {
+  WebSocketGateway,
+  SubscribeMessage,
+  ConnectedSocket,
+  OnGatewayConnection,
+  OnGatewayDisconnect,
+} from '@nestjs/websockets';
+import { UseGuards, Logger } from '@nestjs/common';
 import { Socket } from 'socket.io';
 import { WsJwtAuthGuard } from '../auth/guards/ws-jwt-auth.guard';
+import { WsThrottlerGuard } from '../common/guards/ws-throttler.guard';
+import { wsManager } from '../common/utils/websocket.utils';
+import { JwtService } from '@nestjs/jwt';
 
 @WebSocketGateway({ namespace: '/notifications' })
-export class NotificationsGateway {
+@UseGuards(WsThrottlerGuard)
+export class NotificationsGateway implements OnGatewayConnection, OnGatewayDisconnect {
+  private readonly logger = new Logger(NotificationsGateway.name);
+
+  constructor(private readonly jwtService: JwtService) {}
+
+  async handleConnection(client: Socket) {
+    try {
+      const token =
+        client.handshake.auth?.token || client.handshake.headers?.authorization?.split(' ')[1];
+      if (!token) {
+        throw new Error('No token provided');
+      }
+
+      const payload = await this.jwtService.verifyAsync(token);
+      const userId = payload.sub || payload.id;
+
+      const registered = wsManager.registerConnection(userId, client);
+      if (!registered) {
+        // Connection rejected by global limits
+        return;
+      }
+      this.logger.log(`Client connected: ${client.id}`);
+    } catch (_error) {
+      client.emit('error', { message: 'Unauthorized' });
+      client.disconnect(true);
+    }
+  }
+
+  handleDisconnect(client: Socket) {
+    wsManager.cleanupSocket(client);
+    this.logger.log(`Client disconnected: ${client.id}`);
+  }
+
   @UseGuards(WsJwtAuthGuard)
   @SubscribeMessage('subscribe_notifications')
   async handleSubscribe(@ConnectedSocket() client: Socket) {

src/health/health.controller.ts

@@ -4,7 +4,6 @@ import Redis from 'ioredis';
 import { SkipThrottle } from '@nestjs/throttler';
 import { HealthService } from './health.service';
 
-
 @Version(VERSION_NEUTRAL)
 @SkipThrottle()
 @Controller('health')

src/health/health.service.ts

@@ -1,6 +1,6 @@
 import { Injectable, Logger } from '@nestjs/common';
 import { DataSource } from 'typeorm';
-import Redis from 'ioredis';
+import { Redis } from 'ioredis';
 import * as fs from 'fs';
 import * as _path from 'path';
 import axios from 'axios';
@@ -182,7 +182,7 @@ export class HealthService {
   }> {
     const startTime = Date.now();
     try {
-      const pong = await redis.ping();
+      const pong = await (redis as any).ping();
       const responseTime = Date.now() - startTime;
       return {
         status: pong === 'PONG' && responseTime < 500 ? 'up' : 'degraded',

src/main.ts

@@ -111,7 +111,9 @@ async function bootstrapWorker() {
   // ─── Swagger ──────────────────────────────────────────────────────────────
   const config = new DocumentBuilder()
     .setTitle('TeachLink API')
-    .setDescription(`The TeachLink API documentation - Unified System. ${API_VERSIONING_DOCUMENTATION}`)
+    .setDescription(
+      `The TeachLink API documentation - Unified System. ${API_VERSIONING_DOCUMENTATION}`,
+    )
     .setVersion('1.0')
     .addBearerAuth()
     .addTag('gamification', 'Gamification and user rewards')

src/notifications/notifications.gateway.ts

@@ -11,13 +11,16 @@ import { Server, Socket } from 'socket.io';
 import { UseGuards, Logger } from '@nestjs/common';
 import { WsJwtAuthGuard } from '../auth/guards/ws-jwt-auth.guard';
 import { Notification } from './entities/notification.entity';
+import { wsManager } from '../common/utils/websocket.utils';
+import { WsThrottlerGuard } from '../common/guards/ws-throttler.guard';
 
 @WebSocketGateway({
   cors: {
     origin: '*',
   },
   namespace: 'notifications',
 })
+@UseGuards(WsThrottlerGuard)
 export class NotificationsGateway implements OnGatewayConnection, OnGatewayDisconnect {
   @WebSocketServer()
   server: Server;
@@ -26,10 +29,20 @@ export class NotificationsGateway implements OnGatewayConnection, OnGatewayDisco
   private userSockets: Map<string, Set<string>> = new Map();
 
   async handleConnection(client: Socket) {
+    // Optionally we can get userId here from token if required, but currently it handles subscribe via message.
+    // However, for connection limit we can temporarily register with client id if no token or if handled at subscribe.
+    // Let's enforce global connection limits here.
+    if (wsManager.getTotalConnections() >= 5000) {
+      // Same as MAX_GLOBAL_CONNECTIONS
+      client.emit('error', { message: 'Server is at maximum capacity' });
+      client.disconnect(true);
+      return;
+    }
     this.logger.log(`Client connected: ${client.id}`);
   }
 
   handleDisconnect(client: Socket) {
+    wsManager.cleanupSocket(client);
     this.logger.log(`Client disconnected: ${client.id}`);
     this.removeSocketId(client.id);
   }
@@ -48,6 +61,11 @@ export class NotificationsGateway implements OnGatewayConnection, OnGatewayDisco
       return;
     }
 
+    const registered = wsManager.registerConnection(userId, client);
+    if (!registered) {
+      return { status: 'error', message: 'Connection limit reached' };
+    }
+
     this.addSocketId(userId, client.id);
     client.join(`user:${userId}`);
     this.logger.log(`User ${userId} subscribed to notifications via socket ${client.id}`);