ambient-ci 0.14.0

//! Run QEMU.

use std::{
    ffi::{OsStr, OsString},
    path::{Path, PathBuf},
    process::Command,
};

use bytesize::MIB;
use clingwrap::runner::{CommandError, CommandRunner};
use tempfile::{tempdir_in, TempDir};

use crate::{
    cloud_init::LocalDataStore,
    config::Config,
    qemu_utils::create_cow_image,
    runlog::{RunLog, RunLogError, RunLogSource},
    util::{copy_file_rw, create_file},
    vdrive::{VirtualDrive, VirtualDriveError},
};

/// Path in VM to executor drive.
pub const EXECUTOR_DRIVE: &str = "/dev/vdb";
/// Path in VM to source drive.
pub const SOURCE_DRIVE: &str = "/dev/vdc";
/// Path in VM to artifact drive.
pub const ARTIFACT_DRIVE: &str = "/dev/vdd";
/// Path in VM to cache drive.
pub const CACHE_DRIVE: &str = "/dev/vde";
/// Path in VM to dependencies drive.
pub const DEPS_DRIVE: &str = "/dev/vdf";

/// Path in VM to workspace root directory.
pub const WORKSPACE_DIR: &str = "/ci";
/// Path in VM to source directory.
pub const SOURCE_DIR: &str = "/ci/src";
/// Path in VM to dependencies directory.
pub const DEPS_DIR: &str = "/ci/deps";
/// Path in VM to cache directory.
pub const CACHE_DIR: &str = "/ci/cache";
/// Path in VM to artifactsdirectory.
pub const ARTIFACTS_DIR: &str = "/ci/artifacts";

/// Run QEMU.
#[derive(Default)]
pub struct QemuRunner<'a> {
    config: Option<&'a Config>,
    base_image: Option<PathBuf>,
    cow_image: Option<PathBuf>,
    cloud_init: Option<LocalDataStore>,
    executor: Option<&'a VirtualDrive>,
    source: Option<&'a VirtualDrive>,
    dependencies: Option<&'a VirtualDrive>,
    cache: Option<&'a VirtualDrive>,
    artifacts: Option<&'a VirtualDrive>,
    console_log: Option<PathBuf>,
    raw_log: Option<PathBuf>,
    network: bool,
    uefi: bool,
}

impl<'a> QemuRunner<'a> {
    /// Set configuration.
    pub fn config(mut self, value: &'a Config) -> Self {
        self.config = Some(value);
        self
    }

    /// Set base image.
    pub fn base_image(mut self, filename: &Path) -> Self {
        self.base_image = Some(filename.into());
        self
    }

    /// Set copy-on-write image.
    pub fn cow_image(mut self, filename: &Path) -> Self {
        self.cow_image = Some(filename.into());
        self
    }

    /// Set `cloud-init` data store.
    pub fn cloud_init(mut self, ds: &LocalDataStore) -> Self {
        self.cloud_init = Some(ds.clone());
        self
    }

    /// Set executor drive.
    pub fn executor(mut self, value: &'a VirtualDrive) -> Self {
        self.executor = Some(value);
        self
    }

    /// Set source drive.
    pub fn source(mut self, value: &'a VirtualDrive) -> Self {
        self.source = Some(value);
        self
    }

    /// Set dependencies drive.
    pub fn dependencies(mut self, value: &'a VirtualDrive) -> Self {
        self.dependencies = Some(value);
        self
    }

    /// Set cache drive.
    pub fn cache(mut self, value: &'a VirtualDrive) -> Self {
        self.cache = Some(value);
        self
    }

    /// Set artifacts drive.
    pub fn artifacts(mut self, value: &'a VirtualDrive) -> Self {
        self.artifacts = Some(value);
        self
    }

    /// Set console log file.
    pub fn console_log(mut self, value: &'a Path) -> Self {
        self.console_log = Some(value.into());
        self
    }

    /// Set run log file.
    pub fn raw_log(mut self, value: &'a Path) -> Self {
        self.raw_log = Some(value.into());
        self
    }

    /// Allow network?
    pub fn network(mut self, network: bool) -> Self {
        self.network = network;
        self
    }

    /// Use UEFI?
    pub fn uefi(mut self, uefi: bool) -> Self {
        self.uefi = uefi;
        self
    }

    /// Run QEMU.
    pub fn run(&self, source: RunLogSource, runlog: &mut RunLog) -> Result<i32, QemuError> {
        let config = self.config.ok_or(QemuError::Missing("config"))?;
        let image = self.base_image.clone().ok_or(QemuError::Missing("image"))?;
        let cloud_init = self
            .cloud_init
            .clone()
            .ok_or(QemuError::Missing("cloud_init"))?;
        let executor_drive = self.executor.ok_or(QemuError::Missing("executor_drive"))?;

        let raw_log = self.raw_log.clone().ok_or(QemuError::Missing("raw_log"))?;
        let console_log = self
            .console_log
            .clone()
            .ok_or(QemuError::Missing("console_log"))?;

        let qemu = Qemu {
            kvm_binary: config.kvm_binary(),
            ovmf_code: config.ovmf_code_file(),
            ovmf_vars: config.ovmf_vars_file(),
            image,
            cow_image: self.cow_image.clone(),
            tmpdir: tempdir_in(config.tmpdir()).map_err(QemuError::TempDir)?,
            console_log,
            raw_log: raw_log.clone(),
            cloud_init,
            executor: executor_drive.clone(),
            source: self.source,
            artifact: self.artifacts,
            dependencies: self.dependencies,
            cache: self.cache,
            cpus: config.cpus(),
            memory: config.memory().as_u64(),
            network: self.network,
        };
        let result = qemu.run(source, runlog);
        if result.is_err() {
            if let Ok(data) = std::fs::read(&raw_log) {
                eprintln!(
                    "========================= raw log:\n{}=========================\n",
                    String::from_utf8_lossy(&data)
                );
            }
        }
        result
    }
}

