ostool 0.21.0

//! Runtime artifact preparation helpers.
//!
//! The build pipeline may produce an executable in Cargo's artifact directory
//! or receive a custom ELF path. This module normalizes that input into the ELF
//! and optional BIN files consumed by QEMU, U-Boot, TFTP, and board runners.

use std::{
    fs,
    path::{Path, PathBuf},
};

use anyhow::{Context, anyhow};
use colored::Colorize;
use object::{Architecture, Object};

use crate::{process::ProcessContext, utils::PathResultExt};

/// Options controlling how an input ELF is prepared for runtime use.
pub(crate) struct RuntimeArtifactOptions {
    /// Input ELF path produced by Cargo or supplied by a custom build.
    pub(crate) elf_path: PathBuf,
    /// Whether to also produce a raw BIN image.
    pub(crate) to_bin: bool,
    /// Optional directory for the generated BIN image.
    pub(crate) bin_dir: Option<PathBuf>,
    /// Whether to preserve debug information when producing a BIN image.
    pub(crate) debug: bool,
    /// Cargo artifact directory reported by `--message-format=json`, when known.
    pub(crate) cargo_artifact_dir: Option<PathBuf>,
    /// Whether to copy the input ELF into a stripped runtime `.elf` file first.
    pub(crate) strip_elf: bool,
    /// Objcopy executable used for ELF/BIN materialization.
    pub(crate) objcopy_program: PathBuf,
}

/// Runtime artifacts prepared from a single input ELF.
pub(crate) struct PreparedRuntimeArtifacts {
    elf: PathBuf,
    bin: Option<PathBuf>,
    cargo_artifact_dir: Option<PathBuf>,
    runtime_artifact_dir: Option<PathBuf>,
    arch: Option<Architecture>,
}

impl PreparedRuntimeArtifacts {
    /// Returns the runtime ELF path.
    pub(crate) fn elf(&self) -> &Path {
        &self.elf
    }

    /// Returns the runtime BIN path, when one was generated.
    pub(crate) fn bin(&self) -> Option<&Path> {
        self.bin.as_deref()
    }

    /// Returns the Cargo artifact directory that produced the input ELF.
    pub(crate) fn cargo_artifact_dir(&self) -> Option<&Path> {
        self.cargo_artifact_dir.as_deref()
    }

    /// Returns the directory containing the artifact consumed by runners.
    pub(crate) fn runtime_artifact_dir(&self) -> Option<&Path> {
        self.runtime_artifact_dir.as_deref()
    }

    /// Returns the architecture detected from the input ELF.
    pub(crate) fn arch(&self) -> Option<Architecture> {
        self.arch
    }
}

/// Prepares runtime ELF/BIN artifacts from a Cargo or custom-build ELF.
pub(crate) fn prepare_runtime_artifacts(
    context: &ProcessContext,
    options: RuntimeArtifactOptions,
) -> anyhow::Result<PreparedRuntimeArtifacts> {
    let input_elf = options
        .elf_path
        .canonicalize()
        .with_path("failed to canonicalize file", &options.elf_path)?;
    let input_dir = input_elf
        .parent()
        .ok_or_else(|| anyhow!("invalid ELF file path: {}", input_elf.display()))?
        .to_path_buf();
    let arch = detect_architecture(&input_elf)?;

    let runtime_elf = if options.strip_elf {
        strip_runtime_elf(context, &options.objcopy_program, &input_elf, arch)?
    } else {
        input_elf
    };

    let mut prepared = PreparedRuntimeArtifacts {
        elf: runtime_elf.clone(),
        bin: None,
        cargo_artifact_dir: options.cargo_artifact_dir.or(Some(input_dir)),
        runtime_artifact_dir: runtime_elf.parent().map(PathBuf::from),
        arch: Some(arch),
    };

    if options.to_bin {
        let bin_path = convert_runtime_bin(
            context,
            &options.objcopy_program,
            &runtime_elf,
            options.bin_dir,
            options.debug,
        )?;
        prepared.runtime_artifact_dir = bin_path.parent().map(PathBuf::from);
        prepared.bin = Some(bin_path);
    }

    Ok(prepared)
}

fn detect_architecture(path: &Path) -> anyhow::Result<Architecture> {
    let binary_data = fs::read(path).with_path("failed to read ELF file", path)?;
    let file = object::File::parse(binary_data.as_slice())
        .with_context(|| format!("failed to parse ELF file: {}", path.display()))?;
    Ok(file.architecture())
}

fn strip_runtime_elf(
    context: &ProcessContext,
    objcopy_program: &Path,
    elf_path: &Path,
    arch: Architecture,
) -> anyhow::Result<PathBuf> {
    let stripped_elf_path = elf_path.with_file_name(
        elf_path
            .file_stem()
            .ok_or_else(|| anyhow!("invalid ELF file path: {}", elf_path.display()))?
            .to_string_lossy()
            .to_string()
            + ".elf",
    );
    println!(
        "{}",
        format!(
            "Stripping ELF file...\r\n  original elf: {}\r\n  stripped elf: {}",
            elf_path.display(),
            stripped_elf_path.display()
        )
        .bold()
        .purple()
    );

    let mut objcopy = crate::process::command(objcopy_program, context);
    objcopy.arg(format!(
        "--binary-architecture={}",
        format!("{arch:?}").to_lowercase()
    ));
    objcopy.arg(elf_path);
    objcopy.arg(&stripped_elf_path);
    objcopy.run()?;

    Ok(stripped_elf_path)
}

