NestJS+Redis实现手写一个限流器

import { Injectable } from '@nestjs/common';
import * as Redis from 'ioredis';
import { ConfigService } from '@app/common';
import { createHash } from 'crypto'
import { v4 as uuidv4 } from 'uuid';

const rateLimitScript = ""// 后面单独列出

@Injectable()
export class RedisService {
    private readonly redisClient: Redis.Redis;
    private luaScript: any;
    constructor(
        private readonly configService: ConfigService,
    ) {

        const self = this;
        const connConfig = this.configService.get("redisService")
        this.redisClient = new Redis.Redis(connConfig)
        this.luaScript = {
            rateLimit: {
                script: rateLimitScript,
                hash: self.hashStr(rateLimitScript)
            },
        }
        
    }

    private hashStr(value: string) {
        return createHash("sha1").update(value).digest('hex')
    }

    async rateLimit(opts: any): Promise<boolean> {

        const { key, limit, windowSize } = opts;
        const uuid = uuidv4()
        let result;
        const { script, hash } = this.luaScript.rateLimit
  
        try {

            const shaResult = await this.redisClient.evalsha(hash, 1, key, limit, windowSize, uuid)
            result = shaResult

        } catch (error) {

            const shaResult = await this.redisClient.eval(script, 1, key, limit, windowSize, uuid)
            result = shaResult

        }
        return result == 1
    }
}

--传入四个参数 分别是key,限制次数,时间范围,唯一值
local key = KEYS[1]
local limit = tonumber(ARGV[1])
local windowSize = tonumber(ARGV[2]) --单位毫秒
local uuid = ARGV[3] -- 唯一值是为了防止zset 重复

-- 使用redis 来获取时间，防止多进程生成相似的边界导致超频。时间单位是微秒
local date = redis.call("time")
local now = tonumber(date[1]) * 1000000 + tonumber(date[2])
local startTime = now - windowSize * 1000
local endTime = now +  1000000

-- 计算过期时间 时间单位是秒
local expireSec = tonumber(math.ceil(windowSize / 1000)) + 1

-- 统计当前zset数组里的数据,超出范围则返回0,
-- 否则做3件事,然后返回1
-- 1、向数组里增加新值
-- 2、删除数组中开始时间之前的数据,防止数组过大
-- 3、给数组续过期时间
local count = tonumber(redis.call('zcount', key, startTime, endTime))

if count + 1 > limit then
    return 0
else
    redis.call('zadd', key, now, uuid)
    redis.call('zremrangebyscore', key, 0, startTime - 100000)
    redis.call('expire', key, expireSec)
    return 1
end

import { SetMetadata } from '@nestjs/common';

export interface rateLimitOptions {
    keyPrefix: string,
    limit: number,
    windowSize: number
}

export const RateLimit = (options: rateLimitOptions): MethodDecorator => SetMetadata('rateLimit', options)

import { Injectable, ExecutionContext, CanActivate } from "@nestjs/common";
import { Reflector } from '@nestjs/core';
import { RedisService, rateLimitOptions } from "@app/common";
import { BusinessException } from "@app/common";

@Injectable()
export class RateLimitGuard implements CanActivate {
    constructor(
        private reflector: Reflector,
        private redisService: RedisService
    ) { }

    private getIpFromRequest(request: { ip: string }): string {
        return request.ip?.match(/\b\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}\b/)?.[0]
    }

    async canActivate(context: ExecutionContext) {

         // 通过反射拿到前面设置的值
        const rateLimitConfig = this.reflector.get<rateLimitOptions>("rateLimit", context.getHandler());
           if (!rateLimitConfig) {
			 // 当前接口如果没设置参数则定义默认参数
            const cMethod = this.reflector.get("method", context.getHandler());// 是GET,POST 等http method
            const cPath = this.reflector.get("path", context.getHandler());// 接口的具体路径
            rateLimitConfig = {
                keyPrefix: cMethod + ":" + cPath,
                limit: 1,
                windowSize: 5000
            }
        }

        const { keyPrefix, limit, windowSize } = rateLimitConfig
        const request = context.switchToHttp().getRequest();
        const ip = this.getIpFromRequest(request)
        const key = keyPrefix + ":" + ip

        const isPass = await this.redisService.rateLimit({
            key,
            limit,
            windowSize
        })
        if (!isPass) {
            // 返回自定义的错误
            throw new BusinessException("RATE_LIMIT_EXCEEDED_LIMIT")
        }
        return true
    }
}

@Controller('user')
export class UserController {
  constructor(private readonly userService: UserService) { }
  
  @Public()
  @RateLimit({ keyPrefix: "login", limit: 3, windowSize: 1000 })
  @UseGuards(RateLimitGuard)
  @Post('register')
  register(@Body() createUserDto: CreateUserDto) {
    return this.userService.register(createUserDto);
   }
}

@Module({
  providers: [
    {
      provide: APP_GUARD,
      useClass: RateLimitGuard
    }
  ],
})