smartvm/ts/classes.smartvm.ts

import * as plugins from './plugins.js';
import type { ISmartVMOptions, IMicroVMConfig } from './interfaces/index.js';
import { SmartVMError } from './interfaces/index.js';
import { ImageManager } from './classes.imagemanager.js';
import { NetworkManager } from './classes.networkmanager.js';
import { MicroVM } from './classes.microvm.js';

/**
 * Top-level orchestrator for creating and managing Firecracker MicroVMs.
 */
export class SmartVM {
  private options: ISmartVMOptions;
  public imageManager: ImageManager;
  public networkManager: NetworkManager;
  private activeVMs: Map<string, MicroVM> = new Map();
  private smartExitInstance: InstanceType<typeof plugins.smartexit.SmartExit>;
  private firecrackerVersion: string | null = null;
  private firecrackerBinaryPath: string | null = null;

  constructor(options: ISmartVMOptions = {}) {
    this.options = {
      dataDir: options.dataDir || '/tmp/.smartvm',
      arch: options.arch || 'x86_64',
      bridgeName: options.bridgeName || 'svbr0',
      subnet: options.subnet || '172.30.0.0/24',
      ...options,
    };

    this.imageManager = new ImageManager(this.options.dataDir!, this.options.arch);
    this.networkManager = new NetworkManager({
      bridgeName: this.options.bridgeName,
      subnet: this.options.subnet,
    });

    // If a custom binary path is provided, use it directly
    if (this.options.firecrackerBinaryPath) {
      this.firecrackerBinaryPath = this.options.firecrackerBinaryPath;
    }

    // Register global cleanup
    this.smartExitInstance = new plugins.smartexit.SmartExit({ silent: true });
    this.smartExitInstance.addCleanupFunction(async () => {
      await this.cleanup();
    });
  }

  /**
   * Ensure the Firecracker binary is available.
   * Downloads it if not present.
   */
  public async ensureBinary(): Promise<string> {
    // If custom binary path is set, just verify it exists
    if (this.firecrackerBinaryPath) {
      try {
        await plugins.fs.promises.access(this.firecrackerBinaryPath);
        return this.firecrackerBinaryPath;
      } catch {
        // File doesn't exist, fall through to error
      }
      throw new SmartVMError(
        `Firecracker binary not found at ${this.firecrackerBinaryPath}`,
        'BINARY_NOT_FOUND',
      );
    }

    // Ensure data directories exist
    await this.imageManager.ensureDirectories();

    // Determine version
    let version = this.options.firecrackerVersion;
    if (!version) {
      version = await this.imageManager.getLatestVersion();
    }
    this.firecrackerVersion = version;

    // Check if binary exists
    if (await this.imageManager.hasBinary(version)) {
      this.firecrackerBinaryPath = this.imageManager.getFirecrackerPath(version);
      return this.firecrackerBinaryPath;
    }

    // Download the binary
    this.firecrackerBinaryPath = await this.imageManager.downloadFirecracker(version);
    return this.firecrackerBinaryPath;
  }

  /**
   * Create a new MicroVM with the given configuration.
   * Returns the MicroVM instance (not yet started).
   */
  public async createVM(config: IMicroVMConfig): Promise<MicroVM> {
    // Ensure binary is available
    if (!this.firecrackerBinaryPath) {
      await this.ensureBinary();
    }

    // Generate VM ID if not provided
    const vmId = config.id || plugins.smartunique.uuid4();

    // Generate socket path
    const socketPath = this.imageManager.getSocketPath(vmId);

    // Create MicroVM instance
    const vm = new MicroVM(
      vmId,
      config,
      this.firecrackerBinaryPath!,
      socketPath,
      this.networkManager,
    );

    // Register in active VMs
    this.activeVMs.set(vmId, vm);

    return vm;
  }

  /**
   * Get an active VM by ID.
   */
  public getVM(id: string): MicroVM | undefined {
    return this.activeVMs.get(id);
  }

  /**
   * List all active VM IDs.
   */
  public listVMs(): string[] {
    return Array.from(this.activeVMs.keys());
  }

  /**
   * Get the number of active VMs.
   */
  public get vmCount(): number {
    return this.activeVMs.size;
  }

  /**
   * Stop all running VMs.
   */
  public async stopAll(): Promise<void> {
    const stopPromises: Promise<void>[] = [];
    for (const vm of this.activeVMs.values()) {
      if (vm.state === 'running' || vm.state === 'paused') {
        stopPromises.push(
          vm.stop().catch((err) => {
            console.error(`Failed to stop VM ${vm.id}: ${err.message}`);
          }),
        );
      }
    }
    await Promise.all(stopPromises);
  }

  /**
   * Clean up all VMs and networking resources.
   */
  public async cleanup(): Promise<void> {
    // Clean up all VMs
    const cleanupPromises: Promise<void>[] = [];
    for (const vm of this.activeVMs.values()) {
      cleanupPromises.push(
        vm.cleanup().catch((err) => {
          console.error(`Failed to clean up VM ${vm.id}: ${err.message}`);
        }),
      );
    }
    await Promise.all(cleanupPromises);
    this.activeVMs.clear();

    // Clean up networking
    await this.networkManager.cleanup();
  }

  /**
   * Remove a VM from the active list (after it's been cleaned up).
   */
  public removeVM(id: string): boolean {
    return this.activeVMs.delete(id);
  }
}