/// A QEMU runner.
#[derive(Debug)]
struct Qemu<'a> {
    kvm_binary: PathBuf,
    ovmf_code: PathBuf,
    ovmf_vars: PathBuf,
    image: PathBuf,
    cow_image: Option<PathBuf>,
    cloud_init: LocalDataStore,
    tmpdir: TempDir,
    console_log: PathBuf,
    raw_log: PathBuf,
    executor: VirtualDrive,
    source: Option<&'a VirtualDrive>,
    artifact: Option<&'a VirtualDrive>,
    dependencies: Option<&'a VirtualDrive>,
    cache: Option<&'a VirtualDrive>,
    cpus: usize,
    memory: u64,
    network: bool,
}

impl Qemu<'_> {
    /// Run QEMU in the specified way.
    #[allow(clippy::unwrap_used)]
    fn run(&self, source: RunLogSource, runlog: &mut RunLog) -> Result<i32, QemuError> {
        let tmp = tempdir_in(&self.tmpdir).map_err(QemuError::TempDir)?;

        let cow_image = self
            .cow_image
            .clone()
            .unwrap_or_else(|| tmp.path().join("vm.qcow2"));
        let iso = tmp.path().join("cloud_init.iso");
        let vars = tmp.path().join("vars.fd");

        create_cow_image(&self.image, &cow_image).map_err(QemuError::COW)?;

        copy_file_rw(&self.ovmf_vars, &vars)
            .map_err(|e| QemuError::Copy(self.ovmf_vars.to_path_buf(), Box::new(e)))?;
        assert!(!std::fs::metadata(&vars).unwrap().permissions().readonly());

        self.cloud_init
            .iso(&iso)
            .map_err(|err| QemuError::Iso(iso.clone(), err))?;

        create_file(&self.console_log).map_err(QemuError::CreateFile)?;

        let raw_log_filename = create_file(&self.raw_log).map_err(QemuError::CreateFile)?;

        let cpus = format!("cpus={}", self.cpus);
        let memory = format!("{}", self.memory / MIB);
        let mut args = QemuArgs::default()
            .with_valued_arg("-m", &memory)
            .with_valued_arg("-smp", &cpus)
            .with_valued_arg("-cpu", "kvm64")
            .with_valued_arg("-machine", "type=q35,accel=kvm,usb=off")
            .with_valued_arg("-uuid", "a85c9de7-edc0-4e54-bead-112e5733582c")
            .with_valued_arg("-boot", "strict=on")
            .with_valued_arg("-name", "ambient-ci-vm")
            .with_valued_arg("-rtc", "base=utc,driftfix=slew")
            .with_valued_arg("-display", "none")
            .with_valued_arg("-device", "virtio-rng-pci")
            .with_valued_arg("-serial", &format!("file:{}", self.console_log.display())) // ttyS0
            .with_valued_arg("-serial", &format!("file:{}", raw_log_filename.display())) // ttyS1
            .with_qcow2(&cow_image.to_string_lossy())
            .with_raw(self.executor.filename(), true)
            .with_valued_arg("-cdrom", &iso.display().to_string());

        args = args.with_ipflash(0, &self.ovmf_code.to_string_lossy(), true);
        args = args.with_ipflash(1, &vars.to_string_lossy(), false);

        if let Some(drive) = self.source {
            args = args.with_raw(drive.filename(), true);
        }
        if let Some(drive) = self.artifact {
            args = args.with_raw(drive.filename(), false);
        }
        if let Some(drive) = self.cache {
            args = args.with_raw(drive.filename(), false);
        }
        if let Some(drive) = self.dependencies {
            args = args.with_raw(drive.filename(), true);
        }
        if self.network {
            args = args.with_valued_arg("-nic", "user,model=virtio");
        }
        args = args.with_arg("-nodefaults").with_arg("-no-user-config");

        let mut cmd = Command::new(&self.kvm_binary);
        cmd.args(args.iter());

        runlog.start_qemu(source, &cmd);
        let runner = CommandRunner::new(cmd);
        let result = runner.execute();
        let (vm_runlog, exit) = Self::parse_raw_log(&raw_log_filename)?;
        for msg in vm_runlog.msgs() {
            runlog.write(msg);
        }
        match &result {
            Ok(output) => {
                runlog.qemu_succeeded(source, output);
                Ok(exit)
            }
            Err(CommandError::KilledBySignal { .. }) => {
                let err = result.unwrap_err();
                runlog.qemu_failed(source, &err);
                Err(QemuError::Qemu(err))
            }
            Err(CommandError::CommandFailed { exit_code, .. }) => {
                let exit_code = *exit_code;
                let err = result.unwrap_err();
                runlog.qemu_failed(source, &err);
                Ok(exit_code)
            }
            Err(_) => {
                let err = result.unwrap_err();
                runlog.qemu_failed(source, &err);
                Err(QemuError::Qemu(err))
            }
        }
    }

    fn parse_raw_log(filename: &Path) -> Result<(RunLog, i32), QemuError> {
        const BEGIN: &str = "====================== BEGIN ======================";
        const EXIT: &str = "\nEXIT CODE: ";

        let raw = std::fs::read(filename).map_err(|e| QemuError::ReadLog(filename.into(), e))?;
        let runlog = RunLog::from_raw(raw.clone()).map_err(QemuError::ParseRaw)?;

        let log = String::from_utf8_lossy(&raw);
        if let Some((_, log)) = log.split_once(BEGIN) {
            if let Some((_, rest)) = log.split_once(EXIT) {
                if let Some((exit, _)) = rest.split_once('\n') {
                    let exit = exit.trim();
                    let exit = exit
                        .parse::<i32>()
                        .or(Err(QemuError::ParseExit(exit.to_string())))?;
                    Ok((runlog, exit))
                } else {
                    Err(QemuError::BadExitCode)
                }
            } else {
                Err(QemuError::NoBeginMarker)
            }
        } else {
            Err(QemuError::NoExit)
        }
    }
}