fn convert_runtime_bin(
    context: &ProcessContext,
    objcopy_program: &Path,
    elf_path: &Path,
    bin_dir: Option<PathBuf>,
    debug: bool,
) -> anyhow::Result<PathBuf> {
    let elf_path = elf_path
        .canonicalize()
        .with_path("failed to canonicalize file", elf_path)?;
    let bin_name = elf_path
        .file_stem()
        .ok_or_else(|| anyhow!("invalid ELF file path: {}", elf_path.display()))?
        .to_string_lossy()
        .to_string()
        + ".bin";

    let bin_path = if let Some(bin_dir) = bin_dir {
        bin_dir.join(bin_name)
    } else {
        elf_path.with_file_name(bin_name)
    };

    if let Some(parent) = bin_path.parent() {
        fs::create_dir_all(parent).with_path("failed to create directory", parent)?;
    }

    println!(
        "{}",
        format!(
            "Converting ELF to BIN format...\r\n  elf: {}\r\n  bin: {}",
            elf_path.display(),
            bin_path.display()
        )
        .bold()
        .purple()
    );

    let mut objcopy = crate::process::command(objcopy_program, context);

    if !debug {
        objcopy.arg("--strip-all");
    }

    objcopy
        .arg("-O")
        .arg("binary")
        .arg(&elf_path)
        .arg(&bin_path);
    objcopy.run()?;

    Ok(bin_path)
}

#[cfg(test)]
mod tests {
    use std::{fs, path::PathBuf};

    use crate::{
        artifact::runtime::{RuntimeArtifactOptions, prepare_runtime_artifacts},
        process::ProcessContext,
        project::{resolve_project_layout, variables::VariableScope},
    };

    fn process_context(root: &std::path::Path) -> ProcessContext {
        fs::write(
            root.join("Cargo.toml"),
            "[package]\nname = \"sample\"\nversion = \"0.1.0\"\nedition = \"2024\"\n",
        )
        .unwrap();
        fs::create_dir_all(root.join("src")).unwrap();
        fs::write(root.join("src/lib.rs"), "").unwrap();

        let layout = resolve_project_layout(Some(root.to_path_buf())).unwrap();
        let scope = VariableScope::for_package(&layout, root.to_path_buf());
        ProcessContext::new(root.to_path_buf(), root.to_path_buf(), scope, None)
    }

    fn fake_objcopy(root: &std::path::Path) -> PathBuf {
        let script = root.join("fake-rust-objcopy");
        fs::write(
            &script,
            "#!/bin/sh\nlast=\"\"\nprev=\"\"\nfor arg in \"$@\"; do prev=\"$last\"; last=\"$arg\"; done\ncp \"$prev\" \"$last\"\n",
        )
        .unwrap();
        #[cfg(unix)]
        {
            use std::os::unix::fs::PermissionsExt;
            let mut permissions = fs::metadata(&script).unwrap().permissions();
            permissions.set_mode(0o755);
            fs::set_permissions(&script, permissions).unwrap();
        }
        script
    }

    #[test]
    fn prepares_stripped_elf_without_mutating_tool_state() {
        let temp = tempfile::tempdir().unwrap();
        let context = process_context(temp.path());
        let source = std::env::current_exe().unwrap();
        let input = temp.path().join("sample");
        fs::copy(&source, &input).unwrap();
        let objcopy_program = fake_objcopy(temp.path());

        let prepared = prepare_runtime_artifacts(
            &context,
            RuntimeArtifactOptions {
                elf_path: input.clone(),
                to_bin: false,
                bin_dir: None,
                debug: false,
                cargo_artifact_dir: Some(temp.path().join("target/debug")),
                strip_elf: true,
                objcopy_program,
            },
        )
        .unwrap();

        let expected_elf = input.with_file_name("sample.elf");
        assert_eq!(prepared.elf(), expected_elf);
        assert!(prepared.bin().is_none());
        assert_eq!(
            prepared.cargo_artifact_dir(),
            Some(temp.path().join("target/debug").as_path())
        );
        assert_eq!(prepared.runtime_artifact_dir(), Some(temp.path()));
        assert!(prepared.arch().is_some());
        assert!(expected_elf.exists());
    }

    #[test]
    fn prepares_optional_bin_in_custom_output_dir() {
        let temp = tempfile::tempdir().unwrap();
        let context = process_context(temp.path());
        let source = std::env::current_exe().unwrap();
        let input = temp.path().join("sample");
        fs::copy(&source, &input).unwrap();
        let bin_dir = temp.path().join("bin-out");
        let objcopy_program = fake_objcopy(temp.path());

        let prepared = prepare_runtime_artifacts(
            &context,
            RuntimeArtifactOptions {
                elf_path: input,
                to_bin: true,
                bin_dir: Some(bin_dir.clone()),
                debug: false,
                cargo_artifact_dir: None,
                strip_elf: true,
                objcopy_program,
            },
        )
        .unwrap();

        let expected_bin = bin_dir.join("sample.bin");
        assert_eq!(prepared.bin(), Some(expected_bin.as_path()));
        assert_eq!(prepared.runtime_artifact_dir(), Some(bin_dir.as_path()));
        assert!(expected_bin.exists());
    }
}