navflow-proxy-server/src/tunnelManager.js

Dynamic WebSocket proxy server for NavFlow

/**
 * Tunnel Manager - Handle active tunnel connections
 */

const { v4: uuidv4 } = require('uuid');
const { generateTunnelId, generatePassword } = require('./utils');

class TunnelManager {
  constructor() {
    this.tunnels = new Map(); // tunnelId -> tunnel object
    this.clients = new Map(); // tunnelId -> Set of client WebSockets
    this.connectionCount = 0;
  }

  /**
   * Register a new tunnel connection
   * @param {WebSocket} websocket 
   * @returns {Object} Tunnel credentials { tunnelId, password }
   */
  registerTunnel(websocket) {
    const tunnelId = generateTunnelId();
    const password = generatePassword();
    const connectionId = uuidv4();

    const tunnel = {
      id: tunnelId,
      password: password,
      connectionId: connectionId,
      websocket: websocket,
      connected: true,
      createdAt: new Date(),
      lastActivity: new Date(),
      requestCount: 0
    };

    this.tunnels.set(tunnelId, tunnel);
    this.clients.set(tunnelId, new Set()); // Initialize empty client set
    this.connectionCount++;

    console.log(`[TunnelManager] New tunnel registered: ${tunnelId} (password: ${password})`);
    console.log(`[TunnelManager] Active tunnels: ${this.tunnels.size}`);

    // Set up WebSocket event handlers
    websocket.on('close', () => {
      this.removeTunnel(tunnelId);
    });

    websocket.on('error', (error) => {
      console.error(`[TunnelManager] WebSocket error for tunnel ${tunnelId}:`, error);
      this.removeTunnel(tunnelId);
    });

    websocket.on('pong', () => {
      // Update last activity on pong response
      if (this.tunnels.has(tunnelId)) {
        this.tunnels.get(tunnelId).lastActivity = new Date();
      }
    });

    return { tunnelId, password, connectionId };
  }

  /**
   * Get tunnel by ID
   * @param {string} tunnelId 
   * @returns {Object|null} Tunnel object or null
   */
  getTunnel(tunnelId) {
    return this.tunnels.get(tunnelId) || null;
  }

  /**
   * Remove tunnel connection
   * @param {string} tunnelId 
   */
  removeTunnel(tunnelId) {
    const tunnel = this.tunnels.get(tunnelId);
    if (tunnel) {
      console.log(`[TunnelManager] Removing tunnel: ${tunnelId}`);
      
      // Close all client connections for this tunnel
      const clients = this.clients.get(tunnelId);
      if (clients) {
        for (const clientWs of clients) {
          if (clientWs.readyState === 1) {
            clientWs.close();
          }
        }
        this.clients.delete(tunnelId);
      }
      
      // Close WebSocket if still open
      if (tunnel.websocket && tunnel.websocket.readyState === 1) {
        tunnel.websocket.close();
      }
      
      this.tunnels.delete(tunnelId);
      console.log(`[TunnelManager] Active tunnels: ${this.tunnels.size}`);
    }
  }

  /**
   * Validate tunnel authentication
   * @param {string} tunnelId 
   * @param {string} password 
   * @returns {boolean}
   */
  validateAuth(tunnelId, password) {
    const tunnel = this.getTunnel(tunnelId);
    if (!tunnel || !tunnel.connected) {
      return false;
    }
    return tunnel.password === password;
  }

  /**
   * Update tunnel activity and increment request count
   * @param {string} tunnelId 
   */
  updateActivity(tunnelId) {
    const tunnel = this.getTunnel(tunnelId);
    if (tunnel) {
      tunnel.lastActivity = new Date();
      tunnel.requestCount++;
    }
  }

  /**
   * Send message through tunnel WebSocket (deprecated - use new sendToTunnel)
   * @param {string} tunnelId 
   * @param {Object} message 
   * @returns {Promise<Object>} Response from tunnel
   */
  async sendToTunnelDeprecated(tunnelId, message) {
    return this.sendToTunnelOriginal(tunnelId, message);
  }

  /**
   * Get statistics about active tunnels
   * @returns {Object}
   */
  getStats() {
    const tunnelStats = Array.from(this.tunnels.values()).map(tunnel => ({
      id: tunnel.id,
      connected: tunnel.connected,
      createdAt: tunnel.createdAt,
      lastActivity: tunnel.lastActivity,
      requestCount: tunnel.requestCount
    }));

    return {
      totalTunnels: this.tunnels.size,
      connectionCount: this.connectionCount,
      tunnels: tunnelStats
    };
  }

