Backport timewindow function to v9

README.md

@@ -143,7 +143,7 @@ await fastify.register(import('@fastify/rate-limit'), {
 - `global` : indicates if the plugin should apply rate limiting to all routes within the encapsulation scope.
 - `max`: maximum number of requests a single client can perform inside a timeWindow. It can be an async function with the signature `async (request, key) => {}` where `request` is the Fastify request object and `key` is the value generated by the `keyGenerator`. The function **must** return a number.
 - `ban`: maximum number of 429 responses to return to a client before returning 403 responses. When the ban limit is exceeded, the context argument that is passed to `errorResponseBuilder` will have its `ban` property set to `true`. **Note:** `0` can also be passed to directly return 403 responses when a client exceeds the `max` limit.
-- `timeWindow:` the duration of the time window. It can be expressed in milliseconds or as a string (in the [`ms`](https://github.com/zeit/ms) format)
+- `timeWindow:` the duration of the time window. It can be expressed in milliseconds, as a string (in the [`ms`](https://github.com/zeit/ms) format), or as an async function with the signature `async (request, key) => {}` where `request` is the Fastify request object and `key` is the value generated by the `keyGenerator`. The function **must** return a number.- `cache`: this plugin internally uses a lru cache to handle the clients, you can change the size of the cache with this option
 - `cache`: this plugin internally uses a lru cache to handle the clients, you can change the size of the cache with this option
 - `allowList`: array of string of ips to exclude from rate limiting. It can be a sync or async function with the signature `(request, key) => {}` where `request` is the Fastify request object and `key` is the value generated by the `keyGenerator`. If the function return a truthy value, the request will be excluded from the rate limit.
 - `redis`: by default, this plugin uses an in-memory store, but if an application runs on multiple servers, an external store will be needed. This plugin requires the use of [`ioredis`](https://github.com/redis/ioredis).<br> **Note:** the [default settings](https://github.com/redis/ioredis/blob/v4.16.0/API.md#new_Redis_new) of an ioredis instance are not optimal for rate limiting. We recommend customizing the `connectTimeout` and `maxRetriesPerRequest` parameters as shown in the [`example`](https://github.com/fastify/fastify-rate-limit/tree/master/example/example.js).

index.js

@@ -64,11 +64,18 @@ async function fastifyRateLimit (fastify, settings) {
     : defaultMax
 
   // Global time window
-  globalParams.timeWindow = typeof settings.timeWindow === 'string'
-    ? ms.parse(settings.timeWindow)
-    : typeof settings.timeWindow === 'number' && Number.isFinite(settings.timeWindow) && settings.timeWindow >= 0
-      ? Math.trunc(settings.timeWindow)
-      : defaultTimeWindow
+  const twType = typeof settings.timeWindow
+  globalParams.timeWindow = defaultTimeWindow
+  if (twType === 'function') {
+    globalParams.timeWindow = settings.timeWindow
+  } else if (twType === 'string') {
+    globalParams.timeWindow = ms.parse(settings.timeWindow)
+  } else if (
+    twType === 'number' &&
+    Number.isFinite(settings.timeWindow) && settings.timeWindow >= 0
+  ) {
+    globalParams.timeWindow = Math.trunc(settings.timeWindow)
+  }
 
   globalParams.hook = settings.hook || defaultHook
   globalParams.allowList = settings.allowList || settings.whitelist || null
@@ -147,7 +154,7 @@ function mergeParams (...params) {
     result.timeWindow = ms.parse(result.timeWindow)
   } else if (typeof result.timeWindow === 'number' && Number.isFinite(result.timeWindow) && result.timeWindow >= 0) {
     result.timeWindow = Math.trunc(result.timeWindow)
-  } else {
+  } else if (typeof result.timeWindow !== 'function') {
     result.timeWindow = defaultTimeWindow
   }
 
@@ -180,7 +187,6 @@ function addRouteRateHook (pluginComponent, params, routeOptions) {
 
 function rateLimitRequestHandler (pluginComponent, params) {
   const { rateLimitRan, store } = pluginComponent
-  const timeWindowString = ms.format(params.timeWindow, true)
 
   return async (req, res) => {
     if (req[rateLimitRan]) {
@@ -204,6 +210,7 @@ function rateLimitRequestHandler (pluginComponent, params) {
     }
 
     const max = typeof params.max === 'number' ? params.max : await params.max(req, key)
+    const timeWindow = typeof params.timeWindow === 'number' ? params.timeWindow : await params.timeWindow(req, key)
     let current = 0
     let ttl = 0
     let timeLeftInSeconds = 0
@@ -214,7 +221,7 @@ function rateLimitRequestHandler (pluginComponent, params) {
       const res = await new Promise((resolve, reject) => {
         store.incr(key, (err, res) => {
           err ? reject(err) : resolve(res)
-        }, max, params.ban)
+        }, timeWindow, max, params.ban)
       })
 
       current = res.current
@@ -250,7 +257,7 @@ function rateLimitRequestHandler (pluginComponent, params) {
       ban,
       max,
       ttl,
-      after: timeWindowString
+      after: ms.format(timeWindow, true)
     }
 
     if (ban) {

store/LocalStore.js

@@ -8,17 +8,17 @@ function LocalStore (cache = 5000, timeWindow, continueExceeding) {
   this.continueExceeding = continueExceeding
 }
 
-LocalStore.prototype.incr = function (ip, cb, max, ban) {
+LocalStore.prototype.incr = function (ip, cb, timeWindow, max, ban) {
   const nowInMs = Date.now()
   let current = this.lru.get(ip)
 
   if (!current) {
     // Item doesn't exist
-    current = { current: 1, ttl: this.timeWindow, ban: false, iterationStartMs: nowInMs }
-  } else if (current.iterationStartMs + this.timeWindow <= nowInMs) {
+    current = { current: 1, ttl: timeWindow, ban: false, iterationStartMs: nowInMs }
+  } else if (current.iterationStartMs + timeWindow <= nowInMs) {
     // Item has expired
     current.current = 1
-    current.ttl = this.timeWindow
+    current.ttl = timeWindow
     current.ban = false
     current.iterationStartMs = nowInMs
   } else {
@@ -27,10 +27,10 @@ LocalStore.prototype.incr = function (ip, cb, max, ban) {
 
     // Reset TLL if max has been exceeded and `continueExceeding` is enabled
     if (this.continueExceeding && current.current > max) {
-      current.ttl = this.timeWindow
+      current.ttl = timeWindow
       current.iterationStartMs = nowInMs
     } else {
-      current.ttl = this.timeWindow - (nowInMs - current.iterationStartMs)
+      current.ttl = timeWindow - (nowInMs - current.iterationStartMs)
     }
   }
 

store/RedisStore.js

@@ -41,8 +41,8 @@ function RedisStore (redis, timeWindow, continueExceeding, key) {
   }
 }
 
-RedisStore.prototype.incr = function (ip, cb, max, ban) {
-  this.redis.rateLimit(this.key + ip, this.timeWindow, max, ban, this.continueExceeding, (err, result) => {
+RedisStore.prototype.incr = function (ip, cb, timeWindow, max, ban) {
+  this.redis.rateLimit(this.key + ip, timeWindow, max, ban, this.continueExceeding, (err, result) => {
     err ? cb(err, null) : cb(null, { current: result[0], ttl: result[1], ban: result[2] })
   })
 }

test/global-rate-limit.test.js

@@ -148,6 +148,54 @@ test('When passing NaN to the timeWindow property then the timeWindow should be
   })
 })
 
+test('With function timeWindow', async t => {
+  t.plan(15)
+  t.context.clock = FakeTimers.install()
+  const fastify = Fastify()
+  await fastify.register(rateLimit, { max: 2, timeWindow: (_, __) => 1000 })
+
+  fastify.get('/', async (req, reply) => 'hello!')
+
+  let res
+
+  res = await fastify.inject('/')
+
+  t.equal(res.statusCode, 200)
+  t.equal(res.headers['x-ratelimit-limit'], '2')
+  t.equal(res.headers['x-ratelimit-remaining'], '1')
+
+  res = await fastify.inject('/')
+
+  t.equal(res.statusCode, 200)
+  t.equal(res.headers['x-ratelimit-limit'], '2')
+  t.equal(res.headers['x-ratelimit-remaining'], '0')
+
+  res = await fastify.inject('/')
+
+  t.equal(res.statusCode, 429)
+  t.equal(res.headers['content-type'], 'application/json; charset=utf-8')
+  t.equal(res.headers['x-ratelimit-limit'], '2')
+  t.equal(res.headers['x-ratelimit-remaining'], '0')
+  t.equal(res.headers['retry-after'], '1')
+  t.same({
+    statusCode: 429,
+    error: 'Too Many Requests',
+    message: 'Rate limit exceeded, retry in 1 second'
+  }, JSON.parse(res.payload))
+
+  t.context.clock.tick(1100)
+
+  res = await fastify.inject('/')
+
+  t.equal(res.statusCode, 200)
+  t.equal(res.headers['x-ratelimit-limit'], '2')
+  t.equal(res.headers['x-ratelimit-remaining'], '1')
+
+  t.teardown(() => {
+    t.context.clock.uninstall()
+  })
+})
+
 test('With ips allowList, allowed ips should not result in rate limiting', async t => {
   t.plan(3)
   const fastify = Fastify()

test/route-rate-limit.test.js

@@ -92,7 +92,69 @@ test('With text timeWindow', async t => {
   await fastify.register(rateLimit, { global: false })
 
   fastify.get('/', {
-    config: defaultRouteConfig
+    config: {
+      rateLimit: {
+        max: 2,
+        timeWindow: '1s'
+      },
+      someOtherPlugin: {
+        someValue: 1
+      }
+    }
+  }, async (req, reply) => 'hello!')
+
+  let res
+
+  res = await fastify.inject('/')
+  t.equal(res.statusCode, 200)
+  t.equal(res.headers['x-ratelimit-limit'], '2')
+  t.equal(res.headers['x-ratelimit-remaining'], '1')
+
+  res = await fastify.inject('/')
+  t.equal(res.statusCode, 200)
+  t.equal(res.headers['x-ratelimit-limit'], '2')
+  t.equal(res.headers['x-ratelimit-remaining'], '0')
+
+  res = await fastify.inject('/')
+  t.equal(res.statusCode, 429)
+  t.equal(res.headers['content-type'], 'application/json; charset=utf-8')
+  t.equal(res.headers['x-ratelimit-limit'], '2')
+  t.equal(res.headers['x-ratelimit-remaining'], '0')
+  t.equal(res.headers['retry-after'], '1')
+  t.same(JSON.parse(res.payload), {
+    statusCode: 429,
+    error: 'Too Many Requests',
+    message: 'Rate limit exceeded, retry in 1 second'
+  })
+
+  t.context.clock.tick(1100)
+
+  res = await fastify.inject('/')
+  t.equal(res.statusCode, 200)
+  t.equal(res.headers['x-ratelimit-limit'], '2')
+  t.equal(res.headers['x-ratelimit-remaining'], '1')
+
+  t.teardown(() => {
+    t.context.clock.uninstall()
+  })
+})
+
+test('With function timeWindow', async t => {
+  t.plan(15)
+  t.context.clock = FakeTimers.install()
+  const fastify = Fastify()
+  await fastify.register(rateLimit, { global: false })
+
+  fastify.get('/', {
+    config: {
+      rateLimit: {
+        max: 2,
+        timeWindow: (_, __) => 1000
+      },
+      someOtherPlugin: {
+        someValue: 1
+      }
+    }
   }, async (req, reply) => 'hello!')
 
   let res

types/index.d.ts

@@ -97,10 +97,14 @@ declare namespace fastifyRateLimit {
 
   export interface RateLimitOptions {
     max?:
-    | number
-    | ((req: FastifyRequest, key: string) => number)
-    | ((req: FastifyRequest, key: string) => Promise<number>);
-    timeWindow?: number | string;
+        | number
+        | ((req: FastifyRequest, key: string) => number)
+        | ((req: FastifyRequest, key: string) => Promise<number>);
+    timeWindow?:
+        | number
+        | string
+        | ((req: FastifyRequest, key: string) => number)
+        | ((req: FastifyRequest, key: string) => Promise<number>);
     hook?: RateLimitHook;
     cache?: number;
     store?: FastifyRateLimitStoreCtor;