Merge pull request #1711 from sdg9670/issue/1709

feat: upgrade @nestjs/throttler to v6, add new throttling features

Author: kkoomen

package.json

@@ -49,7 +49,7 @@
     "@nestjs/platform-ws": ">=10.2.4",
     "@nestjs/schematics": ">=10.0.2",
     "@nestjs/testing": ">=10.2.4",
-    "@nestjs/throttler": "^5.0.0",
+    "@nestjs/throttler": "^6.0.0",
     "@nestjs/websockets": ">=10.2.4",
     "@types/express": "^4.17.17",
     "@types/jest": "29.5.12",
@@ -84,7 +84,7 @@
   "peerDependencies": {
     "@nestjs/common": "^7.0.0 || ^8.0.0 || ^9.0.0 || ^10.0.0",
     "@nestjs/core": "^7.0.0 || ^8.0.0 || ^9.0.0 || ^10.0.0",
-    "@nestjs/throttler": ">=5.0.0",
+    "@nestjs/throttler": ">=6.0.0",
     "ioredis": ">=5.0.0",
     "reflect-metadata": "^0.2.1"
   },

src/throttler-storage-redis.interface.ts

@@ -1,4 +1,4 @@
-import { ThrottlerStorageRecord } from '@nestjs/throttler/dist/throttler-storage-record.interface';
+import { ThrottlerStorage } from '@nestjs/throttler';
 import Redis, { Cluster } from 'ioredis';
 
 export interface ThrottlerStorageRedis {
@@ -11,7 +11,7 @@ export interface ThrottlerStorageRedis {
    * Increment the amount of requests for a given record. The record will
    * automatically be removed from the storage once its TTL has been reached.
    */
-  increment(key: string, ttl: number): Promise<ThrottlerStorageRecord>;
+  increment: ThrottlerStorage['increment'];
 }
 
 export const ThrottlerStorageRedis = Symbol('ThrottlerStorageRedis');

src/throttler-storage-redis.service.ts

@@ -31,38 +31,72 @@ export class ThrottlerStorageRedisService implements ThrottlerStorageRedis, OnMo
     // Credits to wyattjoh for the fast implementation you see below.
     // https://github.com/wyattjoh/rate-limit-redis/blob/main/src/lib.ts
     return `
-      local totalHits = redis.call("INCR", KEYS[1])
-      local timeToExpire = redis.call("PTTL", KEYS[1])
-      if timeToExpire <= 0
-        then
-          redis.call("PEXPIRE", KEYS[1], tonumber(ARGV[1]))
-          timeToExpire = tonumber(ARGV[1])
-        end
-      return { totalHits, timeToExpire }
+      local hitKey = KEYS[1]
+      local blockKey = KEYS[2]
+      local throttlerName = ARGV[1]
+      local ttl = tonumber(ARGV[2])
+      local limit = tonumber(ARGV[3])
+      local blockDuration = tonumber(ARGV[4])
+
+      local totalHits = redis.call('INCR', hitKey)
+      local timeToExpire = redis.call('PTTL', hitKey)
+      
+      if timeToExpire <= 0 then
+        redis.call('PEXPIRE', hitKey, ttl)
+        timeToExpire = ttl
+      end
+
+      local isBlocked = redis.call('GET', blockKey)
+      local timeToBlockExpire = 0
+
+      if isBlocked then
+        timeToBlockExpire = redis.call('PTTL', blockKey)
+      elseif totalHits > limit then
+        redis.call('SET', blockKey, 1, 'PX', blockDuration)
+        isBlocked = '1'
+        timeToBlockExpire = blockDuration
+      end
+
+      if isBlocked and timeToBlockExpire <= 0 then
+        redis.call('DEL', blockKey)
+        redis.call('SET', hitKey, 1, 'PX', ttl)
+        totalHits = 1
+        timeToExpire = ttl
+        isBlocked = false
+      end
+
+      return { totalHits, timeToExpire, isBlocked and 1 or 0, timeToBlockExpire }
     `
       .replace(/^\s+/gm, '')
       .trim();
   }
 
-  async increment(key: string, ttl: number): Promise<ThrottlerStorageRecord> {
-    // Use EVAL instead of EVALSHA to support both redis instances and clusters.
+  async increment(
+    key: string,
+    ttl: number,
+    limit: number,
+    blockDuration: number,
+    throttlerName: string,
+  ): Promise<ThrottlerStorageRecord> {
+    const hitKey = `${this.redis.options.keyPrefix}{${key}:${throttlerName}}:hits`;
+    const blockKey = `${this.redis.options.keyPrefix}{${key}:${throttlerName}}:blocked`;
     const results: number[] = (await this.redis.call(
       'EVAL',
       this.scriptSrc,
-      1,
-      `${this.redis.options.keyPrefix}${key}`,
+      2,
+      hitKey,
+      blockKey,
+      throttlerName,
       ttl,
+      limit,
+      blockDuration,
     )) as number[];
 
     if (!Array.isArray(results)) {
       throw new TypeError(`Expected result to be array of values, got ${results}`);
     }
 
-    if (results.length !== 2) {
-      throw new Error(`Expected 2 values, got ${results.length}`);
-    }
-
-    const [totalHits, timeToExpire] = results;
+    const [totalHits, timeToExpire, isBlocked, timeToBlockExpire] = results;
 
     if (typeof totalHits !== 'number') {
       throw new TypeError('Expected totalHits to be a number');
@@ -72,9 +106,19 @@ export class ThrottlerStorageRedisService implements ThrottlerStorageRedis, OnMo
       throw new TypeError('Expected timeToExpire to be a number');
     }
 
+    if (typeof isBlocked !== 'number') {
+      throw new TypeError('Expected isBlocked to be a number');
+    }
+
+    if (typeof timeToBlockExpire !== 'number') {
+      throw new TypeError('Expected timeToBlockExpire to be a number');
+    }
+
     return {
       totalHits,
       timeToExpire: Math.ceil(timeToExpire / 1000),
+      isBlocked: isBlocked === 1,
+      timeToBlockExpire: Math.ceil(timeToBlockExpire / 1000),
     };
   }
 

yarn.lock

@@ -1168,10 +1168,10 @@
   dependencies:
     tslib "2.6.2"
 
-"@nestjs/throttler@^5.0.0":
-  version "5.1.2"
-  resolved "https://registry.yarnpkg.com/@nestjs/throttler/-/throttler-5.1.2.tgz#dc65634153c8b887329b1cc6061db2e556517dcb"
-  integrity sha512-60MqhSLYUqWOgc38P6C6f76JIpf6mVjly7gpuPBCKtVd0p5e8Fq855j7bJuO4/v25vgaOo1OdVs0U1qtgYioGw==
+"@nestjs/throttler@^6.0.0":
+  version "6.0.0"
+  resolved "https://registry.yarnpkg.com/@nestjs/throttler/-/throttler-6.0.0.tgz#4f602ce5c584dd21b728eee612baf5d8118bac5f"
+  integrity sha512-2mccn9uCMuDiZBG1c+tIHlNIhkyVJAg167wK66GcNNoj6+6qf8fCOxBZbAHUSyOJeoIxcZAhT9wjJRSXqbWY4w==
 
 "@nestjs/websockets@>=10.2.4":
   version "10.3.8"