smartproxy/ts/proxies/smart-proxy/rust-proxy-bridge.ts

import * as plugins from '../../plugins.js';
import { logger } from '../../core/utils/logger.js';
import { RustBinaryLocator } from './rust-binary-locator.js';
import type { IRouteConfig } from './models/route-types.js';
import { ChildProcess, spawn } from 'child_process';
import { createInterface, Interface as ReadlineInterface } from 'readline';

/**
 * Management request sent to the Rust binary via stdin.
 */
interface IManagementRequest {
  id: string;
  method: string;
  params: Record<string, any>;
}

/**
 * Management response received from the Rust binary via stdout.
 */
interface IManagementResponse {
  id: string;
  success: boolean;
  result?: any;
  error?: string;
}

/**
 * Management event received from the Rust binary (unsolicited).
 */
interface IManagementEvent {
  event: string;
  data: any;
}

/**
 * Bridge between TypeScript SmartProxy and the Rust binary.
 * Communicates via JSON-over-stdin/stdout IPC protocol.
 */
export class RustProxyBridge extends plugins.EventEmitter {
  private locator = new RustBinaryLocator();
  private process: ChildProcess | null = null;
  private readline: ReadlineInterface | null = null;
  private pendingRequests = new Map<string, {
    resolve: (value: any) => void;
    reject: (error: Error) => void;
    timer: NodeJS.Timeout;
  }>();
  private requestCounter = 0;
  private isRunning = false;
  private binaryPath: string | null = null;
  private readonly requestTimeoutMs = 30000;

  /**
   * Spawn the Rust binary in management mode.
   * Returns true if the binary was found and spawned successfully.
   */
  public async spawn(): Promise<boolean> {
    this.binaryPath = await this.locator.findBinary();
    if (!this.binaryPath) {
      return false;
    }

    return new Promise<boolean>((resolve) => {
      try {
        this.process = spawn(this.binaryPath!, ['--management'], {
          stdio: ['pipe', 'pipe', 'pipe'],
          env: { ...process.env },
        });

        // Handle stderr (logging from Rust goes here)
        this.process.stderr?.on('data', (data: Buffer) => {
          const lines = data.toString().split('\n').filter(l => l.trim());
          for (const line of lines) {
            logger.log('debug', `[rustproxy] ${line}`, { component: 'rust-bridge' });
          }
        });

        // Handle stdout (JSON IPC)
        this.readline = createInterface({ input: this.process.stdout! });
        this.readline.on('line', (line: string) => {
          this.handleLine(line.trim());
        });

        // Handle process exit
        this.process.on('exit', (code, signal) => {
          logger.log('info', `RustProxy process exited (code=${code}, signal=${signal})`, { component: 'rust-bridge' });
          this.cleanup();
          this.emit('exit', code, signal);
        });

        this.process.on('error', (err) => {
          logger.log('error', `RustProxy process error: ${err.message}`, { component: 'rust-bridge' });
          this.cleanup();
          resolve(false);
        });

        // Wait for the 'ready' event from Rust
        const readyTimeout = setTimeout(() => {
          logger.log('error', 'RustProxy did not send ready event within 10s', { component: 'rust-bridge' });
          this.kill();
          resolve(false);
        }, 10000);

        this.once('management:ready', () => {
          clearTimeout(readyTimeout);
          this.isRunning = true;
          logger.log('info', 'RustProxy bridge connected', { component: 'rust-bridge' });
          resolve(true);
        });
      } catch (err: any) {
        logger.log('error', `Failed to spawn RustProxy: ${err.message}`, { component: 'rust-bridge' });
        resolve(false);
      }
    });
  }

