vmrunner 0.0.5

micro-vm runner for testcases that require root or invasive IO
use std::{
    fs,
    path::{Path, PathBuf},
    process::Command,
    time::Instant,
};

use color_eyre::eyre::{Result, WrapErr, eyre};
use libtest_mimic::{Arguments, Completion, Failed, Trial};
use vmrunner::{RootDevice, RootFsType, RootMount, TestCase};
const GUEST_SMOKE_MARKER: &str = "vmrunner-smoke-ok";
const FEDORA_AARCH64_GUEST_TARGET: &str = "aarch64-unknown-linux-gnu";
const AARCH64_ELF_MACHINE: u16 = 183;

type ImageSource = vmrunner::ExternalQCowFilesystem;

fn main() {
    let args = Arguments::from_args();
    let trials = vec![
        Trial::ignorable_test(
            "ubuntu_vm_smoke",
            || result_to_completion(ubuntu_vm_smoke()),
        )
        .with_kind("vm"),
        Trial::ignorable_test(
            "fedora_vm_smoke",
            || result_to_completion(fedora_vm_smoke()),
        )
        .with_kind("vm"),
        Trial::ignorable_test("fedora_aarch64_sysroot_compile_smoke", || {
            result_to_completion(fedora_aarch64_sysroot_compile_smoke())
        })
        .with_kind("sysroot"),
    ];

    libtest_mimic::run(&args, trials).exit();
}

fn result_to_completion(result: Result<()>) -> Result<Completion, Failed> {
    result
        .map(|()| Completion::Completed)
        .map_err(|error| Failed::from(format!("{error:#}")))
}

fn ubuntu_vm_smoke() -> Result<()> {
    run_linux_vm_smoke("ubuntu_vm_smoke", "ubuntu", host_ubuntu_qcow2()?, |_| {
        Ok(RootMount::new(RootDevice::virtio_first_partition()))
    })
}

fn fedora_vm_smoke() -> Result<()> {
    run_linux_vm_smoke(
        "fedora_vm_smoke",
        "fedora",
        host_fedora_qcow2()?,
        fedora_root_mount,
    )
}

fn fedora_aarch64_sysroot_compile_smoke() -> Result<()> {
    let test_name = "fedora_aarch64_sysroot_compile_smoke";
    let image = fedora_aarch64_qcow2()?;
    let root_qcow2 = timed_step(
        test_name,
        "resolve and mkosi-prepare Fedora aarch64 qcow",
        || vmrunner::ensure_root_qcow2(image.url),
    )?;
    let sysroot_parent = tempfile::Builder::new()
        .prefix("vmrunner-fedora-aarch64-sysroot-")
        .tempdir()
        .wrap_err("create temporary Fedora aarch64 sysroot output parent")?;
    let build_dir = tempfile::Builder::new()
        .prefix("vmrunner-fedora-aarch64-sysroot-build-")
        .tempdir()
        .wrap_err("create temporary Fedora aarch64 sysroot compile directory")?;
    let sysroot = timed_step(test_name, "extract Fedora aarch64 sysroot", || {
        vmrunner_sysroot::linux::extract_linux_sysroot(vmrunner_sysroot::SysrootOptions::new(
            &root_qcow2,
            FEDORA_AARCH64_GUEST_TARGET,
            sysroot_parent.path(),
        ))
        .wrap_err_with(|| {
            format!(
                "extract Fedora aarch64 sysroot from '{}'",
                root_qcow2.display()
            )
        })
    })?;

    timed_step(
        test_name,
        "compile and link Fedora aarch64 sysroot smoke",
        || compile_aarch64_fedora_libc_smoke(&sysroot, build_dir.path()),
    )
}

fn run_linux_vm_smoke(
    test_name: &str,
    guest_name: &str,
    root_qcow2: ImageSource,
    root_mount: impl FnOnce(&Path) -> Result<RootMount>,
) -> Result<()> {
    let root_qcow2 = timed_step(test_name, "resolve and mkosi-prepare root qcow", || {
        vmrunner::ensure_root_qcow2(root_qcow2.url)
    })?;
    let root_mount = root_mount(&root_qcow2)?;
    let guest_target = host_linux_guest_target()?;
    timed_step(test_name, "extract sysroot and build init.krun", || {
        vmrunner::ensure_guest_init(&root_qcow2, Some(guest_target), None)
    })?;

    if timed_step(test_name, "enter host child process", || {
        vmrunner::run_current_test_in_platform_child(test_name)
    })? {
        return Ok(());
    }

    timed_step(test_name, "launch VM and wait for smoke marker", || {
        run_guest_true(guest_name, &root_qcow2, root_mount)
    })
}

