Apply more suggestions from the code review.

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

Author: aalekhpatel07

src/command.rs

@@ -1,4 +1,4 @@
-use crate::escape;
+use crate::escape::escape;
 
 use super::stdio::TryFromChildIo;
 use super::RemoteChild;
@@ -7,7 +7,6 @@ use super::{Error, Session};
 
 use std::borrow::Cow;
 use std::ffi::OsStr;
-use std::os::unix::prelude::OsStrExt;
 use std::process;
 
 #[derive(Debug)]
@@ -90,15 +89,13 @@ pub trait OverSsh {
 
 impl OverSsh for std::process::Command {
     fn over_session<'session>(&self, session: &'session Session) -> Command<'session> {
-        let program_escaped = escape(self.get_program().as_bytes());
-        let mut command = session.command(Cow::Borrowed(program_escaped.as_str()));
+        let program_escaped: Cow<'_, OsStr> = escape(self.get_program());
+        let mut command = session.raw_command(program_escaped);
 
-        self
-        .get_args()
-        .for_each(|arg| {
-            let arg_escaped = escape(arg.as_bytes());
-            command.arg(Cow::Borrowed(arg_escaped.as_str()));
-        });
+        let args = self
+            .get_args()
+            .map(|arg| escape(arg));
+        command.raw_args(args);
         command
     }
 }

src/escape.rs

@@ -4,10 +4,17 @@
 //! [`shell-escape`]: https://crates.io/crates/shell-escape
 //! [`shell-escape::unix`]: https://docs.rs/shell-escape/latest/src/shell_escape/lib.rs.html#101
 
+use std::{
+    borrow::Cow, 
+    ffi::{OsStr, OsString},
+    os::unix::ffi::OsStrExt,
+    os::unix::ffi::OsStringExt,
+};
 
-fn whitelisted(ch: char) -> bool {
-    match ch {
-        'a'..='z' | 'A'..='Z' | '0'..='9' | '-' | '_' | '=' | '/' | ',' | '.' | '+' => true,
+
+fn whitelisted(byte: u8) -> bool {
+    match byte {
+        b'a'..=b'z' | b'A'..=b'Z' | b'0'..=b'9' | b'-' | b'_' | b'=' | b'/' | b',' | b'.' | b'+' => true,
         _ => false,
     }
 }
@@ -20,62 +27,74 @@ fn whitelisted(ch: char) -> bool {
 /// 
 /// [`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));
+pub fn escape(s: &OsStr) -> Cow<'_, OsStr> {
+    let s = s.as_bytes();
+    let all_whitelisted = s.iter().all(|x| whitelisted(*x));
 
     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();
+        return OsString::from_vec(s.to_vec()).into();
     }
 
-    let mut escaped = String::with_capacity(s.len() + 2);
-    escaped.push('\'');
+    let mut escaped = Vec::with_capacity(s.len() + 2);
+    escaped.push(b'\'');
 
     for &b in s {
         match b {
             b'\'' | b'!' => {
-                escaped.push_str("'\\");
-                escaped.push(b as char);
-                escaped.push('\'');
+                escaped.push(b'\'');
+                escaped.push(b'\\');
+                escaped.push(b);
+                escaped.push(b'\'');
             }
-            _ => escaped.push(b as char),
+            _ => escaped.push(b),
         }
     }
-    escaped.push('\'');
-    escaped
+    escaped.push(b'\'');
+    OsString::from_vec(escaped).into()
 }
 
 
 #[cfg(test)]
 mod tests {
     use super::*;
 
+    fn test_escape_case(input: &str, expected: &str) {
+        let input_os_str = OsStr::from_bytes(input.as_bytes());
+        let observed_os_str = escape(input_os_str);
+        let expected_os_str = OsStr::from_bytes(expected.as_bytes());
+        assert_eq!(observed_os_str, expected_os_str);
+    }
+
     // These tests are courtesy of the `shell-escape` crate.
     #[test]
     fn test_escape() {
-        assert_eq!(
-            escape(b"abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789-_=/,.+"), 
+        test_escape_case(
+            "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789-_=/,.+",
             "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789-_=/,.+"
         );
-        assert_eq!(
-            escape(b"--aaa=bbb-ccc"), 
+        test_escape_case(
+            "--aaa=bbb-ccc",
             "--aaa=bbb-ccc"
         );
-        assert_eq!(
-            escape(b"linker=gcc -L/foo -Wl,bar"), 
+        test_escape_case(
+            "linker=gcc -L/foo -Wl,bar",
             r#"'linker=gcc -L/foo -Wl,bar'"#
         );
-        assert_eq!(
-            escape(br#"--features="default""#), 
+        test_escape_case(
+            r#"--features="default""#,
             r#"'--features="default"'"#
         );
-        assert_eq!(
-            escape(br#"'!\$`\\\n "#), 
+        test_escape_case(
+            r#"'!\$`\\\n "#,
             r#"''\'''\!'\$`\\\n '"#
         );
-        assert_eq!(escape(b""), r#"''"#);
-        assert_eq!(escape(b" "), r#"' '"#);
-        assert_eq!(escape(b"\xC4b"), r#"'Äb'"#);
+        test_escape_case(
+            "",
+            r#"''"#
+        );
+        test_escape_case(
+            " ",
+            r#"' '"#
+        );
     }
 }

src/lib.rs

@@ -169,8 +169,7 @@ pub use builder::{KnownHosts, SessionBuilder};
 mod command;
 pub use command::{Command, OverSsh};
 
-mod escape;
-pub use escape::escape;
+pub(crate) mod escape;
 
 mod child;
 pub use child::RemoteChild;