const WebSocket = require('ws'); const config = require('../config'); let wss; let connectedAgent = null; const pendingRequests = new Map(); // requestId → { resolve, reject, timeout } function startWebSocketServer() { wss = new WebSocket.Server({ port: config.ws.port }); wss.on('connection', (ws, req) => { // Better URL parsing and logging let token = null; try { const url = new URL(req.url, `http://${req.headers.host}/netanalyze`); token = url.searchParams.get('token'); console.log(`🔑 Incoming connection from ${req.socket.remoteAddress} | Path: ${req.url} | Token present: ${!!token}`); } catch (e) { console.log('❌ Failed to parse connection URL'); } if (!token || token !== config.ws.token) { console.log(`❌ Token mismatch! Received: "${token}" | Expected: "${config.ws.token ? config.ws.token.substring(0, 8) + '...' : 'MISSING'}"`); ws.close(1008, 'Invalid or missing token'); return; } console.log('✅ Authorized Remote Agent connected'); connectedAgent = ws; ws.on('message', (data) => { try { const response = JSON.parse(data); const requestId = response.requestId; if (requestId && pendingRequests.has(requestId)) { const { resolve, reject, timeout } = pendingRequests.get(requestId); clearTimeout(timeout); pendingRequests.delete(requestId); if (response.error) { const err = new Error(response.error); err.status = response.status; reject(err); } else { resolve(response); } } else { console.log('⚠️ Received response with unknown requestId:', requestId); } } catch (e) { console.error('Failed to parse agent message:', e); } }); ws.on('close', (code, reason) => { console.log(`❌ Remote Agent disconnected | Code: ${code} | Reason: ${reason || 'none'}`); connectedAgent = null; }); // Add ping/pong to keep connection alive ws.isAlive = true; ws.on('pong', () => { ws.isAlive = true; }); }); // Keep connections alive const interval = setInterval(() => { wss.clients.forEach((ws) => { if (ws.isAlive === false) return ws.terminate(); ws.isAlive = false; ws.ping(); }); }, 30000); // every 30 seconds console.log(`WebSocket server running on ws://0.0.0.0:${config.ws.port}`); } async function proxyRequest(requestConfig) { return new Promise((resolve, reject) => { if (!connectedAgent || connectedAgent.readyState !== WebSocket.OPEN) { return reject(new Error('No remote agent connected')); } const requestId = `req_${Date.now()}_${Math.random().toString(36).slice(2)}`; const timeout = setTimeout(() => { pendingRequests.delete(requestId); reject(new Error(`Proxy request timeout after 45s`)); }, 45000); // Increased timeout pendingRequests.set(requestId, { resolve, reject, timeout }); const payload = JSON.stringify({ action: 'proxyRequest', requestId, ...requestConfig }); connectedAgent.send(payload); }); } module.exports = { startWebSocketServer, proxyRequest };