  /**
   * Send a management command to the Rust process and wait for the response.
   */
  public async sendCommand(method: string, params: Record<string, any> = {}): Promise<any> {
    if (!this.process || !this.isRunning) {
      throw new Error('RustProxy bridge is not running');
    }

    const id = `req_${++this.requestCounter}`;
    const request: IManagementRequest = { id, method, params };

    return new Promise<any>((resolve, reject) => {
      const timer = setTimeout(() => {
        this.pendingRequests.delete(id);
        reject(new Error(`RustProxy command '${method}' timed out after ${this.requestTimeoutMs}ms`));
      }, this.requestTimeoutMs);

      this.pendingRequests.set(id, { resolve, reject, timer });

      const json = JSON.stringify(request) + '\n';
      this.process!.stdin!.write(json, (err) => {
        if (err) {
          clearTimeout(timer);
          this.pendingRequests.delete(id);
          reject(new Error(`Failed to write to RustProxy stdin: ${err.message}`));
        }
      });
    });
  }

  // Convenience methods for each management command

  public async startProxy(config: any): Promise<void> {
    await this.sendCommand('start', { config });
  }

  public async stopProxy(): Promise<void> {
    await this.sendCommand('stop');
  }

  public async updateRoutes(routes: IRouteConfig[]): Promise<void> {
    await this.sendCommand('updateRoutes', { routes });
  }

  public async getMetrics(): Promise<any> {
    return this.sendCommand('getMetrics');
  }

  public async getStatistics(): Promise<any> {
    return this.sendCommand('getStatistics');
  }

  public async provisionCertificate(routeName: string): Promise<void> {
    await this.sendCommand('provisionCertificate', { routeName });
  }

  public async renewCertificate(routeName: string): Promise<void> {
    await this.sendCommand('renewCertificate', { routeName });
  }

  public async getCertificateStatus(routeName: string): Promise<any> {
    return this.sendCommand('getCertificateStatus', { routeName });
  }

  public async getListeningPorts(): Promise<number[]> {
    const result = await this.sendCommand('getListeningPorts');
    return result?.ports ?? [];
  }

  public async getNftablesStatus(): Promise<any> {
    return this.sendCommand('getNftablesStatus');
  }

  public async setSocketHandlerRelay(socketPath: string): Promise<void> {
    await this.sendCommand('setSocketHandlerRelay', { socketPath });
  }

  public async addListeningPort(port: number): Promise<void> {
    await this.sendCommand('addListeningPort', { port });
  }

  public async removeListeningPort(port: number): Promise<void> {
    await this.sendCommand('removeListeningPort', { port });
  }

  public async loadCertificate(domain: string, cert: string, key: string, ca?: string): Promise<void> {
    await this.sendCommand('loadCertificate', { domain, cert, key, ca });
  }

  /**
   * Kill the Rust process.
   */
  public kill(): void {
    if (this.process) {
      this.process.kill('SIGTERM');
      // Force kill after 5 seconds
      setTimeout(() => {
        if (this.process) {
          this.process.kill('SIGKILL');
        }
      }, 5000).unref();
    }
  }

  /**
   * Whether the bridge is currently running.
   */
  public get running(): boolean {
    return this.isRunning;
  }

  private handleLine(line: string): void {
    if (!line) return;

    let parsed: any;
    try {
      parsed = JSON.parse(line);
    } catch {
      logger.log('warn', `Non-JSON output from RustProxy: ${line}`, { component: 'rust-bridge' });
      return;
    }

    // Check if it's an event (has 'event' field)
    if ('event' in parsed) {
      const event = parsed as IManagementEvent;
      this.emit(`management:${event.event}`, event.data);
      return;
    }

    // Otherwise it's a response (has 'id' field)
    if ('id' in parsed) {
      const response = parsed as IManagementResponse;
      const pending = this.pendingRequests.get(response.id);
      if (pending) {
        clearTimeout(pending.timer);
        this.pendingRequests.delete(response.id);
        if (response.success) {
          pending.resolve(response.result);
        } else {
          pending.reject(new Error(response.error || 'Unknown error from RustProxy'));
        }
      }
    }
  }

  private cleanup(): void {
    this.isRunning = false;
    this.process = null;

    if (this.readline) {
      this.readline.close();
      this.readline = null;
    }

    // Reject all pending requests
    for (const [id, pending] of this.pendingRequests) {
      clearTimeout(pending.timer);
      pending.reject(new Error('RustProxy process exited'));
    }
    this.pendingRequests.clear();
  }
}