lab/bastion/tests/integration/helpers/pxe-vm.ts

// Create a blank UEFI VM for PXE boot testing.
// Unlike cloud image VMs, these have an empty disk and boot from network.
// Each VM gets a serial console on a TCP port for debugging without network/SSH.

import { execSync, spawnSync, type SpawnSyncReturns } from "node:child_process";
import { existsSync } from "node:fs";
import { join } from "node:path";
import { createConnection } from "node:net";
import { log } from "./libvirt.js";

const IMAGE_DIR = "/var/lib/libvirt/images";
const IS_ROOT = process.getuid?.() === 0;

function run(cmd: string, opts?: { timeout?: number }): string {
  const full = IS_ROOT ? cmd : `sudo ${cmd}`;
  return execSync(full, { encoding: "utf-8", stdio: "pipe", timeout: opts?.timeout ?? 60_000 });
}

function virsh(...args: string[]): SpawnSyncReturns<string> {
  const cmd = IS_ROOT ? "virsh" : "sudo";
  const finalArgs = IS_ROOT ? args : ["virsh", ...args];
  return spawnSync(cmd, finalArgs, { encoding: "utf-8", stdio: "pipe", timeout: 30_000 });
}

export interface PxeVmConfig {
  name: string;
  memory: number;       // MB
  vcpus: number;
  diskSize: number;     // GB
  network: string;      // libvirt network name
  /**
   * Serial console TCP port. Defaults to 4555 for backwards compatibility with
   * the single-VM tests, but it MUST be unique per concurrently-running VM:
   * qemu fails with "Failed to find an available port: Address already in use"
   * and the domain never starts.
   */
  serialPort?: number;
  arch?: "x86_64" | "aarch64";
  /**
   * Extra NICs enumerated BEFORE the PXE NIC, on a network with no route to
   * the bastion (defaults to libvirt's "default").
   *
   * Real multi-NIC boxes expose a class of bug a single-NIC VM cannot: an
   * initramfs that picks "the first connected interface" grabs one of these
   * instead of the NIC that PXE booted, and then cannot reach the bastion.
   * Defaults to 0 (single NIC).
   */
  decoyNics?: number;
  decoyNetwork?: string;
  /**
   * Extra blank data disks, attached after the OS disk. Each gets a fixed
   * virtio serial, so inside the guest it has a stable
   * /dev/disk/by-id/virtio-<serial> path -- the moral equivalent of the
   * nvme-..._<serial> links the dedicated-data-disk feature requires on real
   * hardware. virtio-blk truncates serials at 20 characters, so keep them
   * short or the by-id link quietly won't match.
   */
  dataDisks?: { serial: string; sizeGb: number }[];
}

