netanalyzer/services/websocket.js

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