feat: enhance WebSocket connection handling with token validation and update response caching

Author: Artuomka

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);

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);

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');
 }

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,
 		);