/** Create a blank UEFI VM that PXE boots from the network. */
export function createPxeVm(config: PxeVmConfig): void {
  destroyPxeVm(config.name);

  const arch = config.arch ?? "x86_64";
  log(`Creating PXE VM: ${config.name} (${arch}, ${config.memory}MB RAM, ${config.vcpus} vCPU, ${config.diskSize}GB disk)`);

  // Create blank disk
  const diskPath = join(IMAGE_DIR, `${config.name}.qcow2`);
  run(`qemu-img create -f qcow2 "${diskPath}" ${config.diskSize}G`);

  const dataDiskArgs: string[] = [];
  for (const dd of config.dataDisks ?? []) {
    if (dd.serial.length > 20) throw new Error(`virtio serial '${dd.serial}' exceeds 20 chars -- the guest by-id link would be truncated`);
    const p = join(IMAGE_DIR, `${config.name}-${dd.serial}.qcow2`);
    run(`qemu-img create -f qcow2 "${p}" ${dd.sizeGb}G`);
    dataDiskArgs.push(`--disk=path=${p},format=qcow2,bus=virtio,serial=${dd.serial}`);
  }

  // UEFI firmware paths (Fedora)
  if (arch === "aarch64") {
    const aavmf = "/usr/share/edk2/aarch64/QEMU_EFI.fd";
    if (!existsSync(aavmf)) {
      throw new Error(`AAVMF firmware not found at ${aavmf}. Install: sudo dnf install edk2-aarch64`);
    }
  } else {
    const ovmf = "/usr/share/edk2/ovmf/OVMF_CODE.fd";
    if (!existsSync(ovmf)) {
      throw new Error(`OVMF firmware not found at ${ovmf}. Install: sudo dnf install edk2-ovmf`);
    }
  }

  const virtInstallArgs = [
    "virt-install",
    `--name=${config.name}`,
    `--memory=${config.memory}`,
    `--vcpus=${config.vcpus}`,
    `--disk=path=${diskPath},format=qcow2,bus=virtio`,
    ...dataDiskArgs,
    // Decoys first so they enumerate ahead of the PXE NIC. They are up and
    // carry a lease, but have no route to the bastion.
    ...Array.from({ length: config.decoyNics ?? 0 }, () =>
      `--network=network=${config.decoyNetwork ?? "default"},model=virtio`),
    `--network=network=${config.network},model=virtio`,
    // UEFI firmware — required for PXE boot in modern mode
    `--boot=uefi,network,hd`,
    // No OS to install — PXE provides everything
    "--os-variant=generic",
    "--noautoconsole",
    "--wait=0",
    // Graphics for debugging (VNC, connect with virt-viewer if needed)
    "--graphics=vnc,listen=127.0.0.1",
    // Serial console via TCP — allows exec without network/SSH
    // Connect: socat - TCP:127.0.0.1:<serialPort>
    `--serial=tcp,host=127.0.0.1:${config.serialPort ?? 4555},mode=bind,protocol=telnet`,
  ];

  if (arch === "aarch64") {
    virtInstallArgs.push("--arch=aarch64", "--machine=virt");
  }

  log(`Running: virt-install --name=${config.name} --boot=uefi,network ...`);
  run(virtInstallArgs.join(" "), { timeout: 30_000 });
  log(`PXE VM ${config.name} created (serial: telnet 127.0.0.1 ${config.serialPort ?? 4555})`);
}

/** Destroy a PXE VM and clean up its disk. */
export function destroyPxeVm(name: string): void {
  const result = virsh("dominfo", name);
  if (result.status !== 0) return;

  log(`Destroying PXE VM: ${name}`);
  virsh("destroy", name);
  virsh("undefine", name, "--remove-all-storage", "--nvram");
}

/** Get the MAC address of a VM's first NIC. */
export function getVmMac(name: string, network?: string): string | null {
  const result = virsh("domiflist", name);
  if (result.status !== 0) return null;
  // Output format: Interface  Type  Source  Model  MAC
  // With decoy NICs present, match the line for the PXE network so we return
  // the NIC that actually boots rather than whichever is listed first.
  const lines = result.stdout.split("\n");
  const candidates = network === undefined
    ? lines
    : lines.filter((l) => l.split(/\s+/).includes(network));
  for (const line of candidates.length > 0 ? candidates : lines) {
    const m = line.match(/([0-9a-f]{2}:[0-9a-f]{2}:[0-9a-f]{2}:[0-9a-f]{2}:[0-9a-f]{2}:[0-9a-f]{2})/i);
    if (m) return m[1].toLowerCase();
  }
  return null;
}

/** Reboot a VM (force off + start). */
export function rebootPxeVm(name: string): void {
  log(`Rebooting PXE VM: ${name}`);
  virsh("destroy", name);
  // Brief pause to let resources release
  spawnSync("sleep", ["2"]);
  virsh("start", name);
  log(`PXE VM ${name} restarted`);
}

/**
 * Read raw output from the VM's serial console (telnet TCP port).
 * Returns the last N lines. Useful for diagnostics when SSH isn't available.
 */
