| 1 | import type { |
| 2 | CanActivate, |
| 3 | ExecutionContext, |
| 4 | } from '@nestjs/common' |
| 5 | import { |
| 6 | HttpException, |
| 7 | HttpStatus, |
| 8 | Injectable, |
| 9 | Logger, |
| 10 | } from '@nestjs/common' |
| 11 | import { Reflector } from '@nestjs/core' |
| 12 | import { ServerRedisService } from '../redis' |
| 13 | |
| 14 | export const RATE_LIMIT_KEY = 'rate_limit' |
| 15 | |
| 16 | export interface RateLimitOptions { |
| 17 | /** |
| 18 | * 时间窗口（秒） |
| 19 | */ |
| 20 | ttl: number |
| 21 | /** |
| 22 | * 时间窗口内允许的最大请求次数 |
| 23 | */ |
| 24 | limit: number |
| 25 | /** |
| 26 | * 自定义键生成器 |
| 27 | */ |
| 28 | keyGenerator?: (req: any) => string |
| 29 | } |
| 30 | |
| 31 | /** |
| 32 | * 装饰器：设置接口速率限制 |
| 33 | */ |
| 34 | export function RateLimit(options: RateLimitOptions) { |
| 35 | return (target: any, propertyKey?: string, descriptor?: PropertyDescriptor) => { |
| 36 | if (propertyKey && descriptor) { |
| 37 | Reflect.defineMetadata(RATE_LIMIT_KEY, options, descriptor.value) |
| 38 | } |
| 39 | else { |
| 40 | Reflect.defineMetadata(RATE_LIMIT_KEY, options, target) |
| 41 | } |
| 42 | return descriptor |
| 43 | } |
| 44 | } |
| 45 | |
| 46 | @Injectable() |
| 47 | export class RateLimitGuard implements CanActivate { |
| 48 | private readonly logger = new Logger(RateLimitGuard.name) |
| 49 | private readonly reflector = new Reflector() |
| 50 | |
| 51 | constructor(private readonly redisService: ServerRedisService) {} |
| 52 | |
| 53 | async canActivate(context: ExecutionContext): Promise<boolean> { |
| 54 | const rateLimitOptions = this.reflector.getAllAndOverride<RateLimitOptions | undefined>( |
| 55 | RATE_LIMIT_KEY, |
| 56 | [context.getHandler(), context.getClass()], |
| 57 | ) |
| 58 | |
| 59 | if (!rateLimitOptions) { |
| 60 | return true |
| 61 | } |
| 62 | |
| 63 | const request = context.switchToHttp().getRequest() |
| 64 | const { ttl, limit, keyGenerator } = rateLimitOptions |
| 65 | |
| 66 | // 生成限流键 |
| 67 | const key = keyGenerator |
| 68 | ? keyGenerator(request) |
| 69 | : this.getDefaultKey(request) |
| 70 | |
| 71 | try { |
| 72 | const count = await this.redisService.incrementRateLimit(key, ttl) |
| 73 | |
| 74 | // 设置响应头 |
| 75 | const response = context.switchToHttp().getResponse() |
| 76 | response.setHeader('X-RateLimit-Limit', limit.toString()) |
| 77 | response.setHeader('X-RateLimit-Reset', (Date.now() + ttl * 1000).toString()) |
| 78 | |
| 79 | // 检查是否超过限制 |
| 80 | if (count > limit) { |
| 81 | response.setHeader('X-RateLimit-Remaining', '0') |
| 82 | this.logger.warn(`Rate limit exceeded for key: ${key}, count: ${count}, limit: ${limit}`) |
| 83 | throw new HttpException( |
| 84 | { |
| 85 | code: HttpStatus.TOO_MANY_REQUESTS, |
| 86 | message: 'Too many requests', |
| 87 | data: { ttl }, |
| 88 | }, |
| 89 | HttpStatus.TOO_MANY_REQUESTS, |
| 90 | ) |
| 91 | } |
| 92 | |
| 93 | response.setHeader('X-RateLimit-Remaining', (limit - count).toString()) |
| 94 | |
| 95 | return true |
| 96 | } |
| 97 | catch (error) { |
| 98 | if (error instanceof HttpException) { |
| 99 | throw error |
| 100 | } |
| 101 | this.logger.fatal(error, `Rate limit check failed`) |
| 102 | // 如果Redis出错，允许请求通过（优雅降级） |
| 103 | return true |
| 104 | } |
| 105 | } |
| 106 | |
| 107 | private getDefaultKey(request: any): string { |
| 108 | // 优先使用用户ID，其次使用IP地址 |
| 109 | const userId = request.user?.id |
| 110 | const ip = this.getClientIp(request) |
| 111 | const path = request.route?.path || request.url |
| 112 | |
| 113 | return userId ? `user:${userId}:${path}` : `ip:${ip}:${path}` |
| 114 | } |
| 115 | |
| 116 | private getClientIp(request: any): string { |
| 117 | return ( |
| 118 | request.headers['x-forwarded-for']?.split(',')[0] |
| 119 | || request.headers['x-real-ip'] |
| 120 | || request.connection?.remoteAddress |
| 121 | || request.socket?.remoteAddress |
| 122 | || 'unknown' |
| 123 | ) |
| 124 | } |
| 125 | } |
| 126 |