netanalyzer/services/websocket.js

const WebSocket = require('ws');
const config = require('../config');
const logger = require('../utils/logger');

const PROXY_REQUEST_TIMEOUT_MS = 45000;
const PING_INTERVAL_MS = 30000;
const MAX_PENDING_REQUESTS = 200;

let wss;
let connectedAgent = null;
let pingInterval = null;
const pendingRequests = new Map(); // requestId → { resolve, reject, timeout }

/**
 * Pull a bearer token from the standard `Authorization` header first
 * (preferred — keeps the token out of access logs) and fall back to the
 * legacy `?token=` query parameter for backward compatibility.
 */
function extractToken(req) {
  const auth = req.headers['authorization'];
  if (auth && /^Bearer\s+/i.test(auth)) {
    return auth.replace(/^Bearer\s+/i, '').trim();
  }
  try {
    const url = new URL(req.url, `http://${req.headers.host || 'localhost'}`);
    return url.searchParams.get('token');
  } catch (_e) {
    return null;
  }
}

function rejectPending(reason) {
  for (const { reject, timeout } of pendingRequests.values()) {
    clearTimeout(timeout);
    reject(new Error(reason));
  }
  pendingRequests.clear();
}

function handleAgentMessage(data) {
  let response;
  try {
    response = JSON.parse(data);
  } catch (e) {
    logger.error('Failed to parse agent message', { error: e.message });
    return;
  }

  const { requestId } = response;
  if (!requestId || !pendingRequests.has(requestId)) {
    logger.warn('Received response with unknown requestId', { requestId });
    return;
  }

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

function startWebSocketServer() {
  wss = new WebSocket.Server({ port: config.ws.port });

  wss.on('error', err => {
    logger.error('WebSocket server error', { error: err.message });
  });

  wss.on('connection', (ws, req) => {
    const token = extractToken(req);
    const expected = config.ws.token;
    const remote = req.socket.remoteAddress;

    if (!expected) {
      logger.error('WS_TOKEN not configured; refusing connection', { remote });
      ws.close(1011, 'Server misconfigured');
      return;
    }

    if (!token || token !== expected) {
      logger.warn('Rejected agent connection: token mismatch', {
        remote,
        tokenPresent: !!token,
      });
      ws.close(1008, 'Invalid or missing token');
      return;
    }

    if (connectedAgent && connectedAgent.readyState === WebSocket.OPEN) {
      logger.warn('Replacing previously connected remote agent', { remote });
      try {
        connectedAgent.close(1013, 'Replaced by newer agent');
      } catch (_e) {
        // ignore
      }
      rejectPending('Remote agent replaced; in-flight requests dropped');
    }

    logger.info('Remote agent connected', { remote });
    connectedAgent = ws;
    ws.isAlive = true;

    ws.on('message', handleAgentMessage);
    ws.on('pong', () => {
      ws.isAlive = true;
    });
    ws.on('error', err => {
      logger.error('Remote agent socket error', { error: err.message });
    });
    ws.on('close', (code, reason) => {
      logger.info('Remote agent disconnected', {
        code,
        reason: reason?.toString() || 'none',
      });
      if (connectedAgent === ws) {
        connectedAgent = null;
        rejectPending('Remote agent disconnected');
      }
    });
  });

  pingInterval = setInterval(() => {
    wss.clients.forEach(ws => {
      if (ws.isAlive === false) return ws.terminate();
      ws.isAlive = false;
      ws.ping();
    });
  }, PING_INTERVAL_MS);

  logger.info('WebSocket server listening', { port: config.ws.port });
}

function stopWebSocketServer() {
  if (pingInterval) {
    clearInterval(pingInterval);
    pingInterval = null;
  }

  if (connectedAgent) {
    try {
      connectedAgent.close();
    } catch (_e) {
      // ignore errors during shutdown
    }
    connectedAgent = null;
  }

  if (wss) {
    logger.info('Shutting down WebSocket server');
    wss.close(() => logger.info('WebSocket server closed'));
    wss = null;
  }

  rejectPending('Server shutting down');
}

function isAgentConnected() {
  return !!connectedAgent && connectedAgent.readyState === WebSocket.OPEN;
}

class AgentNotConnectedError extends Error {
  constructor(message = 'No remote agent connected') {
    super(message);
    this.name = 'AgentNotConnectedError';
    this.code = 'AGENT_NOT_CONNECTED';
  }
}

async function proxyRequest(requestConfig) {
  return new Promise((resolve, reject) => {
    if (!isAgentConnected()) {
      return reject(new AgentNotConnectedError());
    }

    if (pendingRequests.size >= MAX_PENDING_REQUESTS) {
      return reject(new Error('Too many in-flight proxy requests'));
    }

    const requestId = `req_${Date.now()}_${Math.random().toString(36).slice(2)}`;

    const timeout = setTimeout(() => {
      pendingRequests.delete(requestId);
      reject(new Error(`Proxy request timeout after ${PROXY_REQUEST_TIMEOUT_MS / 1000}s`));
    }, PROXY_REQUEST_TIMEOUT_MS);

    pendingRequests.set(requestId, { resolve, reject, timeout });

    const payload = JSON.stringify({
      action: 'proxyRequest',
      requestId,
      ...requestConfig,
    });

    connectedAgent.send(payload, err => {
      if (err) {
        clearTimeout(timeout);
        pendingRequests.delete(requestId);
        reject(err);
      }
    });
  });
}

module.exports = {
  startWebSocketServer,
  stopWebSocketServer,
  proxyRequest,
  isAgentConnected,
  AgentNotConnectedError,
};