Repository: davidlattimore/wild (rust)

Issue:
WILD_SAVE_DIR/WILD_SAVE_BASE should use Wild's arguments, not the process ones
Works fine when the binary is just the linker, but when integrating Wild with some other binary, wrong arguments are saved because of this line: https://github.com/davidlattimore/wild/blob/9644a9f7054c72b5a3f1cd31bac216f217e54056/libwild/src/save_dir.rs#L88

For example, when creating interposer for `cc` binary with integrated Wild, `run-with` will contain CC's arguments which is not desired.

Pull request description:
save-dir: Handle @filename and thin archives
Issue #575
Fixes #681

Patch:
diff --git a/libwild/src/args.rs b/libwild/src/args.rs
index 4ad6be8..23194d7 100644
--- a/libwild/src/args.rs
+++ b/libwild/src/args.rs
@@ -289,7 +289,7 @@ impl Default for Args {
 }
 
 // Parse the supplied input arguments, which should not include the program name.
-pub(crate) fn parse<S: AsRef<str>, I: Iterator<Item = S>>(mut input: I) -> Result<Args> {
+pub(crate) fn parse<F: Fn() -> I, S: AsRef<str>, I: Iterator<Item = S>>(input: F) -> Result<Args> {
     let mut args = Args {
         files_per_group: std::env::var(FILES_PER_GROUP_ENV)
             .ok()
@@ -300,7 +300,9 @@ pub(crate) fn parse<S: AsRef<str>, I: Iterator<Item = S>>(mut input: I) -> Resul
 
     let mut unrecognised = Vec::new();
 
-    let mut save_dir = SaveDir::new()?;
+    let mut save_dir = SaveDir::new(&input)?;
+
+    let mut input = input();
 
     let mut modifier_stack = vec![Modifiers::default()];
 
@@ -654,11 +656,12 @@ const fn default_target_arch() -> Architecture {
 fn parse_from_argument_file(path: &Path) -> Result<Args> {
     let contents = std::fs::read_to_string(path)
         .with_context(|| format!("Failed to read arguments from file `{}`", path.display()))?;
-    parse(arguments_from_string(&contents)?.into_iter())
+    let arguments = arguments_from_string(&contents)?;
+    parse(|| arguments.iter())
 }
 
 impl Args {
-    pub fn parse<S: AsRef<str>, I: Iterator<Item = S>>(input: I) -> Result<Args> {
+    pub fn parse<F: Fn() -> I, S: AsRef<str>, I: Iterator<Item = S>>(input: F) -> Result<Args> {
         parse(input)
     }
 
@@ -1032,7 +1035,7 @@ mod tests {
 
     #[test]
     fn test_parse() {
-        let args = super::parse(INPUT1.iter()).unwrap();
+        let args = super::parse(|| INPUT1.iter()).unwrap();
         assert!(args.is_relocatable());
         assert_eq!(
             args.inputs
diff --git a/libwild/src/save_dir.rs b/libwild/src/save_dir.rs
index 8126f0f..63e576a 100644
--- a/libwild/src/save_dir.rs
+++ b/libwild/src/save_dir.rs
@@ -1,7 +1,11 @@
 //! Support for saving inputs for later use.
 
+use crate::archive::ArchiveEntry;
+use crate::archive::ArchiveIterator;
+use crate::bail;
 use crate::error::Context as _;
 use crate::error::Result;
+use crate::file_kind::FileKind;
 use foldhash::HashMap as FoldHashMap;
 use foldhash::HashMapExt as _;
 use std::io::BufWriter;
@@ -9,10 +13,7 @@ use std::io::Write;
 use std::path::Path;
 use std::path::PathBuf;
 
-pub(crate) struct SaveDir {
-    dir: Option<PathBuf>,
-    copied_paths: FoldHashMap<String, String>,
-}
+pub(crate) struct SaveDir(Option<SaveDirState>);
 
 const SAVE_DIR_ENV: &str = "WILD_SAVE_DIR";
 const SAVE_BASE_ENV: &str = "WILD_SAVE_BASE";
@@ -20,60 +21,100 @@ const SKIP_LINKING_ENV: &str = "WILD_SAVE_SKIP_LINKING";
 
 const PRELUDE: &str = include_str!("save-dir-prelude.sh");
 
+struct SaveDirState {
+    dir: PathBuf,
+    copied_paths: FoldHashMap<String, PathBuf>,
+    args: Vec<String>,
+}
+
 impl SaveDir {
-    pub(crate) fn new() -> Result<Self> {
-        if let Ok(dir) = std::env::var(SAVE_DIR_ENV) {
-            let dir = PathBuf::from(dir);
-            if dir.exists() {
-                std::fs::remove_dir_all(&dir).with_context(|| {
-                    format!(
-                        "Failed to delete `{}`. If you're running multiple link commands \
-                         concurrently, try using {} instead.",
-                        dir.display(),
-                        SAVE_BASE_ENV
-                    )
-                })?;
-            }
-            std::fs::create_dir_all(&dir).with_context(|| {
-                format!(
-                    "Failed to create directory `{}` specified by {SAVE_DIR_ENV}",
-                    dir.display()
-                )
-            })?;
-            return Ok(Self::with_dir(Some(dir)));
+    pub(crate) fn new<F: Fn() -> I, S: AsRef<str>, I: Iterator<Item = S>>(
+        args: &F,
+    ) -> Result<Self> {
+        let Some(dir) = save_dir_from_env()? else {
+            return Ok(Self(None));
+        };
+
+        Ok(Self(Some(SaveDirState::new(
+            dir,
+            args().map(|s| s.as_ref().to_owned()).collect(),
+        ))))
+    }
+
+    pub(crate) fn finish(&self) -> Result {
+        if let Some(state) = self.0.as_ref() {
+            state.finish()?;
+        }
+        Ok(())
+    }
+
+    pub(crate) fn handle_file(&mut self, arg: &str) -> Result {
+        if let Some(state) = self.0.as_mut() {
+            state.handle_file(arg)?;
         }
-        if let Ok(dir) = std::env::var(SAVE_BASE_ENV) {
-            let dir = PathBuf::from(dir);
-            std::fs::create_dir_all(&dir).with_context(|| {
+        Ok(())
+    }
+}
+
+fn save_dir_from_env() -> Result<Option<PathBuf>> {
+    if let Ok(d) = std::env::var(SAVE_DIR_ENV) {
+        let dir = PathBuf::from(d);
+
+        if dir.exists() {
+            std::fs::remove_dir_all(&dir).with_context(|| {
                 format!(
-                    "Failed to create directory `{}` specified by {SAVE_BASE_ENV}",
-                    dir.display()
+                    "Failed to delete `{}`. If you're running multiple link commands \
+                         concurrently, try using {} instead.",
+                    dir.display(),
+                    SAVE_BASE_ENV
                 )
             })?;
-            let mut counter = 0;
-            loop {
-                let subdir = dir.join(counter.to_string());
-                if std::fs::create_dir(&subdir).is_ok() {
-                    return Ok(Self::with_dir(Some(subdir)));
-                }
-                counter += 1;
+        }
+
+        std::fs::create_dir_all(&dir).with_context(|| {
+            format!(
+                "Failed to create directory `{}` specified by {SAVE_DIR_ENV}",
+                dir.display()
+            )
+        })?;
+
+        return Ok(Some(dir));
+    }
+
+    if let Ok(d) = std::env::var(SAVE_BASE_ENV) {
+        let base = PathBuf::from(d);
+        std::fs::create_dir_all(&base).with_context(|| {
+            format!(
+                "Failed to create directory `{}` specified by {SAVE_BASE_ENV}",
+                base.display()
+            )
+        })?;
+
+        let mut counter = 0;
+
+        loop {
+            let subdir = base.join(counter.to_string());
+            if std::fs::create_dir(&subdir).is_ok() {
+                return Ok(Some(subdir));
             }
+            counter += 1;
         }
-        Ok(Self::with_dir(None))
+    } else {
+        Ok(None)
     }
+}
 
-    fn with_dir(dir: Option<PathBuf>) -> Self {
-        SaveDir {
+impl SaveDirState {
+    fn new(dir: PathBuf, args: Vec<String>) -> Self {
+        SaveDirState {
             dir,
             copied_paths: FoldHashMap::new(),
+            args,
         }
     }
 
-    pub(crate) fn finish(&self) -> Result {
-        let Some(dir) = self.dir.as_ref() else {
-            return Ok(());
-        };
-        let run_with_file = dir.join("run-with");
+    fn finish(&self) -> Result {
+        let run_with_file = self.dir.join("run-with");
         self.write_args_file(&run_with_file)
             .with_context(|| format!("Failed to write `{}`", run_with_file.display()))?;
         if std::env::var(SKIP_LINKING_ENV).is_ok() {
@@ -85,36 +126,38 @@ impl SaveDir {
     fn write_args_file(&self, run_file: &Path) -> Result {
         let mut file = std::fs::File::create(run_file)?;
         let mut out = BufWriter::new(&mut file);
-        let mut args = std::env::args();
         out.write_all(PRELUDE.as_bytes())?;
-        args.next();
-        self.write_args(args, &mut out)?;
+        self.write_args(&mut out)?;
         drop(out);
         crate::fs::make_executable(&file)?;
         Ok(())
     }
 
-    fn write_args(&self, args: std::env::Args, out: &mut BufWriter<&mut std::fs::File>) -> Result {
+    fn write_args(&self, out: &mut BufWriter<&mut std::fs::File>) -> Result {
         let mut is_output_file = false;
         let mut original_output_file = None;
-        for arg in args {
+
+        for arg in &self.args {
             out.write_all(b" \\\n  ")?;
+
             if is_output_file {
                 out.write_all(b"$OUT")?;
                 is_output_file = false;
                 original_output_file = Some(arg);
                 continue;
             }
+
             is_output_file = arg == "-o";
+
             let maybe_path = if let Some(eq_index) = arg.find('=') {
                 out.write_all(&arg.as_bytes()[..=eq_index])?;
                 &arg[eq_index + 1..]
             } else {
                 arg.as_str()
             };
+
             if let Some(copied) = self.copied_paths.get(maybe_path) {
-                out.write_all(b"$D/")?;
-                out.write_all(copied.as_bytes())?;
+                write_copied_file_arg(out, copied)?;
             } else {
                 for b in maybe_path.bytes() {
                     if b" $\\".contains(&b) {
@@ -124,6 +167,7 @@ impl SaveDir {
                 }
             }
         }
+
         if let Some(orig) = original_output_file {
             out.write_all(b"\n# Original output file: ")?;
             out.write_all(orig.as_bytes())?;
@@ -131,57 +175,107 @@ impl SaveDir {
         Ok(())
     }
 
-    pub(crate) fn handle_file(&mut self, arg: &str) -> Result {
-        let Some(dir) = self.dir.as_ref() else {
-            return Ok(());
-        };
+    /// To save disk space, we first attempt to hard link the file. If that fails, then just
+    /// copy it. Returns where we copied the file to.
+    fn copy_file(&self, source_path: &PathBuf) -> Result<PathBuf> {
+        let dest_path = unique_dest_path(&self.dir, source_path)?;
+
+        copy_file_to(source_path, &dest_path)?;
+
+        Ok(dest_path)
+    }
+
+    fn handle_file(&mut self, arg: &str) -> Result {
         let source_path = std::fs::canonicalize(Path::new(arg))?;
-        if source_path.is_dir() {
+
+        if source_path.is_dir() || self.copied_paths.contains_key(arg) {
             return Ok(());
         }
-        if let Some(dest_path) = unique_dest_path(dir, &source_path) {
-            // To save disk space, we first attempt to hard link the file. If that fails, then just
-            // copy it.
-            if std::fs::hard_link(&source_path, &dest_path).is_err() {
-                std::fs::copy(&source_path, &dest_path).with_context(|| {
-                    format!(
-                        "Failed to copy `{}` to `{}`",
-                        source_path.display(),
-                        dest_path.display()
-                    )
-                })?;
+
+        let file_bytes = std::fs::read(&source_path)?;
+        let file_kind = FileKind::identify_bytes(&file_bytes)?;
+        if file_kind == FileKind::ThinArchive {
+            self.handle_thin_archive(&source_path, &file_bytes)?;
+        }
+
+        let copied = self.copy_file(&source_path)?;
+
+        self.copied_paths.insert(arg.to_owned(), copied);
+
+        Ok(())
+    }
+
+    fn handle_thin_archive(&self, path: &Path, file_bytes: &[u8]) -> Result {
+        let parent_path = path.parent().unwrap();
+        let mut extended_filenames = None;
+
+        for entry in ArchiveIterator::from_archive_bytes(file_bytes)? {
+            match entry? {
+                ArchiveEntry::Filenames(t) => extended_filenames = Some(t),
+                ArchiveEntry::Thin(entry) => {
+                    let entry_path = entry.identifier(extended_filenames).as_path();
+                    if entry_path.is_absolute() {
+                        bail!(
+                            "Thin archive `{}` contained absolute path `{}`",
+                            path.display(),
+                            entry_path.display()
+                        );
+                    }
+                    let absolute_entry_path = parent_path.join(entry_path);
+
+                    if let Some(relative_dir) = entry_path.parent() {
+                        let directory = self.dir.join(relative_dir);
+                        std::fs::create_dir_all(&directory).with_context(|| {
+                            format!("Failed to create directory `{}`", directory.display())
+                        })?;
+                    }
+
+                    let dest_path = self.dir.join(entry_path);
+                    copy_file_to(&absolute_entry_path, &dest_path)?;
+                }
+                _ => {}
             }
-            self.copied_paths.insert(
-                arg.to_owned(),
-                dest_path
-                    .file_name()
-                    .unwrap()
-                    .to_str()
-                    .context("Path is not valid UTF-8")?
-                    .to_owned(),
-            );
         }
 
         Ok(())
     }
 }
 
+fn copy_file_to(source_path: &PathBuf, dest_path: &PathBuf) -> Result<(), crate::error::Error> {
+    if std::fs::hard_link(source_path, dest_path).is_err() {
+        std::fs::copy(source_path, dest_path).with_context(|| {
+            format!(
+                "Failed to copy `{}` to `{}`",
+                source_path.display(),
+                dest_path.display()
+            )
+        })?;
+    }
+
+    Ok(())
+}
+
+fn write_copied_file_arg(out: &mut BufWriter<&mut std::fs::File>, path: &Path) -> Result {
+    let file_name = path.file_name().context("Invalid copied file name")?;
+    out.write_all(b"$D/")?;
+    out.write_all(file_name.as_encoded_bytes())?;
+    Ok(())
+}
+
 /// Return the full path to a new filename in `dir` that if possible has the same filename as that
 /// of `path`.
-fn unique_dest_path(dir: &Path, path: &Path) -> Option<PathBuf> {
-    let file_name = path.file_name()?;
+fn unique_dest_path(dir: &Path, path: &Path) -> Result<PathBuf> {
+    let file_name = path.file_name().context("Missing file_name")?;
     let raw_dest_path = dir.join(file_name);
     let mut result = raw_dest_path.clone();
     let mut seq = 0;
     while result.exists() {
-        if let Some(n) = sequence_path(&raw_dest_path, seq) {
-            result = n;
-        } else {
-            return None;
-        }
+        result = sequence_path(&raw_dest_path, seq)
+            .with_context(|| format!("Invalid path `{}`", raw_dest_path.display()))?;
+
         seq += 1;
     }
-    Some(result)
+    Ok(result)
 }
 
 fn sequence_path(path: &Path, sequence: i32) -> Option<PathBuf> {
diff --git a/wild/src/main.rs b/wild/src/main.rs
index 3f7db18..2b7c23d 100644
--- a/wild/src/main.rs
+++ b/wild/src/main.rs
@@ -10,7 +10,7 @@ fn main() -> libwild::error::Result {
     #[cfg(feature = "dhat")]
     let _profiler = dhat::Profiler::new_heap();
 
-    let args = libwild::Args::parse(std::env::args().skip(1))?;
+    let args = libwild::Args::parse(|| std::env::args().skip(1))?;
 
     if args.should_fork() {
         // Safety: We haven't spawned any threads yet.


Test patch:
diff --git a/wild/tests/integration_tests.rs b/wild/tests/integration_tests.rs
index 012306c..bd19199 100644
--- a/wild/tests/integration_tests.rs
+++ b/wild/tests/integration_tests.rs
@@ -1672,7 +1672,7 @@ impl LinkCommand {
                 .context("Linker args must be valid utf-8")?;
 
             let linker = libwild::Linker::new();
-            let parsed_args = libwild::Args::parse(args.iter())?;
+            let parsed_args = libwild::Args::parse(|| args.iter())?;
 
             // This call is expected to error for all but the first call.
             let _ = libwild::setup_tracing(&parsed_args);