feat(collaboration): add WebSocket message size limit validation to p…

.env.example

@@ -326,6 +326,10 @@ REQUEST_BODY_LIMIT=1mb
 # Max file upload size (in bytes, default: 10MB)
 FILE_UPLOAD_MAX_BYTES=10485760
 
+# Max WebSocket message payload size (in bytes, default: 65536 = 64KB)
+# Applies to both Socket.IO transport level and application-level validation
+WS_MAX_PAYLOAD_BYTES=65536
+
 # HTTP request timeout (milliseconds)
 REQUEST_TIMEOUT=30000
 

src/collaboration/collaboration.gateway.ts

@@ -13,6 +13,7 @@ import { JoinSessionDto, CollaborativeOperationDto, SyncRequestDto } from './dto
 import { OtCrdtService, Operation } from './ot-crdt.service';
 import { PresenceService } from './presence.service';
 import { ChangeHistoryService } from './change-history.service';
+import { WsPayloadSizeGuardService } from './guards/ws-payload-size-guard.service';
 
 @WebSocketGateway({ namespace: '/collaboration', cors: { origin: '*' } })
 export class CollaborationGateway implements OnGatewayDisconnect {
@@ -27,6 +28,7 @@ export class CollaborationGateway implements OnGatewayDisconnect {
     private readonly otCrdt: OtCrdtService,
     private readonly presence: PresenceService,
     private readonly history: ChangeHistoryService,
+    private readonly payloadSizeGuard: WsPayloadSizeGuardService,
   ) {}
 
   handleDisconnect(client: Socket): void {
@@ -44,6 +46,8 @@ export class CollaborationGateway implements OnGatewayDisconnect {
 
   @SubscribeMessage(COLLABORATION_EVENTS.JOIN_SESSION)
   handleJoin(@MessageBody() dto: JoinSessionDto, @ConnectedSocket() client: Socket) {
+    this.payloadSizeGuard.validate(dto);
+
     client.join(dto.sessionId);
     this.socketMap.set(client.id, { sessionId: dto.sessionId, userId: dto.userId });
     const presenceInfo = this.presence.join(dto.sessionId, dto.userId);
@@ -70,6 +74,8 @@ export class CollaborationGateway implements OnGatewayDisconnect {
     @MessageBody() dto: CollaborativeOperationDto,
     @ConnectedSocket() client: Socket,
   ) {
+    this.payloadSizeGuard.validate(dto);
+
     const incomingOp = dto.operation as Operation;
     const revision = this.otCrdt.nextRevision(dto.sessionId);
     const op: Operation = { ...incomingOp, sessionId: dto.sessionId, userId: dto.userId, revision };
@@ -101,6 +107,8 @@ export class CollaborationGateway implements OnGatewayDisconnect {
 
   @SubscribeMessage(COLLABORATION_EVENTS.REQUEST_SYNC)
   handleSync(@MessageBody() dto: SyncRequestDto) {
+    this.payloadSizeGuard.validate(dto);
+
     const revision = this.otCrdt.currentRevision(dto.sessionId);
     const history = this.history.getLatest(dto.sessionId);
 
@@ -112,6 +120,8 @@ export class CollaborationGateway implements OnGatewayDisconnect {
 
   @SubscribeMessage(COLLABORATION_EVENTS.RESOLVE_CONFLICT)
   handleConflict(@MessageBody() body: { op1: Operation; op2: Operation; sessionId: string }) {
+    this.payloadSizeGuard.validate(body);
+
     const resolved = this.otCrdt.resolveConflict(body.op1, body.op2);
     this.server.to(body.sessionId).emit(COLLABORATION_EVENTS.CONFLICT_RESOLVED, { resolved });
     return { event: COLLABORATION_EVENTS.CONFLICT_RESOLVED, data: { resolved } };

src/collaboration/collaboration.module.ts

@@ -1,11 +1,20 @@
 import { Module } from '@nestjs/common';
+import { ConfigModule } from '@nestjs/config';
 import { OtCrdtService } from './ot-crdt.service';
 import { PresenceService } from './presence.service';
 import { ChangeHistoryService } from './change-history.service';
 import { CollaborationGateway } from './collaboration.gateway';
+import { WsPayloadSizeGuardService } from './guards/ws-payload-size-guard.service';
 
 @Module({
-  providers: [OtCrdtService, PresenceService, ChangeHistoryService, CollaborationGateway],
+  imports: [ConfigModule],
+  providers: [
+    OtCrdtService,
+    PresenceService,
+    ChangeHistoryService,
+    CollaborationGateway,
+    WsPayloadSizeGuardService,
+  ],
   exports: [OtCrdtService, PresenceService, ChangeHistoryService],
 })
 export class CollaborationModule {}

src/collaboration/guards/ws-payload-size-guard.service.spec.ts

@@ -0,0 +1,244 @@
+import { Test, TestingModule } from '@nestjs/testing';
+import { ConfigService } from '@nestjs/config';
+import { WsException } from '@nestjs/websockets';
+import {
+  WsPayloadSizeGuardService,
+  WS_PAYLOAD_TOO_LARGE_CODE,
+  DEFAULT_WS_MAX_PAYLOAD_BYTES,
+} from './ws-payload-size-guard.service';
+
+describe('WsPayloadSizeGuardService', () => {
+  // ---------------------------------------------------------------------------
+  // Tests with default limit (64KB)
+  // ---------------------------------------------------------------------------
+  describe('with default limit', () => {
+    let service: WsPayloadSizeGuardService;
+
+    beforeEach(async () => {
+      const module: TestingModule = await Test.createTestingModule({
+        providers: [
+          WsPayloadSizeGuardService,
+          {
+            provide: ConfigService,
+            useValue: {
+              get: jest.fn().mockReturnValue(undefined),
+            },
+          },
+        ],
+      }).compile();
+
+      service = module.get<WsPayloadSizeGuardService>(WsPayloadSizeGuardService);
+    });
+
+    it('should use default limit of 64KB', () => {
+      expect(service.getMaxPayloadBytes()).toBe(DEFAULT_WS_MAX_PAYLOAD_BYTES);
+      expect(service.getMaxPayloadBytes()).toBe(65_536);
+    });
+
+    it('should accept a small payload', () => {
+      const payload = { sessionId: 'session-1', userId: 'user-1', data: 'hello' };
+      expect(() => service.validate(payload)).not.toThrow();
+    });
+
+    it('should accept an empty object', () => {
+      expect(() => service.validate({})).not.toThrow();
+    });
+
+    it('should accept a payload just under the limit', () => {
+      // Create a payload that is just under 64KB
+      const padding = 'x'.repeat(60_000);
+      const payload = { data: padding };
+      expect(() => service.validate(payload)).not.toThrow();
+    });
+
+    it('should reject a payload exceeding 64KB', () => {
+      // Create a payload that is well over 64KB
+      const padding = 'x'.repeat(70_000);
+      const payload = { data: padding };
+
+      expect(() => service.validate(payload)).toThrow(WsException);
+
+      try {
+        service.validate(payload);
+      } catch (error) {
+        expect(error).toBeInstanceOf(WsException);
+        const wsError = error as WsException;
+        const errorPayload = wsError.getError() as { code: string; message: string };
+        expect(errorPayload.code).toBe(WS_PAYLOAD_TOO_LARGE_CODE);
+        expect(errorPayload.message).toContain('exceeds the maximum allowed size');
+      }
+    });
+
+    it('should reject a large nested object', () => {
+      // Build a deeply nested object that exceeds 64KB when serialized
+      const largeArray = Array.from({ length: 5000 }, (_, i) => ({
+        id: i,
+        content: `This is element number ${i} with some additional padding text to increase size`,
+        nested: { a: 'value', b: i * 100 },
+      }));
+      const payload = { operations: largeArray };
+
+      expect(() => service.validate(payload)).toThrow(WsException);
+    });
+
+    it('should include PAYLOAD_TOO_LARGE code in the error', () => {
+      const padding = 'x'.repeat(70_000);
+      const payload = { data: padding };
+
+      try {
+        service.validate(payload);
+        fail('Expected WsException to be thrown');
+      } catch (error) {
+        const wsError = error as WsException;
+        const errorPayload = wsError.getError() as { code: string; message: string };
+        expect(errorPayload.code).toBe('PAYLOAD_TOO_LARGE');
+      }
+    });
+
+    it('should include byte sizes in the error message', () => {
+      const padding = 'x'.repeat(70_000);
+      const payload = { data: padding };
+
+      try {
+        service.validate(payload);
+        fail('Expected WsException to be thrown');
+      } catch (error) {
+        const wsError = error as WsException;
+        const errorPayload = wsError.getError() as { code: string; message: string };
+        expect(errorPayload.message).toMatch(/\d+ bytes exceeds/);
+        expect(errorPayload.message).toMatch(/maximum allowed size of \d+ bytes/);
+      }
+    });
+  });
+
+  // ---------------------------------------------------------------------------
+  // Tests with custom limit
+  // ---------------------------------------------------------------------------
+  describe('with custom limit', () => {
+    let service: WsPayloadSizeGuardService;
+    const customLimit = 1_024; // 1KB
+
+    beforeEach(async () => {
+      const module: TestingModule = await Test.createTestingModule({
+        providers: [
+          WsPayloadSizeGuardService,
+          {
+            provide: ConfigService,
+            useValue: {
+              get: jest.fn((key: string) => {
+                if (key === 'WS_MAX_PAYLOAD_BYTES') return customLimit;
+                return undefined;
+              }),
+            },
+          },
+        ],
+      }).compile();
+
+      service = module.get<WsPayloadSizeGuardService>(WsPayloadSizeGuardService);
+    });
+
+    it('should use the configured limit', () => {
+      expect(service.getMaxPayloadBytes()).toBe(customLimit);
+    });
+
+    it('should accept a payload under the custom limit', () => {
+      const payload = { key: 'value' };
+      expect(() => service.validate(payload)).not.toThrow();
+    });
+
+    it('should reject a payload over the custom 1KB limit', () => {
+      const padding = 'x'.repeat(2_000);
+      const payload = { data: padding };
+
+      expect(() => service.validate(payload)).toThrow(WsException);
+    });
+  });
+
+  // ---------------------------------------------------------------------------
+  // Tests with realistic collaboration payloads
+  // ---------------------------------------------------------------------------
+  describe('with realistic collaboration payloads', () => {
+    let service: WsPayloadSizeGuardService;
+
+    beforeEach(async () => {
+      const module: TestingModule = await Test.createTestingModule({
+        providers: [
+          WsPayloadSizeGuardService,
+          {
+            provide: ConfigService,
+            useValue: {
+              get: jest.fn().mockReturnValue(undefined), // Use default 64KB
+            },
+          },
+        ],
+      }).compile();
+
+      service = module.get<WsPayloadSizeGuardService>(WsPayloadSizeGuardService);
+    });
+
+    it('should accept a normal collaborative operation', () => {
+      const payload = {
+        sessionId: 'session-abc-123',
+        userId: 'user-456',
+        resourceType: 'document',
+        operation: {
+          type: 'insert',
+          position: 42,
+          content: 'Hello, this is a normal collaborative edit.',
+          revision: 10,
+        },
+      };
+      expect(() => service.validate(payload)).not.toThrow();
+    });
+
+    it('should accept a join-session message', () => {
+      const payload = {
+        sessionId: 'session-abc-123',
+        userId: 'user-456',
+        resourceType: 'document',
+      };
+      expect(() => service.validate(payload)).not.toThrow();
+    });
+
+    it('should accept a sync-request message', () => {
+      const payload = {
+        sessionId: 'session-abc-123',
+        userId: 'user-456',
+      };
+      expect(() => service.validate(payload)).not.toThrow();
+    });
+
+    it('should reject a malicious megabyte-scale operation payload', () => {
+      const maliciousPayload = {
+        sessionId: 'session-abc-123',
+        userId: 'attacker',
+        resourceType: 'document',
+        operation: {
+          type: 'insert',
+          position: 0,
+          // 1MB of content — clearly malicious
+          content: 'A'.repeat(1_048_576),
+        },
+      };
+      expect(() => service.validate(maliciousPayload)).toThrow(WsException);
+    });
+
+    it('should reject a payload with excessive metadata', () => {
+      const payload = {
+        sessionId: 'session-abc-123',
+        userId: 'user-456',
+        resourceType: 'document',
+        operation: {
+          type: 'insert',
+          position: 0,
+          content: 'small content',
+          // Attacker stuffs massive metadata
+          metadata: Object.fromEntries(
+            Array.from({ length: 5000 }, (_, i) => [`key-${i}`, `value-${'x'.repeat(20)}-${i}`]),
+          ),
+        },
+      };
+      expect(() => service.validate(payload)).toThrow(WsException);
+    });
+  });
+});