From b496bf3147d2cd873d24371be02cb7ea5dbd8621 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Iv=C3=A1n=20Ovejero?= <ivov.src@gmail.com>
Date: Fri, 8 Nov 2024 10:09:24 +0100
Subject: [PATCH] fix(core): Ensure task runner server closes websocket
 connection correctly (#11633)

---
 .../__tests__/task-runner-server.test.ts      | 67 +++++++++++++++++++
 .../cli/src/runners/task-runner-server.ts     |  2 +-
 2 files changed, 68 insertions(+), 1 deletion(-)
 create mode 100644 packages/cli/src/runners/__tests__/task-runner-server.test.ts

diff --git a/packages/cli/src/runners/__tests__/task-runner-server.test.ts b/packages/cli/src/runners/__tests__/task-runner-server.test.ts
new file mode 100644
index 0000000000..ae25cd1231
--- /dev/null
+++ b/packages/cli/src/runners/__tests__/task-runner-server.test.ts
@@ -0,0 +1,67 @@
+import type { GlobalConfig } from '@n8n/config';
+import { mock } from 'jest-mock-extended';
+import { ServerResponse } from 'node:http';
+import type WebSocket from 'ws';
+
+import type { TaskRunnerAuthController } from '@/runners/auth/task-runner-auth.controller';
+import { TaskRunnerServer } from '@/runners/task-runner-server';
+
+import type { TaskRunnerServerInitRequest } from '../runner-types';
+
+describe('TaskRunnerServer', () => {
+	describe('handleUpgradeRequest', () => {
+		it('should close WebSocket when response status code is > 200', () => {
+			const ws = mock<WebSocket>();
+			const request = mock<TaskRunnerServerInitRequest>({
+				url: '/runners/_ws',
+				ws,
+			});
+
+			const server = new TaskRunnerServer(
+				mock(),
+				mock<GlobalConfig>({ taskRunners: { path: '/runners' } }),
+				mock<TaskRunnerAuthController>(),
+				mock(),
+			);
+
+			// @ts-expect-error Private property
+			server.handleUpgradeRequest(request, mock(), Buffer.from(''));
+
+			const response = new ServerResponse(request);
+			response.writeHead = (statusCode) => {
+				if (statusCode > 200) ws.close();
+				return response;
+			};
+
+			response.writeHead(401);
+			expect(ws.close).toHaveBeenCalledWith(); // no args
+		});
+
+		it('should not close WebSocket when response status code is 200', () => {
+			const ws = mock<WebSocket>();
+			const request = mock<TaskRunnerServerInitRequest>({
+				url: '/runners/_ws',
+				ws,
+			});
+
+			const server = new TaskRunnerServer(
+				mock(),
+				mock<GlobalConfig>({ taskRunners: { path: '/runners' } }),
+				mock<TaskRunnerAuthController>(),
+				mock(),
+			);
+
+			// @ts-expect-error Private property
+			server.handleUpgradeRequest(request, mock(), Buffer.from(''));
+
+			const response = new ServerResponse(request);
+			response.writeHead = (statusCode) => {
+				if (statusCode > 200) ws.close();
+				return response;
+			};
+
+			response.writeHead(200);
+			expect(ws.close).not.toHaveBeenCalled();
+		});
+	});
+});
diff --git a/packages/cli/src/runners/task-runner-server.ts b/packages/cli/src/runners/task-runner-server.ts
index 6dd0fd5919..56c56e02ae 100644
--- a/packages/cli/src/runners/task-runner-server.ts
+++ b/packages/cli/src/runners/task-runner-server.ts
@@ -181,7 +181,7 @@ export class TaskRunnerServer {
 
 			const response = new ServerResponse(request);
 			response.writeHead = (statusCode) => {
-				if (statusCode > 200) ws.close(100);
+				if (statusCode > 200) ws.close();
 				return response;
 			};