chore(server): clean up throttler (#6326)

packages/backend/server/src/fundamentals/cache/index.ts

@@ -1,11 +1,12 @@
 import { Global, Module } from '@nestjs/common';
 
 import { Cache, SessionCache } from './instances';
+import { CacheInterceptor } from './interceptor';
 
 @Global()
 @Module({
-  providers: [Cache, SessionCache],
-  exports: [Cache, SessionCache],
+  providers: [Cache, SessionCache, CacheInterceptor],
+  exports: [Cache, SessionCache, CacheInterceptor],
 })
 export class CacheModule {}
 export { Cache, SessionCache };

packages/backend/server/src/fundamentals/index.ts

@@ -27,7 +27,7 @@ export {
 export type { PrismaTransaction } from './prisma';
 export * from './storage';
 export { type StorageProvider, StorageProviderFactory } from './storage';
-export { AuthThrottlerGuard, CloudThrottlerGuard, Throttle } from './throttler';
+export { CloudThrottlerGuard, Throttle } from './throttler';
 export {
   getRequestFromHost,
   getRequestResponseFromContext,

packages/backend/server/src/fundamentals/throttler/decorators.ts

@@ -0,0 +1,38 @@
+import { applyDecorators, SetMetadata } from '@nestjs/common';
+import { SkipThrottle, Throttle as RawThrottle } from '@nestjs/throttler';
+
+export type Throttlers = 'default' | 'strict';
+export const THROTTLER_PROTECTED = 'affine_throttler:protected';
+
+/**
+ * Choose what throttler to use
+ *
+ * If a Controller or Query do not protected behind a Throttler,
+ * it will never be rate limited.
+ *
+ * - Ease: 120 calls within 60 seconds
+ * - Strict: 10 calls within 60 seconds
+ *
+ * @example
+ *
+ * \@Throttle()
+ * \@Throttle('strict')
+ *
+ * // the config call be override by the second parameter,
+ * // and the call count will be calculated separately
+ * \@Throttle('default', { limit: 10, ttl: 10 })
+ *
+ */
+export function Throttle(
+  type: Throttlers = 'default',
+  override: { limit?: number; ttl?: number } = {}
+): MethodDecorator & ClassDecorator {
+  return applyDecorators(
+    SetMetadata(THROTTLER_PROTECTED, type),
+    RawThrottle({
+      [type]: override,
+    })
+  );
+}
+
+export { SkipThrottle };

packages/backend/server/src/fundamentals/throttler/index.ts

@@ -1,15 +1,20 @@
 import { ExecutionContext, Global, Injectable, Module } from '@nestjs/common';
+import { Reflector } from '@nestjs/core';
 import {
-  Throttle,
+  InjectThrottlerOptions,
+  InjectThrottlerStorage,
   ThrottlerGuard,
   ThrottlerModule,
-  ThrottlerModuleOptions,
+  type ThrottlerModuleOptions,
+  ThrottlerOptions,
   ThrottlerOptionsFactory,
   ThrottlerStorageService,
 } from '@nestjs/throttler';
+import type { Request } from 'express';
 
 import { Config } from '../config';
 import { getRequestResponseFromContext } from '../utils/request';
+import { THROTTLER_PROTECTED, Throttlers } from './decorators';
 
 @Injectable()
 export class ThrottlerStorage extends ThrottlerStorageService {}
@@ -25,13 +30,16 @@ class CustomOptionsFactory implements ThrottlerOptionsFactory {
     const options: ThrottlerModuleOptions = {
       throttlers: [
         {
+          name: 'default',
           ttl: this.config.rateLimiter.ttl * 1000,
           limit: this.config.rateLimiter.limit,
         },
+        {
+          name: 'strict',
+          ttl: this.config.rateLimiter.ttl * 1000,
+          limit: 20,
+        },
       ],
-      skipIf: () => {
-        return !this.config.node.prod || this.config.affine.canary;
-      },
       storage: this.storage,
     };
 
@@ -39,6 +47,132 @@ class CustomOptionsFactory implements ThrottlerOptionsFactory {
   }
 }
 
+@Injectable()
+export class CloudThrottlerGuard extends ThrottlerGuard {
+  constructor(
+    @InjectThrottlerOptions() options: ThrottlerModuleOptions,
+    @InjectThrottlerStorage() storageService: ThrottlerStorage,
+    reflector: Reflector,
+    private readonly config: Config
+  ) {
+    super(options, storageService, reflector);
+  }
+
+  override getRequestResponse(context: ExecutionContext) {
+    return getRequestResponseFromContext(context) as any;
+  }
+
+  override getTracker(req: Request): Promise<string> {
+    return Promise.resolve(
+      //           ↓ prefer session id if available
+      `throttler:${req.sid ?? req.get('CF-Connecting-IP') ?? req.get('CF-ray') ?? req.ip}`
+      // ^ throttler prefix make the key in store recognizable
+    );
+  }
+
+  override generateKey(
+    context: ExecutionContext,
+    tracker: string,
+    throttler: string
+  ) {
+    if (tracker.endsWith(';custom')) {
+      return `${tracker};${throttler}:${context.getClass().name}.${context.getHandler().name}`;
+    }
+
+    return `${tracker};${throttler}`;
+  }
+
+  override async handleRequest(
+    context: ExecutionContext,
+    limit: number,
+    ttl: number,
+    throttlerOptions: ThrottlerOptions
+  ) {
+    // give it 'default' if no throttler is specified,
+    // so the unauthenticated users visits will always hit default throttler
+    // authenticated users will directly bypass unprotected APIs in [CloudThrottlerGuard.canActivate]
+    const throttler = this.getSpecifiedThrottler(context) ?? 'default';
+
+    // by pass unmatched throttlers
+    if (throttlerOptions.name !== throttler) {
+      return true;
+    }
+
+    const { req, res } = this.getRequestResponse(context);
+    const ignoreUserAgents =
+      throttlerOptions.ignoreUserAgents ?? this.commonOptions.ignoreUserAgents;
+    if (Array.isArray(ignoreUserAgents)) {
+      for (const pattern of ignoreUserAgents) {
+        const ua = req.headers['user-agent'];
+        if (ua && pattern.test(ua)) {
+          return true;
+        }
+      }
+    }
+
+    let tracker = await this.getTracker(req);
+
+    if (this.config.node.dev) {
+      limit = Number.MAX_SAFE_INTEGER;
+    } else {
+      // custom limit or ttl APIs will be treated standalone
+      if (limit !== throttlerOptions.limit || ttl !== throttlerOptions.ttl) {
+        tracker += ';custom';
+      }
+    }
+
+    const key = this.generateKey(
+      context,
+      tracker,
+      throttlerOptions.name ?? 'default'
+    );
+    const { timeToExpire, totalHits } = await this.storageService.increment(
+      key,
+      ttl
+    );
+
+    if (totalHits > limit) {
+      res.header('Retry-After', timeToExpire.toString());
+      await this.throwThrottlingException(context, {
+        limit,
+        ttl,
+        key,
+        tracker,
+        totalHits,
+        timeToExpire,
+      });
+    }
+
+    res.header(`${this.headerPrefix}-Limit`, limit.toString());
+    res.header(
+      `${this.headerPrefix}-Remaining`,
+      (limit - totalHits).toString()
+    );
+    res.header(`${this.headerPrefix}-Reset`, timeToExpire.toString());
+    return true;
+  }
+
+  override async canActivate(context: ExecutionContext): Promise<boolean> {
+    const { req } = this.getRequestResponse(context);
+
+    const throttler = this.getSpecifiedThrottler(context);
+
+    // if user is logged in, bypass non-protected handlers
+    if (!throttler && req.user) {
+      return true;
+    }
+
+    return super.canActivate(context);
+  }
+
+  getSpecifiedThrottler(context: ExecutionContext) {
+    return this.reflector.getAllAndOverride<Throttlers | undefined>(
+      THROTTLER_PROTECTED,
+      [context.getHandler(), context.getClass()]
+    );
+  }
+}
+
 @Global()
 @Module({
   imports: [
@@ -46,46 +180,9 @@ class CustomOptionsFactory implements ThrottlerOptionsFactory {
       useClass: CustomOptionsFactory,
     }),
   ],
-  providers: [ThrottlerStorage],
-  exports: [ThrottlerStorage],
+  providers: [ThrottlerStorage, CloudThrottlerGuard],
+  exports: [ThrottlerStorage, CloudThrottlerGuard],
 })
 export class RateLimiterModule {}
 
-@Injectable()
-export class CloudThrottlerGuard extends ThrottlerGuard {
-  override getRequestResponse(context: ExecutionContext) {
-    return getRequestResponseFromContext(context) as any;
-  }
-
-  protected override getTracker(req: Record<string, any>): Promise<string> {
-    return Promise.resolve(
-      req?.get('CF-Connecting-IP') ?? req?.get('CF-ray') ?? req?.ip
-    );
-  }
-}
-
-@Injectable()
-export class AuthThrottlerGuard extends CloudThrottlerGuard {
-  override async handleRequest(
-    context: ExecutionContext,
-    limit: number,
-    ttl: number
-  ): Promise<boolean> {
-    const { req } = this.getRequestResponse(context);
-
-    if (req?.url === '/api/auth/session') {
-      // relax throttle for session auto renew
-      return super.handleRequest(context, limit * 20, ttl, {
-        ttl: ttl * 20,
-        limit: limit * 20,
-      });
-    }
-
-    return super.handleRequest(context, limit, ttl, {
-      ttl,
-      limit,
-    });
-  }
-}
-
-export { Throttle };
+export * from './decorators';