cargo-hyperlight 0.1.12

use std::collections::{BTreeMap, HashMap};
use std::env;
use std::ffi::{OsStr, OsString};
use std::fmt::Debug;
use std::path::{Path, PathBuf};
use std::process::Command as StdCommand;

use anyhow::{Context, Result};
use os_str_bytes::OsStrBytesExt;

use crate::CargoCommandExt;
use crate::cargo_cmd::{CargoBinary, CargoCmd as _, find_cargo, merge_env};
use crate::cli::{Args, Warning};

/// A process builder for cargo commands, providing a similar API to `std::process::Command`.
///
/// `Command` is a wrapper around `std::process::Command` specifically designed for
/// executing cargo commands targeting [hyperlight](https://github.com/hyperlight-dev/hyperlight)
/// guest code.
/// Before executing the desired command, `Command` takes care of setting up the
/// appropriate environment. It:
/// * creates a custom rust target for hyperlight guest code
/// * creates a sysroot with Rust's libs core and alloc
/// * finds the appropriate compiler and archiver for any C dependencies
/// * sets up necessary environment variables for `cc-rs` and `bindgen` to work correctly.
///
/// # Examples
///
/// Basic usage:
///
/// ```rust,no_run
/// use cargo_hyperlight::cargo;
///
/// let mut command = cargo().unwrap();
/// command.arg("build").arg("--release");
/// command.exec(); // This will replace the current process
/// ```
///
/// Setting environment variables and working directory:
///
/// ```rust
/// use cargo_hyperlight::cargo;
///
/// let mut command = cargo().unwrap();
/// command
///     .current_dir("/path/to/project")
///     .env("CARGO_TARGET_DIR", "/custom/target")
///     .args(["build", "--release"]);
/// ```
#[derive(Clone)]
pub struct Command {
    cargo: CargoBinary,
    /// Arguments to pass to the cargo program
    args: Vec<OsString>,
    /// Environment variable mappings to set for the child process
    inherit_envs: bool,
    inherit_cargo_envs: bool,
    envs: BTreeMap<OsString, Option<OsString>>,
    // Working directory for the child process
    current_dir: Option<PathBuf>,
}

impl Debug for Command {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let args = self.build_args_infallible();
        let mut cmd = self.command();
        cmd.populate_from_args(&args, false);

        write!(f, "env ")?;
        if let Some(current_dir) = &self.current_dir {
            write!(f, "-C {current_dir:?} ")?;
        }
        if !self.inherit_envs {
            write!(f, "-i ")?;
        }
        for (k, v) in cmd.get_envs() {
            if v.is_none() {
                write!(f, "-u {} ", k.to_string_lossy())?;
            }
        }
        for (k, v) in cmd.get_envs() {
            if let Some(v) = v {
                write!(f, "{}={:?} ", k.to_string_lossy(), v)?;
            }
        }
        write!(f, "{:?} ", self.get_program())?;
        for arg in &self.args {
            write!(f, "{arg:?} ")?;
        }
        writeln!(f)
    }
}

impl Command {
    /// Constructs a new `Command` for launching the cargo program.
    ///
    /// The value of the `CARGO` environment variable is used if it is set; otherwise, the
    /// default `cargo` from the system PATH is used.
    /// If `RUSTUP_TOOLCHAIN` is set in the environment, it is also propagated to the
    /// child process to ensure correct functioning of the rustup wrappers.
    ///
    /// The default configuration is:
    /// - No arguments to the program
    /// - Inherits the current process's environment
    /// - Inherits the current process's working directory
    ///
    /// # Errors
    ///
    /// This function will return an error if:
    /// - If the `CARGO` environment variable is set but it specifies an invalid path
    /// - If the `CARGO` environment variable is not set and the `cargo` program cannot be found in the system PATH
    ///
    /// # Examples
    ///
    /// Basic usage:
    ///
    /// ```rust
    /// use cargo_hyperlight::cargo;
    ///
    /// let command = cargo().unwrap();
    /// ```
    pub(crate) fn new() -> Result<Self> {
        let cargo = find_cargo()?;
        Ok(Self {
            cargo,
            args: Vec::new(),
            envs: BTreeMap::new(),
            inherit_envs: true,
            inherit_cargo_envs: true,
            current_dir: None,
        })
    }

