smartvm/ts/classes.smartvm.ts

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

function getErrorMessage(err: unknown): string {
  return err instanceof Error ? err.message : String(err);
}

/**
 * Top-level orchestrator for creating and managing Firecracker MicroVMs.
 */
export class SmartVM {
  private options: ISmartVMOptions;
  public imageManager: ImageManager;
  public baseImageManager: BaseImageManager;
  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',
      runtimeDir: options.runtimeDir || this.getDefaultRuntimeDir(),
      ephemeralWritableDrives: options.ephemeralWritableDrives ?? true,
      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.baseImageManager = new BaseImageManager({
      arch: this.options.arch,
      cacheDir: this.options.baseImageCacheDir,
      maxStoredBaseImages: this.options.maxStoredBaseImages,
      hostedManifestUrl: this.options.baseImageManifestUrl,
      hostedManifestPath: this.options.baseImageManifestPath,
    });
    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();
    });
  }

  private getDefaultRuntimeDir(): string {
    const tmpfsDir = '/dev/shm';
    try {
      if (plugins.fs.existsSync(tmpfsDir) && plugins.fs.statSync(tmpfsDir).isDirectory()) {
        return plugins.path.join(tmpfsDir, '.smartvm', 'runtime');
      }
    } catch {
      // Fall back to os.tmpdir() below.
    }
    return plugins.path.join(plugins.os.tmpdir(), '.smartvm', 'runtime');
  }

  public getRuntimeDir(): string {
    return this.options.runtimeDir!;
  }

  private sanitizePathPart(value: string): string {
    const sanitized = value.replace(/[^a-zA-Z0-9._-]/g, '_');
    if (!sanitized || sanitized === '.' || sanitized === '..') {
      return 'item';
    }
    return sanitized;
  }

  /**
   * 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();

    // Keep per-VM runtime artifacts in tmpfs by default.
    const vmRuntimeDir = plugins.path.join(this.options.runtimeDir!, this.sanitizePathPart(vmId));
    const socketPath = plugins.path.join(vmRuntimeDir, 'firecracker.sock');

    // Create MicroVM instance
    const vm = new MicroVM(
      vmId,
      config,
      this.firecrackerBinaryPath!,
      socketPath,
      this.networkManager,
      {
        runtimeDir: this.options.runtimeDir,
        ephemeralWritableDrives: this.options.ephemeralWritableDrives,
      },
    );

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

    return vm;
  }

  /**
   * Ensure a Firecracker CI base image bundle is available locally.
   */
  public async ensureBaseImage(options: IEnsureBaseImageOptions = {}): Promise<IBaseImageBundle> {
    return this.baseImageManager.ensureBaseImage(options);
  }

  /**
   * 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}: ${getErrorMessage(err)}`);
          }),
        );
      }
    }
    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}: ${getErrorMessage(err)}`);
        }),
      );
    }
    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);
  }
}