lab/bastion/tests/integration/helpers/provision-harness.ts

// Reusable bastion + VM provisioning harness for partition-layout tests.
//
// pxe-provision.test.ts grew a ~200-line beforeAll that stands up a real
// bastion, caches Fedora boot artifacts, and drives one VM through discovery
// and install. Testing a partition MATRIX needs that same sequence for several
// VMs and several roles, so it lives here instead of being copy-pasted per
// topology.

import { existsSync, mkdirSync, copyFileSync, symlinkSync, writeFileSync } from "node:fs";
import { execSync } from "node:child_process";
import { join } from "node:path";
import { tmpdir } from "node:os";
import type { BastionConfig } from "../../../src/shared/src/types/config.js";
import type { Role } from "../../../src/shared/src/types/state.js";
import { log, waitForSsh } from "./libvirt.js";
import { ensurePxeNetwork, deleteNftablesRejectRules, PXE_NETWORK_NAME, PXE_BRIDGE, PXE_GATEWAY, PXE_SUBNET } from "./pxe-network.js";
import { createPxeVm, destroyPxeVm, getVmMac, rebootPxeVm } from "./pxe-vm.js";
import { sshExec } from "./ssh.js";

export const DISCOVERY_TIMEOUT_MS = 6 * 60_000;
export const INSTALL_TIMEOUT_MS = 60 * 60_000;
export const SSH_TIMEOUT_MS = 12 * 60_000;

export const sleep = (ms: number) => new Promise((r) => setTimeout(r, ms));

export interface Harness {
  config: BastionConfig;
  close: () => Promise<void>;
  httpPort: number;
  sshUser: string;
  sshKeyPath: string;
}

async function pollUntil<T>(url: string, check: (d: T) => boolean, timeoutMs: number, intervalMs = 5000): Promise<T> {
  const start = Date.now();
  let lastErr = "";
  while (Date.now() - start < timeoutMs) {
    try {
      const res = await fetch(url);
      if (res.ok) {
        const data = (await res.json()) as T;
        if (check(data)) return data;
      }
    } catch (e) { lastErr = e instanceof Error ? e.message : String(e); }
    await sleep(intervalMs);
  }
  throw new Error(`Timeout after ${timeoutMs}ms polling ${url}${lastErr ? ` (last error: ${lastErr})` : ""}`);
}

/**
 * Stand up a real bastion on the isolated PXE network.
 *
 * FEDORA_MIRROR is honoured (config.ts now reads it) -- point it at a local
 * mirror before running a multi-VM matrix, or every install pulls the full
 * package set from download.fedoraproject.org.
 */
