wsServer.js

Go to the documentation of this file.

// File operation request/response map
const fileListRequests = new Map(); // requestId → { resolve, reject, timeout }
 
// Utility to send file-list request to agent and await response
/**
 *
 * @param agentId
 * @param path
 */
function requestAgentFileList(agentId, path) {
  return new Promise((resolve, reject) => {
    const ws = agentConnections.get(agentId);
    if (!ws || ws.readyState !== 1) return reject(new Error('Agent not connected'));
    const requestId = 'filelist-' + Date.now() + '-' + Math.random().toString(36).slice(2);
    // Store resolver
    fileListRequests.set(requestId, {
      resolve,
      reject,
      timeout: setTimeout(() => {
        fileListRequests.delete(requestId);
        reject(new Error('Timeout waiting for agent file list'));
      }, 10000)
    });
    // Send request
    ws.send(JSON.stringify({ type: 'file-list', requestId, path }));
  });
}
// backend/wsServer.js
const { WebSocketServer } = require('ws');
const url = require('url');
 
/**
 * agentConnections:  agent_uuid → ws
 * viewerConnections: sessionId → ws
 * activeRdpSessions: sessionId → { agentId, viewer, agent, createdAt }
 */
const agentConnections = new Map();
const viewerConnections = new Map();
const activeRdpSessions = new Map();
 
/**
 *
 * @param server
 */
