smartvm/ts/classes.microvm.ts

import * as plugins from './plugins.js';
import type {
  TVMState,
  IMicroVMConfig,
  ISnapshotCreateParams,
  ISnapshotLoadParams,
  ITapDevice,
} from './interfaces/index.js';
import { SmartVMError } from './interfaces/index.js';
import { VMConfig } from './classes.vmconfig.js';
import { SocketClient } from './classes.socketclient.js';
import { FirecrackerProcess } from './classes.firecrackerprocess.js';
import { NetworkManager } from './classes.networkmanager.js';

/**
 * Represents a single Firecracker MicroVM with full lifecycle management.
 * State machine: created → configuring → running → paused → stopped
 */
export class MicroVM {
  public readonly id: string;
  public state: TVMState = 'created';

  private vmConfig: VMConfig;
  private process: FirecrackerProcess | null = null;
  private socketClient: SocketClient | null = null;
  private networkManager: NetworkManager;
  private binaryPath: string;
  private socketPath: string;
  private tapDevices: ITapDevice[] = [];

  constructor(
    id: string,
    config: IMicroVMConfig,
    binaryPath: string,
    socketPath: string,
    networkManager: NetworkManager,
  ) {
    this.id = id;
    this.vmConfig = new VMConfig(config);
    this.binaryPath = binaryPath;
    this.socketPath = socketPath;
    this.networkManager = networkManager;
  }

  /**
   * Assert that the VM is in one of the expected states.
   */
  private assertState(expected: TVMState[], operation: string): void {
    if (!expected.includes(this.state)) {
      throw new SmartVMError(
        `Cannot ${operation}: VM is in state '${this.state}', expected one of [${expected.join(', ')}]`,
        'INVALID_STATE',
      );
    }
  }

  /**
   * Start the MicroVM.
   * Validates config, starts the Firecracker process, applies pre-boot configuration, and boots the VM.
   */
  public async start(): Promise<void> {
    this.assertState(['created'], 'start');

    // Validate configuration
    const validation = this.vmConfig.validate();
    if (!validation.valid) {
      throw new SmartVMError(
        `Invalid VM configuration: ${validation.errors.join('; ')}`,
        'INVALID_CONFIG',
      );
    }

    this.state = 'configuring';

    try {
      // Start the Firecracker process
      this.process = new FirecrackerProcess({
        binaryPath: this.binaryPath,
        socketPath: this.socketPath,
      });
      await this.process.start();
      this.socketClient = this.process.socketClient;

      // Apply pre-boot configuration in order

      // 1. Logger (optional, must be first)
      const loggerPayload = this.vmConfig.toLoggerPayload();
      if (loggerPayload) {
        await this.apiPut('/logger', loggerPayload);
      }

      // 2. Metrics (optional)
      const metricsPayload = this.vmConfig.toMetricsPayload();
      if (metricsPayload) {
        await this.apiPut('/metrics', metricsPayload);
      }

      // 3. Machine config
      await this.apiPut('/machine-config', this.vmConfig.toMachineConfigPayload());

      // 4. Boot source
      await this.apiPut('/boot-source', this.vmConfig.toBootSourcePayload());

      // 5. Drives
      if (this.vmConfig.config.drives) {
        for (const drive of this.vmConfig.config.drives) {
          const payload = this.vmConfig.toDrivePayload(drive);
          await this.apiPut(`/drives/${drive.driveId}`, payload);
        }
      }

      // 6. Network interfaces
      if (this.vmConfig.config.networkInterfaces) {
        for (const iface of this.vmConfig.config.networkInterfaces) {
          // Create TAP device if hostDevName not manually specified
          if (!iface.hostDevName) {
            const tap = await this.networkManager.createTapDevice(this.id, iface.ifaceId);
            this.tapDevices.push(tap);
            iface.hostDevName = tap.tapName;
            if (!iface.guestMac) {
              iface.guestMac = tap.mac;
            }
          }
          const payload = this.vmConfig.toNetworkInterfacePayload(iface);
          await this.apiPut(`/network-interfaces/${iface.ifaceId}`, payload);
        }
      }

      // 7. Vsock (optional)
      const vsockPayload = this.vmConfig.toVsockPayload();
      if (vsockPayload) {
        await this.apiPut('/vsock', vsockPayload);
      }

      // 8. Balloon (optional)
      const balloonPayload = this.vmConfig.toBalloonPayload();
      if (balloonPayload) {
        await this.apiPut('/balloon', balloonPayload);
      }

      // 9. MMDS config (optional)
      const mmdsPayload = this.vmConfig.toMmdsConfigPayload();
      if (mmdsPayload) {
        await this.apiPut('/mmds/config', mmdsPayload);
      }

      // Boot the VM
      await this.apiPut('/actions', { action_type: 'InstanceStart' });

      this.state = 'running';
    } catch (err) {
      this.state = 'error';
      // Clean up on failure
      await this.cleanup();
      if (err instanceof SmartVMError) {
        throw err;
      }
      throw new SmartVMError(
        `Failed to start VM ${this.id}: ${err.message}`,
        'START_FAILED',
      );
    }
  }

  /**
   * Pause the running VM.
   */
  public async pause(): Promise<void> {
    this.assertState(['running'], 'pause');
    await this.apiPatch('/vm', { state: 'Paused' });
    this.state = 'paused';
  }