  /**
   * Cleanup inactive tunnels (run periodically)
   */
  cleanupInactiveTunnels() {
    const now = new Date();
    const timeout = 5 * 60 * 1000; // 5 minutes

    for (const [tunnelId, tunnel] of this.tunnels.entries()) {
      if (now - tunnel.lastActivity > timeout) {
        console.log(`[TunnelManager] Cleaning up inactive tunnel: ${tunnelId}`);
        this.removeTunnel(tunnelId);
      }
    }
  }

  /**
   * Send ping to all tunnels to keep connections alive
   */
  pingTunnels() {
    for (const [tunnelId, tunnel] of this.tunnels.entries()) {
      if (tunnel.websocket && tunnel.websocket.readyState === 1) {
        tunnel.websocket.ping();
      }
    }
  }

  /**
   * Authenticate client connection
   * @param {string} tunnelId 
   * @param {string} password 
   * @returns {boolean}
   */
  authenticateClient(tunnelId, password) {
    return this.validateAuth(tunnelId, password);
  }

  /**
   * Add client connection to tunnel
   * @param {string} tunnelId 
   * @param {WebSocket} clientWs 
   */
  addClient(tunnelId, clientWs) {
    const clients = this.clients.get(tunnelId);
    if (clients) {
      clients.add(clientWs);
      console.log(`[TunnelManager] Client added to tunnel ${tunnelId}. Total clients: ${clients.size}`);
    }
  }

  /**
   * Remove client connection from tunnel
   * @param {string} tunnelId 
   * @param {WebSocket} clientWs 
   */
  removeClient(tunnelId, clientWs) {
    const clients = this.clients.get(tunnelId);
    if (clients) {
      clients.delete(clientWs);
      console.log(`[TunnelManager] Client removed from tunnel ${tunnelId}. Total clients: ${clients.size}`);
    }
  }

  /**
   * Broadcast message to all clients connected to a tunnel
   * @param {string} tunnelId 
   * @param {Object} message 
   */
  broadcastToClients(tunnelId, message) {
    const clients = this.clients.get(tunnelId);
    if (clients) {
      const messageStr = JSON.stringify(message);
      for (const clientWs of clients) {
        if (clientWs.readyState === 1) {
          clientWs.send(messageStr);
        }
      }
      console.log(`[TunnelManager] Broadcasted message to ${clients.size} clients for tunnel ${tunnelId}`);
    }
  }

  /**
   * Send message to tunnel (enhanced for WebRTC signaling)
   * @param {string} tunnelId 
   * @param {Object} message 
   */
  sendToTunnel(tunnelId, message) {
    const tunnel = this.getTunnel(tunnelId);
    if (!tunnel || !tunnel.connected || tunnel.websocket.readyState !== 1) {
      console.error(`[TunnelManager] Tunnel ${tunnelId} not available for sending message`);
      return;
    }

    // For WebRTC signaling, we don't need to wait for a response
    if (message.type && message.type.startsWith('webrtc_')) {
      tunnel.websocket.send(JSON.stringify(message));
      console.log(`[TunnelManager] Sent ${message.type} to tunnel ${tunnelId}`);
      return;
    }

    // For other messages, use the original implementation
    return this.sendToTunnelOriginal(tunnelId, message);
  }

  /**
   * Original sendToTunnel method for request-response pattern
   * @param {string} tunnelId 
   * @param {Object} message 
   * @returns {Promise<Object>}
   */
  async sendToTunnelOriginal(tunnelId, message) {
    const tunnel = this.getTunnel(tunnelId);
    if (!tunnel || !tunnel.connected || tunnel.websocket.readyState !== 1) {
      throw new Error('Tunnel not available');
    }

    return new Promise((resolve, reject) => {
      const requestId = uuidv4();
      const requestMessage = {
        id: requestId,
        ...message
      };

      // Set up response handler
      const responseHandler = (data) => {
        try {
          const response = JSON.parse(data);
          if (response.id === requestId) {
            tunnel.websocket.removeListener('message', responseHandler);
            resolve(response);
          }
        } catch (error) {
          // Ignore non-JSON messages or messages for other requests
        }
      };

      // Set up timeout - increased for flow execution
      const timeout = setTimeout(() => {
        tunnel.websocket.removeListener('message', responseHandler);
        reject(new Error('Tunnel request timeout'));
      }, 120000); // 120 second timeout for flow execution

      tunnel.websocket.on('message', responseHandler);

      // Send the request
      tunnel.websocket.send(JSON.stringify(requestMessage));
      
      // Clear timeout when response is received
      responseHandler.originalResolve = resolve;
      const originalResolve = resolve;
      resolve = (data) => {
        clearTimeout(timeout);
        originalResolve(data);
      };
    });
  }
}

module.exports = TunnelManager;