From b80b05df944ca40d19de9df8aa75485594f675fa Mon Sep 17 00:00:00 2001 From: Aniket Sinha Date: Thu, 26 Sep 2024 22:37:54 +0530 Subject: [PATCH] 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 --- README.md | 25 ++++ index.js | 19 ++- test/group-rate-limit.test.js | 252 ++++++++++++++++++++++++++++++++++ 3 files changed, 294 insertions(+), 2 deletions(-) create mode 100644 test/group-rate-limit.test.js diff --git a/README.md b/README.md index 13b9f56..a33e767 100644 --- 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 index 3895395..378fe90 100644 --- 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 new file mode 100644 index 0000000..038c767 --- /dev/null +++ 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') + } +})