export async function startTestBastion(opts: {
  httpPort: number;
  sshUser: string;
  sshPubKey: string;
  sshKeyPath: string;
  label: string;
  /**
   * Leave the post-provision k3s auto-install ON (the default) or suppress it.
   *
   * Any test that forms a cluster ITSELF must pass false. The auto-install runs
   * a bare `INSTALL_K3S_EXEC=server|agent` with no server URL and no token, so
   * three combo nodes become three separate single-node clusters -- and it
   * stops k3s and deletes CNI state minutes after an install lands, landing on
   * top of whatever the test is doing by then.
   */
  autoK3s?: boolean;
}): Promise<Harness> {
  ensurePxeNetwork();

  const testDir = join(tmpdir(), `lab-${opts.label}-${Date.now()}`);
  for (const sub of ["tftp", "http", "logs"]) mkdirSync(join(testDir, sub), { recursive: true });

  const { createApp } = await import("../../../src/bastion/src/server.js");
  const { loadConfig } = await import("../../../src/bastion/src/config.js");
  const { generateDnsmasqConf, startDnsmasq, stopDnsmasq } = await import("../../../src/bastion/src/services/dnsmasq.js");
  const { generateDiscoverKickstart } = await import("../../../src/bastion/src/services/kickstart-generator.js");
  const { renderBootIpxe } = await import("../../../src/bastion/src/templates/boot.ipxe.js");

  const config = loadConfig({
    bastionDir: testDir,
    httpPort: opts.httpPort,
    iface: "virbr-pxe",
    serverIp: PXE_GATEWAY,
    network: `${PXE_SUBNET}.0`,
    gateway: PXE_GATEWAY,
    dhcpMode: "full",
    dhcpRangeStart: `${PXE_SUBNET}.100`,
    dhcpRangeEnd: `${PXE_SUBNET}.200`,
    domain: "pxe-test.local",
    sshKeys: [opts.sshPubKey],
    adminUser: opts.sshUser,
    ...(opts.autoK3s === undefined ? {} : { autoK3s: opts.autoK3s }),
  });

  const ipxeSrc = "/usr/share/ipxe/ipxe-snponly-x86_64.efi";
  if (!existsSync(ipxeSrc)) throw new Error(`iPXE not found: ${ipxeSrc}. Install: sudo dnf install ipxe-bootimgs-x86`);
  const ipxeDest = join(config.tftpDir, "ipxe.efi");
  copyFileSync(ipxeSrc, ipxeDest);

  const cacheDir = "/var/lib/libvirt/images/lab-pxe-cache";
  execSync(`mkdir -p "${cacheDir}"`, { stdio: "pipe" });
  const kernel = join(cacheDir, `vmlinuz-${config.fedoraVersion}`);
  const initrd = join(cacheDir, `initrd-${config.fedoraVersion}.img`);
  if (!existsSync(kernel)) {
    log(`Downloading Fedora ${config.fedoraVersion} kernel from ${config.fedoraMirror}...`);
    execSync(`curl -sS -L -f -o "${kernel}" "${config.fedoraMirror}/images/pxeboot/vmlinuz"`, { stdio: "inherit", timeout: 300_000 });
  }
  if (!existsSync(initrd)) {
    log(`Downloading Fedora ${config.fedoraVersion} initrd...`);
    execSync(`curl -sS -L -f -o "${initrd}" "${config.fedoraMirror}/images/pxeboot/initrd.img"`, { stdio: "inherit", timeout: 300_000 });
  }
  copyFileSync(kernel, join(config.httpDir, "vmlinuz"));
  copyFileSync(initrd, join(config.httpDir, "initrd.img"));
  try { symlinkSync(ipxeDest, join(config.httpDir, "ipxe.efi")); } catch { /* exists */ }

  writeFileSync(join(config.httpDir, "discover.ks"), generateDiscoverKickstart(config));
  writeFileSync(join(config.httpDir, "boot.ipxe"), renderBootIpxe({ serverIp: config.serverIp, httpPort: config.httpPort }));
  generateDnsmasqConf(config);

  // Kill dnsmasq processes left behind by EARLIER TEST RUNS.
  //
  // startTestBastion's close() shuts down the HTTP app but never stopped
  // dnsmasq, so every run leaked one. A leaked dnsmasq keeps :67 on the bridge
  // and keeps answering DHCP -- with the PREVIOUS run's config, pointing PXE
  // clients at a bastion port that is no longer listening. The next run then
  // sees perfect DHCP (offer, request, ack, even the right bootfile name) and
  // no machine ever reaches it. That is not a failure mode anyone guesses from
  // "discovery timed out".
  killOrphanedDnsmasq();

  const { app, syslog } = createApp(config);
  await app.listen({ port: config.httpPort, host: "0.0.0.0" });
  syslog.start();
  // NOT awaited, deliberately.
  //
  // startDnsmasq returns the execa child handle, and because the function is
  // `async` the promise it returns ADOPTS that thenable -- which settles only
  // when dnsmasq EXITS. With --no-daemon that is never, so `await` here hangs
  // the harness silently and forever, with the log frozen right after
  // "TFTP root is ...". (Learned by doing exactly that.) The rejection still
  // has to be handled or an early bind failure becomes an unhandled rejection.
  startDnsmasq(config).catch((err) => log(`dnsmasq: ${err instanceof Error ? err.message : String(err)}`));
  await sleep(1000);

  // Prove DHCP is actually listening before handing the harness back.
  //
  // startDnsmasq's rejection is swallowed into a log line, so a dnsmasq that
  // failed to bind looks exactly like one that started. The symptom then
  // arrives six minutes later as four VMs timing out in discovery, which
  // reads as a PXE or kickstart problem rather than "there was never any
  // DHCP". Fail here, immediately, with the real reason.
  await assertDhcpListening(join(testDir, "dnsmasq.conf"));
  log(`Bastion up on :${opts.httpPort} (mirror: ${config.fedoraMirror}, autoK3s=${config.autoK3s})`);

  return {
    config,
    httpPort: opts.httpPort,
    sshUser: opts.sshUser,
    sshKeyPath: opts.sshKeyPath,
    close: async () => {
      // Stop OUR dnsmasq. close() never used to, which is how a leaked one
      // ended up owning :67 and serving the PREVIOUS run's TFTP root and
      // bastion port to the NEXT run's VMs -- they booted toward a dead port
      // and simply never appeared, while DHCP looked perfect.
      try { stopDnsmasq(); } catch { /* already gone */ }
      await app.close();
      // Belt and braces: if that did not stop it, sweep by config path.
      killOrphanedDnsmasq();
    },
  };
}

