从nestjs-rate-limiter学写一个拦截器

无意在github看到 nestjs-rate-limiter 的issue中有一个关于自定义error body的特性 #60 , 之前用到过这个库, 感觉不能自定义error body太伤了, 毕竟大多数后台应用都有一套自己定义的返回值消息格式, 而不是基于标准的web HttpException , 于是自己提交了一个pull request来实现该需求, 看到源码时候突发奇想有了本文, 何不写个博客来复习下拦截器interceptor呢 ？

开始

先看一下基本的 nestjs-rate-limter 使用方法, 比如我们想对接口做访问频率控制, 1分钟只能访问1次, 我们需要注册一个全局的rate limter拦截器(通常在App Module中)

@Module({
  imports: [	 
    RateLimiterModule.register({
      for: 'Express',
      type: 'Memory',
      points: 1,
      duration: 60, // 60秒访问1次
        })],
  providers: [
    {
      provide: APP_INTERCEPTOR,
      useClass: RateLimiterInterceptor,
    }
 ]})

既然是全局的拦截器, 可以使用 useGlobalInterceptors 来注册全局拦截器吗？

const app = await NestFactory.create(ApplicationModule);
app.useGlobalInterceptors(new LoggingInterceptor());

答案: 不可以！ 因为使用useGlobalInterceptors注册就无法拿到 RateLimiterModule.register 的参数了。

只需上面的配置, 我们就已经对接口做了访问频率限制。下面让我们来看看RateLimiterInterceptor这个拦截器是如何编写的。

Typescript

Nest.js对TS支持非常好, 第三方库一般都会被包成一个"TS完备"的Nest.js模块, 开发起来体验非常好, 因为有各种代码提示和API说明。所以开发拦截器的第一步我们先要定义TS类型, 让其他开发者使用我们的库会感觉非常舒服。 从前面 register 的使用来看, 它是一个静态方法, 并接收一个对象作为配置参数。

export class RateLimiterModule {
	static register(options: RateLimiterOptions = defaultRateLimiterOptions): DynamicModule {
		return {
			module: RateLimiterModule,
			providers: [{ provide: 'RATE_LIMITER_OPTIONS', useValue: options }]
		}
	}

配置参数是 RateLimiterOptions 类型, 定义这个类型很关键, 因为一般从配置参数就可以大致看出这个包所提供的功能, 下面是 RateLimiterOptions 的定义

export interface RateLimiterOptions {
	for?: 'Express' | 'Fastify' | 'Microservice' | 'ExpressGraphql' | 'FastifyGraphql'
	type?: 'Memory' | 'Redis' | 'Memcache' | 'Postgres' | 'MySQL' | 'Mongo'
	keyPrefix?: string
	points?: number
	pointsConsumed?: number
	inmemoryBlockDuration?: number
	duration?: number
	blockDuration?: number
	inmemoryBlockOnConsum

编写一个拦截器

拦截器有全局的和非全局的, 全局的就如我们前面介绍的注册方法, 在根模块的 provider 中增加一个对象, provide 属性为固定的 APP_INTERCEPTOR （由@nestjs/core导出）, useClass 就是一个拦截器:

拦截器就是使用 @Injectable 装饰并且实现 NestInterceptor 的普通类

  providers: [
    {
      provide: APP_INTERCEPTOR,
      useClass: RateLimiterInterceptor,
    }]

比如 nest-rate-limiter

@Injectable()
export class RateLimiterInterceptor implements NestInterceptor {
  	...
		async intercept(context: ExecutionContext, next: CallHandler): Promise<any> {
      // 拦截到请求后, 这里实现对请求进行限制的逻辑
      
    	return next.handle() // 最后转移控制权
    }
		...
}

这么看来拦截器很像middleware呀？确实！ 实现接口限流也可以使用 express 的 express-rate-limit 包来实现, 它就是使用middleware实现的。

如果想局部使用 RateLimiterInterceptor 拦截器的话, 可以在路由上使用 UseInterceptors 

@UseInterceptors(RateLimiterInterceptor)
@Get('/login')
public async login() {
    console.log('hello');
}

这样, 这个拦截器只对 /login 路由生效。

nest-rate-limiter 实现细节

rate-limiter拦截器如何获取到配置的参数?

我们在根模块提供了一个配置对象

RateLimiterModule.register({
      for: 'Express',

那在拦截器中如何获取这个配置参数呢? 使用注入的方式

class RateLimiterModule {
	static register(options: RateLimiterOptions = defaultRateLimiterOptions): DynamicModule {
		return {
			module: RateLimiterModule,
      // 接收到配置对象options后, 给它唯一标识RATE_LIMITER_OPTIONS, 为了注入使用
			providers: [{ provide: 'RATE_LIMITER_OPTIONS', useValue: options }] 
		}
	}

在拦截器中可以直接将 RATE_LIMITER_OPTIONS 注入到拦截器中

@Injectable()
export class RateLimiterInterceptor implements NestInterceptor {
	constructor(@Inject('RATE_LIMITER_OPTIONS') private options: RateLimiterOptions) {}
  
  // 使用@Inject注入配置后 我们就可以获取到传入的options了, 比如 for:express

如何实现对某个接口进行自定义限制？

如果我们想设置一个全局的拦截器并配置了60秒访问1次, 但对 /login 接口单独设置一个限流规则比如1分钟访问10次该怎么办呢? nest-rate-limiter实现这个需求非常简单, 那他是如何实现的?

@RateLimit({ points: 10, duration: 60 })
@Get('/login')

先定义一个RateLimit的装饰器

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

RateLimit使用 SetMetadata 来定义来一个元数据, key是 rateLimit , options就是传入的 { points: 10, duration: 60 } 。

定义了元数据后, 可以在拦截器中通过 Reflector 类来获取到元数据

@Injectable()
export class RateLimiterInterceptor implements NestInterceptor {
  // 注入Reflector
	constructor(@Inject('Reflector') private readonly reflector: Reflector) {}
  async intercept(context: ExecutionContext, next: CallHandler): Promise<any> {
    // 获取到 { points: 10, duration: 60 }
		const reflectedOptions = this.reflector.get<RateLimiterOptions>('rateLimit', context.getHandler())
  }
}

逻辑限流是怎样的?

1. 根据配置生成 rateLimiter 

let rateLimiter: RateLimiterAbstract = this.rateLimiters.get(libraryArguments.keyPrefix)

	if (!rateLimiter) {
			switch (this.spesificOptions?.type || this.options.type) {
				case 'Memory':
					rateLimiter = new RateLimiterMemory(libraryArguments)
					Logger.log(`Rate Limiter started with ${limiterOptions.keyPrefix} key prefix`, 'RateLimiterMemory')
					break;
			...

nest-rate-limiter只是将rate-limiter-flexible以拦截器的方式包装成nest.js的模块, 限流逻辑是由底层的 rate-limiter-flexible 实现的, 比如第6行的 RateLimiterMemory类。

2. 生成的 rateLimiter 提供了对IP访问管控的方法比如 consume 、 block 。每当某个IP请求访问一次, 就调用一次 consume 方法来消耗配置的 points 如此往复
3. 当请求次数过多, 将所有的points消耗光了, 则抛出一个 HttpException 异常, 状态为429

写在最后的

希望通过对nest-rate-limiter的源码分析, 大家可以理解nest.js拦截器的相关概念, 并且可以编写自己的拦截器去解决实际开发中遇到的问题。

都看到这里了, 点个赞吧！