    /// Adds an argument to pass to the cargo program.
    ///
    /// Only one argument can be passed per use. So instead of:
    ///
    /// ```no_run
    /// # let mut command = cargo_hyperlight::cargo().unwrap();
    /// command.arg("--features some_feature");
    /// ```
    ///
    /// usage would be:
    ///
    /// ```no_run
    /// # let mut command = cargo_hyperlight::cargo().unwrap();
    /// command.arg("--features").arg("some_feature");
    /// ```
    ///
    /// To pass multiple arguments see [`args`].
    ///
    /// [`args`]: Command::args
    ///
    /// Note that the argument is not shell-escaped, so if you pass an argument like
    /// `"hello world"`, it will be passed as a single argument with the literal
    /// `hello world`, not as two arguments `hello` and `world`.
    ///
    /// # Examples
    ///
    /// Basic usage:
    ///
    /// ```no_run
    /// use cargo_hyperlight::cargo;
    ///
    /// cargo()
    ///     .unwrap()
    ///     .arg("build")
    ///     .arg("--release")
    ///     .exec();
    /// ```
    pub fn arg(&mut self, arg: impl AsRef<OsStr>) -> &mut Self {
        self.args.push(arg.as_ref().to_os_string());
        self
    }

    /// Adds multiple arguments to pass to the cargo program.
    ///
    /// To pass a single argument see [`arg`].
    ///
    /// [`arg`]: Command::arg
    ///
    /// Note that the arguments are not shell-escaped, so if you pass an argument
    /// like `"hello world"`, it will be passed as a single argument with the
    /// literal `hello world`, not as two arguments `hello` and `world`.
    ///
    /// # Examples
    ///
    /// Basic usage:
    ///
    /// ```no_run
    /// use cargo_hyperlight::cargo;
    ///
    /// cargo()
    ///     .unwrap()
    ///     .args(["build", "--release"])
    ///     .exec();
    /// ```
    pub fn args(&mut self, args: impl IntoIterator<Item = impl AsRef<OsStr>>) -> &mut Self {
        for arg in args {
            self.arg(arg);
        }
        self
    }

    /// Sets the working directory for the child process.
    ///
    /// # Examples
    ///
    /// Basic usage:
    ///
    /// ```no_run
    /// use cargo_hyperlight::cargo;
    ///
    /// cargo()
    ///     .unwrap()
    ///     .current_dir("path/to/project")
    ///     .arg("build")
    ///     .exec();
    /// ```
    ///
    /// [`canonicalize`]: std::fs::canonicalize
    pub fn current_dir(&mut self, dir: impl AsRef<Path>) -> &mut Self {
        self.current_dir = Some(dir.as_ref().to_path_buf());
        self
    }

    /// Inserts or updates an explicit environment variable mapping.
    ///
    /// This method allows you to add an environment variable mapping to the spawned process
    /// or overwrite a variable if it already exists.
    ///
    /// Child processes will inherit environment variables from their parent process by
    /// default. Environment variables explicitly set using [`env`] take precedence
    /// over inherited variables. You can disable environment variable inheritance entirely
    /// using [`env_clear`] or for a single key using [`env_remove`].
    ///
    /// Note that environment variable names are case-insensitive (but
    /// case-preserving) on Windows and case-sensitive on all other platforms.
    ///
    /// # Examples
    ///
    /// Basic usage:
    ///
    /// ```no_run
    /// use cargo_hyperlight::cargo;
    ///
    /// cargo()
    ///     .unwrap()
    ///     .env("CARGO_TARGET_DIR", "/path/to/target")
    ///     .arg("build")
    ///     .exec();
    /// ```
    ///
    /// [`env`]: Command::env
    /// [`env_clear`]: Command::env_clear
    /// [`env_remove`]: Command::env_remove
    pub fn env(&mut self, key: impl AsRef<OsStr>, value: impl AsRef<OsStr>) -> &mut Self {
        self.envs
            .insert(key.as_ref().to_owned(), Some(value.as_ref().to_owned()));
        self
    }

    /// Clears all environment variables that will be set for the child process.
    ///
    /// This method will remove all environment variables from the child process,
    /// including those that would normally be inherited from the parent process.
    /// Environment variables can be added back individually using [`env`].
    ///
    /// If `RUSTUP_TOOLCHAIN` was set in the parent process, it will be preserved.
    ///
    /// # Examples
    ///
    /// Basic usage:
    ///
    /// ```no_run
    /// use cargo_hyperlight::cargo;
    ///
    /// cargo()
    ///     .unwrap()
    ///     .env_clear()
    ///     .env("CARGO_TARGET_DIR", "/path/to/target")
    ///     .arg("build")
    ///     .exec();
    /// ```
    ///
    /// [`env`]: Command::env
    pub fn env_clear(&mut self) -> &mut Self {
        self.inherit_envs = false;
        self.envs.clear();
        self
    }

