From 500f5c1f127697bfbe683e0278f6dd8be32e0bb5 Mon Sep 17 00:00:00 2001
From: Samuel Moelius <samuel.moelius@trailofbits.com>
Date: Fri, 5 Apr 2024 08:01:09 -0400
Subject: [PATCH] Split args before performing replacements

---
 README.md          |  2 ++
 src/util/common.rs | 65 +++++++++++++++++++++++++++++++++++++++++++---
 2 files changed, 64 insertions(+), 3 deletions(-)

diff --git a/README.md b/README.md
index 1c3af37..c90a286 100644
--- a/README.md
+++ b/README.md
@@ -44,6 +44,8 @@ Note that the above environment variables are read **when the build script is li
 - `{VAR}` is replaced with the value of environment variable `VAR`.
 - `{{` is replaced with `{`.
 - `}}` is replaced with `}`.
+- `\` followed by a whitespace character is replaced with that whitespace character.
+- `\\` is replaced with `\`.
 
 ## How `build-wrap` works
 
diff --git a/src/util/common.rs b/src/util/common.rs
index 16cfbe7..5ca951f 100644
--- a/src/util/common.rs
+++ b/src/util/common.rs
@@ -61,15 +61,18 @@ fn unpack_and_exec(bytes: &[u8]) -> Result<()> {
     // They will cause the wrapped build script to be rerun, however.
     let cmd =
         option_env!("BUILD_WRAP_CMD").ok_or_else(|| anyhow!("`BUILD_WRAP_CMD` is undefined"))?;
-    let expanded_cmd = __expand_cmd(cmd, &temp_path)?;
-    let expanded_args = expanded_cmd.split_ascii_whitespace().collect::<Vec<_>>();
+    let args = split_escaped(cmd)?;
+    let expanded_args = args
+        .into_iter()
+        .map(|arg| __expand_cmd(&arg, &temp_path))
+        .collect::<Result<Vec<_>>>()?;
     eprintln!("expanded `BUILD_WRAP_CMD`: {:#?}", &expanded_args);
     ensure!(
         !expanded_args.is_empty(),
         "expanded `BUILD_WRAP_CMD` is empty or all whitespace"
     );
 
-    let mut command = Command::new(expanded_args[0]);
+    let mut command = Command::new(&expanded_args[0]);
     command.args(&expanded_args[1..]);
     let _: Output = exec(command, true)?;
 
@@ -91,6 +94,62 @@ impl<T: AsRef<Path>> ToUtf8 for T {
     }
 }
 
+fn split_escaped(mut s: &str) -> Result<Vec<String>> {
+    let mut v = vec![String::new()];
+
+    while let Some(i) = s.find(|c: char| c.is_ascii_whitespace() || c == '\\') {
+        debug_assert!(!v.is_empty());
+        // smoelius: Only the last string in `v` can be empty.
+        debug_assert!(v
+            .iter()
+            .position(String::is_empty)
+            .map_or(true, |i| i == v.len() - 1));
+
+        let c = s.as_bytes()[i];
+
+        v.last_mut().unwrap().push_str(&s[..i]);
+
+        s = &s[i + 1..];
+
+        // smoelius: `i` shouldn't be needed anymore.
+        #[allow(unused_variables, clippy::let_unit_value)]
+        let i = ();
+
+        if c.is_ascii_whitespace() {
+            if !v.last().unwrap().is_empty() {
+                v.push(String::new());
+            }
+            continue;
+        }
+
+        // smoelius: If the previous `if` fails, then `c` must be a backslash.
+        if !s.is_empty() {
+            let c = s.as_bytes()[0];
+            // smoelius: Verify that `c` is a legally escapable character before subslicing `s`.
+            ensure!(
+                c.is_ascii_whitespace() || c == b'\\',
+                "illegally escaped character"
+            );
+            s = &s[1..];
+            v.last_mut().unwrap().push(c as char);
+            continue;
+        }
+
+        bail!("trailing backslash");
+    }
+
+    // smoelius: Push whatever is left.
+    v.last_mut().unwrap().push_str(s);
+
+    if v.last().unwrap().is_empty() {
+        v.pop();
+    }
+
+    debug_assert!(!v.iter().any(String::is_empty));
+
+    Ok(v)
+}
+
 // smoelius: `__expand_cmd` is `pub` simply to allow testing it in an integration test. It is not
 // meant to be used outside of this module.
 pub fn __expand_cmd(mut cmd: &str, build_script_path: &Path) -> Result<String> {