export async function readSerialLog(
  port: number,
  opts: { lastLines?: number; timeoutMs?: number } = {},
): Promise<string> {
  const { lastLines = 50, timeoutMs = 10_000 } = opts;
  return new Promise((resolve) => {
    const sock = createConnection({ host: "127.0.0.1", port });
    let buf = "";
    const timer = setTimeout(() => { sock.destroy(); resolve(buf); }, timeoutMs);
    sock.on("data", (d: Buffer) => { buf += d.toString(); });
    sock.on("error", () => { clearTimeout(timer); resolve(`(connection error) ${buf}`); });
    sock.on("close", () => { clearTimeout(timer); resolve(buf); });
    // Send a newline to trigger any buffered output / prompt
    setTimeout(() => sock.write("\r\n"), 500);
  }).then((raw: unknown) => {
    const lines = (raw as string).split("\n").map(l => l.trimEnd()).filter(Boolean);
    return lines.slice(-lastLines).join("\n");
  });
}

/**
 * Execute a command on the VM's serial console via socat.
 * Requires auto-login root shell on the serial port.
 */
export function serialExec(
  port: number,
  command: string,
  timeoutMs = 15_000,
): string {
  const marker = `__END_${Date.now()}__`;
  // Use socat to handle telnet negotiation properly
  const input = `\r\n${command}; echo '${marker}'\r\n`;
  const result = spawnSync("bash", ["-c",
    `echo -e '${input.replace(/'/g, "\\'")}' | socat -T${Math.ceil(timeoutMs / 1000)} - TCP:127.0.0.1:${port} 2>/dev/null`
  ], { encoding: "utf-8", stdio: "pipe", timeout: timeoutMs + 5000 });
  const output = result.stdout ?? "";
  const markerIdx = output.indexOf(marker);
  if (markerIdx < 0) return `(no marker) ${output.slice(-500)}`;
  // Get lines between command echo and marker
  const before = output.substring(0, markerIdx);
  const lines = before.split("\n");
  // Skip everything up to and including the command echo line
  const cmdIdx = lines.findIndex(l => l.includes(command.substring(0, 20)));
  return lines.slice(cmdIdx >= 0 ? cmdIdx + 1 : 1).join("\n").trim();
}

export interface IsoVmConfig {
  name: string;
  memory: number;       // MB
  vcpus: number;
  diskSize: number;     // GB
  network: string;      // libvirt network name
  isoPath: string;      // path to boot ISO
  arch?: "x86_64" | "aarch64";
}

/** Create a UEFI VM that boots from a CD-ROM ISO (not PXE). */
export function createIsoVm(config: IsoVmConfig): void {
  destroyPxeVm(config.name);

  const arch = config.arch ?? "x86_64";
  log(`Creating ISO boot VM: ${config.name} (${arch}, ${config.memory}MB RAM, ${config.vcpus} vCPU, ${config.diskSize}GB disk)`);

  // Create blank disk
  const diskPath = join(IMAGE_DIR, `${config.name}.qcow2`);
  run(`qemu-img create -f qcow2 "${diskPath}" ${config.diskSize}G`);

  const virtInstallArgs = [
    "virt-install",
    `--name=${config.name}`,
    `--memory=${config.memory}`,
    `--vcpus=${config.vcpus}`,
    `--disk=path=${diskPath},format=qcow2,bus=virtio`,
    // Boot ISO as CD-ROM
    `--disk=path=${config.isoPath},device=cdrom,readonly=on`,
    `--network=network=${config.network},model=virtio`,
    // UEFI firmware, boot from cdrom (not network)
    "--boot=uefi,cdrom",
    "--os-variant=generic",
    "--noautoconsole",
    "--wait=0",
    "--graphics=vnc,listen=127.0.0.1",
    // Serial console via TCP (port 4556 to avoid conflict with PXE VM)
    "--serial=tcp,host=127.0.0.1:4556,mode=bind,protocol=telnet",
  ];

  if (arch === "aarch64") {
    virtInstallArgs.push("--arch=aarch64", "--machine=virt");
  }

  log(`Running: virt-install --name=${config.name} --boot=uefi,cdrom ...`);
  run(virtInstallArgs.join(" "), { timeout: 60_000 });
  log(`ISO boot VM ${config.name} created (serial: telnet 127.0.0.1 4556)`);
}