smartproxy/ts/proxies/smart-proxy/port-manager.ts

import * as plugins from '../../plugins.js';
import type { ISmartProxyOptions } from './models/interfaces.js';
import { RouteConnectionHandler } from './route-connection-handler.js';

/**
 * PortManager handles the dynamic creation and removal of port listeners
 * 
 * This class provides methods to add and remove listening ports at runtime,
 * allowing SmartProxy to adapt to configuration changes without requiring
 * a full restart.
 */
export class PortManager {
  private servers: Map<number, plugins.net.Server> = new Map();
  private settings: ISmartProxyOptions;
  private routeConnectionHandler: RouteConnectionHandler;
  private isShuttingDown: boolean = false;

  /**
   * Create a new PortManager
   * 
   * @param settings The SmartProxy settings
   * @param routeConnectionHandler The handler for new connections
   */
  constructor(
    settings: ISmartProxyOptions,
    routeConnectionHandler: RouteConnectionHandler
  ) {
    this.settings = settings;
    this.routeConnectionHandler = routeConnectionHandler;
  }

  /**
   * Start listening on a specific port
   * 
   * @param port The port number to listen on
   * @returns Promise that resolves when the server is listening or rejects on error
   */
  public async addPort(port: number): Promise<void> {
    // Check if we're already listening on this port
    if (this.servers.has(port)) {
      console.log(`PortManager: Already listening on port ${port}`);
      return;
    }

    // Create a server for this port
    const server = plugins.net.createServer((socket) => {
      // Check if shutting down
      if (this.isShuttingDown) {
        socket.end();
        socket.destroy();
        return;
      }
      
      // Delegate to route connection handler
      this.routeConnectionHandler.handleConnection(socket);
    }).on('error', (err: Error) => {
      console.log(`Server Error on port ${port}: ${err.message}`);
    });
    
    // Start listening on the port
    return new Promise<void>((resolve, reject) => {
      server.listen(port, () => {
        const isHttpProxyPort = this.settings.useHttpProxy?.includes(port);
        console.log(
          `SmartProxy -> OK: Now listening on port ${port}${
            isHttpProxyPort ? ' (HttpProxy forwarding enabled)' : ''
          }`
        );
        
        // Store the server reference
        this.servers.set(port, server);
        resolve();
      }).on('error', (err) => {
        console.log(`Failed to listen on port ${port}: ${err.message}`);
        reject(err);
      });
    });
  }

  /**
   * Stop listening on a specific port
   * 
   * @param port The port to stop listening on
   * @returns Promise that resolves when the server is closed
   */
  public async removePort(port: number): Promise<void> {
    // Get the server for this port
    const server = this.servers.get(port);
    if (!server) {
      console.log(`PortManager: Not listening on port ${port}`);
      return;
    }

    // Close the server
    return new Promise<void>((resolve) => {
      server.close((err) => {
        if (err) {
          console.log(`Error closing server on port ${port}: ${err.message}`);
        } else {
          console.log(`SmartProxy -> Stopped listening on port ${port}`);
        }
        
        // Remove the server reference
        this.servers.delete(port);
        resolve();
      });
    });
  }

  /**
   * Add multiple ports at once
   * 
   * @param ports Array of ports to add
   * @returns Promise that resolves when all servers are listening
   */
  public async addPorts(ports: number[]): Promise<void> {
    const uniquePorts = [...new Set(ports)];
    await Promise.all(uniquePorts.map(port => this.addPort(port)));
  }

  /**
   * Remove multiple ports at once
   * 
   * @param ports Array of ports to remove
   * @returns Promise that resolves when all servers are closed
   */
  public async removePorts(ports: number[]): Promise<void> {
    const uniquePorts = [...new Set(ports)];
    await Promise.all(uniquePorts.map(port => this.removePort(port)));
  }

  /**
   * Update listening ports to match the provided list
   * 
   * This will add any ports that aren't currently listening,
   * and remove any ports that are no longer needed.
   * 
   * @param ports Array of ports that should be listening
   * @returns Promise that resolves when all operations are complete
   */
  public async updatePorts(ports: number[]): Promise<void> {
    const targetPorts = new Set(ports);
    const currentPorts = new Set(this.servers.keys());
    
    // Find ports to add and remove
    const portsToAdd = ports.filter(port => !currentPorts.has(port));
    const portsToRemove = Array.from(currentPorts).filter(port => !targetPorts.has(port));
    
    // Log the changes
    if (portsToAdd.length > 0) {
      console.log(`PortManager: Adding new listeners for ports: ${portsToAdd.join(', ')}`);
    }
    
    if (portsToRemove.length > 0) {
      console.log(`PortManager: Removing listeners for ports: ${portsToRemove.join(', ')}`);
    }
    
    // Add and remove ports
    await this.removePorts(portsToRemove);
    await this.addPorts(portsToAdd);
  }

  /**
   * Get all ports that are currently listening
   * 
   * @returns Array of port numbers
   */
  public getListeningPorts(): number[] {
    return Array.from(this.servers.keys());
  }

  /**
   * Mark the port manager as shutting down
   */
  public setShuttingDown(isShuttingDown: boolean): void {
    this.isShuttingDown = isShuttingDown;
  }

  /**
   * Close all listening servers
   * 
   * @returns Promise that resolves when all servers are closed
   */
  public async closeAll(): Promise<void> {
    const allPorts = Array.from(this.servers.keys());
    await this.removePorts(allPorts);
  }

  /**
   * Get all server instances (for testing or debugging)
   */
  public getServers(): Map<number, plugins.net.Server> {
    return new Map(this.servers);
  }
}