Adds rate limits to assistant route

FlowFuse · Jul 16, 2024 · bbbebe0 · bbbebe0
1 parent 10b7e1f
commit bbbebe0
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 2 deletions.
diff --git a/forge/routes/api/assistant.js b/forge/routes/api/assistant.js
@@ -28,6 +28,18 @@ module.exports = async function (app) {
      * use an alternative means of accessing it.
     */
     app.post('/:method', {
+        config: {
+            rateLimit: app.config.rate_limits
+                ? {
+                    hook: 'preHandler', // apply the rate as a preHandler so that session is available
+                    max: 5, // max requests per window
+                    timeWindow: 30000, // 30 seconds
+                    keyGenerator: (request) => {
+                        return request.session?.ownerId || request.ip
+                    }
+                }
+                : false
+        },
         schema: {
             hide: true, // dont show in swagger
             body: {

diff --git a/test/unit/forge/routes/api/rateLimits/rateLimits_spec.js b/test/unit/forge/routes/api/rateLimits/rateLimits_spec.js
@@ -162,6 +162,7 @@ describe('Endpoint Rate Limiting', () => {
             { url: '/account/register', method: 'POST', shouldLimit: true },
             { url: '/account/forgot_password', method: 'POST', shouldLimit: true },
             { url: '/account/reset_password/:token', method: 'POST', shouldLimit: true },
+            { url: '/api/v1/assistant/:method', method: 'POST', shouldLimit: true, customLimits: true },
             // routes that are never rate limited
             { url: '/api/comms/auth/client', method: 'POST', shouldLimit: false },
             { url: '/api/comms/auth/acl', method: 'POST', shouldLimit: false },
@@ -281,7 +282,11 @@ describe('Endpoint Rate Limiting', () => {
                     it(`Route ${route.method} ${route.url} should be rate limited`, async function () {
                         const routeConfig = route.fastifyRoute.config
                         routeConfig.should.have.property('rateLimit').and.be.an.Object()
-                        if (routeConfig.rateLimit.hard) {
+                        if (route.customLimits) {
+                            // should have one of the following properties: max, timeWindow, keyGenerator
+                            const hasKeys = Object.keys(routeConfig.rateLimit).some((key) => ['max', 'timeWindow', 'keyGenerator'].includes(key))
+                            should(hasKeys).be.true()
+                        } else if (routeConfig.rateLimit.hard) {
                             routeConfig.rateLimit.should.have.property('max')
                             routeConfig.rateLimit.max.should.be.equalOneOf(5, 2)
                             routeConfig.rateLimit.should.have.property('timeWindow')
@@ -395,7 +400,11 @@ describe('Endpoint Rate Limiting', () => {
                     it(`Route ${route.method} ${route.url} should be rate limited`, async function () {
                         const routeConfig = route.fastifyRoute.config
                         routeConfig.should.have.property('rateLimit').and.be.an.Object()
-                        if (routeConfig.rateLimit.hard) {
+                        if (route.customLimits) {
+                            // should have one of the following properties: max, timeWindow, keyGenerator
+                            const hasKeys = Object.keys(routeConfig.rateLimit).some((key) => ['max', 'timeWindow', 'keyGenerator'].includes(key))
+                            should(hasKeys).be.true()
+                        } else if (routeConfig.rateLimit.hard) {
                             routeConfig.rateLimit.should.have.property('max')
                             routeConfig.rateLimit.max.should.be.equalOneOf(5, 2)
                             routeConfig.rateLimit.should.have.property('timeWindow')