feat: properly handle sudo with a command builder

Cargo.toml

@@ -62,6 +62,7 @@ addr2line = "0.25"
 gimli = "0.32"
 open = "5.3.2"
 tabled = "0.17"
+shell-words = "1.1.0"
 
 [target.'cfg(target_os = "linux")'.dependencies]
 procfs = "0.17.0"

src/run/runner/helpers/apt.rs

@@ -96,10 +96,10 @@ pub fn install(system_info: &SystemInfo, packages: &[&str]) -> Result<()> {
 
     info!("Installing packages: {}", packages.join(", "));
 
-    run_with_sudo(&["apt-get", "update"])?;
-    let mut install_cmd = vec!["apt-get", "install", "-y", "--allow-downgrades"];
-    install_cmd.extend_from_slice(packages);
-    run_with_sudo(&install_cmd)?;
+    run_with_sudo("apt-get", ["update"])?;
+    let mut install_argv = vec!["install", "-y", "--allow-downgrades"];
+    install_argv.extend_from_slice(packages);
+    run_with_sudo("apt-get", &install_argv)?;
 
     debug!("Packages installed successfully");
     Ok(())
@@ -155,9 +155,7 @@ fn restore_from_cache(system_info: &SystemInfo, cache_dir: &Path) -> Result<()>
         .to_str()
         .ok_or_else(|| anyhow!("Invalid cache directory path"))?;
 
-    let copy_cmd = format!("cp -r {cache_dir_str}/* /");
-
-    run_with_sudo(&["bash", "-c", &copy_cmd])?;
+    run_with_sudo("cp", ["-r", &format!("{cache_dir_str}/*"), "/"])?;
 
     debug!("Cache restored successfully");
     Ok(())

src/run/runner/helpers/command.rs

@@ -0,0 +1,188 @@
+use std::{
+    collections::BTreeMap,
+    ffi::{OsStr, OsString},
+    process::Command,
+};
+
+#[derive(Debug)]
+pub struct CommandBuilder {
+    program: OsString,
+    argv: Vec<OsString>,
+    envs: BTreeMap<OsString, OsString>,
+    cwd: Option<OsString>,
+}
+
+impl CommandBuilder {
+    pub fn new<S: AsRef<OsStr>>(program: S) -> Self {
+        Self {
+            program: program.as_ref().to_owned(),
+            argv: Vec::new(),
+            envs: BTreeMap::new(),
+            cwd: None,
+        }
+    }
+
+    pub fn build(self) -> Command {
+        let mut command = Command::new(&self.program);
+        command.args(&self.argv);
+        command.envs(&self.envs);
+        if let Some(cwd) = self.cwd {
+            command.current_dir(cwd);
+        }
+        command
+    }
+
+    pub fn arg<S: AsRef<OsStr>>(&mut self, arg: S) -> &mut Self {
+        self.argv.push(arg.as_ref().to_owned());
+        self
+    }
+
+    pub fn args<I, S>(&mut self, args: I) -> &mut Self
+    where
+        I: IntoIterator<Item = S>,
+        S: AsRef<OsStr>,
+    {
+        for arg in args {
+            self.arg(arg.as_ref());
+        }
+        self
+    }
+
+    pub fn current_dir<D>(&mut self, dir: D)
+    where
+        D: AsRef<OsStr>,
+    {
+        self.cwd = Some(dir.as_ref().to_owned());
+    }
+
+    pub fn wrap<S, I, T>(&mut self, wrapper: S, wrapper_args: I) -> &mut Self
+    where
+        S: AsRef<OsStr>,
+        I: IntoIterator<Item = T>,
+        T: AsRef<OsStr>,
+    {
+        let mut new_argv = Vec::new();
+
+        // Add wrapper arguments first
+        for arg in wrapper_args {
+            new_argv.push(arg.as_ref().to_owned());
+        }
+
+        // Add the current program
+        new_argv.push(self.program.clone());
+
+        // Add the current arguments
+        new_argv.extend(self.argv.iter().cloned());
+
+        // Update program to wrapper and argv to the new argument list
+        self.program = wrapper.as_ref().to_owned();
+        self.argv = new_argv;
+        self
+    }
+
+    pub fn wrap_with(&mut self, wrapper_cmd: CommandBuilder) -> &mut Self {
+        let mut new_argv = Vec::new();
+
+        // Add wrapper command arguments first
+        new_argv.extend(wrapper_cmd.argv.iter().cloned());
+
+        // Add the current program
+        new_argv.push(self.program.clone());
+
+        // Add the current arguments
+        new_argv.extend(self.argv.iter().cloned());
+
+        // Update program to wrapper and argv to the new argument list
+        self.program = wrapper_cmd.program;
+        self.argv = new_argv;
+        // Update cwd if wrapper has it set
+        self.cwd = wrapper_cmd.cwd.or(self.cwd.take());
+        // Merge environment variables, with wrapper's envs taking precedence
+        self.envs.extend(wrapper_cmd.envs);
+        self
+    }
+
+    /// Returns the command line as a string for debugging/testing purposes
+    pub fn as_command_line(&self) -> String {
+        let mut parts: Vec<String> = vec![self.program.to_string_lossy().into_owned()];
+        parts.extend(
+            self.argv
+                .iter()
+                .map(|arg| arg.to_string_lossy().into_owned()),
+        );
+        shell_words::join(parts)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_wrap_with_args() {
+        let mut builder = CommandBuilder::new("ls");
+        builder.arg("-la").wrap("sudo", ["-n"]);
+        assert_eq!(builder.as_command_line(), "sudo -n ls -la");
+    }
+
+    #[test]
+    fn test_wrap_without_args() {
+        let mut builder = CommandBuilder::new("ls");
+        builder.arg("-la").wrap("time", [] as [&str; 0]);
+        assert_eq!(builder.as_command_line(), "time ls -la");
+    }
+
+    #[test]
+    fn test_multiple_wraps() {
+        let mut builder = CommandBuilder::new("valgrind");
+        builder
+            .arg("my-program")
+            .wrap("setarch", ["x86_64", "-R"])
+            .wrap("sudo", ["-n"]);
+        assert_eq!(
+            builder.as_command_line(),
+            "sudo -n setarch x86_64 -R valgrind my-program"
+        );
+    }
+
+    #[test]
+    fn test_wrap_with_spaces() {
+        let mut builder = CommandBuilder::new("echo");
+        builder.arg("hello world").wrap("bash", ["-c"]);
+        assert_eq!(builder.as_command_line(), "bash -c echo 'hello world'");
+    }
+
+    #[test]
+    fn test_wrap_and_build() {
+        let mut builder = CommandBuilder::new("ls");
+        builder.arg("-la").wrap("sudo", ["-n"]);
+
+        let cmd = builder.build();
+        assert_eq!(cmd.get_program(), "sudo");
+
+        let args: Vec<_> = cmd.get_args().collect();
+        assert_eq!(args, vec!["-n", "ls", "-la"]);
+    }
+
+    #[test]
+    fn test_wrap_with_builder() {
+        let mut wrapper = CommandBuilder::new("sudo");
+        wrapper.arg("-n");
+
+        let mut builder = CommandBuilder::new("ls");
+        builder.arg("-la").wrap_with(wrapper);
+
+        assert_eq!(builder.as_command_line(), "sudo -n ls -la");
+    }
+
+    #[test]
+    fn test_wrap_with_preserves_env() {
+        let mut wrapper = CommandBuilder::new("env");
+        wrapper.arg("FOO=bar");
+
+        let mut builder = CommandBuilder::new("ls");
+        builder.arg("-la").wrap_with(wrapper);
+
+        assert_eq!(builder.as_command_line(), "env 'FOO=bar' ls -la");
+    }
+}

src/run/runner/helpers/mod.rs

@@ -1,4 +1,5 @@
 pub mod apt;
+pub mod command;
 pub mod env;
 pub mod get_bench_command;
 pub mod introspected_golang;

src/run/runner/helpers/run_with_sudo.rs

@@ -1,5 +1,8 @@
-use crate::{local_logger::suspend_progress_bar, prelude::*};
+use crate::{
+    local_logger::suspend_progress_bar, prelude::*, run::runner::helpers::command::CommandBuilder,
+};
 use std::{
+    ffi::OsStr,
     io::IsTerminal,
     process::{Command, Stdio},
 };
@@ -35,9 +38,7 @@ fn validate_sudo_access() -> Result<()> {
 
     if needs_password {
         suspend_progress_bar(|| {
-            info!(
-                "Sudo privileges are required to continue. Please enter your password if prompted."
-            );
+            info!("Sudo privileges are required to continue. Please enter your password.");
 
             // Validate and cache sudo credentials
             let auth_status = Command::new("sudo")
@@ -57,47 +58,47 @@ fn validate_sudo_access() -> Result<()> {
     Ok(())
 }
 
-/// Creates the base sudo command after validating sudo access
-pub fn validated_sudo_command() -> Result<Command> {
-    validate_sudo_access()?;
-    let mut cmd = Command::new("sudo");
-    // Password prompt should not appear here since it has already been validated
-    cmd.arg("--non-interactive");
-    Ok(cmd)
-}
-
-/// Build a command wrapped with sudo if possible
-pub fn build_command_with_sudo(command_args: &[&str]) -> Result<Command> {
-    let command_str = command_args.join(" ");
+/// Wrap with sudo if not running as root
+pub fn wrap_with_sudo(mut cmd_builder: CommandBuilder) -> Result<CommandBuilder> {
     if is_root_user() {
-        debug!("Running command without sudo: {command_str}");
-        let mut c = Command::new(command_args[0]);
-        c.args(&command_args[1..]);
-        Ok(c)
+        Ok(cmd_builder)
     } else if is_sudo_available() {
-        debug!("Sudo is required for command: {command_str}");
-        let mut c = validated_sudo_command()?;
-        c.args(command_args);
-        Ok(c)
+        debug!("Wrapping with sudo: {}", cmd_builder.as_command_line());
+        validate_sudo_access()?;
+        cmd_builder.wrap(
+            "sudo",
+            // Password prompt should not appear here since it has already been validated
+            ["--non-interactive"],
+        );
+        Ok(cmd_builder)
     } else {
-        bail!("Sudo is not available to run the command: {command_str}");
+        bail!(
+            "Sudo is not available to run the command: {}",
+            cmd_builder.as_command_line()
+        );
     }
 }
 
 /// Run a command with sudo after validating sudo access
-pub fn run_with_sudo(command_args: &[&str]) -> Result<()> {
-    let command_str = command_args.join(" ");
-    debug!("Running command with sudo: {command_str}");
-    let mut cmd = build_command_with_sudo(command_args)?;
+pub fn run_with_sudo<S, I, T>(program: S, argv: I) -> Result<()>
+where
+    S: AsRef<OsStr>,
+    I: IntoIterator<Item = T>,
+    T: AsRef<OsStr>,
+{
+    let mut builder = CommandBuilder::new(program);
+    builder.args(argv);
+    debug!("Running command with sudo: {}", builder.as_command_line());
+    let mut cmd = wrap_with_sudo(builder)?.build();
     let output = cmd
         .stdout(Stdio::piped())
         .output()
-        .map_err(|_| anyhow!("Failed to execute command with sudo: {command_str}"))?;
+        .map_err(|_| anyhow!("Failed to execute command with sudo: {cmd:?}"))?;
 
     if !output.status.success() {
         info!("stdout: {}", String::from_utf8_lossy(&output.stdout));
         error!("stderr: {}", String::from_utf8_lossy(&output.stderr));
-        bail!("Failed to execute command with sudo: {command_str}");
+        bail!("Failed to execute command with sudo: {cmd:?}");
     }
 
     Ok(())