| 123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108 |
- import {
- CanActivate,
- ExecutionContext,
- Injectable,
- Logger,
- HttpException,
- HttpStatus,
- } from '@nestjs/common';
- import type { FastifyRequest } from 'fastify';
- interface RateLimitBucket {
- count: number;
- resetAt: number;
- }
- @Injectable()
- export class RateLimitGuard implements CanActivate {
- private readonly logger = new Logger(RateLimitGuard.name);
- private readonly buckets = new Map<string, RateLimitBucket>();
- private readonly limit = 10; // Max requests
- private readonly windowMs = 60_000; // 1 minute window
- private readonly cleanupInterval: NodeJS.Timeout;
- constructor() {
- // Cleanup expired buckets every 5 minutes
- this.cleanupInterval = setInterval(() => {
- this.cleanup();
- }, 5 * 60_000);
- }
- canActivate(context: ExecutionContext): boolean {
- const request = context.switchToHttp().getRequest<FastifyRequest>();
- const ip = this.getClientIp(request);
- const endpoint = `${request.method}:${request.url.split('?')[0]}`;
- const key = `${ip}:${endpoint}`;
- const now = Date.now();
- let bucket = this.buckets.get(key);
- if (!bucket || now > bucket.resetAt) {
- // Create new bucket or reset expired one
- bucket = {
- count: 1,
- resetAt: now + this.windowMs,
- };
- this.buckets.set(key, bucket);
- return true;
- }
- bucket.count++;
- if (bucket.count > this.limit) {
- const retryAfter = Math.ceil((bucket.resetAt - now) / 1000);
- this.logger.warn(
- `Rate limit exceeded for ${ip} on ${endpoint} (${bucket.count}/${this.limit})`,
- );
- throw new HttpException(
- {
- statusCode: 429,
- message: `Too many requests. Please try again in ${retryAfter} seconds.`,
- code: 'RATE_LIMITED',
- retryAfter,
- },
- HttpStatus.TOO_MANY_REQUESTS,
- );
- }
- return true;
- }
- private getClientIp(request: FastifyRequest): string {
- // Check common headers for real IP
- const forwardedFor = request.headers['x-forwarded-for'];
- if (forwardedFor) {
- const ips = Array.isArray(forwardedFor) ? forwardedFor[0] : forwardedFor;
- return ips.split(',')[0].trim();
- }
- const realIp = request.headers['x-real-ip'];
- if (realIp) {
- return Array.isArray(realIp) ? realIp[0] : realIp;
- }
- return request.ip || 'unknown';
- }
- private cleanup() {
- const now = Date.now();
- let cleaned = 0;
- for (const [key, bucket] of this.buckets.entries()) {
- if (now > bucket.resetAt) {
- this.buckets.delete(key);
- cleaned++;
- }
- }
- if (cleaned > 0) {
- this.logger.debug(`Cleaned up ${cleaned} expired rate limit buckets`);
- }
- }
- onModuleDestroy() {
- if (this.cleanupInterval) {
- clearInterval(this.cleanupInterval);
- }
- }
- }
|