vmexec 0.7.1

use std::fs;
use std::path::{Path, PathBuf};
use std::process::Command;
use std::time::Duration;

use eyre::{Context, OptionExt, Result, bail};

use dir_lock::DirLock;
use directories::ProjectDirs;
use tokio::io::AsyncWriteExt;
use tokio::task;
use tracing::{debug, instrument, trace, warn};
use walkdir::WalkDir;

use crate::error::Error;
use crate::vms::CID_FILENAME;

pub(crate) static HEX_ALPHABET: [char; 16] = [
    '1', '2', '3', '4', '5', '6', '7', '8', '9', '0', 'a', 'b', 'c', 'd', 'e', 'f',
];

/// Ensure that a required directory exists
pub(crate) fn ensure_directory(purpose: &str, path: &Path) -> Result<()> {
    if !path.exists() {
        debug!("{purpose} dir {path:?} doesn't exist yet, creating");
        fs::create_dir_all(path).wrap_err(format!("Creating {purpose} dir {path:?}"))?;
    }
    Ok(())
}

/// Ensure that a required directory exists
pub(crate) async fn ensure_directory_async(purpose: &str, path: &Path) -> Result<()> {
    let purpose = purpose.to_owned();
    let path = path.to_owned();
    task::spawn_blocking(move || ensure_directory(&purpose, &path)).await?
}

#[derive(Clone, Debug)]
pub struct VmexecDirs {
    /// Dir containing cached stuff like kernels and VM images (usually ~/.cache/vmexec/)
    pub cache_dir: PathBuf,

    /// Dir containing the downloaded VM images (usually ~/.cache/vmexec/images/)
    pub images_dir: PathBuf,

    /// Dir containing secrets (usually ~/.local/state/vmexec/secrets/)
    pub secrets_dir: PathBuf,

    /// Dir containing all VMs (usually ~/.local/state/vmexec/vms/)
    pub vms_dir: PathBuf,

    /// Dir used for locking between vmexec processes via [`DirLock`] (usually ~/.local/state/vmexec/lockdir/)
    pub lock_dir: PathBuf,
}

impl VmexecDirs {
    pub fn new() -> Result<Self> {
        let project_dir =
            ProjectDirs::from("", "", "vmexec").ok_or_eyre("Couldn't get project dir")?;

        let cache_dir = project_dir.cache_dir().to_path_buf();
        ensure_directory("cache", &cache_dir)?;

        let images_dir = cache_dir.join("images");
        ensure_directory("images", &images_dir)?;

        let data_dir = project_dir
            .state_dir()
            .ok_or_eyre("Couldn't get state dir")?
            .to_path_buf();
        ensure_directory("data", &data_dir)?;

        let secrets_dir = data_dir.join("secrets");
        ensure_directory("secrets", &secrets_dir)?;

        let vms_dir = data_dir.join("vms");
        ensure_directory("vms", &vms_dir)?;

        let lock_dir = data_dir.join("lockdir");

        Ok(Self {
            cache_dir,
            images_dir,
            secrets_dir,
            vms_dir,
            lock_dir,
        })
    }
}

/// Path escaping, like `systemd-escape --path`.
///
/// From https://github.com/lucab/libsystemd-rs/blob/b43fa5e3b5eca3e6aa16a6c2fad87220dc0ad7a0/src/unit.rs
pub(crate) fn escape_path(path: &str) -> String {
    let trimmed = path.trim_matches('/');
    if trimmed.is_empty() {
        return "-".to_string();
    }

    let mut slash_seq = false;
    let parts: Vec<String> = trimmed
        .bytes()
        .filter(|b| {
            let is_slash = *b == b'/';
            let res = !(is_slash && slash_seq);
            slash_seq = is_slash;
            res
        })
        .enumerate()
        .map(|(n, b)| escape_byte(b, n))
        .collect();
    parts.join("")
}

