diff --git a/src/collaboration/collaboration.gateway.ts b/src/collaboration/collaboration.gateway.ts index 6d98a20e..8a60cb52 100644 --- a/src/collaboration/collaboration.gateway.ts +++ b/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() diff --git a/src/collaboration/collaboration.module.ts b/src/collaboration/collaboration.module.ts index cc1c9bb8..9c8932c3 100644 --- a/src/collaboration/collaboration.module.ts +++ b/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], }) diff --git a/src/collaboration/guards/ws-jwt-auth.guard.spec.ts b/src/collaboration/guards/ws-jwt-auth.guard.spec.ts new file mode 100644 index 00000000..9e43e0f0 --- /dev/null +++ b/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 = {}) => ({ + id: 'socket-123', + handshake: { auth: {}, query: {}, ...overrides }, + data: {} as Record, +}); + +const makeContext = (client: ReturnType): ExecutionContext => + ({ + switchToWs: () => ({ getClient: () => client }), + }) as unknown as ExecutionContext; + +describe('WsJwtAuthGuard', () => { + let guard: WsJwtAuthGuard; + let jwtService: jest.Mocked; + + beforeEach(() => { + jwtService = { verify: jest.fn() } as unknown as jest.Mocked; + 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)); + }); +}); diff --git a/src/collaboration/guards/ws-jwt-auth.guard.ts b/src/collaboration/guards/ws-jwt-auth.guard.ts new file mode 100644 index 00000000..29de90f3 --- /dev/null +++ b/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(); + const token = + (client.handshake.auth as Record)?.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'); + } + } +} diff --git a/src/session/session.module.ts b/src/session/session.module.ts index ca2ed7da..be9ad268 100644 --- a/src/session/session.module.ts +++ b/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], }) diff --git a/src/session/session.service.spec.ts b/src/session/session.service.spec.ts index 50c74643..700ed47d 100644 --- a/src/session/session.service.spec.ts +++ b/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(), }; diff --git a/src/session/session.service.ts b/src/session/session.service.ts index 966a17f8..6031a503 100644 --- a/src/session/session.service.ts +++ b/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 { + 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 { + await this.redis.zadd(`user:sessions:${userId}`, Date.now(), sid); + } + + async removeSessionFromUserIndex(userId: string, sid: string): Promise { + await this.redis.zrem(`user:sessions:${userId}`, sid); + } + + async getUserSessionIds(userId: string): Promise { + return this.redis.zrange(`user:sessions:${userId}`, 0, -1); } /** diff --git a/src/session/tasks/session-cleanup.task.spec.ts b/src/session/tasks/session-cleanup.task.spec.ts new file mode 100644 index 00000000..d2b8918c --- /dev/null +++ b/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(); + }); +}); diff --git a/src/session/tasks/session-cleanup.task.ts b/src/session/tasks/session-cleanup.task.ts new file mode 100644 index 00000000..e9da1cf3 --- /dev/null +++ b/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 { + 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); + } + } +}