Apply more suggestions from the code review.

Signed-off-by: Aalekh Patel <aalekh.gwpeck.7998@icloud.com>

Author: aalekhpatel07

src/command.rs

@@ -1,10 +1,13 @@
+use crate::escape;
+
 use super::stdio::TryFromChildIo;
 use super::RemoteChild;
 use super::Stdio;
 use super::{Error, Session};
 
 use std::borrow::Cow;
 use std::ffi::OsStr;
+use std::os::unix::prelude::OsStrExt;
 use std::process;
 
 #[derive(Debug)]
@@ -57,15 +60,10 @@ pub trait OverSsh {
     /// **Note**: The command to be executed on the remote machine does not include 
     /// any environment variables or the current directory. They must be
     /// set explicitly, if desired.
-    fn over_session<'session>(&self, session: &'session Session) -> crate::Command<'session>;
-}
-
-impl OverSsh for std::process::Command {
-    /// Given a session, convert a `std::process::Command` into an `openssh::Command` 
-    /// that can be executed over that session.
-    /// **Note**: The command to be executed on the remote machine does not include 
-    /// any environment variables or the current directory. They must be
-    /// set explicitly, if desired.
+    /// 
+    /// # Example
+    /// Consider the implementation of `OverSsh` for `std::process::Command`:
+    /// 
     /// ```no_run
     /// # #[tokio::main(flavor = "current_thread")]
     /// async fn main() -> Result<(), Box<dyn std::error::Error>> {
@@ -87,42 +85,26 @@ impl OverSsh for std::process::Command {
     /// }
     /// 
     /// ```
+    fn over_session<'session>(&self, session: &'session Session) -> crate::Command<'session>;
+}
+
+impl OverSsh for std::process::Command {
     fn over_session<'session>(&self, session: &'session Session) -> Command<'session> {
-        let mut command = session.command(self.get_program().to_string_lossy());
-        command.raw_args(self.get_args());
+        let program_escaped = escape(self.get_program().as_bytes());
+        let mut command = session.command(Cow::Borrowed(program_escaped.as_str()));
+        
+        self
+        .get_args()
+        .for_each(|arg| {
+            let arg_escaped = escape(arg.as_bytes());
+            command.arg(Cow::Borrowed(arg_escaped.as_str()));
+        });
         command
     }
 }
 
 
 impl OverSsh for tokio::process::Command {
-    /// Given a session, convert a `tokio::process::Command` into an `openssh::Command` 
-    /// that can be executed over that session.
-    /// 
-    /// **Note**: The command to be executed on the remote machine does not include 
-    /// any environment variables or the current directory. They must be
-    /// set explicitly, if desired.
-    /// ```no_run
-    /// # #[tokio::main(flavor = "current_thread")]
-    /// # async fn main() -> Result<(), Box<dyn std::error::Error>> {
-    ///     use tokio::process::Command;
-    ///     use openssh::{Session, KnownHosts, OverSsh};
-
-    ///     let session = Session::connect_mux("me@ssh.example.com", KnownHosts::Strict).await?;
-    ///     let ls = 
-    ///         Command::new("ls")
-    ///         .arg("-l")
-    ///         .arg("-a")
-    ///         .arg("-h")
-    ///         .over_session(&session)
-    ///         .output()
-    ///         .await?;
-    /// 
-    ///     assert!(String::from_utf8(ls.stdout).unwrap().contains("total"));
-    /// #   Ok(())
-    /// # }
-    /// 
-    /// ```
     fn over_session<'session>(&self, session: &'session Session) -> Command<'session> {
         self.as_std().over_session(session)
     }
@@ -138,15 +120,6 @@ where
     }
 }
 
-impl<S> OverSsh for &mut S
-where
-    S: OverSsh
-{
-    fn over_session<'session>(&self, session: &'session Session) -> Command<'session> {
-        <S as OverSsh>::over_session(self, session)
-    }
-}
-
 impl<S> OverSsh for Box<S>
 where
     S: OverSsh

src/escape.rs

@@ -0,0 +1,81 @@
+//! Escape characters that may have special meaning in a shell, including spaces.
+//! This is a modified version of the [`shell-escape::unix`] module of [`shell-escape`] crate.
+//! 
+//! [`shell-escape`]: https://crates.io/crates/shell-escape
+//! [`shell-escape::unix`]: https://docs.rs/shell-escape/latest/src/shell_escape/lib.rs.html#101
+
+
+fn whitelisted(ch: char) -> bool {
+    match ch {
+        'a'..='z' | 'A'..='Z' | '0'..='9' | '-' | '_' | '=' | '/' | ',' | '.' | '+' => true,
+        _ => false,
+    }
+}
+
+/// Escape characters that may have special meaning in a shell, including spaces.
+/// 
+/// **Note**: This function is an adaptation of [`shell-escape::unix::escape`].
+/// This function exists only for type compatibility and the implementation is
+/// almost exactly the same as [`shell-escape::unix::escape`].
+/// 
+/// [`shell-escape::unix::escape`]: https://docs.rs/shell-escape/latest/src/shell_escape/lib.rs.html#101
+/// 
+pub fn escape(s: &[u8]) -> String {
+    let all_whitelisted = s.iter().all(|x| whitelisted(*x as char));
+
+    if !s.is_empty() && all_whitelisted {
+        // All bytes are whitelisted and valid single-byte UTF-8, 
+        // so we can build the original string and return as is.
+        return String::from_utf8(s.to_vec()).unwrap();
+    }
+
+    let mut escaped = String::with_capacity(s.len() + 2);
+    escaped.push('\'');
+
+    for &b in s {
+        match b {
+            b'\'' | b'!' => {
+                escaped.push_str("'\\");
+                escaped.push(b as char);
+                escaped.push('\'');
+            }
+            _ => escaped.push(b as char),
+        }
+    }
+    escaped.push('\'');
+    escaped
+}
+
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    // These tests are courtesy of the `shell-escape` crate.
+    #[test]
+    fn test_escape() {
+        assert_eq!(
+            escape(b"abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789-_=/,.+"), 
+            "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789-_=/,.+"
+        );
+        assert_eq!(
+            escape(b"--aaa=bbb-ccc"), 
+            "--aaa=bbb-ccc"
+        );
+        assert_eq!(
+            escape(b"linker=gcc -L/foo -Wl,bar"), 
+            r#"'linker=gcc -L/foo -Wl,bar'"#
+        );
+        assert_eq!(
+            escape(br#"--features="default""#), 
+            r#"'--features="default"'"#
+        );
+        assert_eq!(
+            escape(br#"'!\$`\\\n "#), 
+            r#"''\'''\!'\$`\\\n '"#
+        );
+        assert_eq!(escape(b""), r#"''"#);
+        assert_eq!(escape(b" "), r#"' '"#);
+        assert_eq!(escape(b"\xC4b"), r#"'Äb'"#);
+    }
+}

src/lib.rs

@@ -169,6 +169,9 @@ pub use builder::{KnownHosts, SessionBuilder};
 mod command;
 pub use command::{Command, OverSsh};
 
+mod escape;
+pub use escape::escape;
+
 mod child;
 pub use child::RemoteChild;