/**
 * Kill any dnsmasq started by a previous (or current) lab test run.
 *
 * Matched on the config path, which is always /tmp/lab-<label>-<ts>/dnsmasq.conf
 * -- so this never touches libvirt's own dnsmasq for the default network, or
 * anything else on the host.
 */
function killOrphanedDnsmasq(): void {
  try {
    // Match on the process NAME first (pgrep -x), then inspect each one's
    // arguments. A `pgrep -f 'dnsmasq.*/tmp/lab-'` matches the shell running
    // that very pattern, so the naive version can kill its own shell -- and
    // libvirt's dnsmasq for the default network must never be touched either.
    const pids = execSync("pgrep -x dnsmasq 2>/dev/null || true", { encoding: "utf-8" })
      .trim().split("\n").filter((x) => /^\d+$/.test(x.trim()));
    let killed = 0;
    for (const pid of pids) {
      let args = "";
      try {
        args = execSync(`tr '\\0' ' ' < /proc/${pid.trim()}/cmdline 2>/dev/null || true`, { encoding: "utf-8" });
      } catch { continue; }
      if (!args.includes("--conf-file=/tmp/lab-")) continue;
      log(`Killing orphaned test dnsmasq (pid ${pid.trim()}) — it would answer DHCP with stale config`);
      try { execSync(`sudo kill -9 ${pid.trim()} 2>/dev/null || kill -9 ${pid.trim()} 2>/dev/null || true`); } catch { /* gone */ }
      killed++;
    }
    // Give the socket a moment to be released before we bind it.
    if (killed > 0) execSync("sleep 1");
  } catch { /* nothing to kill */ }
}

/**
 * Is something serving DHCP on the PXE bridge?
 *
 * `ss -lunp` shows UDP listeners; dnsmasq binds port 67. Checked on the bridge
 * address specifically, because a dnsmasq bound to the wrong interface is the
 * failure being guarded against.
 */
async function assertDhcpListening(confPath: string, timeoutMs = 20_000): Promise<void> {
  const start = Date.now();
  let last = "";
  while (Date.now() - start < timeoutMs) {
    try {
      // Check for OUR dnsmasq by its config path, not merely for something on
      // :67. A leaked dnsmasq from a previous run also holds :67 and also
      // answers DHCP -- with the wrong bastion port -- so a port check alone
      // reports success for precisely the situation it is meant to catch.
      const pids = execSync("pgrep -x dnsmasq 2>/dev/null || true", { encoding: "utf-8" })
        .trim().split("\n").filter((x) => /^\d+$/.test(x.trim()));
      last = pids.filter((pid) => {
        try {
          return execSync(`tr '\\0' ' ' < /proc/${pid.trim()}/cmdline 2>/dev/null || true`, { encoding: "utf-8" })
            .includes(confPath);
        } catch { return false; }
      }).join(",");
      if (last !== "") {
        const ports = execSync("ss -lun 2>/dev/null || true", { encoding: "utf-8" });
        if (/:67\b/.test(ports)) {
          log("dnsmasq (ours) is serving DHCP on :67");
          return;
        }
      }
    } catch { /* retry */ }
    await sleep(1000);
  }
  throw new Error(
    `dnsmasq for ${confPath} is not serving DHCP after ${timeoutMs / 1000}s. ` +
    `It most likely failed to bind to ${PXE_BRIDGE} -- check the bastion log for ` +
    `"error binding DHCP socket", and check for a leaked dnsmasq from an earlier run ` +
    `(pgrep -af 'dnsmasq.*--conf-file=/tmp/lab-'). Without this every VM sits in ` +
    `discovery until it times out, six minutes later, looking like a PXE problem.`,
  );
}

