smartproxy/ts/forwarding/config/domain-manager.ts

import * as plugins from '../../plugins.js';
import type { DomainConfig } from './domain-config.js';
import { ForwardingHandler } from '../handlers/base-handler.js';
import { ForwardingHandlerEvents } from './forwarding-types.js';
import { ForwardingHandlerFactory } from '../factory/forwarding-factory.js';

/**
 * Events emitted by the DomainManager
 */
export enum DomainManagerEvents {
  DOMAIN_ADDED = 'domain-added',
  DOMAIN_REMOVED = 'domain-removed',
  DOMAIN_MATCHED = 'domain-matched',
  DOMAIN_MATCH_FAILED = 'domain-match-failed',
  CERTIFICATE_NEEDED = 'certificate-needed',
  CERTIFICATE_LOADED = 'certificate-loaded',
  ERROR = 'error'
}

/**
 * Manages domains and their forwarding handlers
 */
export class DomainManager extends plugins.EventEmitter {
  private domainConfigs: DomainConfig[] = [];
  private domainHandlers: Map<string, ForwardingHandler> = new Map();
  
  /**
   * Create a new DomainManager
   * @param initialDomains Optional initial domain configurations
   */
  constructor(initialDomains?: DomainConfig[]) {
    super();
    
    if (initialDomains) {
      this.setDomainConfigs(initialDomains);
    }
  }
  
  /**
   * Set or replace all domain configurations
   * @param configs Array of domain configurations
   */
  public async setDomainConfigs(configs: DomainConfig[]): Promise<void> {
    // Clear existing handlers
    this.domainHandlers.clear();
    
    // Store new configurations
    this.domainConfigs = [...configs];
    
    // Initialize handlers for each domain
    for (const config of this.domainConfigs) {
      await this.createHandlersForDomain(config);
    }
  }
  
  /**
   * Add a new domain configuration
   * @param config The domain configuration to add
   */
  public async addDomainConfig(config: DomainConfig): Promise<void> {
    // Check if any of these domains already exist
    for (const domain of config.domains) {
      if (this.domainHandlers.has(domain)) {
        // Remove existing handler for this domain
        this.domainHandlers.delete(domain);
      }
    }
    
    // Add the new configuration
    this.domainConfigs.push(config);
    
    // Create handlers for the new domain
    await this.createHandlersForDomain(config);
    
    this.emit(DomainManagerEvents.DOMAIN_ADDED, {
      domains: config.domains,
      forwardingType: config.forwarding.type
    });
  }
  
  /**
   * Remove a domain configuration
   * @param domain The domain to remove
   * @returns True if the domain was found and removed
   */
  public removeDomainConfig(domain: string): boolean {
    // Find the config that includes this domain
    const index = this.domainConfigs.findIndex(config => 
      config.domains.includes(domain)
    );
    
    if (index === -1) {
      return false;
    }
    
    // Get the config
    const config = this.domainConfigs[index];
    
    // Remove all handlers for this config
    for (const domainName of config.domains) {
      this.domainHandlers.delete(domainName);
    }
    
    // Remove the config
    this.domainConfigs.splice(index, 1);
    
    this.emit(DomainManagerEvents.DOMAIN_REMOVED, {
      domains: config.domains
    });
    
    return true;
  }
  
  /**
   * Find the handler for a domain
   * @param domain The domain to find a handler for
   * @returns The handler or undefined if no match
   */
  public findHandlerForDomain(domain: string): ForwardingHandler | undefined {
    // Try exact match
    if (this.domainHandlers.has(domain)) {
      return this.domainHandlers.get(domain);
    }
    
    // Try wildcard matches
    const wildcardHandler = this.findWildcardHandler(domain);
    if (wildcardHandler) {
      return wildcardHandler;
    }
    
    // No match found
    return undefined;
  }
  