fn escape_byte(b: u8, index: usize) -> String {
    let c = char::from(b);
    match c {
        '/' => '-'.to_string(),
        ':' | '_' | '0'..='9' | 'a'..='z' | 'A'..='Z' => c.to_string(),
        '.' if index > 0 => c.to_string(),
        _ => format!(r#"\x{b:02x}"#),
    }
}

/// Get a random unused CID to use with vsock
///
/// The way this works is that every vm dir inside `vms_dir` contains its own CID. We then look
/// at all the CIDs in all vm dirs to get the current list of CIDs that are in-use and just pick
/// the next free one.
///
/// The caller must hold the vmexec lock (hence the `_lock` parameter) so that multiple instances
/// of `vmexec` do not race each other.
#[instrument(skip(_lock))]
pub(crate) fn create_free_cid(_lock: &DirLock, vms_dir: &Path, vm_dir: &Path) -> Result<u32> {
    let mut cids = vec![];

    for entry in WalkDir::new(vms_dir) {
        let entry = entry?;
        let filename = entry.file_name();
        if filename.to_string_lossy() == CID_FILENAME {
            trace!("Found CID file at {:?}", entry.path());
            let cid = fs::read_to_string(entry.path())?;
            cids.push(cid.parse::<u32>()?);
        }
    }

    // Get the next CID.
    cids.sort();
    let cid = if let Some(last_cid) = cids.iter().next_back() {
        last_cid + 1
    } else {
        // We get here if the current list of CIDs is empty. So we'll just start with some
        // arbitrary CID.
        10
    };

    debug!("Our new CID: {cid}");
    fs::write(vm_dir.join(CID_FILENAME), cid.to_string())?;

    Ok(cid)
}

/// Take the inter-process lock that for state-mutating operations
///
/// Blocks until the lock is free.
pub(crate) fn lock_state_dir(dirs: &VmexecDirs) -> Result<DirLock> {
    let lock_dir = &dirs.lock_dir;
    if lock_dir.exists() {
        warn!("Already locked by another vmexec operation, waiting...");
    }
    trace!("Trying to lock {lock_dir:?}");
    let lock = DirLock::new_sync(lock_dir)?;
    Ok(lock)
}

#[derive(Clone, Debug)]
pub struct ExecutablePaths {
    pub qemu_path: PathBuf,
    pub virtiofsd_path: PathBuf,
    pub virt_copy_out_path: PathBuf,
}

/// Check whether necessary tools are installed and return their paths
pub fn find_required_tools() -> Result<ExecutablePaths, Error> {
    // Find QEMUU
    let qemu_path = which::which_global("qemu-system-x86_64").map_err(|_| Error::MissingTool {
        tool: "qemu-system-x86_64",
        searched: "PATH",
    })?;

    // Find virtiofsd
    let virtiofsd_path =
        which::which_in("virtiofsd", Some("/usr/lib:/usr/libexec"), "/").map_err(|_| {
            Error::MissingTool {
                tool: "virtiofsd",
                searched: "/usr/lib or /usr/libexec",
            }
        })?;

    // Find virt-copy-out
    let virt_copy_out_path =
        which::which_global("virt-copy-out").map_err(|_| Error::MissingTool {
            tool: "virt-copy-out",
            searched: "PATH",
        })?;

    check_unshare()?;

    Ok(ExecutablePaths {
        qemu_path,
        virtiofsd_path,
        virt_copy_out_path,
    })
}

/// Check whether unshare is working as expected
fn check_unshare() -> Result<()> {
    let unshare_output = Command::new("unshare")
        .env("LC_ALL", "C")
        .arg("-r")
        .arg("id")
        .output()?;
    let unshare_stdout = std::str::from_utf8(&unshare_output.stdout)?;
    let unshare_stderr = std::str::from_utf8(&unshare_output.stderr)?;
    if !unshare_output.status.success() {
        bail!(
            "Test command 'unshare -r id' didn't exit succesfully, stdout: {unshare_stdout}, stderr: {unshare_stderr}"
        );
    }
    if !unshare_stdout.starts_with("uid=0(root) gid=0(root) groups=0(root)") {
        bail!(
            "Expected output to start with 'unshare -r id' to report 'uid=0(root) gid=0(root) groups=0(root)' but got: {unshare_stdout}"
        );
    }

    Ok(())
}

// Wrap flush in a retry loop
//
// This is required because somehow flush sometimes will just not be ready yet and fail.
pub(crate) async fn safe_flush<W: AsyncWriteExt + std::marker::Unpin>(
    writer: &mut W,
) -> std::io::Result<()> {
    use std::io::ErrorKind;

    loop {
        match writer.flush().await {
            Err(ref e) if e.kind() == ErrorKind::WouldBlock => {
                // Wait a bit and retry
                tokio::time::sleep(Duration::from_millis(1)).await;
                continue;
            }
            result => return result,
        }
    }
}