From af98f5b8ceadb1cfbd97da8777e4cfdf4822da5d Mon Sep 17 00:00:00 2001
From: David Knaack <davidkna@users.noreply.github.com>
Date: Tue, 28 Dec 2021 06:56:06 +0100
Subject: [PATCH] fix: set cwd for command execution (#3309)

---
 src/context.rs | 23 +++++++++++++++-----
 src/utils.rs   | 59 +++++++++++++++++++++-----------------------------
 2 files changed, 42 insertions(+), 40 deletions(-)

diff --git a/src/context.rs b/src/context.rs
index e0e8ce25..a67fdc06 100644
--- a/src/context.rs
+++ b/src/context.rs
@@ -1,7 +1,7 @@
 use crate::config::{RootModuleConfig, StarshipConfig};
 use crate::configs::StarshipRootConfig;
 use crate::module::Module;
-use crate::utils::{exec_cmd, CommandOutput};
+use crate::utils::{create_command, exec_timeout, CommandOutput};
 
 use crate::modules;
 use crate::utils::{self, home_dir};
@@ -316,16 +316,27 @@ impl<'a> Context<'a> {
         cmd: T,
         args: &[U],
     ) -> Option<CommandOutput> {
+        log::trace!(
+            "Executing command {:?} with args {:?} from context",
+            cmd,
+            args
+        );
         #[cfg(test)]
         {
             let command = crate::utils::display_command(&cmd, args);
-            if let Some(output) = self.cmd.get(command.as_str()) {
-                return output.clone();
+            if let Some(output) = self
+                .cmd
+                .get(command.as_str())
+                .cloned()
+                .or_else(|| crate::utils::mock_cmd(&cmd, args))
+            {
+                return output;
             }
         }
-        exec_cmd(
-            &cmd,
-            args,
+        let mut cmd = create_command(cmd).ok()?;
+        cmd.args(args).current_dir(&self.current_dir);
+        exec_timeout(
+            &mut cmd,
             Duration::from_millis(self.root_config.command_timeout),
         )
     }
diff --git a/src/utils.rs b/src/utils.rs
index 6dc27a87..8e05ac27 100644
--- a/src/utils.rs
+++ b/src/utils.rs
@@ -38,7 +38,7 @@ pub fn get_command_string_output(command: CommandOutput) -> String {
 /// This function also initialises std{err,out,in} to protect against processes changing the console mode
 pub fn create_command<T: AsRef<OsStr>>(binary_name: T) -> Result<Command> {
     let binary_name = binary_name.as_ref();
-    log::trace!("Creating Command struct with binary name {:?}", binary_name);
+    log::trace!("Creating Command for binary {:?}", binary_name);
 
     let full_path = match which::which(binary_name) {
         Ok(full_path) => {
@@ -85,23 +85,26 @@ pub fn display_command<T: AsRef<OsStr> + Debug, U: AsRef<OsStr> + Debug>(
 }
 
 /// Execute a command and return the output on stdout and stderr if successful
-#[cfg(not(test))]
 pub fn exec_cmd<T: AsRef<OsStr> + Debug, U: AsRef<OsStr> + Debug>(
     cmd: T,
     args: &[U],
     time_limit: Duration,
 ) -> Option<CommandOutput> {
+    log::trace!("Executing command {:?} with args {:?}", cmd, args);
+    #[cfg(test)]
+    if let Some(o) = mock_cmd(&cmd, args) {
+        return o;
+    }
     internal_exec_cmd(cmd, args, time_limit)
 }
 
 #[cfg(test)]
-pub fn exec_cmd<T: AsRef<OsStr> + Debug, U: AsRef<OsStr> + Debug>(
+pub fn mock_cmd<T: AsRef<OsStr> + Debug, U: AsRef<OsStr> + Debug>(
     cmd: T,
     args: &[U],
-    time_limit: Duration,
-) -> Option<CommandOutput> {
+) -> Option<Option<CommandOutput>> {
     let command = display_command(&cmd, args);
-    match command.as_str() {
+    let out = match command.as_str() {
         "cobc -version" => Some(CommandOutput {
             stdout: String::from("\
 cobc (GnuCOBOL) 3.1.2.0
@@ -313,9 +316,9 @@ CMake suite maintained and supported by Kitware (kitware.com/cmake).\n",
             stdout: String::from("22.1.3\n"),
             stderr: String::default(),
         }),
-        // If we don't have a mocked command fall back to executing the command
-        _ => internal_exec_cmd(&cmd, args, time_limit),
-    }
+        _ => return None,
+    };
+    Some(out)
 }
 
 /// Wraps ANSI color escape sequences in the shell-appropriate wrappers.
@@ -376,36 +379,20 @@ fn internal_exec_cmd<T: AsRef<OsStr> + Debug, U: AsRef<OsStr> + Debug>(
     args: &[U],
     time_limit: Duration,
 ) -> Option<CommandOutput> {
-    log::trace!("Executing command {:?} with args {:?}", cmd, args);
-
-    let full_path = match which::which(&cmd) {
-        Ok(full_path) => {
-            log::trace!("Using {:?} as {:?}", full_path, cmd);
-            full_path
-        }
-        Err(error) => {
-            log::trace!("Unable to find {:?} in PATH, {:?}", cmd, error);
-            return None;
-        }
-    };
+    let mut cmd = create_command(cmd).ok()?;
+    cmd.args(args);
+    exec_timeout(&mut cmd, time_limit)
+}
 
+pub fn exec_timeout(cmd: &mut Command, time_limit: Duration) -> Option<CommandOutput> {
     let start = Instant::now();
-
-    #[allow(clippy::disallowed_method)]
-    let process = match Command::new(full_path)
-        .args(args)
-        .stderr(Stdio::piped())
-        .stdout(Stdio::piped())
-        .stdin(Stdio::null())
-        .spawn()
-    {
+    let process = match cmd.spawn() {
         Ok(process) => process,
         Err(error) => {
-            log::info!("Unable to run {:?}, {:?}", cmd, error);
+            log::info!("Unable to run {:?}, {:?}", cmd.get_program(), error);
             return None;
         }
     };
-
     match process.with_output_timeout(time_limit).terminating().wait() {
         Ok(Some(output)) => {
             let stdout_string = match String::from_utf8(output.stdout) {
@@ -441,12 +428,16 @@ fn internal_exec_cmd<T: AsRef<OsStr> + Debug, U: AsRef<OsStr> + Debug>(
             })
         }
         Ok(None) => {
-            log::warn!("Executing command {:?} timed out.", cmd);
+            log::warn!("Executing command {:?} timed out.", cmd.get_program());
             log::warn!("You can set command_timeout in your config to a higher value to allow longer-running commands to keep executing.");
             None
         }
         Err(error) => {
-            log::info!("Executing command {:?} failed by: {:?}", cmd, error);
+            log::info!(
+                "Executing command {:?} failed by: {:?}",
+                cmd.get_program(),
+                error
+            );
             None
         }
     }