From 31b34d9ea2bce7247cecded1dfef1e05308f3f0c Mon Sep 17 00:00:00 2001
From: Artem Niehrieiev <r.tem@tuta.io>
Date: Tue, 2 Jun 2026 13:02:14 +0000
Subject: [PATCH] feat: enhance WebSocket connection handling with token
 validation and update response caching

---
 autoadmin-ws-server/src/handlers/command.ts        |  2 +-
 autoadmin-ws-server/src/handlers/websocket.ts      | 14 ++++++++++++++
 autoadmin-ws-server/src/services/response-cache.ts |  3 +++
 .../data-access-object-agent.ts                    |  5 ++---
 4 files changed, 20 insertions(+), 4 deletions(-)

diff --git a/autoadmin-ws-server/src/handlers/command.ts b/autoadmin-ws-server/src/handlers/command.ts
index 1d9d3ad72..64284661f 100644
--- a/autoadmin-ws-server/src/handlers/command.ts
+++ b/autoadmin-ws-server/src/handlers/command.ts
@@ -42,7 +42,7 @@ export async function executeCommand(c: Context): Promise<Response> {
 			resolve(c.json({ error: message }, status as 400));
 		};
 
-		cacheResponse(resId, handleResolve, handleReject, handleSendError);
+		cacheResponse(resId, handleResolve, handleReject, handleSendError, connectionToken.token);
 
 		try {
 			sendCommandToClient(connectionToken.token, body, resId);
diff --git a/autoadmin-ws-server/src/handlers/websocket.ts b/autoadmin-ws-server/src/handlers/websocket.ts
index aba519fab..ee38702a8 100644
--- a/autoadmin-ws-server/src/handlers/websocket.ts
+++ b/autoadmin-ws-server/src/handlers/websocket.ts
@@ -9,6 +9,10 @@ import { validateConnectionToken } from '../services/token-validator.js';
 import { hashToken } from '../utils/crypto.js';
 import { logger } from '../utils/logger.js';
 
+interface AuthenticatedSocket extends WebSocket {
+	agentToken?: string;
+}
+
 export function setupWebSocketServer(server: Server): WebSocketServer {
 	const wss = new WebSocketServer({ server });
 
@@ -46,12 +50,22 @@ export function setupWebSocketServer(server: Server): WebSocketServer {
 				);
 				connectionToken = hashedToken;
 				data.connectionToken = connectionToken;
+				(ws as AuthenticatedSocket).agentToken = hashedToken;
 			}
 
 			if (operationType === COMMAND_TYPE.dataFromAgent && resId) {
 				const cachedResponse = responseCache.get(resId);
 
 				if (cachedResponse) {
+					const socketToken = (ws as AuthenticatedSocket).agentToken;
+					if (!socketToken || socketToken !== cachedResponse.routedToken) {
+						logger.warn(
+							{ resId, authenticated: !!socketToken },
+							'Discarding dataFromAgent from a socket not bound to the routed connection token',
+						);
+						return;
+					}
+
 					logger.debug({ resId }, 'Received data from agent');
 					cachedResponse.resolve(rawMessage.toString());
 					responseCache.delete(resId);
diff --git a/autoadmin-ws-server/src/services/response-cache.ts b/autoadmin-ws-server/src/services/response-cache.ts
index 6fbd13e86..044629da5 100644
--- a/autoadmin-ws-server/src/services/response-cache.ts
+++ b/autoadmin-ws-server/src/services/response-cache.ts
@@ -7,6 +7,7 @@ interface CachedResponse {
 	reject: (error: Error) => void;
 	sendError: (status: number, message: string) => void;
 	createdAt: Date;
+	routedToken: string;
 }
 
 export const responseCache = new LRUCache<string, CachedResponse>({
@@ -23,12 +24,14 @@ export function cacheResponse(
 	resolve: (data: string) => void,
 	reject: (error: Error) => void,
 	sendError: (status: number, message: string) => void,
+	routedToken: string,
 ): void {
 	responseCache.set(resId, {
 		resolve,
 		reject,
 		sendError,
 		createdAt: new Date(),
+		routedToken,
 	});
 	logger.debug({ resId }, 'Response cached');
 }
diff --git a/shared-code/src/data-access-layer/data-access-objects/data-access-object-agent.ts b/shared-code/src/data-access-layer/data-access-objects/data-access-object-agent.ts
index 096163cbb..cb52fd77a 100644
--- a/shared-code/src/data-access-layer/data-access-objects/data-access-object-agent.ts
+++ b/shared-code/src/data-access-layer/data-access-objects/data-access-object-agent.ts
@@ -821,13 +821,12 @@ export class DataAccessObjectAgent implements IDataAccessObjectAgent {
 	}
 
 	private generateJWT(connectionToken: string): string {
-		const exp = new Date();
-		exp.setDate(exp.getDate() + 60);
+		const exp = Math.floor(Date.now() / 1000) + 120;
 		const secret = process.env.JWT_SECRET;
 		return jwt.sign(
 			{
 				token: connectionToken,
-				exp: Math.floor(exp.getTime() / 1000),
+				exp,
 			},
 			secret,
 		);