fireclaw/src/vm.ts

import { type ChildProcess } from "node:child_process";
import { mkdirSync } from "node:fs";
import { join } from "node:path";
import { randomBytes } from "node:crypto";
import { CONFIG } from "./config.js";
import type { VMConfig, RunResult, RunOptions } from "./types.js";
import * as api from "./firecracker-api.js";
import { allocateIp, releaseIp, deleteTap } from "./network.js";
import {
  ensureBaseImage,
  ensureSshKeypair,
  createRunCopy,
  injectSshKey,
  deleteRunCopy,
} from "./rootfs.js";
import { waitForSsh, execCommand } from "./ssh.js";
import { registerVm, unregisterVm } from "./cleanup.js";
import { snapshotExists } from "./snapshot.js";
import {
  setupNetwork,
  spawnFirecracker,
  bootVM,
  killFirecracker,
} from "./firecracker-vm.js";

function log(verbose: boolean, msg: string) {
  if (verbose) process.stderr.write(`[fireclaw] ${msg}\n`);
}

export class VMInstance {
  private config: VMConfig;
  private process: ChildProcess | null = null;
  private octet = 0;

  constructor(config: VMConfig) {
    this.config = config;
  }

  static async run(
    command: string,
    opts: RunOptions = {}
  ): Promise<RunResult> {
    if (!opts.noSnapshot && snapshotExists()) {
      return VMInstance.runFromSnapshot(command, opts);
    }
    return VMInstance.runColdBoot(command, opts);
  }

  private static async runFromSnapshot(
    command: string,
    opts: RunOptions
  ): Promise<RunResult> {
    const id = `fc-snap-${randomBytes(3).toString("hex")}`;
    const verbose = opts.verbose ?? false;
    const timeoutMs = opts.timeout ?? CONFIG.vm.defaultTimeoutMs;
    const snap = CONFIG.snapshot;

    mkdirSync(CONFIG.socketDir, { recursive: true });

    const config: VMConfig = {
      id,
      guestIp: snap.ip,
      tapDevice: snap.tapDevice,
      socketPath: join(CONFIG.socketDir, `${id}.sock`),
      rootfsPath: "",
      timeoutMs,
      verbose,
    };

    const vm = new VMInstance(config);
    vm.octet = 0;
    registerVm(vm);

    try {
      log(verbose, `VM ${id}: restoring from snapshot...`);
      setupNetwork(snap.tapDevice);

      vm.process = await spawnFirecracker(config.socketPath);
      await api.putSnapshotLoad(
        config.socketPath,
        snap.statePath,
        snap.memPath
      );
      await api.patchVm(config.socketPath, "Resumed");

      log(verbose, `VM ${id}: resumed, waiting for SSH...`);
      await waitForSsh(snap.ip);

      log(verbose, `VM ${id}: executing command...`);
      const result = await execCommand(snap.ip, command, timeoutMs, verbose);

      log(
        verbose,
        `VM ${id}: done (exit=${result.exitCode}, ${result.durationMs}ms)`
      );
      return result;
    } finally {
      await vm.destroy();
      unregisterVm(vm);
    }
  }

  private static async runColdBoot(
    command: string,
    opts: RunOptions
  ): Promise<RunResult> {
    const id = `fc-${randomBytes(3).toString("hex")}`;
    const verbose = opts.verbose ?? false;
    const timeoutMs = opts.timeout ?? CONFIG.vm.defaultTimeoutMs;

    ensureBaseImage();
    ensureSshKeypair();

    const { ip, octet } = allocateIp();
    const tapDevice = `fctap${octet}`;

    const config: VMConfig = {
      id,
      guestIp: ip,
      tapDevice,
      socketPath: join(CONFIG.socketDir, `${id}.sock`),
      rootfsPath: "",
      timeoutMs,
      verbose,
    };

    const vm = new VMInstance(config);
    vm.octet = octet;
    registerVm(vm);

    try {
      log(verbose, `VM ${id}: preparing rootfs...`);
      config.rootfsPath = createRunCopy(id);
      injectSshKey(config.rootfsPath);

      log(verbose, `VM ${id}: creating tap ${tapDevice}...`);
      setupNetwork(tapDevice);

      log(verbose, `VM ${id}: booting...`);
      vm.process = await spawnFirecracker(config.socketPath);
      await bootVM({
        socketPath: config.socketPath,
        rootfsPath: config.rootfsPath,
        tapDevice,
        ip,
        octet,
        vcpu: opts.vcpu,
        mem: opts.mem,
      });

      log(verbose, `VM ${id}: waiting for SSH at ${ip}...`);
      await waitForSsh(ip);

      log(verbose, `VM ${id}: executing command...`);
      const result = await execCommand(ip, command, timeoutMs, verbose);

      log(
        verbose,
        `VM ${id}: done (exit=${result.exitCode}, ${result.durationMs}ms)`
      );
      return result;
    } finally {
      await vm.destroy();
      unregisterVm(vm);
    }
  }

  async destroy() {
    const { config } = this;
    log(config.verbose, `VM ${config.id}: cleaning up...`);

    await killFirecracker(this.process, config.socketPath);
    deleteTap(config.tapDevice);

    if (this.octet > 0) {
      releaseIp(this.octet);
    }

    if (config.rootfsPath) {
      deleteRunCopy(config.rootfsPath);
    }
  }
}