From a4a2adb0f82ae31373e161c4820c79031823a165 Mon Sep 17 00:00:00 2001
From: David Knaack <davidkna@users.noreply.github.com>
Date: Sat, 3 Jul 2021 23:30:27 +0200
Subject: [PATCH] fix(init): improve starship path escaping (#2848)

---
 src/init/mod.rs | 52 ++++++++++++++++++++++++++++++++++++++-----------
 1 file changed, 41 insertions(+), 11 deletions(-)

diff --git a/src/init/mod.rs b/src/init/mod.rs
index 190e67a2..ed9c6153 100644
--- a/src/init/mod.rs
+++ b/src/init/mod.rs
@@ -33,8 +33,17 @@ impl StarshipPath {
             .ok_or_else(|| io::Error::new(io::ErrorKind::Other, "can't convert to str"))?;
         Ok(current_exe)
     }
+
+    /// Returns POSIX quoted path to starship binary
     fn sprint(&self) -> io::Result<String> {
-        self.str_path().map(|s| s.replace("\"", "\"'\"'\""))
+        self.str_path().map(|p| shell_words::quote(p).into_owned())
+    }
+
+    /// PowerShell specific path escaping
+    fn sprint_pwsh(&self) -> io::Result<String> {
+        self.str_path()
+            .map(|s| s.replace("'", "''"))
+            .map(|s| format!("'{}'", s))
     }
     fn sprint_posix(&self) -> io::Result<String> {
         // On non-Windows platform, return directly.
@@ -71,7 +80,7 @@ impl StarshipPath {
                 str_path
             }
         };
-        Ok(posix_path.replace("\"", "\"'\"'\""))
+        Ok(shell_words::quote(posix_path).into_owned())
     }
 }
 
@@ -138,25 +147,25 @@ pub fn init_stub(shell_name: &str) -> io::Result<()> {
             starship.sprint_posix()?
         ),
         "zsh" => print!(
-            r#"source <("{}" init zsh --print-full-init)"#,
+            r#"source <({} init zsh --print-full-init)"#,
             starship.sprint_posix()?
         ),
         "fish" => print!(
             // Fish does process substitution with pipes and psub instead of bash syntax
-            r#"source ("{}" init fish --print-full-init | psub)"#,
+            r#"source ({} init fish --print-full-init | psub)"#,
             starship.sprint_posix()?
         ),
         "powershell" => print!(
-            r#"Invoke-Expression (& "{}" init powershell --print-full-init | Out-String)"#,
-            starship.sprint()?
+            r#"Invoke-Expression (& {} init powershell --print-full-init | Out-String)"#,
+            starship.sprint_pwsh()?
         ),
         "ion" => print!("eval $({} init ion --print-full-init)", starship.sprint()?),
         "elvish" => print!(
-            r#"eval ("{}" init elvish --print-full-init | slurp)"#,
+            r#"eval ({} init elvish --print-full-init | slurp)"#,
             starship.sprint_posix()?
         ),
         "tcsh" => print!(
-            r#"eval `("{}" init tcsh --print-full-init)`"#,
+            r#"eval `({} init tcsh --print-full-init)`"#,
             starship.sprint_posix()?
         ),
         _ => {
@@ -190,7 +199,7 @@ pub fn init_main(shell_name: &str) -> io::Result<()> {
         "bash" => print_script(BASH_INIT, &starship_path.sprint_posix()?),
         "zsh" => print_script(ZSH_INIT, &starship_path.sprint_posix()?),
         "fish" => print_script(FISH_INIT, &starship_path.sprint_posix()?),
-        "powershell" => print_script(PWSH_INIT, &starship_path.sprint()?),
+        "powershell" => print_script(PWSH_INIT, &starship_path.sprint_pwsh()?),
         "ion" => print_script(ION_INIT, &starship_path.sprint()?),
         "elvish" => print_script(ELVISH_INIT, &starship_path.sprint_posix()?),
         "tcsh" => print_script(TCSH_INIT, &starship_path.sprint_posix()?),
@@ -206,8 +215,7 @@ pub fn init_main(shell_name: &str) -> io::Result<()> {
 }
 
 fn print_script(script: &str, path: &str) {
-    let starship_path_string = format!("\"{}\"", path);
-    let script = script.replace("::STARSHIP::", &starship_path_string);
+    let script = script.replace("::STARSHIP::", path);
     print!("{}", script);
 }
 
@@ -239,3 +247,25 @@ const ION_INIT: &str = include_str!("starship.ion");
 const ELVISH_INIT: &str = include_str!("starship.elv");
 
 const TCSH_INIT: &str = include_str!("starship.tcsh");
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    #[test]
+    fn escape_pwsh() -> io::Result<()> {
+        let starship_path = StarshipPath {
+            native_path: PathBuf::from(r"C:\starship.exe"),
+        };
+        assert_eq!(starship_path.sprint_pwsh()?, r"'C:\starship.exe'");
+        Ok(())
+    }
+
+    #[test]
+    fn escape_tick_pwsh() -> io::Result<()> {
+        let starship_path = StarshipPath {
+            native_path: PathBuf::from(r"C:\'starship.exe"),
+        };
+        assert_eq!(starship_path.sprint_pwsh()?, r"'C:\''starship.exe'");
+        Ok(())
+    }
+}