feat: add WsJwtAuthGuard to CollaborationGateway and periodic session index cleanup

src/collaboration/collaboration.gateway.ts

@@ -1,4 +1,4 @@
-import { Logger } from '@nestjs/common';
+import { Logger, UseGuards } from '@nestjs/common';
 import {
   WebSocketGateway,
   WebSocketServer,
@@ -7,6 +7,7 @@ import {
   ConnectedSocket,
   OnGatewayDisconnect,
 } from '@nestjs/websockets';
+import { WsJwtAuthGuard } from './guards/ws-jwt-auth.guard';
 import { Server, Socket } from 'socket.io';
 import { COLLABORATION_EVENTS } from './constants/collaboration-events.constants';
 import { JoinSessionDto, CollaborativeOperationDto, SyncRequestDto } from './dto/websocket.dto';
@@ -15,6 +16,7 @@ import { PresenceService } from './presence.service';
 import { ChangeHistoryService } from './change-history.service';
 import { WsPayloadSizeGuardService } from './guards/ws-payload-size-guard.service';
 
+@UseGuards(WsJwtAuthGuard)
 @WebSocketGateway({ namespace: '/collaboration', cors: { origin: '*' } })
 export class CollaborationGateway implements OnGatewayDisconnect {
   @WebSocketServer()

src/collaboration/collaboration.module.ts

@@ -1,19 +1,28 @@
 import { Module } from '@nestjs/common';
 import { ConfigModule } from '@nestjs/config';
+import { JwtModule } from '@nestjs/jwt';
 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';
+import { WsJwtAuthGuard } from './guards/ws-jwt-auth.guard';
 
 @Module({
-  imports: [ConfigModule],
+  imports: [
+    ConfigModule,
+    JwtModule.register({
+      secret: process.env.JWT_SECRET || 'default-jwt-secret',
+      signOptions: { expiresIn: (process.env.JWT_EXPIRES_IN || '15m') as any },
+    }),
+  ],
   providers: [
     OtCrdtService,
     PresenceService,
     ChangeHistoryService,
     CollaborationGateway,
     WsPayloadSizeGuardService,
+    WsJwtAuthGuard,
   ],
   exports: [OtCrdtService, PresenceService, ChangeHistoryService],
 })

src/collaboration/guards/ws-jwt-auth.guard.spec.ts

@@ -0,0 +1,66 @@
+import { ExecutionContext } from '@nestjs/common';
+import { JwtService } from '@nestjs/jwt';
+import { WsException } from '@nestjs/websockets';
+import { WsJwtAuthGuard } from './ws-jwt-auth.guard';
+
+const makeClient = (overrides: Record<string, unknown> = {}) => ({
+  id: 'socket-123',
+  handshake: { auth: {}, query: {}, ...overrides },
+  data: {} as Record<string, unknown>,
+});
+
+const makeContext = (client: ReturnType<typeof makeClient>): ExecutionContext =>
+  ({
+    switchToWs: () => ({ getClient: () => client }),
+  }) as unknown as ExecutionContext;
+
+describe('WsJwtAuthGuard', () => {
+  let guard: WsJwtAuthGuard;
+  let jwtService: jest.Mocked<JwtService>;
+
+  beforeEach(() => {
+    jwtService = { verify: jest.fn() } as unknown as jest.Mocked<JwtService>;
+    guard = new WsJwtAuthGuard(jwtService);
+  });
+
+  it('throws WsException when no token is provided', () => {
+    const client = makeClient();
+    expect(() => guard.canActivate(makeContext(client))).toThrow(WsException);
+    expect(() => guard.canActivate(makeContext(client))).toThrow('missing token');
+  });
+
+  it('throws WsException when token verification fails', () => {
+    jwtService.verify.mockImplementation(() => {
+      throw new Error('jwt expired');
+    });
+    const client = makeClient({ auth: { token: 'bad.token.here' } });
+    expect(() => guard.canActivate(makeContext(client))).toThrow(WsException);
+    expect(() => guard.canActivate(makeContext(client))).toThrow('invalid token');
+  });
+
+  it('returns true and sets client.data.user when token is valid (auth object)', () => {
+    const payload = { sub: 'user-1', email: 'test@example.com' };
+    jwtService.verify.mockReturnValue(payload as any);
+    const client = makeClient({ auth: { token: 'valid.token' } });
+    const result = guard.canActivate(makeContext(client));
+    expect(result).toBe(true);
+    expect(client.data.user).toEqual(payload);
+  });
+
+  it('returns true and sets client.data.user when token is in query param', () => {
+    const payload = { sub: 'user-2', email: 'other@example.com' };
+    jwtService.verify.mockReturnValue(payload as any);
+    const client = makeClient({ query: { token: 'valid.query.token' } });
+    const result = guard.canActivate(makeContext(client));
+    expect(result).toBe(true);
+    expect(client.data.user).toEqual(payload);
+  });
+
+  it('prefers auth.token over query.token', () => {
+    const payload = { sub: 'user-3' };
+    jwtService.verify.mockReturnValue(payload as any);
+    const client = makeClient({ auth: { token: 'auth-token' }, query: { token: 'query-token' } });
+    guard.canActivate(makeContext(client));
+    expect(jwtService.verify).toHaveBeenCalledWith('auth-token', expect.any(Object));
+  });
+});

src/collaboration/guards/ws-jwt-auth.guard.ts

@@ -0,0 +1,38 @@
+import { CanActivate, ExecutionContext, Injectable, Logger } from '@nestjs/common';
+import { JwtService } from '@nestjs/jwt';
+import { WsException } from '@nestjs/websockets';
+import { Socket } from 'socket.io';
+
+/**
+ * Guards WebSocket connections by validating a JWT supplied in the
+ * handshake auth object or as a query parameter.
+ */
+@Injectable()
+export class WsJwtAuthGuard implements CanActivate {
+  private readonly logger = new Logger(WsJwtAuthGuard.name);
+
+  constructor(private readonly jwtService: JwtService) {}
+
+  canActivate(context: ExecutionContext): boolean {
+    const client: Socket = context.switchToWs().getClient<Socket>();
+    const token =
+      (client.handshake.auth as Record<string, string>)?.token ??
+      (client.handshake.query?.token as string);
+
+    if (!token) {
+      this.logger.warn(`WS connection rejected: no token (socketId=${client.id})`);
+      throw new WsException('Unauthorized: missing token');
+    }
+
+    try {
+      const payload = this.jwtService.verify(token, {
+        secret: process.env.JWT_SECRET || 'default-jwt-secret',
+      });
+      client.data.user = payload;
+      return true;
+    } catch {
+      this.logger.warn(`WS connection rejected: invalid token (socketId=${client.id})`);
+      throw new WsException('Unauthorized: invalid token');
+    }
+  }
+}

src/session/session.module.ts

@@ -3,6 +3,7 @@ import { ConfigModule, ConfigService } from '@nestjs/config';
 import { getSharedRedisClient } from '../config/cache.config';
 import { SESSION_REDIS_CLIENT } from './session.constants';
 import { SessionService } from './session.service';
+import { SessionCleanupTask } from './tasks/session-cleanup.task';
 
 /**
  * Registers the session module.
@@ -18,6 +19,7 @@ import { SessionService } from './session.service';
         getSharedRedisClient(configService),
     },
     SessionService,
+    SessionCleanupTask,
   ],
   exports: [SESSION_REDIS_CLIENT, SessionService],
 })

src/session/session.service.spec.ts

@@ -10,6 +10,9 @@ const mockRedis = {
   expire: jest.fn(),
   eval: jest.fn(),
   multi: jest.fn(),
+  zadd: jest.fn().mockResolvedValue(1),
+  zrem: jest.fn().mockResolvedValue(1),
+  zrange: jest.fn().mockResolvedValue([]),
   status: 'ready',
   quit: jest.fn(),
 };

src/session/session.service.ts

@@ -81,6 +81,7 @@ export class SessionService implements OnModuleDestroy {
       'EX',
       this.sessionTtlSeconds,
     );
+    await this.addSessionToUserIndex(userId, sid);
     return sid;
   }
 
@@ -131,7 +132,23 @@ export class SessionService implements OnModuleDestroy {
    * @param sid The sid.
    */
   async removeSession(sid: string): Promise<void> {
+    const session = await this.getSession(sid);
     await this.redis.del(this.sessionKey(sid));
+    if (session) {
+      await this.removeSessionFromUserIndex(session.userId, sid);
+    }
+  }
+
+  async addSessionToUserIndex(userId: string, sid: string): Promise<void> {
+    await this.redis.zadd(`user:sessions:${userId}`, Date.now(), sid);
+  }
+
+  async removeSessionFromUserIndex(userId: string, sid: string): Promise<void> {
+    await this.redis.zrem(`user:sessions:${userId}`, sid);
+  }
+
+  async getUserSessionIds(userId: string): Promise<string[]> {
+    return this.redis.zrange(`user:sessions:${userId}`, 0, -1);
   }
 
   /**

src/session/tasks/session-cleanup.task.spec.ts

@@ -0,0 +1,84 @@
+import { Test, TestingModule } from '@nestjs/testing';
+import { SessionCleanupTask } from './session-cleanup.task';
+import { SessionService } from '../session.service';
+import { SESSION_REDIS_CLIENT } from '../session.constants';
+
+const mockSessionService = {
+  getUserSessionIds: jest.fn(),
+  getSession: jest.fn(),
+  removeSessionFromUserIndex: jest.fn(),
+};
+
+const mockRedis = {
+  scan: jest.fn(),
+};
+
+describe('SessionCleanupTask', () => {
+  let task: SessionCleanupTask;
+
+  beforeEach(async () => {
+    const module: TestingModule = await Test.createTestingModule({
+      providers: [
+        SessionCleanupTask,
+        { provide: SessionService, useValue: mockSessionService },
+        { provide: SESSION_REDIS_CLIENT, useValue: mockRedis },
+      ],
+    }).compile();
+
+    task = module.get(SessionCleanupTask);
+  });
+
+  afterEach(() => jest.clearAllMocks());
+
+  it('does nothing when no user session keys exist', async () => {
+    mockRedis.scan.mockResolvedValueOnce(['0', []]);
+
+    await task.handleCleanup();
+
+    expect(mockSessionService.getUserSessionIds).not.toHaveBeenCalled();
+    expect(mockSessionService.removeSessionFromUserIndex).not.toHaveBeenCalled();
+  });
+
+  it('removes stale session ids from the index', async () => {
+    mockRedis.scan.mockResolvedValueOnce(['0', ['user:sessions:user-1']]);
+    mockSessionService.getUserSessionIds.mockResolvedValueOnce(['sid-expired', 'sid-live']);
+    mockSessionService.getSession
+      .mockResolvedValueOnce(null) // sid-expired → gone
+      .mockResolvedValueOnce({ sid: 'sid-live', userId: 'user-1' }); // sid-live → still there
+
+    await task.handleCleanup();
+
+    expect(mockSessionService.removeSessionFromUserIndex).toHaveBeenCalledTimes(1);
+    expect(mockSessionService.removeSessionFromUserIndex).toHaveBeenCalledWith(
+      'user-1',
+      'sid-expired',
+    );
+  });
+
+  it('keeps live session ids in the index', async () => {
+    mockRedis.scan.mockResolvedValueOnce(['0', ['user:sessions:user-2']]);
+    mockSessionService.getUserSessionIds.mockResolvedValueOnce(['sid-active']);
+    mockSessionService.getSession.mockResolvedValueOnce({ sid: 'sid-active', userId: 'user-2' });
+
+    await task.handleCleanup();
+
+    expect(mockSessionService.removeSessionFromUserIndex).not.toHaveBeenCalled();
+  });
+
+  it('handles multiple cursor pages', async () => {
+    mockRedis.scan
+      .mockResolvedValueOnce(['42', ['user:sessions:user-3']])
+      .mockResolvedValueOnce(['0', ['user:sessions:user-4']]);
+    mockSessionService.getUserSessionIds.mockResolvedValue([]);
+
+    await task.handleCleanup();
+
+    expect(mockRedis.scan).toHaveBeenCalledTimes(2);
+  });
+
+  it('does not throw when cleanup encounters an error', async () => {
+    mockRedis.scan.mockRejectedValueOnce(new Error('Redis unavailable'));
+
+    await expect(task.handleCleanup()).resolves.not.toThrow();
+  });
+});

src/session/tasks/session-cleanup.task.ts

@@ -0,0 +1,67 @@
+import { Inject, Injectable, Logger } from '@nestjs/common';
+import { Cron, CronExpression } from '@nestjs/schedule';
+import { Counter, Registry } from 'prom-client';
+import Redis from 'ioredis';
+import { SessionService } from '../session.service';
+import { SESSION_REDIS_CLIENT } from '../session.constants';
+
+/**
+ * Periodic job that removes stale members from user session sorted-set indexes.
+ * Primary session keys expire via Redis TTL; this job reconciles the secondary
+ * index so it does not grow unbounded.
+ */
+@Injectable()
+export class SessionCleanupTask {
+  private readonly logger = new Logger(SessionCleanupTask.name);
+  private readonly removedCounter: Counter;
+
+  constructor(
+    private readonly sessionService: SessionService,
+    @Inject(SESSION_REDIS_CLIENT) private readonly redis: Redis,
+  ) {
+    const registry = new Registry();
+    this.removedCounter = new Counter({
+      name: 'session_cleanup_removed_total',
+      help: 'Total stale session index entries removed by cleanup job',
+      registers: [registry],
+    });
+  }
+
+  @Cron(CronExpression.EVERY_HOUR)
+  async handleCleanup(): Promise<void> {
+    this.logger.log('Starting expired session index cleanup...');
+    let removedTotal = 0;
+
+    try {
+      let cursor = '0';
+      do {
+        const [nextCursor, keys] = await this.redis.scan(
+          cursor,
+          'MATCH',
+          'user:sessions:*',
+          'COUNT',
+          100,
+        );
+        cursor = nextCursor;
+
+        for (const key of keys) {
+          const userId = key.replace('user:sessions:', '');
+          const sids = await this.sessionService.getUserSessionIds(userId);
+
+          for (const sid of sids) {
+            const exists = await this.sessionService.getSession(sid);
+            if (!exists) {
+              await this.sessionService.removeSessionFromUserIndex(userId, sid);
+              removedTotal++;
+            }
+          }
+        }
+      } while (cursor !== '0');
+
+      this.removedCounter.inc(removedTotal);
+      this.logger.log(`Cleanup complete. Removed ${removedTotal} stale session index entries.`);
+    } catch (error) {
+      this.logger.error('Session index cleanup failed:', error);
+    }
+  }
+}