#[derive(Debug, Default)]
struct QemuArgs {
    args: Vec<OsString>,
}

impl QemuArgs {
    fn with_arg(mut self, arg: &str) -> Self {
        self.args.push(arg.into());
        self
    }

    fn with_valued_arg(mut self, arg: &str, value: &str) -> Self {
        self.args.push(arg.into());
        self.args.push(value.into());
        self
    }

    fn with_ipflash(mut self, unit: usize, path: &str, readonly: bool) -> Self {
        self.args.push("-drive".into());
        self.args.push(
            format!(
                "if=pflash,format=raw,unit={},file={}{}",
                unit,
                path,
                if readonly { ",readonly=on" } else { "" },
            )
            .into(),
        );
        self
    }

    fn with_qcow2(mut self, path: &str) -> Self {
        self.args.push("-drive".into());
        self.args
            .push(format!("format=qcow2,if=virtio,file={path}").into());
        self
    }

    fn with_raw(mut self, path: &Path, readonly: bool) -> Self {
        self.args.push("-drive".into());
        self.args.push(
            format!(
                "format=raw,if=virtio,file={}{}",
                path.display(),
                if readonly { ",readonly=on" } else { "" },
            )
            .into(),
        );
        self
    }

    fn iter(&self) -> impl Iterator<Item = &OsStr> {
        self.args.iter().map(|s| s.as_os_str())
    }
}

/// Possible errors from running Qemu.
#[allow(missing_docs)]
#[derive(Debug, thiserror::Error)]
pub enum QemuError {
    #[error("missing field in QemuRunner: {0}")]
    Missing(&'static str),

    #[error("failed to create a temporary directory")]
    TempDir(#[source] std::io::Error),

    #[error("failed to copy to temporary directory: {0}")]
    Copy(PathBuf, #[source] Box<crate::util::UtilError>),

    #[error("failed to read log file {0}")]
    ReadLog(PathBuf, #[source] std::io::Error),

    #[error("failed to parse raw log for JSON Lines")]
    ParseRaw(#[source] RunLogError),

    #[error("failed to create a tar archive from {0}")]
    Tar(PathBuf, #[source] Box<VirtualDriveError>),

    #[error("failed to extract cache drive to {0}")]
    ExtractCache(PathBuf, #[source] Box<VirtualDriveError>),

    #[error("failed to read temporary file for logging")]
    TemporaryLog(#[source] std::io::Error),

    #[error("run log lacks exit code of run")]
    NoExit,

    #[error("failed to get length of file {0}")]
    Metadata(PathBuf, #[source] std::io::Error),

    #[error("failed to set length of file to {0}: {1}")]
    SetLen(u64, PathBuf, #[source] std::io::Error),

    #[error("failed to run actions in QEMU")]
    QemuFailed(i32),

    #[error("failed to create cloud-init ISO file {0}")]
    Iso(PathBuf, #[source] crate::cloud_init::CloudInitError),

    #[error(transparent)]
    COW(#[from] crate::qemu_utils::QemuUtilError),

    #[error("failed to run QEMU")]
    Qemu(#[source] CommandError),

    #[error(transparent)]
    CreateFile(#[from] crate::util::UtilError),

    #[error("failed to parse exist code {0:?}")]
    ParseExit(String),

    #[error("run log from QEMU does not have a BEGIN marker")]
    NoBeginMarker,

    #[error("run log from QEMU has malformed exit code marker")]
    BadExitCode,
}