/**
 * Wait for the installer to actually start, rebooting into PXE if it does not.
 *
 * "Started" means the bastion has seen anything at all for this MAC beyond the
 * queued state -- a dispatch, a progress line, a status. A VM that never gets
 * that far is stuck before Anaconda, and only a reboot will shift it.
 */
async function waitForInstallStart(
  base: string,
  mac: string,
  name: string,
  attemptMs = 6 * 60_000,
  attempts = 3,
): Promise<void> {
  for (let attempt = 1; attempt <= attempts; attempt++) {
    try {
      // NOT `status === "installing"`. A machine that has merely been QUEUED
      // already reports that -- the bastion sets it the moment the install is
      // dispatched, before anything has booted. Using it would match instantly
      // and this retry would never fire.
      //
      // The honest signals that Anaconda is alive are `progress` leaving
      // "queued", or log lines arriving over syslog.
      await pollUntil<{ status?: string; progress?: string; log_total?: number }>(
        `${base}/api/logs/${encodeURIComponent(mac)}`,
        (d) => d.status === "installed"
          || (d.progress !== undefined && d.progress !== "queued")
          || (d.log_total ?? 0) > 0,
        attemptMs, 5_000);
      if (attempt > 1) log(`[${name}] PXE boot took on attempt ${attempt}`);
      return;
    } catch {
      if (attempt === attempts) {
        throw new Error(
          `[${name}] installer never started after ${attempts} PXE boots ` +
          `(${attemptMs / 60000} min each). The VM is stuck before Anaconda — ` +
          `check TFTP and the iPXE chainload.`);
      }
      log(`[${name}] no installer activity in ${attemptMs / 60000} min; rebooting into PXE (attempt ${attempt + 1}/${attempts})`);
      rebootPxeVm(name);
      deleteNftablesRejectRules();
    }
  }
}

export interface ProvisionedVm { name: string; mac: string; ip: string }

/**
 * Serial ports are a global resource on the host, so hand out a distinct one
 * per VM. Without this only one VM can run at a time -- qemu refuses to start
 * the second with "Address already in use".
 */
let nextSerialPort = 4600;

/** Create a blank VM, let it discover, queue the install, wait for SSH. */
export async function provisionVm(h: Harness, spec: {
  name: string;
  role: Role;
  schemaVersion?: 1 | 2;
  diskGb: number;
  memory?: number;
  vcpus?: number;
}): Promise<ProvisionedVm> {
  const base = `http://${PXE_GATEWAY}:${h.httpPort}`;
  log(`[${spec.name}] creating VM (${spec.diskGb}GB disk, role=${spec.role}, schema=v${spec.schemaVersion ?? 1})`);
  createPxeVm({
    name: spec.name,
    memory: spec.memory ?? 4096,
    vcpus: spec.vcpus ?? 4,
    diskSize: spec.diskGb,
    network: PXE_NETWORK_NAME,
    serialPort: nextSerialPort++,
  });
  deleteNftablesRejectRules();

  const mac = getVmMac(spec.name);
  if (!mac) throw new Error(`[${spec.name}] no MAC`);

  log(`[${spec.name}] waiting for discovery (${mac})...`);
  await pollUntil<{ discovered: Record<string, unknown> }>(
    `${base}/api/machines`, (d) => mac in (d.discovered ?? {}), DISCOVERY_TIMEOUT_MS);

  log(`[${spec.name}] queueing install...`);
  const res = await fetch(`${base}/api/install`, {
    method: "POST",
    headers: { "Content-Type": "application/json" },
    body: JSON.stringify({
      mac, hostname: spec.name, disk: "", role: spec.role,
      ...(spec.schemaVersion ? { schemaVersion: spec.schemaVersion } : {}),
    }),
  });
  if (!res.ok) throw new Error(`[${spec.name}] install queue rejected: ${res.status} ${await res.text()}`);

  rebootPxeVm(spec.name);
  deleteNftablesRejectRules();

  // Make sure the PXE boot actually TOOK before settling in for the install.
  //
  // A UEFI VM that fails to chainload just retries DHCP quietly, for ever.
  // Observed: one VM of four sat in a DHCP loop for an HOUR with no TFTP
  // fetch, then began installing seconds after the 60-minute install timeout
  // had already fired -- so the run failed on a machine that was, by then,
  // working. TFTP is UDP and four VMs contend for it, so this is a question of
  // when rather than if.
  //
  // Rebooting is cheap and a boot that has not started within a few minutes is
  // not going to start on its own.
  await waitForInstallStart(base, mac, spec.name);

  log(`[${spec.name}] installing (up to ${INSTALL_TIMEOUT_MS / 60000} min)...`);
  await pollUntil<{ status?: string }>(
    `${base}/api/logs/${encodeURIComponent(mac)}`, (d) => d.status === "installed", INSTALL_TIMEOUT_MS, 10_000);

  // The installed record carries the DHCP address the machine ended up with.
  const machines = await pollUntil<{ installed: Record<string, { ip?: string }> }>(
    `${base}/api/machines`, (d) => Boolean(d.installed?.[mac]?.ip), 2 * 60_000);
  const ip = machines.installed[mac]!.ip!;

  try {
    if (execSync(`sudo virsh domstate ${spec.name}`, { encoding: "utf-8" }).trim() === "shut off") {
      execSync(`sudo virsh start ${spec.name}`, { stdio: "pipe" });
    }
  } catch { /* already running */ }
  deleteNftablesRejectRules();

  log(`[${spec.name}] waiting for SSH at ${ip}...`);
  await waitForSsh(ip, h.sshUser, SSH_TIMEOUT_MS, h.sshKeyPath);
  log(`[${spec.name}] up at ${ip}`);
  return { name: spec.name, mac, ip };
}

