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

import * as plugins from '../../plugins.js';
import { NfTablesProxy } from '../nftables-proxy/nftables-proxy.js';
import type { 
  NfTableProxyOptions, 
  PortRange,
  NfTablesStatus
} from '../nftables-proxy/models/interfaces.js';
import type {
  IRouteConfig,
  TPortRange,
  INfTablesOptions
} from './models/route-types.js';
import type { ISmartProxyOptions } from './models/interfaces.js';

/**
 * Manages NFTables rules based on SmartProxy route configurations
 * 
 * This class bridges the gap between SmartProxy routes and the NFTablesProxy,
 * allowing high-performance kernel-level packet forwarding for routes that 
 * specify NFTables as their forwarding engine.
 */
export class NFTablesManager {
  private rulesMap: Map<string, NfTablesProxy> = new Map();
  
  /**
   * Creates a new NFTablesManager
   * 
   * @param options The SmartProxy options
   */
  constructor(private options: ISmartProxyOptions) {}
  
  /**
   * Provision NFTables rules for a route
   * 
   * @param route The route configuration
   * @returns A promise that resolves to true if successful, false otherwise
   */
  public async provisionRoute(route: IRouteConfig): Promise<boolean> {
    // Generate a unique ID for this route
    const routeId = this.generateRouteId(route);
    
    // Skip if route doesn't use NFTables
    if (route.action.forwardingEngine !== 'nftables') {
      return true;
    }
    
    // Create NFTables options from route configuration
    const nftOptions = this.createNfTablesOptions(route);
    
    // Create and start an NFTablesProxy instance
    const proxy = new NfTablesProxy(nftOptions);
    
    try {
      await proxy.start();
      this.rulesMap.set(routeId, proxy);
      return true;
    } catch (err) {
      console.error(`Failed to provision NFTables rules for route ${route.name || 'unnamed'}: ${err.message}`);
      return false;
    }
  }
  
  /**
   * Remove NFTables rules for a route
   * 
   * @param route The route configuration
   * @returns A promise that resolves to true if successful, false otherwise
   */
  public async deprovisionRoute(route: IRouteConfig): Promise<boolean> {
    const routeId = this.generateRouteId(route);
    
    const proxy = this.rulesMap.get(routeId);
    if (!proxy) {
      return true; // Nothing to remove
    }
    
    try {
      await proxy.stop();
      this.rulesMap.delete(routeId);
      return true;
    } catch (err) {
      console.error(`Failed to deprovision NFTables rules for route ${route.name || 'unnamed'}: ${err.message}`);
      return false;
    }
  }
  
  /**
   * Update NFTables rules when route changes
   * 
   * @param oldRoute The previous route configuration
   * @param newRoute The new route configuration
   * @returns A promise that resolves to true if successful, false otherwise
   */
  public async updateRoute(oldRoute: IRouteConfig, newRoute: IRouteConfig): Promise<boolean> {
    // Remove old rules and add new ones
    await this.deprovisionRoute(oldRoute);
    return this.provisionRoute(newRoute);
  }
  
  /**
   * Generate a unique ID for a route
   * 
   * @param route The route configuration
   * @returns A unique ID string
   */
  private generateRouteId(route: IRouteConfig): string {
    // Generate a unique ID based on route properties
    // Include the route name, match criteria, and a timestamp
    const matchStr = JSON.stringify({
      ports: route.match.ports,
      domains: route.match.domains
    });
    
    return `${route.name || 'unnamed'}-${matchStr}-${route.id || Date.now().toString()}`;
  }
  
