All files / middleware rate-limiter.ts

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
import { RateLimiterMemory, RateLimiterRedis } from 'rate-limiter-flexible';
import { redisService } from '../services/redis.service';
 
interface RateLimitOptions {
  points: number;        // 限流点数
  duration: number;      // 时间窗口（秒）
  blockDuration?: number; // 封锁时间（秒）
  keyGenerator?: (ctx: any) => string; // 自定义键生成器
}
 
// 内存限流器（Redis 不可用时使用）
const memoryLimiters = new Map<string, RateLimiterMemory>();
 
// Redis 限流器
const redisLimiters = new Map<string, RateLimiterRedis>();
 
/**
 * 获取或创建限流器
 */
function getLimiter(name: string, options: RateLimitOptions) {
  if (redisService.isAvailable()) {
    if (!redisLimiters.has(name)) {
      const limiter = new RateLimiterRedis({
        storeClient: redisService['client'],
        keyPrefix: `rate_limit:${name}:`,
        points: options.points,
        duration: options.duration,
        blockDuration: options.blockDuration || options.duration * 2,
      });
      redisLimiters.set(name, limiter);
    }
    return redisLimiters.get(name)!;
  } else {
    if (!memoryLimiters.has(name)) {
      const limiter = new RateLimiterMemory({
        points: options.points,
        duration: options.duration,
      });
      memoryLimiters.set(name, limiter);
    }
    return memoryLimiters.get(name)!;
  }
}
 
/**
 * 创建限流中间件
 * @param options 限流配置
 */
export function createRateLimiter(options: RateLimitOptions) {
  const limiter = getLimiter('default', options);
 
  return async (ctx: any, next: any) => {
    try {
      const key = options.keyGenerator 
        ? options.keyGenerator(ctx) 
        : ctx.ip;
 
      await limiter.consume(key, 1);
      return await next();
    } catch (rejRes: any) {
      const retrySecs = Math.round(rejRes.msBeforeNext / 1000) || 1;
      
      ctx.set('Retry-After', String(retrySecs));
      ctx.status = 429;
      ctx.body = {
        code: 429,
        message: '请求过于频繁，请稍后再试',
        retryAfter: retrySecs,
      };
    }
  };
}
 
/**
 * API 全局限流（每分钟 100 次）
 */
export const apiRateLimiter = createRateLimiter({
  points: 100,
  duration: 60,
  keyGenerator: (ctx) => `api:${ctx.ip}`,
});
 
/**
 * 登录接口限流（每分钟 5 次）
 */
export const loginRateLimiter = createRateLimiter({
  points: 5,
  duration: 60,
  blockDuration: 300, // 封锁 5 分钟
  keyGenerator: (ctx) => `login:${ctx.ip}`,
});
 
/**
 * 验证码发送限流（每分钟 1 次，每小时 5 次）
 */
export const smsRateLimiter = createRateLimiter({
  points: 1,
  duration: 60,
  blockDuration: 3600,
  keyGenerator: (ctx) => `sms:${ctx.request.body?.phone || ctx.ip}`,
});
 
/**
 * TTS 生成限流（根据用户等级）
 */
export const ttsRateLimiter = createRateLimiter({
  points: 20,
  duration: 60,
  keyGenerator: (ctx) => `tts:${ctx.state.user?.id || ctx.ip}`,
});
 
/**
 * 文件上传限流（每分钟 10 次）
 */
export const uploadRateLimiter = createRateLimiter({
  points: 10,
  duration: 60,
  keyGenerator: (ctx) => `upload:${ctx.state.user?.id || ctx.ip}`,
});