From 491cb605e3c93d7a261bb0cef0d38f2ddc3affe8 Mon Sep 17 00:00:00 2001
From: Tomi Turtiainen <10324676+tomi@users.noreply.github.com>
Date: Wed, 8 Jan 2025 10:06:36 +0200
Subject: [PATCH] fix(core): Use rate limiter for task runner endpoints
 (#12486)

---
 .../src/task-runners/task-runner-server.ts    | 10 +++++++++
 .../runners/task-runner-server.test.ts        | 22 +++++++++++++++++++
 2 files changed, 32 insertions(+)

diff --git a/packages/cli/src/task-runners/task-runner-server.ts b/packages/cli/src/task-runners/task-runner-server.ts
index 6e68c4fb32..80679f8c41 100644
--- a/packages/cli/src/task-runners/task-runner-server.ts
+++ b/packages/cli/src/task-runners/task-runner-server.ts
@@ -2,6 +2,7 @@ import { GlobalConfig } from '@n8n/config';
 import { Service } from '@n8n/di';
 import compression from 'compression';
 import express from 'express';
+import { rateLimit as expressRateLimit } from 'express-rate-limit';
 import { Logger } from 'n8n-core';
 import * as a from 'node:assert/strict';
 import { randomBytes } from 'node:crypto';
@@ -147,8 +148,16 @@ export class TaskRunnerServer {
 	}
 
 	private configureRoutes() {
+		const createRateLimiter = () =>
+			expressRateLimit({
+				windowMs: 1000,
+				limit: 5,
+				message: { message: 'Too many requests' },
+			});
+
 		this.app.use(
 			this.upgradeEndpoint,
+			createRateLimiter(),
 			// eslint-disable-next-line @typescript-eslint/unbound-method
 			this.taskRunnerAuthController.authMiddleware,
 			(req: TaskRunnerServerInitRequest, res: TaskRunnerServerInitResponse) =>
@@ -158,6 +167,7 @@ export class TaskRunnerServer {
 		const authEndpoint = `${this.getEndpointBasePath()}/auth`;
 		this.app.post(
 			authEndpoint,
+			createRateLimiter(),
 			send(async (req) => await this.taskRunnerAuthController.createGrantToken(req)),
 		);
 
diff --git a/packages/cli/test/integration/runners/task-runner-server.test.ts b/packages/cli/test/integration/runners/task-runner-server.test.ts
index 11d77b53fc..6088af3525 100644
--- a/packages/cli/test/integration/runners/task-runner-server.test.ts
+++ b/packages/cli/test/integration/runners/task-runner-server.test.ts
@@ -19,4 +19,26 @@ describe('TaskRunnerServer', () => {
 			await agent.get('/healthz').expect(200);
 		});
 	});
+
+	describe('/runners/_ws', () => {
+		it('should return 429 when too many requests are made', async () => {
+			await agent.post('/runners/_ws').send({}).expect(401);
+			await agent.post('/runners/_ws').send({}).expect(401);
+			await agent.post('/runners/_ws').send({}).expect(401);
+			await agent.post('/runners/_ws').send({}).expect(401);
+			await agent.post('/runners/_ws').send({}).expect(401);
+			await agent.post('/runners/_ws').send({}).expect(429);
+		});
+	});
+
+	describe('/runners/auth', () => {
+		it('should return 429 when too many requests are made', async () => {
+			await agent.post('/runners/auth').send({ token: 'invalid' }).expect(403);
+			await agent.post('/runners/auth').send({ token: 'invalid' }).expect(403);
+			await agent.post('/runners/auth').send({ token: 'invalid' }).expect(403);
+			await agent.post('/runners/auth').send({ token: 'invalid' }).expect(403);
+			await agent.post('/runners/auth').send({ token: 'invalid' }).expect(403);
+			await agent.post('/runners/auth').send({ token: 'invalid' }).expect(429);
+		});
+	});
 });