  /**
   * Create NFTablesProxy options from a route configuration
   * 
   * @param route The route configuration
   * @returns NFTableProxyOptions object
   */
  private createNfTablesOptions(route: IRouteConfig): NfTableProxyOptions {
    const { action } = route;
    
    // Ensure we have a target
    if (!action.target) {
      throw new Error('Route must have a target to use NFTables forwarding');
    }
    
    // Convert port specifications
    const fromPorts = this.expandPortRange(route.match.ports);
    
    // Determine target port
    let toPorts: number | PortRange | Array<number | PortRange>;
    
    if (action.target.port === 'preserve') {
      // 'preserve' means use the same ports as the source
      toPorts = fromPorts;
    } else if (typeof action.target.port === 'function') {
      // For function-based ports, we can't determine at setup time
      // Use the "preserve" approach and let NFTables handle it
      toPorts = fromPorts;
    } else {
      toPorts = action.target.port;
    }
    
    // Determine target host
    let toHost: string;
    if (typeof action.target.host === 'function') {
      // Can't determine at setup time, use localhost as a placeholder
      // and rely on run-time handling
      toHost = 'localhost';
    } else if (Array.isArray(action.target.host)) {
      // Use first host for now - NFTables will do simple round-robin  
      toHost = action.target.host[0];
    } else {
      toHost = action.target.host;
    }
    
    // Create options
    const options: NfTableProxyOptions = {
      fromPort: fromPorts,
      toPort: toPorts,
      toHost: toHost,
      protocol: action.nftables?.protocol || 'tcp',
      preserveSourceIP: action.nftables?.preserveSourceIP !== undefined ? 
                        action.nftables.preserveSourceIP : 
                        this.options.preserveSourceIP,
      useIPSets: action.nftables?.useIPSets !== false,
      useAdvancedNAT: action.nftables?.useAdvancedNAT,
      enableLogging: this.options.enableDetailedLogging,
      deleteOnExit: true,
      tableName: action.nftables?.tableName || 'smartproxy'
    };
    
    // Add security-related options
    if (route.security?.ipAllowList?.length) {
      options.ipAllowList = route.security.ipAllowList;
    }
    
    if (route.security?.ipBlockList?.length) {
      options.ipBlockList = route.security.ipBlockList;
    }
    
    // Add QoS options
    if (action.nftables?.maxRate || action.nftables?.priority) {
      options.qos = {
        enabled: true,
        maxRate: action.nftables.maxRate,
        priority: action.nftables.priority
      };
    }
    
    return options;
  }
  
  /**
   * Expand port range specifications
   * 
   * @param ports The port range specification
   * @returns Expanded port range
   */
  private expandPortRange(ports: TPortRange): number | PortRange | Array<number | PortRange> {
    // Process different port specifications
    if (typeof ports === 'number') {
      return ports;
    } else if (Array.isArray(ports)) {
      const result: Array<number | PortRange> = [];
      
      for (const item of ports) {
        if (typeof item === 'number') {
          result.push(item);
        } else if ('from' in item && 'to' in item) {
          result.push({ from: item.from, to: item.to });
        }
      }
      
      return result;
    } else if (typeof ports === 'object' && ports !== null && 'from' in ports && 'to' in ports) {
      return { from: (ports as any).from, to: (ports as any).to };
    }
    
    // Fallback to port 80 if something went wrong
    console.warn('Invalid port range specification, using port 80 as fallback');
    return 80;
  }
  
  /**
   * Get status of all managed rules
   * 
   * @returns A promise that resolves to a record of NFTables status objects
   */
  public async getStatus(): Promise<Record<string, NfTablesStatus>> {
    const result: Record<string, NfTablesStatus> = {};
    
    for (const [routeId, proxy] of this.rulesMap.entries()) {
      result[routeId] = await proxy.getStatus();
    }
    
    return result;
  }
  
  /**
   * Check if a route is currently provisioned
   * 
   * @param route The route configuration
   * @returns True if the route is provisioned, false otherwise
   */
  public isRouteProvisioned(route: IRouteConfig): boolean {
    const routeId = this.generateRouteId(route);
    return this.rulesMap.has(routeId);
  }
  
  /**
   * Stop all NFTables rules
   * 
   * @returns A promise that resolves when all rules have been stopped
   */
  public async stop(): Promise<void> {
    // Stop all NFTables proxies
    const stopPromises = Array.from(this.rulesMap.values()).map(proxy => proxy.stop());
    await Promise.all(stopPromises);
    
    this.rulesMap.clear();
  }
}