/**
 * Provision several VMs at once.
 *
 * Installs are ~12 minutes each and almost entirely I/O and network wait, so
 * four in sequence is an hour that four in parallel is twenty minutes. They are
 * staggered because the expensive, contended part is the start: PXE, DHCP and
 * the first package fetches all land at once otherwise.
 *
 * Rejects only after every VM has settled, so a single failure does not leave
 * the others half-built and unreferenced -- the caller still gets the names it
 * needs to tear down.
 */
export async function provisionVms(
  h: Harness,
  specs: Parameters<typeof provisionVm>[1][],
  opts: { staggerMs?: number } = {},
): Promise<ProvisionedVm[]> {
  const stagger = opts.staggerMs ?? 60_000;
  const settled = await Promise.allSettled(
    specs.map(async (spec, i) => {
      await sleep(i * stagger);
      return provisionVm(h, spec);
    }),
  );

  const failures = settled
    .map((r, i) => (r.status === "rejected" ? `${specs[i]?.name}: ${String(r.reason)}` : null))
    .filter((x): x is string => x !== null);
  if (failures.length > 0) {
    throw new Error(`provisionVms: ${failures.length}/${specs.length} failed\n  ${failures.join("\n  ")}`);
  }
  return settled.map((r) => (r as PromiseFulfilledResult<ProvisionedVm>).value);
}

/** LV name -> size in MiB, read off the installed machine. */
export function readLvs(h: Harness, vm: ProvisionedVm): Record<string, number> {
  const res = sshExec(vm.ip, h.sshUser,
    "sudo lvs labvg --noheadings --units m --nosuffix -o lv_name,lv_size --separator=,",
    { keyPath: h.sshKeyPath });
  if (res.exitCode !== 0) throw new Error(`[${vm.name}] lvs failed: ${res.stderr}`);
  const map: Record<string, number> = {};
  for (const line of res.stdout.split("\n")) {
    const [name, size] = line.trim().split(",");
    if (name && size) map[name] = Math.round(Number(size));
  }
  // Log what actually landed on disk. The VMs are destroyed in afterAll, so
  // without this the only record of the real layout is the pass/fail bit.
  const summary = Object.entries(map)
    .sort(([a], [b]) => a.localeCompare(b))
    .map(([n, m]) => `${n}=${(m / 1024).toFixed(0)}G`)
    .join(" ");
  log(`[${vm.name}] on-disk LVs: ${summary}`);
  return map;
}

/** The VG tag that records which schema this disk was built with. */
export function readSchemaTag(h: Harness, vm: ProvisionedVm): string {
  return sshExec(vm.ip, h.sshUser, "sudo vgs --noheadings -o vg_tags labvg",
    { keyPath: h.sshKeyPath }).stdout.trim();
}