fn run_guest_true(name: &str, root_qcow2: &Path, root_mount: RootMount) -> Result<()> {
    let node = vmrunner::Node::new(&[], "/bin/sh")
        .name(name)
        .args(["-c".to_owned(), format!("echo {GUEST_SMOKE_MARKER}")]);
    let mut test_case = TestCase::new(&[&node])
        .root_qcow2(root_qcow2)
        .guest_init(vmrunner::guest_init_path_for_root_qcow2(root_qcow2))
        .root_device(root_mount.device().clone());
    if let Some(root_fstype) = root_mount.root_fstype() {
        test_case = test_case.root_fstype(root_fstype.clone());
    }
    if let Some(root_options) = root_mount.root_options() {
        test_case = test_case.root_options(root_options.clone());
    }

    let running = futures::executor::block_on(test_case.launch())?;
    let output = futures::executor::block_on(running.wait_with_output())?;
    if output.len() == 1 && output[0].stdout.contains(GUEST_SMOKE_MARKER) {
        Ok(())
    } else {
        Err(eyre!(
            "{name} guest did not print {GUEST_SMOKE_MARKER:?}: {output:?}"
        ))
    }
}

fn host_linux_guest_target() -> Result<&'static str> {
    match std::env::consts::ARCH {
        "aarch64" => Ok("aarch64-unknown-linux-gnu"),
        "x86_64" => Ok("x86_64-unknown-linux-gnu"),
        arch => Err(eyre!(
            "unsupported host arch for Linux VM image test: {arch}"
        )),
    }
}

fn host_ubuntu_qcow2() -> Result<ImageSource> {
    host_qcow2("ubuntu")
}

fn host_fedora_qcow2() -> Result<ImageSource> {
    host_qcow2("fedora")
}

fn fedora_aarch64_qcow2() -> Result<ImageSource> {
    qcow2("fedora", "aarch64")
}

fn host_qcow2(name: &str) -> Result<ImageSource> {
    qcow2(name, std::env::consts::ARCH)
}

fn qcow2(name: &str, arch: &str) -> Result<ImageSource> {
    vmrunner::EXTERNAL_QCOW_FILESYSTEMS
        .iter()
        .copied()
        .find(|image| image.name == name && image.arch == arch)
        .ok_or_else(|| eyre!("unsupported arch for {name} VM image test: {arch}"))
}

fn compile_aarch64_fedora_libc_smoke(sysroot: &Path, build_dir: &Path) -> Result<()> {
    let libc = find_sysroot_file(sysroot, "libc.so.6")
        .wrap_err_with(|| format!("find libc.so.6 under '{}'", sysroot.display()))?;
    let dynamic_linker = find_sysroot_file(sysroot, "ld-linux-aarch64.so.1")
        .wrap_err_with(|| format!("find aarch64 dynamic linker under '{}'", sysroot.display()))?;
    let dynamic_linker_guest_path = sysroot_guest_path(sysroot, &dynamic_linker)?;
    let lib_dir = libc
        .parent()
        .ok_or_else(|| eyre!("libc path '{}' has no parent", libc.display()))?;

    let source = build_dir.join("hello.c");
    let object = build_dir.join("hello.o");
    let binary = build_dir.join("hello");
    fs::write(
        &source,
        r#"extern int puts(const char *);
extern void _exit(int);
void _start(void) {
    puts("hello from Fedora aarch64 sysroot");
    _exit(0);
}
"#,
    )
    .wrap_err_with(|| format!("write sysroot compile smoke source '{}'", source.display()))?;

    run_checked_command(
        Command::new("clang")
            .arg("-target")
            .arg(FEDORA_AARCH64_GUEST_TARGET)
            .arg("--sysroot")
            .arg(sysroot)
            .arg("-c")
            .arg(&source)
            .arg("-o")
            .arg(&object),
        "compile Fedora aarch64 sysroot smoke object",
    )?;

    run_checked_command(
        Command::new("clang")
            .arg("-target")
            .arg(FEDORA_AARCH64_GUEST_TARGET)
            .arg("--sysroot")
            .arg(sysroot)
            .arg("-fuse-ld=lld")
            .arg("-nostdlib")
            .arg(format!("-Wl,--dynamic-linker,{dynamic_linker_guest_path}"))
            .arg(format!("-Wl,-rpath-link,{}", lib_dir.display()))
            .arg("-L")
            .arg(lib_dir)
            .arg(&object)
            .arg("-l:libc.so.6")
            .arg("-o")
            .arg(&binary),
        "link Fedora aarch64 sysroot smoke binary",
    )?;

    validate_elf_machine(&binary, AARCH64_ELF_MACHINE)?;
    Ok(())
}