    /// Clears Cargo build-context environment variables from the child process.
    ///
    /// This method removes the `CARGO_` environment variables that Cargo sets
    /// during build script execution and crate compilation (e.g. `CARGO_PKG_*`,
    /// `CARGO_MANIFEST_*`, `CARGO_CFG_*`, `CARGO_FEATURE_*`, etc.).
    ///
    /// User configuration variables are preserved, including:
    /// - `CARGO_HOME` — Cargo home directory
    /// - `CARGO_REGISTRIES_*` — Private registry index URLs, tokens, and credential providers
    /// - `CARGO_REGISTRY_*` — Default registry and crates.io credentials
    /// - `CARGO_HTTP_*` — HTTP/TLS proxy and timeout settings
    /// - `CARGO_NET_*` — Network configuration (retry, offline, git-fetch-with-cli)
    /// - `CARGO_ALIAS_*` — Command aliases
    /// - `CARGO_TERM_*` — Terminal output settings
    ///
    /// Other environment variables will remain unaffected. Environment variables
    /// can be added back individually using [`env`].
    ///
    /// This is particularly useful when using cargo-hyperlight from a build script
    /// or other cargo-invoked context where `CARGO_` variables may change the behavior
    /// of the cargo command being executed.
    ///
    /// # Examples
    ///
    /// Basic usage:
    ///
    /// ```no_run
    /// use cargo_hyperlight::cargo;
    ///
    /// cargo()
    ///     .unwrap()
    ///     .env_clear_cargo()
    ///     .env("CARGO_TARGET_DIR", "/path/to/target")
    ///     .arg("build")
    ///     .exec();
    /// ```
    ///
    /// [`env`]: Command::env
    pub fn env_clear_cargo(&mut self) -> &mut Self {
        self.inherit_cargo_envs = false;
        self.envs.retain(|k, _| should_preserve_cargo_env(k));
        self
    }

    #[doc(hidden)]
    #[deprecated(note = "use `env_clear_cargo` instead")]
    pub fn env_clear_cargo_vars(&mut self) -> &mut Self {
        self.env_clear_cargo()
    }

    /// Removes an explicitly set environment variable and prevents inheriting
    /// it from a parent process.
    ///
    /// This method will ensure that the specified environment variable is not
    /// present in the spawned process's environment, even if it was present
    /// in the parent process. This serves to "unset" environment variables.
    ///
    /// Note that environment variable names are case-insensitive (but
    /// case-preserving) on Windows and case-sensitive on all other platforms.
    ///
    /// # Examples
    ///
    /// Basic usage:
    ///
    /// ```no_run
    /// use cargo_hyperlight::cargo;
    ///
    /// cargo()
    ///     .unwrap()
    ///     .env_remove("CARGO_TARGET_DIR")
    ///     .arg("build")
    ///     .exec();
    /// ```
    pub fn env_remove(&mut self, key: impl AsRef<OsStr>) -> &mut Self {
        self.envs.insert(key.as_ref().to_owned(), None);
        self
    }

    /// Inserts or updates multiple explicit environment variable mappings.
    ///
    /// This method allows you to add multiple environment variable mappings
    /// to the spawned process or overwrite variables if they already exist.
    /// Environment variables can be passed as a `HashMap` or any other type
    /// implementing `IntoIterator` with the appropriate item type.
    ///
    /// Child processes will inherit environment variables from their parent process by
    /// default. Environment variables explicitly set using [`env`] take precedence
    /// over inherited variables. You can disable environment variable inheritance entirely
    /// using [`env_clear`] or for a single key using [`env_remove`].
    ///
    /// Note that environment variable names are case-insensitive (but
    /// case-preserving) on Windows and case-sensitive on all other platforms.
    ///
    /// # Examples
    ///
    /// Basic usage:
    ///
    /// ```no_run
    /// use std::collections::HashMap;
    /// use cargo_hyperlight::cargo;
    ///
    /// let mut envs = HashMap::new();
    /// envs.insert("CARGO_TARGET_DIR", "/path/to/target");
    /// envs.insert("CARGO_HOME", "/path/to/.cargo");
    ///
    /// cargo()
    ///     .unwrap()
    ///     .envs(&envs)
    ///     .arg("build")
    ///     .exec();
    /// ```
    ///
    /// ```no_run
    /// use cargo_hyperlight::cargo;
    ///
    /// cargo()
    ///     .unwrap()
    ///     .envs([
    ///         ("CARGO_TARGET_DIR", "/path/to/target"),
    ///         ("CARGO_HOME", "/path/to/.cargo"),
    ///     ])
    ///     .arg("build")
    ///     .exec();
    /// ```
    ///
    /// [`env`]: Command::env
    /// [`env_clear`]: Command::env_clear
    /// [`env_remove`]: Command::env_remove
    pub fn envs(
        &mut self,
        envs: impl IntoIterator<Item = (impl AsRef<OsStr>, impl AsRef<OsStr>)>,
    ) -> &mut Self {
        for (k, v) in envs {
            self.env(k, v);
        }
        self
    }

