From 1ca2af83295d8642d37de07fa2c5fa12d05f668a Mon Sep 17 00:00:00 2001
From: Josh Stone <cuviper@gmail.com>
Date: Mon, 25 Mar 2024 11:16:56 -0700
Subject: [PATCH] Support `RUSTC_WRAPPER` and `RUSTC_WORKSPACE_WRAPPER`

We need Cargo 1.55 to know whether the workspace wrapper is applicable,
and even `RUSTC_WRAPPER` is incomplete before then because it could also
be set in config files. The environment does take precedence when set
though, so it should hopefully be fine to use it on earlier versions.
The status quo remains to miss config wrappers before Rust 1.55.
---
 src/lib.rs        | 41 ++++++++++++++++++++++++++++++++++++-
 tests/wrappers.rs | 51 +++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 91 insertions(+), 1 deletion(-)
 create mode 100644 tests/wrappers.rs

diff --git a/src/lib.rs b/src/lib.rs
index f519a1a..69e8b6a 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -83,6 +83,8 @@ mod tests;
 pub struct AutoCfg {
     out_dir: PathBuf,
     rustc: PathBuf,
+    rustc_wrapper: Option<PathBuf>,
+    rustc_workspace_wrapper: Option<PathBuf>,
     rustc_version: Version,
     target: Option<OsString>,
     no_std: bool,
@@ -168,6 +170,8 @@ impl AutoCfg {
 
         let mut ac = AutoCfg {
             rustflags: rustflags(&target, &dir),
+            rustc_wrapper: get_rustc_wrapper(false),
+            rustc_workspace_wrapper: get_rustc_wrapper(true),
             out_dir: dir,
             rustc: rustc,
             rustc_version: rustc_version,
@@ -234,7 +238,18 @@ impl AutoCfg {
         static ID: AtomicUsize = ATOMIC_USIZE_INIT;
 
         let id = ID.fetch_add(1, Ordering::Relaxed);
-        let mut command = Command::new(&self.rustc);
+
+        // Build the command with possible wrappers.
+        let mut rustc = self
+            .rustc_wrapper
+            .iter()
+            .chain(self.rustc_workspace_wrapper.iter())
+            .chain(Some(&self.rustc));
+        let mut command = Command::new(rustc.next().unwrap());
+        for arg in rustc {
+            command.arg(arg);
+        }
+
         command
             .arg("--crate-name")
             .arg(format!("probe{}", id))
@@ -478,3 +493,27 @@ fn rustflags(target: &Option<OsString>, dir: &Path) -> Vec<String> {
 
     Vec::new()
 }
+
+fn get_rustc_wrapper(workspace: bool) -> Option<PathBuf> {
+    // We didn't really know whether the workspace wrapper is applicable until Cargo started
+    // deliberately setting or unsetting it in rust-lang/cargo#9601. We'll use the encoded
+    // rustflags as a proxy for that change for now, but we could instead check version 1.55.
+    if workspace && env::var_os("CARGO_ENCODED_RUSTFLAGS").is_none() {
+        return None;
+    }
+
+    let name = if workspace {
+        "RUSTC_WORKSPACE_WRAPPER"
+    } else {
+        "RUSTC_WRAPPER"
+    };
+
+    if let Some(wrapper) = env::var_os(name) {
+        // NB: `OsStr` didn't get `len` or `is_empty` until 1.9.
+        if wrapper != OsString::new() {
+            return Some(wrapper.into());
+        }
+    }
+
+    None
+}
diff --git a/tests/wrappers.rs b/tests/wrappers.rs
new file mode 100644
index 0000000..560cecd
--- /dev/null
+++ b/tests/wrappers.rs
@@ -0,0 +1,51 @@
+extern crate autocfg;
+
+use std::env;
+
+/// Tests that autocfg uses the RUSTC_WRAPPER and/or RUSTC_WORKSPACE_WRAPPER
+/// environment variables when running rustc.
+#[test]
+#[cfg(unix)] // we're using system binaries as wrappers
+fn test_wrappers() {
+    fn set(name: &str, value: Option<bool>) {
+        match value {
+            Some(true) => env::set_var(name, "/usr/bin/env"),
+            Some(false) => env::set_var(name, "/bin/false"),
+            None => env::remove_var(name),
+        }
+    }
+
+    // Use the same path as this test binary.
+    let dir = env::current_exe().unwrap().parent().unwrap().to_path_buf();
+    env::set_var("OUT_DIR", &format!("{}", dir.display()));
+
+    // This is used as a heuristic to detect rust-lang/cargo#9601.
+    env::set_var("CARGO_ENCODED_RUSTFLAGS", "");
+
+    // No wrapper, a good pass-through wrapper, and a bad wrapper.
+    let variants = [None, Some(true), Some(false)];
+
+    for &workspace in &variants {
+        for &rustc in &variants {
+            set("RUSTC_WRAPPER", rustc);
+            set("RUSTC_WORKSPACE_WRAPPER", workspace);
+
+            let ac = autocfg::AutoCfg::new().unwrap();
+            if rustc == Some(false) || workspace == Some(false) {
+                // Everything should fail with bad wrappers.
+                assert!(!ac.probe_type("usize"));
+            } else {
+                // Try known good and bad types for the wrapped rustc.
+                assert!(ac.probe_type("usize"));
+                assert!(!ac.probe_type("mesize"));
+            }
+        }
+    }
+
+    // Finally, make sure that `RUSTC_WRAPPER` is applied outermost
+    // by using something that doesn't pass through at all.
+    env::set_var("RUSTC_WRAPPER", "/bin/true");
+    env::set_var("RUSTC_WORKSPACE_WRAPPER", "/bin/false");
+    let ac = autocfg::AutoCfg::new().unwrap();
+    assert!(ac.probe_type("mesize")); // anything goes!
+}