  /**
   * Handle a connection for a domain
   * @param domain The domain
   * @param socket The client socket
   * @returns True if the connection was handled
   */
  public handleConnection(domain: string, socket: plugins.net.Socket): boolean {
    const handler = this.findHandlerForDomain(domain);
    
    if (!handler) {
      this.emit(DomainManagerEvents.DOMAIN_MATCH_FAILED, {
        domain,
        remoteAddress: socket.remoteAddress
      });
      return false;
    }
    
    this.emit(DomainManagerEvents.DOMAIN_MATCHED, {
      domain,
      handlerType: handler.constructor.name,
      remoteAddress: socket.remoteAddress
    });
    
    // Handle the connection
    handler.handleConnection(socket);
    return true;
  }
  
  /**
   * Handle an HTTP request for a domain
   * @param domain The domain
   * @param req The HTTP request
   * @param res The HTTP response
   * @returns True if the request was handled
   */
  public handleHttpRequest(domain: string, req: plugins.http.IncomingMessage, res: plugins.http.ServerResponse): boolean {
    const handler = this.findHandlerForDomain(domain);
    
    if (!handler) {
      this.emit(DomainManagerEvents.DOMAIN_MATCH_FAILED, {
        domain,
        remoteAddress: req.socket.remoteAddress
      });
      return false;
    }
    
    this.emit(DomainManagerEvents.DOMAIN_MATCHED, {
      domain,
      handlerType: handler.constructor.name,
      remoteAddress: req.socket.remoteAddress
    });
    
    // Handle the request
    handler.handleHttpRequest(req, res);
    return true;
  }
  
  /**
   * Create handlers for a domain configuration
   * @param config The domain configuration
   */
  private async createHandlersForDomain(config: DomainConfig): Promise<void> {
    try {
      // Create a handler for this forwarding configuration
      const handler = ForwardingHandlerFactory.createHandler(config.forwarding);
      
      // Initialize the handler
      await handler.initialize();
      
      // Set up event forwarding
      this.setupHandlerEvents(handler, config);
      
      // Store the handler for each domain in the config
      for (const domain of config.domains) {
        this.domainHandlers.set(domain, handler);
      }
    } catch (error) {
      this.emit(DomainManagerEvents.ERROR, {
        domains: config.domains,
        error: error instanceof Error ? error.message : String(error)
      });
    }
  }
  
  /**
   * Set up event forwarding from a handler
   * @param handler The handler
   * @param config The domain configuration for this handler
   */
  private setupHandlerEvents(handler: ForwardingHandler, config: DomainConfig): void {
    // Forward relevant events
    handler.on(ForwardingHandlerEvents.CERTIFICATE_NEEDED, (data) => {
      this.emit(DomainManagerEvents.CERTIFICATE_NEEDED, {
        ...data,
        domains: config.domains
      });
    });
    
    handler.on(ForwardingHandlerEvents.CERTIFICATE_LOADED, (data) => {
      this.emit(DomainManagerEvents.CERTIFICATE_LOADED, {
        ...data,
        domains: config.domains
      });
    });
    
    handler.on(ForwardingHandlerEvents.ERROR, (data) => {
      this.emit(DomainManagerEvents.ERROR, {
        ...data,
        domains: config.domains
      });
    });
  }
  
  /**
   * Find a handler for a domain using wildcard matching
   * @param domain The domain to find a handler for
   * @returns The handler or undefined if no match
   */
  private findWildcardHandler(domain: string): ForwardingHandler | undefined {
    // Exact match already checked in findHandlerForDomain
    
    // Try subdomain wildcard (*.example.com)
    if (domain.includes('.')) {
      const parts = domain.split('.');
      if (parts.length > 2) {
        const wildcardDomain = `*.${parts.slice(1).join('.')}`;
        if (this.domainHandlers.has(wildcardDomain)) {
          return this.domainHandlers.get(wildcardDomain);
        }
      }
    }
    
    // Try full wildcard
    if (this.domainHandlers.has('*')) {
      return this.domainHandlers.get('*');
    }
    
    // No match found
    return undefined;
  }
  
  /**
   * Get all domain configurations
   * @returns Array of domain configurations
   */
  public getDomainConfigs(): DomainConfig[] {
    return [...this.domainConfigs];
  }
}