    /// Returns an iterator over the arguments that will be passed to the cargo program.
    ///
    /// This does not include the program name itself (which can be retrieved with
    /// [`get_program`]).
    ///
    /// # Examples
    ///
    /// ```no_run
    /// use cargo_hyperlight::cargo;
    ///
    /// let mut command = cargo().unwrap();
    /// command.arg("build").arg("--release");
    ///
    /// let args: Vec<&std::ffi::OsStr> = command.get_args().collect();
    /// assert_eq!(args, &["build", "--release"]);
    /// ```
    ///
    /// [`get_program`]: Command::get_program
    pub fn get_args(&'_ self) -> impl Iterator<Item = &OsStr> {
        self.args.iter().map(|s| s.as_os_str())
    }

    /// Returns the working directory for the child process.
    ///
    /// This returns `None` if the working directory will not be changed from
    /// the current directory of the parent process.
    ///
    /// # Examples
    ///
    /// ```no_run
    /// use std::path::Path;
    /// use cargo_hyperlight::cargo;
    ///
    /// let mut command = cargo().unwrap();
    /// assert_eq!(command.get_current_dir(), None);
    ///
    /// command.current_dir("/tmp");
    /// assert_eq!(command.get_current_dir(), Some(Path::new("/tmp")));
    /// ```
    pub fn get_current_dir(&self) -> Option<&Path> {
        self.current_dir.as_deref()
    }

    /// Returns an iterator over the environment mappings that will be set for the child process.
    ///
    /// Environment variables explicitly set or unset via [`env`], [`envs`], and
    /// [`env_remove`] can be retrieved with this method.
    ///
    /// Note that this output does not include environment variables inherited from the
    /// parent process.
    ///
    /// Each element is a tuple key/value where `None` means the variable is explicitly
    /// unset in the child process environment.
    ///
    /// # Examples
    ///
    /// ```no_run
    /// use std::ffi::OsStr;
    /// use cargo_hyperlight::cargo;
    ///
    /// let mut command = cargo().unwrap();
    /// command.env("CARGO_HOME", "/path/to/.cargo");
    /// command.env_remove("CARGO_TARGET_DIR");
    ///
    /// for (key, value) in command.get_envs() {
    ///     println!("{key:?} => {value:?}");
    /// }
    /// ```
    ///
    /// [`env`]: Command::env
    /// [`envs`]: Command::envs
    /// [`env_remove`]: Command::env_remove
    pub fn get_envs(&'_ self) -> impl Iterator<Item = (&OsStr, Option<&OsStr>)> {
        self.envs.iter().map(|(k, v)| (k.as_os_str(), v.as_deref()))
    }

    /// Returns the base environment variables for the command.
    ///
    /// This method returns the environment variables that will be inherited
    /// from the current process, taking into account whether [`env_clear`] has been called.
    ///
    /// [`env_clear`]: Command::env_clear
    fn base_env(&self) -> impl Iterator<Item = (OsString, OsString)> {
        env::vars_os().filter(|(k, _)| {
            if !self.inherit_envs {
                false
            } else if !self.inherit_cargo_envs {
                should_preserve_cargo_env(k)
            } else {
                true
            }
        })
    }

    fn resolve_env(&self) -> HashMap<OsString, OsString> {
        merge_env(self.base_env(), self.get_envs())
    }

    fn command(&self) -> StdCommand {
        let mut command = self.cargo.command();
        // Filter out --target and --target-dir from forwarded args since
        // populate_from_args sets them via env vars with the resolved values
        let args: Vec<_> = self.get_args().map(|a| a.to_owned()).collect();
        let mut skip_next = false;
        for arg in args.iter() {
            if skip_next {
                skip_next = false;
                continue;
            }
            let arg_str = arg.to_string_lossy();
            if arg_str == "--target" || arg_str == "--target-dir" {
                skip_next = true; // skip the next arg (the value)
                continue;
            }
            if arg_str.starts_with("--target=") || arg_str.starts_with("--target-dir=") {
                continue;
            }
            if arg_str == "--host" {
                skip_next = true;
                continue;
            }
            if arg_str.starts_with("--host=") {
                continue;
            }
            if arg_str.starts_with("--with-guest-capi") {
                continue;
            }
            command.arg(arg);
        }
        if let Some(cwd) = &self.current_dir {
            command.current_dir(cwd);
        }
        if !self.inherit_envs {
            command.env_clear();
        }
        if !self.inherit_cargo_envs {
            for (k, _) in env::vars_os() {
                if !should_preserve_cargo_env(&k) {
                    command.env_remove(k);
                }
            }
        }
        if let Some(rustup_toolchain) = &self.cargo.rustup_toolchain {
            command.env("RUSTUP_TOOLCHAIN", rustup_toolchain);
        }
        for (k, v) in self.get_envs() {
            match v {
                Some(v) => command.env(k, v),
                None => command.env_remove(k),
            };
        }
        command
    }