fn find_sysroot_file(sysroot: &Path, file_name: &str) -> Result<PathBuf> {
    let mut stack = vec![sysroot.to_path_buf()];
    while let Some(dir) = stack.pop() {
        for entry in fs::read_dir(&dir).wrap_err_with(|| format!("read '{}'", dir.display()))? {
            let entry = entry?;
            let path = entry.path();
            if entry.file_name() == file_name && path.is_file() {
                return Ok(path);
            }
            if entry.file_type()?.is_dir() {
                stack.push(path);
            }
        }
    }
    Err(eyre!(
        "file '{file_name}' not found under sysroot '{}'",
        sysroot.display()
    ))
}

fn sysroot_guest_path(sysroot: &Path, path: &Path) -> Result<String> {
    let relative = path.strip_prefix(sysroot).wrap_err_with(|| {
        format!(
            "path '{}' is not under sysroot '{}'",
            path.display(),
            sysroot.display()
        )
    })?;
    Ok(format!("/{}", relative.display()))
}

fn run_checked_command(command: &mut Command, description: &str) -> Result<()> {
    eprintln!("[distro_vms] start: {description}");
    let started = Instant::now();
    let command_debug = format!("{command:?}");
    let output = command
        .output()
        .wrap_err_with(|| format!("spawn command to {description}: {command_debug}"))?;
    if output.status.success() {
        eprintln!("[distro_vms] done: {description} ({:?})", started.elapsed());
        return Ok(());
    }
    eprintln!(
        "[distro_vms] failed: {description} ({:?})",
        started.elapsed()
    );
    Err(eyre!(
        "command failed while trying to {description}: {command_debug}\nstatus: {}\nstdout:\n{}\nstderr:\n{}",
        output.status,
        String::from_utf8_lossy(&output.stdout),
        String::from_utf8_lossy(&output.stderr),
    ))
}

fn validate_elf_machine(path: &Path, expected_machine: u16) -> Result<()> {
    let header =
        fs::read(path).wrap_err_with(|| format!("read ELF binary '{}'", path.display()))?;
    if header.len() < 20 || &header[0..4] != b"\x7fELF" {
        return Err(eyre!("'{}' is not an ELF binary", path.display()));
    }
    if header[4] != 2 || header[5] != 1 {
        return Err(eyre!(
            "'{}' is not a 64-bit little-endian ELF binary",
            path.display()
        ));
    }
    let machine = u16::from_le_bytes(header[18..20].try_into()?);
    if machine == expected_machine {
        Ok(())
    } else {
        Err(eyre!(
            "ELF binary '{}' has machine {machine}, expected {expected_machine}",
            path.display()
        ))
    }
}

fn timed_step<T>(test_name: &str, label: &str, step: impl FnOnce() -> Result<T>) -> Result<T> {
    eprintln!("[{test_name}] start: {label}");
    let started = Instant::now();
    match step() {
        Ok(value) => {
            eprintln!("[{test_name}] done: {label} ({:?})", started.elapsed());
            Ok(value)
        }
        Err(error) => {
            eprintln!("[{test_name}] failed: {label} ({:?})", started.elapsed());
            Err(error)
        }
    }
}

fn fedora_root_mount(_root_qcow2: &Path) -> Result<RootMount> {
    Ok(RootMount::new(RootDevice::virtio_first_partition()).fstype(RootFsType::Btrfs))
}