From 790fd6a4600b6659efd9b635c88e9b449136bd29 Mon Sep 17 00:00:00 2001 From: Cody Tseng Date: Sat, 12 Aug 2023 12:02:15 +0800 Subject: [PATCH] feat: add rate limiting (#53) --- package-lock.json | 43 +++++++++++++++++++- package.json | 3 +- src/app.module.ts | 6 +++ src/common/exceptions/index.ts | 1 + src/common/exceptions/throttler.exception.ts | 7 ++++ src/common/guards/index.ts | 1 + src/common/guards/ws-throttler.guard.spec.ts | 39 ++++++++++++++++++ src/common/guards/ws-throttler.guard.ts | 23 +++++++++++ src/config/config.spec.ts | 4 ++ src/config/config.ts | 2 + src/config/environment.ts | 9 ++++ src/config/index.ts | 1 + src/config/throttler.config.ts | 8 ++++ src/nostr/nostr.gateway.ts | 4 +- 14 files changed, 147 insertions(+), 4 deletions(-) create mode 100644 src/common/exceptions/throttler.exception.ts create mode 100644 src/common/guards/index.ts create mode 100644 src/common/guards/ws-throttler.guard.spec.ts create mode 100644 src/common/guards/ws-throttler.guard.ts create mode 100644 src/config/throttler.config.ts diff --git a/package-lock.json b/package-lock.json index 1bbd2218..e6663028 100644 --- a/package-lock.json +++ b/package-lock.json @@ -15,6 +15,7 @@ "@nestjs/event-emitter": "^2.0.1", "@nestjs/platform-express": "^10.1.3", "@nestjs/platform-ws": "^10.1.3", + "@nestjs/throttler": "^4.2.1", "@nestjs/typeorm": "^10.0.0", "@nestjs/websockets": "^10.1.3", "@noble/curves": "^1.1.0", @@ -1976,6 +1977,19 @@ } } }, + "node_modules/@nestjs/throttler": { + "version": "4.2.1", + "resolved": "https://registry.npmjs.org/@nestjs/throttler/-/throttler-4.2.1.tgz", + "integrity": "sha512-wVPMuIyr0KdrK1RVVQceWVNesogCm9IgYC1V5EkaTZ+usIE4qxEyzdwU5IqQLgOO/Loiq98MLwReDxazX7i9Uw==", + "dependencies": { + "md5": "^2.2.1" + }, + "peerDependencies": { + "@nestjs/common": "^7.0.0 || ^8.0.0 || ^9.0.0 || ^10.0.0", + "@nestjs/core": "^7.0.0 || ^8.0.0 || ^9.0.0 || ^10.0.0", + "reflect-metadata": "^0.1.13" + } + }, "node_modules/@nestjs/typeorm": { "version": "10.0.0", "resolved": "https://registry.npmjs.org/@nestjs/typeorm/-/typeorm-10.0.0.tgz", @@ -3602,6 +3616,14 @@ "integrity": "sha512-mT8iDcrh03qDGRRmoA2hmBJnxpllMR+0/0qlzjqZES6NdiWDcZkCNAk4rPFZ9Q85r27unkiNNg8ZOiwZXBHwcA==", "dev": true }, + "node_modules/charenc": { + "version": "0.0.2", + "resolved": "https://registry.npmjs.org/charenc/-/charenc-0.0.2.tgz", + "integrity": "sha512-yrLQ/yVUFXkzg7EDQsPieE/53+0RlaWTs+wBrvW36cyilJ2SaDWfl4Yj7MtLTXleV9uEKefbAGUPv2/iWSooRA==", + "engines": { + "node": "*" + } + }, "node_modules/chokidar": { "version": "3.5.3", "resolved": "https://registry.npmjs.org/chokidar/-/chokidar-3.5.3.tgz", @@ -4115,6 +4137,14 @@ "node": ">= 8" } }, + "node_modules/crypt": { + "version": "0.0.2", + "resolved": "https://registry.npmjs.org/crypt/-/crypt-0.0.2.tgz", + "integrity": "sha512-mCxBlsHFYh9C+HVpiEacem8FEBnMXgU9gy4zmNC+SXAZNB/1idgp/aulFJ4FgCi7GPEVbfyng092GqL2k2rmow==", + "engines": { + "node": "*" + } + }, "node_modules/date-fns": { "version": "2.30.0", "resolved": "https://registry.npmjs.org/date-fns/-/date-fns-2.30.0.tgz", @@ -6350,8 +6380,7 @@ "node_modules/is-buffer": { "version": "1.1.6", "resolved": "https://registry.npmjs.org/is-buffer/-/is-buffer-1.1.6.tgz", - "integrity": "sha512-NcdALwpXkTm5Zvvbk7owOUSvVvBKDgKP5/ewfXEznmQFfs4ZRmanOeKBTjRVjka3QFoN6XJ+9F3USqfHqTaU5w==", - "dev": true + "integrity": "sha512-NcdALwpXkTm5Zvvbk7owOUSvVvBKDgKP5/ewfXEznmQFfs4ZRmanOeKBTjRVjka3QFoN6XJ+9F3USqfHqTaU5w==" }, "node_modules/is-core-module": { "version": "2.13.0", @@ -7682,6 +7711,16 @@ "node": ">=0.10.0" } }, + "node_modules/md5": { + "version": "2.3.0", + "resolved": "https://registry.npmjs.org/md5/-/md5-2.3.0.tgz", + "integrity": "sha512-T1GITYmFaKuO91vxyoQMFETst+O71VUPEU3ze5GNzDm0OWdP8v1ziTaAEPUr/3kLsY3Sftgz242A1SetQiDL7g==", + "dependencies": { + "charenc": "0.0.2", + "crypt": "0.0.2", + "is-buffer": "~1.1.6" + } + }, "node_modules/media-typer": { "version": "0.3.0", "resolved": "https://registry.npmjs.org/media-typer/-/media-typer-0.3.0.tgz", diff --git a/package.json b/package.json index 567e719e..b68ea896 100644 --- a/package.json +++ b/package.json @@ -34,6 +34,7 @@ "@nestjs/event-emitter": "^2.0.1", "@nestjs/platform-express": "^10.1.3", "@nestjs/platform-ws": "^10.1.3", + "@nestjs/throttler": "^4.2.1", "@nestjs/typeorm": "^10.0.0", "@nestjs/websockets": "^10.1.3", "@noble/curves": "^1.1.0", @@ -96,4 +97,4 @@ "testEnvironment": "node", "maxWorkers": 1 } -} \ No newline at end of file +} diff --git a/src/app.module.ts b/src/app.module.ts index a005667e..6bae1f9a 100644 --- a/src/app.module.ts +++ b/src/app.module.ts @@ -1,6 +1,7 @@ import { Module } from '@nestjs/common'; import { ConfigModule, ConfigService } from '@nestjs/config'; import { EventEmitterModule } from '@nestjs/event-emitter'; +import { ThrottlerModule } from '@nestjs/throttler'; import { TypeOrmModule } from '@nestjs/typeorm'; import { LoggerModule, PinoLogger } from 'nestjs-pino'; import { loggerModuleFactory } from './common/utils/logger-module-factory'; @@ -38,6 +39,11 @@ import { NostrModule } from './nostr/nostr.module'; }, inject: [ConfigService, PinoLogger], }), + ThrottlerModule.forRootAsync({ + useFactory: (configService: ConfigService) => + configService.get('throttler', { infer: true }), + inject: [ConfigService], + }), NostrModule, ], }) diff --git a/src/common/exceptions/index.ts b/src/common/exceptions/index.ts index 017b8c7f..f4264ca8 100644 --- a/src/common/exceptions/index.ts +++ b/src/common/exceptions/index.ts @@ -1,3 +1,4 @@ export * from './client.exception'; export * from './restricted.exception'; export * from './validation.exception'; +export * from './throttler.exception'; diff --git a/src/common/exceptions/throttler.exception.ts b/src/common/exceptions/throttler.exception.ts new file mode 100644 index 00000000..cf7fb3d8 --- /dev/null +++ b/src/common/exceptions/throttler.exception.ts @@ -0,0 +1,7 @@ +import { ClientException } from './client.exception'; + +export class ThrottlerException extends ClientException { + constructor() { + super('rate-limited: slow down there chief'); + } +} diff --git a/src/common/guards/index.ts b/src/common/guards/index.ts new file mode 100644 index 00000000..d46ca065 --- /dev/null +++ b/src/common/guards/index.ts @@ -0,0 +1 @@ +export * from './ws-throttler.guard'; diff --git a/src/common/guards/ws-throttler.guard.spec.ts b/src/common/guards/ws-throttler.guard.spec.ts new file mode 100644 index 00000000..f23018f7 --- /dev/null +++ b/src/common/guards/ws-throttler.guard.spec.ts @@ -0,0 +1,39 @@ +import { createMock } from '@golevelup/ts-jest'; +import { ExecutionContext } from '@nestjs/common'; +import { Reflector } from '@nestjs/core'; +import { ThrottlerStorageService } from '@nestjs/throttler'; +import { WsThrottlerGuard } from './ws-throttler.guard'; + +describe('WsThrottlerGuard', () => { + const context = createMock({ + getClass: jest.fn().mockReturnValue({ name: 'Test' }), + getHandler: jest.fn().mockReturnValue({ name: 'test' }), + switchToWs: jest.fn().mockReturnValue({ + getClient: jest.fn().mockReturnValue({ id: 'test' }), + }), + }); + + let storageService: ThrottlerStorageService; + + beforeEach(() => { + storageService = new ThrottlerStorageService(); + }); + + afterEach(() => { + storageService.onApplicationShutdown(); + }); + + it('should be fine', async () => { + const guard = new WsThrottlerGuard( + { limit: 2, ttl: 2 }, + storageService, + new Reflector(), + ); + + await expect(guard.canActivate(context)).resolves.toBe(true); + await expect(guard.canActivate(context)).resolves.toBe(true); + await expect(guard.canActivate(context)).rejects.toThrowError( + 'rate-limited: slow down there chief', + ); + }); +}); diff --git a/src/common/guards/ws-throttler.guard.ts b/src/common/guards/ws-throttler.guard.ts new file mode 100644 index 00000000..0b48172c --- /dev/null +++ b/src/common/guards/ws-throttler.guard.ts @@ -0,0 +1,23 @@ +import { ExecutionContext, Injectable } from '@nestjs/common'; +import { ThrottlerGuard } from '@nestjs/throttler'; +import { WebSocket } from 'ws'; +import { ThrottlerException } from '../exceptions'; + +@Injectable() +export class WsThrottlerGuard extends ThrottlerGuard { + protected async handleRequest( + context: ExecutionContext, + limit: number, + ttl: number, + ): Promise { + const client = context.switchToWs().getClient(); + const key = this.generateKey(context, client.id); + const { totalHits } = await this.storageService.increment(key, ttl); + + if (totalHits > limit) { + throw new ThrottlerException(); + } + + return true; + } +} diff --git a/src/config/config.spec.ts b/src/config/config.spec.ts index cb5e249e..cbf0f948 100644 --- a/src/config/config.spec.ts +++ b/src/config/config.spec.ts @@ -17,6 +17,8 @@ describe('config', () => { LOG_LEVEL: 'info', EVENT_CREATED_AT_UPPER_LIMIT: '60', EVENT_ID_MIN_LEADING_ZERO_BITS: '16', + THROTTLER_LIMIT: '100', + THROTTLER_TTL: '1', }), ).toEqual({ DOMAIN: 'localhost', @@ -26,6 +28,8 @@ describe('config', () => { LOG_LEVEL: 'info', EVENT_CREATED_AT_UPPER_LIMIT: 60, EVENT_ID_MIN_LEADING_ZERO_BITS: 16, + THROTTLER_LIMIT: 100, + THROTTLER_TTL: 1, }); }); }); diff --git a/src/config/config.ts b/src/config/config.ts index 2de9ee49..2eace6d3 100644 --- a/src/config/config.ts +++ b/src/config/config.ts @@ -4,6 +4,7 @@ import { limitConfig } from './limit.config'; import { loggerConfig } from './logger.config'; import { meiliSearchConfig } from './meili-search'; import { relayInfoDocConfig } from './relay-info-doc.config'; +import { throttlerConfig } from './throttler.config'; export function config() { const env = validateEnvironment(process.env); @@ -15,6 +16,7 @@ export function config() { limit: limitConfig(env), relayInfoDoc: relayInfoDocConfig(env), logger: loggerConfig(env), + throttler: throttlerConfig(env), }; } export type Config = ReturnType; diff --git a/src/config/environment.ts b/src/config/environment.ts index 3bfeb25a..a36b2fa9 100644 --- a/src/config/environment.ts +++ b/src/config/environment.ts @@ -35,6 +35,15 @@ export const EnvironmentSchema = z.object({ .string() .transform((minLeadingZeroBits) => parseInt(minLeadingZeroBits)) .optional(), + + THROTTLER_LIMIT: z + .string() + .transform((limit) => parseInt(limit)) + .optional(), + THROTTLER_TTL: z + .string() + .transform((ttl) => parseInt(ttl)) + .optional(), }); export type Environment = z.infer; diff --git a/src/config/index.ts b/src/config/index.ts index 569e93fc..a26b299c 100644 --- a/src/config/index.ts +++ b/src/config/index.ts @@ -4,3 +4,4 @@ export * from './environment'; export * from './limit.config'; export * from './meili-search'; export * from './relay-info-doc.config'; +export * from './throttler.config'; diff --git a/src/config/throttler.config.ts b/src/config/throttler.config.ts new file mode 100644 index 00000000..b72aa76d --- /dev/null +++ b/src/config/throttler.config.ts @@ -0,0 +1,8 @@ +import { Environment } from './environment'; + +export function throttlerConfig(env: Environment) { + return { + limit: env.THROTTLER_LIMIT ?? 100, + ttl: env.THROTTLER_TTL ?? 1, + }; +} diff --git a/src/nostr/nostr.gateway.ts b/src/nostr/nostr.gateway.ts index 8e90e9f4..9ccb2e66 100644 --- a/src/nostr/nostr.gateway.ts +++ b/src/nostr/nostr.gateway.ts @@ -1,4 +1,4 @@ -import { UseFilters } from '@nestjs/common'; +import { UseFilters, UseGuards } from '@nestjs/common'; import { ConfigService } from '@nestjs/config'; import { ConnectedSocket, @@ -15,6 +15,7 @@ import { concatWith, filter, from, map, of } from 'rxjs'; import { WebSocket, WebSocketServer } from 'ws'; import { RestrictedException } from '../common/exceptions'; import { WsExceptionFilter } from '../common/filters'; +import { WsThrottlerGuard } from '../common/guards'; import { ZodValidationPipe } from '../common/pipes'; import { Config, LimitConfig } from '../config'; import { MessageType } from './constants'; @@ -40,6 +41,7 @@ import { @WebSocketGateway() @UseFilters(WsExceptionFilter) +@UseGuards(WsThrottlerGuard) export class NostrGateway implements OnGatewayInit, OnGatewayConnection, OnGatewayDisconnect {