  /**
   * Resume a paused VM.
   */
  public async resume(): Promise<void> {
    this.assertState(['paused'], 'resume');
    await this.apiPatch('/vm', { state: 'Resumed' });
    this.state = 'running';
  }

  /**
   * Stop the VM.
   * Sends SendCtrlAltDel first for graceful shutdown, then kills the process if needed.
   */
  public async stop(): Promise<void> {
    this.assertState(['running', 'paused'], 'stop');

    try {
      // Try graceful shutdown via SendCtrlAltDel
      await this.apiPut('/actions', { action_type: 'SendCtrlAltDel' });
      // Wait a bit for the VM to shut down
      await plugins.smartdelay.delayFor(2000);
    } catch {
      // SendCtrlAltDel may fail if the VM is already stopping
    }

    // Force stop the process
    if (this.process) {
      await this.process.stop();
    }

    this.state = 'stopped';
  }

  /**
   * Create a snapshot of the VM.
   */
  public async createSnapshot(params: ISnapshotCreateParams): Promise<void> {
    this.assertState(['paused'], 'createSnapshot');
    await this.apiPut('/snapshot/create', {
      snapshot_path: params.snapshotPath,
      mem_file_path: params.memFilePath,
      snapshot_type: params.snapshotType || 'Full',
    });
  }

  /**
   * Load a snapshot into the VM.
   */
  public async loadSnapshot(params: ISnapshotLoadParams): Promise<void> {
    this.assertState(['created', 'configuring'], 'loadSnapshot');
    await this.apiPut('/snapshot/load', {
      snapshot_path: params.snapshotPath,
      mem_file_path: params.memFilePath,
      enable_diff_snapshots: params.enableDiffSnapshots ?? false,
      resume_vm: params.resumeVm ?? false,
    });
    this.state = params.resumeVm ? 'running' : 'paused';
  }

  /**
   * Set MMDS metadata.
   */
  public async setMetadata(data: Record<string, any>): Promise<void> {
    this.assertState(['running', 'paused'], 'setMetadata');
    await this.apiPut('/mmds', data);
  }

  /**
   * Get MMDS metadata.
   */
  public async getMetadata(): Promise<any> {
    this.assertState(['running', 'paused'], 'getMetadata');
    const response = await this.socketClient!.get('/mmds');
    return response.body;
  }

  /**
   * Update a drive on the running VM (hot-plug path update).
   */
  public async updateDrive(driveId: string, pathOnHost: string): Promise<void> {
    this.assertState(['running', 'paused'], 'updateDrive');
    await this.apiPatch(`/drives/${driveId}`, {
      drive_id: driveId,
      path_on_host: pathOnHost,
    });
  }

  /**
   * Update a network interface rate limiter on the running VM.
   */
  public async updateNetworkInterface(
    ifaceId: string,
    update: Record<string, any>,
  ): Promise<void> {
    this.assertState(['running', 'paused'], 'updateNetworkInterface');
    await this.apiPatch(`/network-interfaces/${ifaceId}`, update);
  }

  /**
   * Update balloon device on the running VM.
   */
  public async updateBalloon(amountMib: number): Promise<void> {
    this.assertState(['running', 'paused'], 'updateBalloon');
    await this.apiPatch('/balloon', { amount_mib: amountMib });
  }

  /**
   * Get VM instance info.
   */
  public async getInfo(): Promise<any> {
    const response = await this.socketClient!.get('/');
    return response.body;
  }

  /**
   * Get Firecracker version info.
   */
  public async getVersion(): Promise<any> {
    const response = await this.socketClient!.get('/version');
    return response.body;
  }

  /**
   * Get the TAP devices associated with this VM.
   */
  public getTapDevices(): ITapDevice[] {
    return [...this.tapDevices];
  }

  /**
   * Get the VMConfig instance for payload inspection.
   */
  public getVMConfig(): VMConfig {
    return this.vmConfig;
  }

  /**
   * Full cleanup: stop process, remove socket, remove TAP devices.
   */
  public async cleanup(): Promise<void> {
    // Stop process
    if (this.process) {
      await this.process.cleanup();
      this.process = null;
    }

    // Remove TAP devices
    for (const tap of this.tapDevices) {
      await this.networkManager.removeTapDevice(tap.tapName);
    }
    this.tapDevices = [];

    this.socketClient = null;
    if (this.state !== 'error') {
      this.state = 'stopped';
    }
  }

  /**
   * Helper: PUT request with error handling.
   */
  private async apiPut(path: string, body: Record<string, any>): Promise<void> {
    if (!this.socketClient) {
      throw new SmartVMError('Socket client not initialized', 'NO_CLIENT');
    }
    const response = await this.socketClient.put(path, body);
    if (!response.ok) {
      throw new SmartVMError(
        `API PUT ${path} failed with status ${response.statusCode}: ${JSON.stringify(response.body)}`,
        'API_ERROR',
        response.statusCode,
        response.body,
      );
    }
  }

  /**
   * Helper: PATCH request with error handling.
   */
  private async apiPatch(path: string, body: Record<string, any>): Promise<void> {
    if (!this.socketClient) {
      throw new SmartVMError('Socket client not initialized', 'NO_CLIENT');
    }
    const response = await this.socketClient.patch(path, body);
    if (!response.ok) {
      throw new SmartVMError(
        `API PATCH ${path} failed with status ${response.statusCode}: ${JSON.stringify(response.body)}`,
        'API_ERROR',
        response.statusCode,
        response.body,
      );
    }
  }
}