Improve root escalation system

Cargo.toml

@@ -43,6 +43,7 @@ serde = { version = "1.0.152", features = ["derive"] }
 serde_json = "1.0.94"
 sha1 = "0.10.5"
 sha2 = "0.10.6"
+shell-words = "1.1.0"
 static_cell = "1.0.0"
 thiserror = "1.0.38"
 tokio = { version = "1.25.0", features = ["full"] }

src/burn/handle.rs

@@ -5,7 +5,6 @@ use rand::distributions::Alphanumeric;
 use rand::distributions::DistString;
 use std::fs::remove_file;
 use std::path::PathBuf;
-use std::process::Command;
 use std::{env, pin::Pin};
 use tokio::io::BufReader;
 use tokio_util::compat::FuturesAsyncReadCompatExt;
@@ -20,10 +19,9 @@ use tokio::{
     process::Child,
 };
 
-use tokio::process::Command as AsyncCommand;
-
 use crate::burn::ipc::read_msg_async;
 use crate::escalation::run_escalate;
+use crate::escalation::Command;
 
 use super::ipc::InitialInfo;
 use super::{
@@ -54,14 +52,17 @@ impl Handle {
 
         let mut socket = ChildSocket::new()?;
 
-        let mut cmd = Command::new(proc);
-        cmd.arg(args).arg(&socket.socket_name).env(BURN_ENV, "1");
+        let cmd = Command {
+            proc: proc.to_string_lossy(),
+            envs: vec![(BURN_ENV.into(), "1".into())],
+            args: vec![args.into(), socket.socket_name.to_string_lossy().into()],
+        };
 
         debug!("Starting child process with command: {:?}", cmd);
         let child = if escalate {
             run_escalate(cmd).await?
         } else {
-            AsyncCommand::from(cmd).spawn()?
+            tokio::process::Command::from(cmd).spawn()?
         };
 
         debug!("Waiting for pipe to be opened...");

src/escalation/darwin.rs

@@ -1,16 +1,12 @@
-use std::process::Command;
+use super::unix::{Command, EscalationMethod};
 
-use tokio::process::Command as AsyncCommand;
-
-use super::unix::EscalationMethod;
-
-pub async fn wrap_osascript_escalation(raw: Command) -> anyhow::Result<tokio::process::Child> {
+pub async fn wrap_osascript_escalation(raw: Command<'_>) -> anyhow::Result<tokio::process::Child> {
     for _ in 0..3 {
         // User-friendly thing that lets you use touch ID if you wanted.
         // https://apple.stackexchange.com/questions/23494/what-option-should-i-give-the-sudo-command-to-have-the-password-asked-through-a
         // We loop because your finger might not be recognized sometimes.
 
-        let result = AsyncCommand::new("osascript")
+        let result = tokio::process::Command::new("osascript")
             .arg("-e")
             .arg("do shell script \"mkdir -p /var/db/sudo/$USER; touch /var/db/sudo/$USER\" with administrator privileges")
             .kill_on_drop(true)
@@ -23,6 +19,6 @@ pub async fn wrap_osascript_escalation(raw: Command) -> anyhow::Result<tokio::pr
         }
     }
 
-    let cmd: Command = EscalationMethod::Sudo.wrap_command(raw).into();
-    Ok(AsyncCommand::from(cmd).kill_on_drop(true).spawn()?)
+    let cmd: tokio::process::Command = EscalationMethod::Sudo.wrap_command(raw).into();
+    Ok(cmd.kill_on_drop(true).spawn()?)
 }

src/escalation/mod.rs

@@ -3,7 +3,7 @@
 mod darwin;
 mod unix;
 
-use std::process::Command;
+pub use self::unix::Command;
 use tokio::process::Command as AsyncCommand;
 
 #[derive(Debug, thiserror::Error)]
@@ -16,16 +16,16 @@ pub enum Error {
 }
 
 #[cfg(target_os = "linux")]
-pub async fn run_escalate(cmd: Command) -> anyhow::Result<tokio::process::Child> {
+pub async fn run_escalate(cmd: Command<'_>) -> anyhow::Result<tokio::process::Child> {
     use self::unix::EscalationMethod;
 
-    let mut cmd: AsyncCommand = EscalationMethod::detect()?.wrap_command(cmd).into();
+    let mut cmd: tokio::process::Command = EscalationMethod::detect()?.wrap_command(&cmd).into();
     cmd.kill_on_drop(true);
     Ok(cmd.spawn()?)
 }
 
 #[cfg(target_os = "macos")]
-pub async fn run_escalate(cmd: Command) -> anyhow::Result<tokio::process::Child> {
+pub async fn run_escalate(cmd: Command<'_>) -> anyhow::Result<tokio::process::Child> {
     use self::darwin::wrap_osascript_escalation;
 
     wrap_osascript_escalation(cmd).await

src/escalation/unix.rs

@@ -1,5 +1,7 @@
+use std::borrow::Cow;
+
 use itertools::Itertools;
-use std::{ffi::OsString, process::Command};
+use shell_words::{join, quote};
 use which::which;
 
 use super::Error;
@@ -14,6 +16,14 @@ pub enum EscalationMethod {
     Su,
 }
 
+/// Command components, backed by copy-on-write storage.
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub struct Command<'a> {
+    pub envs: Vec<(Cow<'a, str>, Cow<'a, str>)>,
+    pub proc: Cow<'a, str>,
+    pub args: Vec<Cow<'a, str>>,
+}
+
 impl EscalationMethod {
     const ALL: [EscalationMethod; 3] = [Self::Sudo, Self::Doas, Self::Su];
 
@@ -38,86 +48,144 @@ impl EscalationMethod {
         }
     }
 
-    pub fn wrap_command<'a>(&self, cmd: Command) -> Command {
-        // Yes this is jank. However, it's good enough for our purposes.
-        let envs: String = cmd
-            .get_envs()
-            .map(|(k, v)| {
-                format!(
-                    "{}={}",
-                    k.to_string_lossy(),
-                    v.unwrap_or(&OsString::from("")).to_string_lossy()
-                )
-            })
-            .join(" ");
-
-        let raw = format!("{envs} {cmd:?}");
+    pub fn wrap_command<'a>(&self, cmd: &Command) -> Command {
+        let raw = cmd.to_string();
 
         match self {
-            Self::Sudo => {
-                let mut cmd = Command::new("sudo");
-                cmd.args(["sh", "-c", &raw]);
-                cmd
-            }
-            Self::Doas => {
-                let mut cmd = Command::new("doas");
-                cmd.args(["sh", "-c", &raw]);
-                cmd
-            }
-            Self::Su => {
-                let mut cmd = Command::new("su");
-                cmd.args(["root", "-c", "sh", "-c", &raw]);
-                cmd
-            }
+            Self::Sudo => Command {
+                envs: vec![],
+                proc: "sudo".into(),
+                args: vec!["sh".into(), "-c".into(), raw.into()],
+            },
+            Self::Doas => Command {
+                envs: vec![],
+                proc: "doas".into(),
+                args: vec!["sh".into(), "-c".into(), raw.into()],
+            },
+            Self::Su => Command {
+                envs: vec![],
+                proc: "su".into(),
+                args: vec![
+                    "root".into(),
+                    "-c".into(),
+                    "sh".into(),
+                    "-c".into(),
+                    raw.into(),
+                ],
+            },
+        }
+    }
+}
+
+impl ToString for Command<'_> {
+    fn to_string(&self) -> String {
+        let args = join([&self.proc].into_iter().chain(self.args.iter()));
+
+        if self.envs.is_empty() {
+            args
+        } else {
+            let envs: String = (self.envs.iter())
+                .map(|(k, v)| format!("{}={}", quote(k), quote(v)))
+                .join(" ");
+
+            format!("{envs} {args}")
         }
     }
 }
 
+impl From<Command<'_>> for std::process::Command {
+    fn from(value: Command<'_>) -> Self {
+        let mut c = std::process::Command::new(value.proc.as_ref());
+        c.args(value.args.iter().map(|a| a.as_ref()));
+        c.envs(value.envs.iter().map(|(k, v)| (k.as_ref(), v.as_ref())));
+        c
+    }
+}
+
+impl From<Command<'_>> for tokio::process::Command {
+    fn from(value: Command<'_>) -> Self {
+        std::process::Command::from(value).into()
+    }
+}
+
 #[cfg(test)]
 mod tests {
-    use std::process::Command;
+    use super::*;
+
+    #[test]
+    fn test_to_string_no_env() {
+        let command = Command {
+            envs: vec![],
+            proc: "foo bar".into(),
+            args: vec![
+                "mrrrrp\\x12 mrp nya nya!".into(),
+                "yip yip".into(),
+                "yip".into(),
+            ],
+        };
+
+        let result = command.to_string();
+
+        assert_eq!(result, "'foo bar' 'mrrrrp\\x12 mrp nya nya!' 'yip yip' yip")
+    }
 
-    use super::EscalationMethod;
+    #[test]
+    fn test_to_string_with_env() {
+        let command = Command {
+            envs: vec![("uwu".into(), "nyaaa aaa!".into())],
+            proc: "foo bar".into(),
+            args: vec![
+                "mrrrrp\\x12 mrp nya nya!".into(),
+                "yip yip".into(),
+                "yip".into(),
+            ],
+        };
+
+        let result = command.to_string();
 
-    fn get_test_command() -> Command {
-        let mut cmd = Command::new("some/proc");
-        cmd.arg("two")
-            .arg("--three")
-            .arg("\"four\"")
-            .env("asdf", "foo");
-        cmd
+        assert_eq!(
+            result,
+            "uwu='nyaaa aaa!' 'foo bar' 'mrrrrp\\x12 mrp nya nya!' 'yip yip' yip"
+        )
+    }
+
+    fn get_test_command() -> Command<'static> {
+        Command {
+            envs: vec![("asdf".into(), "foo".into())],
+            proc: "some/proc".into(),
+            args: vec!["two".into(), "--three".into(), "\"four\"".into()],
+        }
     }
 
     #[test]
     fn test_sudo() {
-        let result = EscalationMethod::Sudo.wrap_command(get_test_command());
+        let result = EscalationMethod::Sudo.wrap_command(&get_test_command());
 
-        let printed = format!("{result:?}");
         assert_eq!(
-            printed,
-            r#""sudo" "sh" "-c" "asdf=foo \"some/proc\" \"two\" \"--three\" \"\\\"four\\\"\"""#
+            result.to_string(),
+            "sudo sh -c 'asdf=foo some/proc two --three '\\''\"four\"'\\'''"
         )
     }
 
     #[test]
     fn test_doas() {
-        let result = EscalationMethod::Doas.wrap_command(get_test_command());
+        let result = EscalationMethod::Doas.wrap_command(&get_test_command());
 
         let printed = format!("{result:?}");
         assert_eq!(
-            printed,
-            r#""doas" "sh" "-c" "asdf=foo \"some/proc\" \"two\" \"--three\" \"\\\"four\\\"\"""#
+            result.to_string(),
+            "doas sh -c 'asdf=foo some/proc two --three '\\''\"four\"'\\'''"
         )
     }
 
     #[test]
     fn test_su() {
-        let result = EscalationMethod::Su.wrap_command(get_test_command());
+        let result = EscalationMethod::Su.wrap_command(&get_test_command());
 
         let printed = format!("{result:?}");
         assert_eq!(
-            printed,
-            r#""su" "root" "-c" "sh" "-c" "asdf=foo \"some/proc\" \"two\" \"--three\" \"\\\"four\\\"\"""#
+            result.to_string(),
+            "su root -c sh -c 'asdf=foo some/proc two --three '\\''\"four\"'\\'''"
         )
     }
 }