export function destroyVms(...names: string[]): void {
  for (const n of names) { try { destroyPxeVm(n); } catch { /* already gone */ } }
}

/**
 * Reinstall a machine that is already installed, the way
 * `labctl provision reprovision` does: queue another install for the same MAC
 * and reboot it into PXE.
 *
 * The schema is deliberately NOT passed. A reinstall must reproduce whatever
 * layout is already on the disk -- %pre reads the lab-schema-v2 LVM tag rather
 * than trusting the caller -- and passing one here would hide a regression in
 * exactly that logic.
 */
export async function reprovisionVm(
  h: Harness,
  vm: ProvisionedVm,
  role: Role,
  /**
   * Only set this to prove it is IGNORED. A reinstall takes its layout from the
   * disk, not from the request; passing a schema here is how you test that.
   */
  requestSchemaVersion?: 1 | 2,
): Promise<ProvisionedVm> {
  const base = `http://${PXE_GATEWAY}:${h.httpPort}`;
  log(`[${vm.name}] reprovisioning (role=${role}, schema: whatever is on disk)...`);

  const res = await fetch(`${base}/api/install`, {
    method: "POST",
    headers: { "Content-Type": "application/json" },
    body: JSON.stringify({
      mac: vm.mac, hostname: vm.name, disk: "", role,
      ...(requestSchemaVersion ? { schemaVersion: requestSchemaVersion } : {}),
    }),
  });
  if (!res.ok) throw new Error(`[${vm.name}] reinstall queue rejected: ${res.status} ${await res.text()}`);

  rebootPxeVm(vm.name);
  deleteNftablesRejectRules();

  log(`[${vm.name}] reinstalling...`);
  await pollUntil<{ status?: string }>(
    `${base}/api/logs/${encodeURIComponent(vm.mac)}`,
    (d) => d.status === "installed", INSTALL_TIMEOUT_MS, 10_000);

  const machines = await pollUntil<{ installed: Record<string, { ip?: string }> }>(
    `${base}/api/machines`, (d) => Boolean(d.installed?.[vm.mac]?.ip), 2 * 60_000);
  const ip = machines.installed[vm.mac]!.ip!;

  try {
    if (execSync(`sudo virsh domstate ${vm.name}`, { encoding: "utf-8" }).trim() === "shut off") {
      execSync(`sudo virsh start ${vm.name}`, { stdio: "pipe" });
    }
  } catch { /* already running */ }
  deleteNftablesRejectRules();

  await waitForSsh(ip, h.sshUser, SSH_TIMEOUT_MS, h.sshKeyPath);
  log(`[${vm.name}] back up at ${ip} after reinstall`);
  return { ...vm, ip };
}

/**
 * Drop a marker file into each mount point that a reinstall is supposed to
 * preserve. Ported from test-reprovision.sh, which proved the idea but only
 * ever ran by hand.
 */
export function writeMarkers(h: Harness, vm: ProvisionedVm, mounts: string[]): string {
  const token = `preserve-${Date.now()}-${Math.random().toString(36).slice(2, 8)}`;
  for (const m of mounts) {
    const r = sshExec(vm.ip, h.sshUser,
      `sudo sh -c 'echo ${token} > ${m}/LAB_PRESERVE_TEST.txt && sync'`,
      { keyPath: h.sshKeyPath });
    if (r.exitCode !== 0) throw new Error(`[${vm.name}] could not write marker in ${m}: ${r.stderr}`);
  }
  log(`[${vm.name}] wrote marker ${token} into ${mounts.join(", ")}`);
  return token;
}

/** Which mount points still carry the marker. */
export function readMarkers(h: Harness, vm: ProvisionedVm, mounts: string[]): Record<string, string | null> {
  const out: Record<string, string | null> = {};
  for (const m of mounts) {
    const r = sshExec(vm.ip, h.sshUser, `sudo cat ${m}/LAB_PRESERVE_TEST.txt 2>/dev/null || true`,
      { keyPath: h.sshKeyPath });
    const v = r.stdout.trim();
    out[m] = v === "" ? null : v;
  }
  log(`[${vm.name}] markers after reinstall: ${JSON.stringify(out)}`);
  return out;
}