Added grouping routes for ratelimit (#380)

* Added grouping routes for ratelimit * Added test * Fixed linting and used sleep in test * Added test to increase coverage * Updated test case for invalid gorupId * Updated test case for invalid gorupId * Updated test cases * Updated test cases * Added test case for no gorup id * Updated tests and groupId value
fastify · Sep 26, 2024 · b80b05d · b80b05d
1 parent 4cd5021
commit b80b05d
Show file tree

Hide file tree

Showing 3 changed files with 294 additions and 2 deletions.
diff --git a/README.md b/README.md
@@ -371,6 +371,31 @@ fastify.get('/public/sub-rated-1', {
 }, (request, reply) => {
   reply.send({ hello: 'from sub-rated-1 ... using default max value ... ' })
 })
+
+// gorup routes and add a rate limit
+fastify.get('/otp/send', {
+  config: {
+    rateLimit: {
+      max: 3,
+      timeWindow: '1 minute',
+      groupId:"OTP"
+    }
+  }
+}, (request, reply) => {
+  reply.send({ hello: 'from ... grouped rate limit' })
+})
+
+fastify.get('/otp/resend', {
+  config: {
+    rateLimit: {
+      max: 3,
+      timeWindow: '1 minute',
+      groupId:"OTP"
+    }
+  }
+}, (request, reply) => {
+  reply.send({ hello: 'from ... grouped rate limit' })
+})
 ```
 
 In the route creation you can override the same settings of the plugin registration plus the following additional options:

diff --git a/index.js b/index.js
@@ -143,7 +143,17 @@ async function fastifyRateLimit (fastify, settings) {
         const newPluginComponent = Object.create(pluginComponent)
         const mergedRateLimitParams = mergeParams(globalParams, routeOptions.config.rateLimit, { routeInfo: routeOptions })
         newPluginComponent.store = pluginComponent.store.child(mergedRateLimitParams)
-        addRouteRateHook(newPluginComponent, mergedRateLimitParams, routeOptions)
+
+        if (routeOptions?.config?.rateLimit?.groupId) {
+          const groupId = routeOptions.config.rateLimit.groupId
+          if (typeof groupId === 'string') {
+            addRouteRateHook(pluginComponent, globalParams, routeOptions)
+          } else {
+            throw new Error('groupId must be a string')
+          }
+        } else {
+          addRouteRateHook(newPluginComponent, mergedRateLimitParams, routeOptions)
+        }
       } else if (routeOptions.config.rateLimit !== false) {
         throw new Error('Unknown value for route rate-limit configuration')
       }
@@ -208,7 +218,12 @@ function rateLimitRequestHandler (pluginComponent, params) {
     req[rateLimitRan] = true
 
     // Retrieve the key from the generator (the global one or the one defined in the endpoint)
-    const key = await params.keyGenerator(req)
+    let key = await params.keyGenerator(req)
+    const groupId = req?.routeOptions?.config?.rateLimit?.groupId
+
+    if (groupId) {
+      key += groupId
+    }
 
     // Don't apply any rate limiting if in the allow list
     if (params.allowList) {

diff --git a/test/group-rate-limit.test.js b/test/group-rate-limit.test.js
@@ -0,0 +1,252 @@
+'use strict'
+const { test } = require('node:test')
+const assert = require('assert')
+const Fastify = require('fastify')
+const rateLimit = require('../index')
+
+const sleep = (ms) => new Promise((resolve) => setTimeout(resolve, ms))
+
+test('GroupId from routeConfig', async (t) => {
+  const fastify = Fastify()
+
+  // Register rate limit plugin with groupId in routeConfig
+  await fastify.register(rateLimit, { max: 2, timeWindow: 500 })
+
+  fastify.get(
+    '/routeWithGroupId',
+    {
+      config: {
+        rateLimit: {
+          max: 2,
+          timeWindow: 500,
+          groupId: 'group1' // groupId specified in routeConfig
+        }
+      }
+    },
+    async (req, reply) => 'hello from route with groupId!'
+  )
+
+  // Test: Request should have the correct groupId in response
+  const res = await fastify.inject({ url: '/routeWithGroupId', method: 'GET' })
+  assert.deepStrictEqual(res.statusCode, 200)
+  assert.deepStrictEqual(res.headers['x-ratelimit-limit'], '2')
+  assert.deepStrictEqual(res.headers['x-ratelimit-remaining'], '1')
+})
+
+test('GroupId from routeOptions', async (t) => {
+  const fastify = Fastify()
+
+  // Register rate limit plugin with groupId in routeOptions
+  await fastify.register(rateLimit, { max: 2, timeWindow: 500 })
+
+  fastify.get(
+    '/routeWithGroupIdFromOptions',
+    {
+      config: {
+        rateLimit: {
+          max: 2,
+          timeWindow: 500
+          // groupId not specified here
+        }
+      }
+    },
+    async (req, reply) => 'hello from route with groupId from options!'
+  )
+
+  // Test: Request should have the correct groupId from routeOptions
+  const res = await fastify.inject({ url: '/routeWithGroupIdFromOptions', method: 'GET' })
+  assert.deepStrictEqual(res.statusCode, 200)
+  assert.deepStrictEqual(res.headers['x-ratelimit-limit'], '2')
+  assert.deepStrictEqual(res.headers['x-ratelimit-remaining'], '1')
+})
+
+test('No groupId provided', async (t) => {
+  const fastify = Fastify()
+
+  // Register rate limit plugin without groupId
+  await fastify.register(rateLimit, { max: 2, timeWindow: 500 })
+
+  // Route without groupId
+  fastify.get(
+    '/noGroupId',
+    {
+      config: {
+        rateLimit: {
+          max: 2,
+          timeWindow: 500
+        }
+      }
+    },
+    async (req, reply) => 'hello from no groupId route!'
+  )
+
+  let res
+
+  // Test without groupId
+  res = await fastify.inject({ url: '/noGroupId', method: 'GET' })
+  assert.deepStrictEqual(res.statusCode, 200)
+  assert.deepStrictEqual(res.headers['x-ratelimit-limit'], '2')
+  assert.deepStrictEqual(res.headers['x-ratelimit-remaining'], '1')
+
+  res = await fastify.inject({ url: '/noGroupId', method: 'GET' })
+  assert.deepStrictEqual(res.statusCode, 200)
+  assert.deepStrictEqual(res.headers['x-ratelimit-limit'], '2')
+  assert.deepStrictEqual(res.headers['x-ratelimit-remaining'], '0')
+
+  res = await fastify.inject({ url: '/noGroupId', method: 'GET' })
+  assert.deepStrictEqual(res.statusCode, 429)
+  assert.deepStrictEqual(
+    res.headers['content-type'],
+    'application/json; charset=utf-8'
+  )
+  assert.deepStrictEqual(res.headers['x-ratelimit-limit'], '2')
+  assert.deepStrictEqual(res.headers['x-ratelimit-remaining'], '0')
+  assert.deepStrictEqual(res.headers['retry-after'], '1')
+  assert.deepStrictEqual(
+    {
+      statusCode: 429,
+      error: 'Too Many Requests',
+      message: 'Rate limit exceeded, retry in 500 ms'
+    },
+    JSON.parse(res.payload)
+  )
+})
+
+test('With multiple routes and custom groupId', async (t) => {
+  const fastify = Fastify()
+
+  // Register rate limit plugin
+  await fastify.register(rateLimit, { max: 2, timeWindow: 500 })
+
+  // Route 1 with groupId 'group1'
+  fastify.get(
+    '/route1',
+    {
+      config: {
+        rateLimit: {
+          max: 2,
+          timeWindow: 500,
+          groupId: 'group1'
+        }
+      }
+    },
+    async (req, reply) => 'hello from route 1!'
+  )
+
+  // Route 2 with groupId 'group2'
+  fastify.get(
+    '/route2',
+    {
+      config: {
+        rateLimit: {
+          max: 2,
+          timeWindow: 500,
+          groupId: 'group2'
+        }
+      }
+    },
+    async (req, reply) => 'hello from route 2!'
+  )
+
+  let res
+
+  // Test Route 1
+  res = await fastify.inject({ url: '/route1', method: 'GET' })
+  assert.deepStrictEqual(res.statusCode, 200)
+  assert.deepStrictEqual(res.headers['x-ratelimit-limit'], '2')
+  assert.deepStrictEqual(res.headers['x-ratelimit-remaining'], '1')
+
+  res = await fastify.inject({ url: '/route1', method: 'GET' })
+  assert.deepStrictEqual(res.statusCode, 200)
+  assert.deepStrictEqual(res.headers['x-ratelimit-limit'], '2')
+  assert.deepStrictEqual(res.headers['x-ratelimit-remaining'], '0')
+
+  res = await fastify.inject({ url: '/route1', method: 'GET' })
+  assert.deepStrictEqual(res.statusCode, 429)
+  assert.deepStrictEqual(
+    res.headers['content-type'],
+    'application/json; charset=utf-8'
+  )
+  assert.deepStrictEqual(res.headers['x-ratelimit-limit'], '2')
+  assert.deepStrictEqual(res.headers['x-ratelimit-remaining'], '0')
+  assert.deepStrictEqual(res.headers['retry-after'], '1')
+  assert.deepStrictEqual(
+    {
+      statusCode: 429,
+      error: 'Too Many Requests',
+      message: 'Rate limit exceeded, retry in 500 ms'
+    },
+    JSON.parse(res.payload)
+  )
+
+  // Test Route 2
+  res = await fastify.inject({ url: '/route2', method: 'GET' })
+  assert.deepStrictEqual(res.statusCode, 200)
+  assert.deepStrictEqual(res.headers['x-ratelimit-limit'], '2')
+  assert.deepStrictEqual(res.headers['x-ratelimit-remaining'], '1')
+
+  res = await fastify.inject({ url: '/route2', method: 'GET' })
+  assert.deepStrictEqual(res.statusCode, 200)
+  assert.deepStrictEqual(res.headers['x-ratelimit-limit'], '2')
+  assert.deepStrictEqual(res.headers['x-ratelimit-remaining'], '0')
+
+  res = await fastify.inject({ url: '/route2', method: 'GET' })
+  assert.deepStrictEqual(res.statusCode, 429)
+  assert.deepStrictEqual(
+    res.headers['content-type'],
+    'application/json; charset=utf-8'
+  )
+  assert.deepStrictEqual(res.headers['x-ratelimit-limit'], '2')
+  assert.deepStrictEqual(res.headers['x-ratelimit-remaining'], '0')
+  assert.deepStrictEqual(res.headers['retry-after'], '1')
+  assert.deepStrictEqual(
+    {
+      statusCode: 429,
+      error: 'Too Many Requests',
+      message: 'Rate limit exceeded, retry in 500 ms'
+    },
+    JSON.parse(res.payload)
+  )
+
+  // Wait for the window to reset
+  await sleep(1000)
+
+  // After reset, Route 1 should succeed again
+  res = await fastify.inject({ url: '/route1', method: 'GET' })
+  assert.deepStrictEqual(res.statusCode, 200)
+  assert.deepStrictEqual(res.headers['x-ratelimit-limit'], '2')
+  assert.deepStrictEqual(res.headers['x-ratelimit-remaining'], '1')
+
+  // Route 2 should also succeed after the reset
+  res = await fastify.inject({ url: '/route2', method: 'GET' })
+  assert.deepStrictEqual(res.statusCode, 200)
+  assert.deepStrictEqual(res.headers['x-ratelimit-limit'], '2')
+  assert.deepStrictEqual(res.headers['x-ratelimit-remaining'], '1')
+})
+
+test('Invalid groupId type', async (t) => {
+  const fastify = Fastify()
+
+  // Register rate limit plugin with a route having an invalid groupId
+  await fastify.register(rateLimit, { max: 2, timeWindow: 1000 })
+
+  try {
+    fastify.get(
+      '/invalidGroupId',
+      {
+        config: {
+          rateLimit: {
+            max: 2,
+            timeWindow: 1000,
+            groupId: 123 // Invalid groupId type
+          }
+        }
+      },
+      async (req, reply) => 'hello with invalid groupId!'
+    )
+    assert.fail('should throw')
+    console.log('HER')
+  } catch (err) {
+    assert.deepStrictEqual(err.message, 'groupId must be a string')
+  }
+})