refactor: block keys via penalize method

src/http_limiter.ts

@@ -143,8 +143,9 @@ export class HttpLimiter<KnownStores extends Record<string, LimiterManagerStoreF
     /**
      * Abort when user has exhausted all the requests.
      *
-     * We still run the "consume" method when remaining count is zero, since
-     * that will trigger the block logic on the key
+     * We still run the "consume" method when consumed is same as
+     * the limit, this will allow the consume method to trigger
+     * the block logic.
      */
     if (limiterResponse && limiterResponse.consumed > limiterResponse.limit) {
       debug('requests exhausted for key "%s"', key)

src/limiter.ts

@@ -39,6 +39,13 @@ export class Limiter implements LimiterStoreContract {
     return this.#store.duration
   }
 
+  /**
+   * The duration (in seconds) for which to block the key
+   */
+  get blockDuration() {
+    return this.#store.blockDuration
+  }
+
   constructor(store: LimiterStoreContract) {
     this.#store = store
   }
@@ -76,8 +83,9 @@ export class Limiter implements LimiterStoreContract {
      * Return early when remaining requests are less than
      * zero.
      *
-     * We still run the "consume" method when remaining count is zero, since
-     * that will trigger the block logic on the key
+     * We still run the "consume" method when consumed is same as
+     * the limit, this will allow the consume method to trigger
+     * the block logic.
      */
     const response = await this.get(key)
     if (response && response.consumed > response.limit) {
@@ -118,14 +126,32 @@ export class Limiter implements LimiterStoreContract {
       return [new E_TOO_MANY_REQUESTS(response), null]
     }
 
+    let callbackResult: T
+    let callbackError: unknown
+
     try {
-      const result = await callback()
-      await this.delete(key)
-      return [null, result]
+      callbackResult = await callback()
     } catch (error) {
-      await this.increment(key)
-      throw error
+      callbackError = error
     }
+
+    /**
+     * Consume one point and block the key if there is
+     * an error.
+     */
+    if (callbackError) {
+      const { consumed, limit } = await this.increment(key)
+      if (consumed >= limit && this.blockDuration) {
+        await this.block(key, this.blockDuration)
+      }
+      throw callbackError
+    }
+
+    /**
+     * Reset key
+     */
+    await this.delete(key)
+    return [null, callbackResult!]
   }
 
   /**

src/stores/bridge.ts

@@ -48,6 +48,13 @@ export default abstract class RateLimiterBridge implements LimiterStoreContract
     return this.rateLimiter.duration
   }
 
+  /**
+   * The duration (in seconds) for which to block the key
+   */
+  get blockDuration() {
+    return this.rateLimiter.blockDuration
+  }
+
   constructor(rateLimiter: RateLimiterStoreAbstract | RateLimiterAbstract) {
     this.rateLimiter = rateLimiter
   }

src/types.ts

@@ -176,6 +176,11 @@ export interface LimiterStoreContract {
    */
   readonly duration: number
 
+  /**
+   * The duration (in seconds) for which to block the key
+   */
+  readonly blockDuration: number
+
   /**
    * Consume 1 request for a given key. An exception is raised
    * when all the requests have already been consumed or if

tests/limiter.spec.ts

@@ -240,11 +240,10 @@ test.group('Limiter', () => {
     }, 'Something went wrong')
     assert.equal(await limiter.remaining('ip_localhost'), 1)
 
-    assert.isTrue(
-      await limiter.penalize('ip_localhost', () => {
-        return true
-      })
-    )
+    const [, result] = await limiter.penalize('ip_localhost', () => {
+      return true
+    })
+    assert.isTrue(result)
 
     assert.isNull(await limiter.get('ip_localhost'))
   })
@@ -290,5 +289,31 @@ test.group('Limiter', () => {
     assert.instanceOf(error, ThrottleException)
     assert.equal(error?.response.remaining, 0)
     assert.equal(await limiter.remaining('ip_localhost'), 0)
+    assert.closeTo(await limiter.availableIn('ip_localhost'), 60, 5)
+  })
+
+  test('block key when all requests have been exhausted', async ({ assert }) => {
+    const redis = createRedis(['rlflx:ip_localhost']).connection()
+    const store = new LimiterRedisStore(redis, {
+      duration: '1 minute',
+      requests: 2,
+      blockDuration: '30 mins',
+    })
+
+    const limiter = new Limiter(store)
+
+    await assert.rejects(async () => {
+      await limiter.penalize('ip_localhost', () => {
+        throw new Error('Something went wrong')
+      })
+    }, 'Something went wrong')
+
+    await assert.rejects(async () => {
+      await limiter.penalize('ip_localhost', () => {
+        throw new Error('Something went wrong')
+      })
+    }, 'Something went wrong')
+
+    assert.closeTo(await limiter.availableIn('ip_localhost'), 60 * 30, 5)
   })
 })