function initWebSocketServer(server) {
  const wss = new WebSocketServer({ noServer: true });
 
  // Handle upgrade requests
  server.on('upgrade', (req, socket, head) => {
    const pathname = req.url.split('?')[0];
 
    if (
      pathname === '/ws' ||
      pathname.startsWith('/api/rds/view/') ||
      pathname.startsWith('/api/rds/control/')
    ) {
      wss.handleUpgrade(req, socket, head, (ws) =>
        wss.emit('connection', ws, req)
      );
    }
  });
 
  // Handle all WebSocket connections through here
  wss.on('connection', (ws, req) => {
    const parsed = url.parse(req.url, true);
    const path = parsed.pathname;
 
    // ────────────────────────────────────────────
    // 1. AGENT CONNECTION (/ws)
    // ────────────────────────────────────────────
    if (path === '/ws') {
      const agent_uuid = parsed.query.agent_uuid;
      if (!agent_uuid) return ws.close();
 
      agentConnections.set(agent_uuid, ws);
      console.log(`[WS] Agent connected: ${agent_uuid}`);
 
      ws.on('close', () => {
        agentConnections.delete(agent_uuid);
        console.log(`[WS] Agent disconnected: ${agent_uuid}`);
 
        // Update last_seen in the database
        try {
          const pool = require('./services/db');
          pool.query(
            'UPDATE agents SET last_seen = NOW() WHERE agent_uuid = $1',
            [agent_uuid]
          ).then(() => {
            console.log(`[WS] Updated last_seen for agent ${agent_uuid}`);
          }).catch(err => {
            console.error(`[WS] Failed to update last_seen for agent ${agent_uuid}:`, err);
          });
        } catch (err) {
          console.error(`[WS] Error updating last_seen for agent ${agent_uuid}:`, err);
        }
 
        // clean up any active RDS sessions belonging to this agent
        for (const [sessionId, session] of activeRdpSessions.entries()) {
          if (session.agentId === agent_uuid) {
            if (session.viewer) session.viewer.close();
            viewerConnections.delete(sessionId);
            activeRdpSessions.delete(sessionId);
            console.log(`[RDS] Cleaned orphaned session ${sessionId}`);
          }
        }
      });
 
      return;
    }
 
    // ────────────────────────────────────────────
    // 2. VIEWER CONNECTION (/api/rds/view/:sessionId)
    // ────────────────────────────────────────────
    if (path.startsWith('/api/rds/view/')) {
      const sessionId = parsed.query.session;
      if (!sessionId) return ws.close();
 
      console.log(`[RDS] Viewer connected to session ${sessionId}`);
      viewerConnections.set(sessionId, ws);
 
      const session = activeRdpSessions.get(sessionId);
      if (session) session.viewer = ws;
 
      ws.on('close', () => {
        viewerConnections.delete(sessionId);
        console.log(`[RDS] Viewer disconnected: ${sessionId}`);
      });
 
      return;
    }
 
    // ────────────────────────────────────────────
    // 3. CONTROL CONNECTION (/api/rds/control/:sessionId)
    // ────────────────────────────────────────────
    if (path.startsWith('/api/rds/control/')) {
      const sessionId = parsed.query.session;
      if (!sessionId) return ws.close();
 
      console.log(`[RDS] Control channel connected to session ${sessionId}`);
 
      const session = activeRdpSessions.get(sessionId);
      if (!session) {
        console.log(`[RDS] Session ${sessionId} not found, closing control channel`);
        return ws.close();
      }
 
      // Store control websocket reference
      session.control = ws;
 
      // Forward control messages (mouse/keyboard) to agent
      ws.on('message', (raw) => {
        try {
          const msg = JSON.parse(raw.toString());
          console.log(`[RDS] Control message:`, msg.type);
 
          // Get agent connection
          const agentWs = agentConnections.get(session.agentId);
          if (agentWs && agentWs.readyState === 1) {
            // Forward to agent with sessionId
            agentWs.send(JSON.stringify({
              ...msg,
              sessionId
            }));
          } else {
            console.log(`[RDS] Agent not connected for session ${sessionId}`);
          }
        } catch (err) {
          console.error('[RDS] Error forwarding control message:', err);
        }
      });
 
      ws.on('close', () => {
        if (session) session.control = null;
        console.log(`[RDS] Control channel disconnected: ${sessionId}`);
      });
 
      return;
    }
 
    // If path doesn't match anything
    ws.close();
  });
 
  // ────────────────────────────────────────────
  // FORWARDING PIPELINE (Viewer → Agent)
  // ────────────────────────────────────────────
  wss.on('connection', (ws, req) => {
    ws.on('message', (raw) => {
      try {
        // Viewer → Agent messages are JSON
        const msg = JSON.parse(raw.toString());
        if (!msg.sessionId) return;
 
        const session = activeRdpSessions.get(msg.sessionId);
        if (!session) return;
 
        const agentWs = session.agent;
 
        if (agentWs && agentWs.readyState === 1) {
          agentWs.send(JSON.stringify(msg));
        }
      } catch {
        // Ignore invalid JSON (JPEG screen frames)
      }
    });
  });
 
  // ────────────────────────────────────────────
  // FORWARDING PIPELINE (Agent → Viewer)
  // ────────────────────────────────────────────
  wss.on('connection', (ws, req) => {
    ws.on('message', (raw) => {
      // Try treat as agent-binary-screen-frame first
      for (const [sessionId, session] of activeRdpSessions.entries()) {
        if (session.agent === ws) {
          const viewerWs = session.viewer;
          if (viewerWs && viewerWs.readyState === 1) {
            // Forward binary frames
            viewerWs.send(raw);
          }
        }
      }
 
      // If this WAS JSON (clipboard events, file chunks, file-list, etc)
      try {
        const msg = JSON.parse(raw.toString());
 
        // Handle agent file-list response
        if (msg.type === 'file-list-response' && msg.requestId) {
          const reqObj = fileListRequests.get(msg.requestId);
          if (reqObj) {
            clearTimeout(reqObj.timeout);
            reqObj.resolve(msg.files || []);
            fileListRequests.delete(msg.requestId);
          }
          return;
        }
 
        // ...existing RDS clipboard/file chunk/session logic...
        if (msg.sessionId) {
          const session = activeRdpSessions.get(msg.sessionId);
          if (!session) return;
          const viewerWs = session.viewer;
          if (viewerWs && viewerWs.readyState === 1) {
            viewerWs.send(JSON.stringify(msg));
          }
        }
      } catch {
        // Not JSON → ignore
      }
    });
  });
 
  console.log('✅ WebSocket subsystem ready (agents + RDS)');
  return { agentConnections, viewerConnections, activeRdpSessions, requestAgentFileList };
}
 
module.exports = { initWebSocketServer, agentConnections, viewerConnections, activeRdpSessions, requestAgentFileList };