    /// Returns the path to the cargo program that will be executed.
    ///
    /// # Examples
    ///
    /// ```no_run
    /// use cargo_hyperlight::cargo;
    ///
    /// let command = cargo().unwrap();
    /// println!("Program: {:?}", command.get_program());
    /// ```
    pub fn get_program(&self) -> &OsStr {
        self.cargo.path.as_os_str()
    }

    pub fn build_args(&self) -> Args {
        // parse the arguments and environment variables
        match Args::parse(
            self.get_args(),
            self.resolve_env(),
            self.get_current_dir(),
            Warning::WARN,
        ) {
            Ok(args) => args,
        }
    }

    fn build_args_infallible(&self) -> Args {
        match Args::parse(
            self.get_args(),
            self.resolve_env(),
            self.get_current_dir(),
            Warning::IGNORE,
        ) {
            Ok(args) => args,
            Err(err) => {
                eprintln!("Failed to parse arguments: {err}");
                std::process::exit(1);
            }
        }
    }

    /// Executes a cargo command as a child process, waiting for it to finish and
    /// collecting its exit status.
    ///
    /// The process stdin, stdout and stderr are inherited from the parent.
    ///
    /// # Examples
    ///
    /// Basic usage:
    ///
    /// ```no_run
    /// use cargo_hyperlight::cargo;
    ///
    /// let result = cargo()
    ///     .unwrap()
    ///     .arg("build")
    ///     .status();
    ///
    /// match result {
    ///     Ok(()) => println!("Cargo command succeeded"),
    ///     Err(e) => println!("Cargo command failed: {}", e),
    /// }
    /// ```
    ///
    /// # Errors
    ///
    /// This method will return an error if:
    /// - The sysroot preparation fails
    /// - The cargo process could not be spawned
    /// - The cargo process returned a non-zero exit status
    pub fn status(&self) -> anyhow::Result<()> {
        let args = self.build_args();

        args.prepare_sysroot()
            .context("Failed to prepare sysroot")?;

        self.command()
            .populate_from_args(&args, false)
            .checked_status()
            .context("Failed to execute cargo")?;
        Ok(())
    }
}

/// Returns `true` if the given environment variable should be preserved
/// when clearing Cargo build-context variables.
///
/// Non-`CARGO_` variables are always preserved. The following `CARGO_`
/// user configuration variables are also preserved:
/// - `CARGO_HOME` — Cargo home directory
/// - `CARGO_REGISTRIES_*` — Private registry index URLs, tokens, and credential providers
/// - `CARGO_REGISTRY_*` — Default registry and crates.io credentials
/// - `CARGO_HTTP_*` — HTTP/TLS proxy and timeout settings
/// - `CARGO_NET_*` — Network configuration (retry, offline, git-fetch-with-cli)
/// - `CARGO_ALIAS_*` — Command aliases
/// - `CARGO_TERM_*` — Terminal output settings
///
/// All other `CARGO_*` variables (e.g. `CARGO_PKG_*`, `CARGO_MANIFEST_*`,
/// `CARGO_CFG_*`, `CARGO_FEATURE_*`, `CARGO_MAKEFLAGS`, etc.) are considered
/// build-context variables set by Cargo during compilation and will be cleared.
///
/// See: <https://doc.rust-lang.org/cargo/reference/environment-variables.html>
fn should_preserve_cargo_env(key: &OsStr) -> bool {
    !key.starts_with("CARGO_")
        || key == "CARGO_HOME"
        || key.starts_with("CARGO_REGISTRIES_")
        || key.starts_with("CARGO_REGISTRY_")
        || key.starts_with("CARGO_HTTP_")
        || key.starts_with("CARGO_NET_")
        || key.starts_with("CARGO_ALIAS_")
        || key.starts_with("CARGO_TERM_")
}