Allow conversion of std::process::Command into a openssh::Command…

… given a `openssh::Session`. (#112)

* Add OverSsh trait to allow conversion of std::process::Command into a openssh::Session.

Signed-off-by: Aalekh Patel <[email protected]>

* Add a doc comment and example usage.

Signed-off-by: Aalekh Patel <[email protected]>

* Add tests for OverSsh's with_session

Signed-off-by: Aalekh Patel <[email protected]>

* Use raw_command and raw_args. Refactor the name to 'over_ssh'.

Signed-off-by: Aalekh Patel <[email protected]>

* Remove a dangling test import. Use better name for test.

Signed-off-by: Aalekh Patel <[email protected]>

* Clean up doc comment.

Signed-off-by: Aalekh Patel <[email protected]>

* Apply suggestions of code review.

Signed-off-by: Aalekh Patel <[email protected]>

* Update src/command.rs

Co-authored-by: Jiahao XU <[email protected]>

* Apply more suggestions from the code review.

Signed-off-by: Aalekh Patel <[email protected]>

* Apply more suggestions from the code review.

Signed-off-by: Aalekh Patel <[email protected]>

* Update src/command.rs

Co-authored-by: Jiahao XU <[email protected]>

* Update src/lib.rs

Co-authored-by: Jiahao XU <[email protected]>

* Update src/escape.rs

Co-authored-by: Jiahao XU <[email protected]>

* Apply more suggestions. Add a test for non-utf8 chars for escape.

* Slight refactor of escape.

Signed-off-by: Aalekh Patel <[email protected]>

* Run cargo clippy --fix and persist all modifications to escape.rs

Signed-off-by: Aalekh Patel <[email protected]>

* Apply rustfmt to command.rs and escape.rs

Signed-off-by: Aalekh Patel <[email protected]>

* Lint over_session test

Signed-off-by: Aalekh Patel <[email protected]>

* Update src/escape.rs

Co-authored-by: Jiahao XU <[email protected]>

* Update src/escape.rs

Co-authored-by: Jiahao XU <[email protected]>

* Update src/escape.rs

Co-authored-by: Jiahao XU <[email protected]>

* Apply more suggestions.

Signed-off-by: Aalekh Patel <[email protected]>

* Rustfmt fix.

Signed-off-by: Aalekh Patel <[email protected]>

* Fix a doctest for overssh.

Signed-off-by: Aalekh Patel <[email protected]>

* Add a test for overssh that requires escaping the argument passed to the command.

Signed-off-by: Aalekh Patel <[email protected]>

* Update overssh require-escaping-arguments test to contain a space.

Signed-off-by: Aalekh Patel <[email protected]>

* Apply rustfmt to tests/openssh.rs

Signed-off-by: Aalekh Patel <[email protected]>

---------

Signed-off-by: Aalekh Patel <[email protected]>
Signed-off-by: Aalekh Patel <[email protected]>
Co-authored-by: Jiahao XU <[email protected]>
Co-authored-by: Aalekh Patel <[email protected]>

src/command.rs

@@ -1,3 +1,5 @@
+use crate::escape::escape;
+
 use super::stdio::TryFromChildIo;
 use super::RemoteChild;
 use super::Stdio;
@@ -49,6 +51,130 @@ macro_rules! delegate {
     }};
 }
 
+/// If a command is `OverSsh` then it can be executed over an SSH session.
+///
+/// Primarily a way to allow `std::process::Command` to be turned directly into an `openssh::Command`.
+pub trait OverSsh {
+    /// Given an ssh session, return a command that can be executed over that ssh session.
+    ///
+    /// ### Notes
+    ///
+    /// The command to be executed on the remote machine should not explicitly
+    /// set environment variables or the current working directory. It errors if the source command
+    /// has environment variables or a current working directory set, since `openssh` doesn't (yet) have
+    /// a method to set environment variables and `ssh` doesn't support setting a current working directory
+    /// outside of `bash/dash/zsh` (which is not always available).
+    ///
+    /// ###  Examples
+    ///
+    /// 1. Consider the implementation of `OverSsh` for `std::process::Command`. Let's build a
+    /// `ls -l -a -h` command and execute it over an SSH session.
+    ///
+    /// ```no_run
+    /// # #[tokio::main(flavor = "current_thread")]
+    /// async fn main() -> Result<(), Box<dyn std::error::Error>> {
+    ///     use std::process::Command;
+    ///     use openssh::{Session, KnownHosts, OverSsh};
+    ///
+    ///     let session = Session::connect_mux("[email protected]", KnownHosts::Strict).await?;
+    ///     let ls =
+    ///         Command::new("ls")
+    ///         .arg("-l")
+    ///         .arg("-a")
+    ///         .arg("-h")
+    ///         .over_ssh(&session)?
+    ///         .output()
+    ///         .await?;
+    ///
+    ///     assert!(String::from_utf8(ls.stdout).unwrap().contains("total"));
+    /// #   Ok(())
+    /// }
+    ///
+    /// ```
+    /// 2. Building a command with environment variables or a current working directory set will
+    /// results in an error.
+    ///
+    /// ```no_run
+    /// # #[tokio::main(flavor = "current_thread")]
+    /// async fn main() -> Result<(), Box<dyn std::error::Error>> {
+    ///     use std::process::Command;
+    ///     use openssh::{Session, KnownHosts, OverSsh};
+    ///
+    ///     let session = Session::connect_mux("[email protected]", KnownHosts::Strict).await?;
+    ///     let echo =
+    ///         Command::new("echo")
+    ///         .env("MY_ENV_VAR", "foo")
+    ///         .arg("$MY_ENV_VAR")
+    ///         .over_ssh(&session);
+    ///     assert!(matches!(echo, Err(openssh::Error::CommandHasEnv)));
+    ///
+    /// #   Ok(())
+    /// }
+    ///
+    /// ```
+    fn over_ssh<'session>(
+        &self,
+        session: &'session Session,
+    ) -> Result<crate::Command<'session>, crate::Error>;
+}
+
+impl OverSsh for std::process::Command {
+    fn over_ssh<'session>(
+        &self,
+        session: &'session Session,
+    ) -> Result<Command<'session>, crate::Error> {
+        // I'd really like `!self.get_envs().is_empty()` here, but that's
+        // behind a `exact_size_is_empty` feature flag.
+        if self.get_envs().len() > 0 {
+            return Err(crate::Error::CommandHasEnv);
+        }
+
+        if self.get_current_dir().is_some() {
+            return Err(crate::Error::CommandHasCwd);
+        }
+
+        let program_escaped: Cow<'_, OsStr> = escape(self.get_program());
+        let mut command = session.raw_command(program_escaped);
+
+        let args = self.get_args().map(escape);
+        command.raw_args(args);
+        Ok(command)
+    }
+}
+
+impl OverSsh for tokio::process::Command {
+    fn over_ssh<'session>(
+        &self,
+        session: &'session Session,
+    ) -> Result<Command<'session>, crate::Error> {
+        self.as_std().over_ssh(session)
+    }
+}
+
+impl<S> OverSsh for &S
+where
+    S: OverSsh,
+{
+    fn over_ssh<'session>(
+        &self,
+        session: &'session Session,
+    ) -> Result<Command<'session>, crate::Error> {
+        <S as OverSsh>::over_ssh(self, session)
+    }
+}
+
+impl<S> OverSsh for &mut S
+where
+    S: OverSsh,
+{
+    fn over_ssh<'session>(
+        &self,
+        session: &'session Session,
+    ) -> Result<Command<'session>, crate::Error> {
+        <S as OverSsh>::over_ssh(self, session)
+    }
+}
+
 /// A remote process builder, providing fine-grained control over how a new remote process should
 /// be spawned.
 ///

src/error.rs

@@ -67,6 +67,16 @@ pub enum Error {
     /// IO Error when creating/reading/writing from ChildStdin, ChildStdout, ChildStderr.
     #[error("failure while accessing standard i/o of remote process")]
     ChildIo(#[source] io::Error),
+
+    /// The command has some env variables that it expects to carry over ssh.
+    /// However, OverSsh does not support passing env variables over ssh.
+    #[error("rejected runing a command over ssh that expects env variables to be carried over to remote.")]
+    CommandHasEnv,
+
+    /// The command expects to be in a specific working directory in remote.
+    /// However, OverSsh does not support setting a working directory for commands to be executed over ssh.
+    #[error("rejected runing a command over ssh that expects a specific working directory to be carried over to remote.")]
+    CommandHasCwd,
 }
 
 #[cfg(feature = "native-mux")]

src/escape.rs

@@ -0,0 +1,91 @@
+//! 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
+
+use std::{
+    borrow::Cow,
+    ffi::{OsStr, OsString},
+    os::unix::ffi::OsStrExt,
+    os::unix::ffi::OsStringExt,
+};
+
+fn allowed(byte: u8) -> bool {
+    matches!(byte, b'a'..=b'z' | b'A'..=b'Z' | b'0'..=b'9' | b'-' | b'_' | b'=' | b'/' | b',' | b'.' | b'+')
+}
+
+/// 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(crate) fn escape(s: &OsStr) -> Cow<'_, OsStr> {
+    let as_bytes = s.as_bytes();
+    let all_allowed = as_bytes.iter().copied().all(allowed);
+
+    if !as_bytes.is_empty() && all_allowed {
+        return Cow::Borrowed(s);
+    }
+
+    let mut escaped = Vec::with_capacity(as_bytes.len() + 2);
+    escaped.push(b'\'');
+
+    for &b in as_bytes {
+        match b {
+            b'\'' | b'!' => {
+                escaped.reserve(4);
+                escaped.push(b'\'');
+                escaped.push(b'\\');
+                escaped.push(b);
+                escaped.push(b'\'');
+            }
+            _ => escaped.push(b),
+        }
+    }
+    escaped.push(b'\'');
+    OsString::from_vec(escaped).into()
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    fn test_escape_case(input: &str, expected: &str) {
+        test_escape_from_bytes(input.as_bytes(), expected.as_bytes())
+    }
+
+    fn test_escape_from_bytes(input: &[u8], expected: &[u8]) {
+        let input_os_str = OsStr::from_bytes(input);
+        let observed_os_str = escape(input_os_str);
+        let expected_os_str = OsStr::from_bytes(expected);
+        assert_eq!(observed_os_str, expected_os_str);
+    }
+
+    // These tests are courtesy of the `shell-escape` crate.
+    #[test]
+    fn test_escape() {
+        test_escape_case(
+            "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789-_=/,.+",
+            "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789-_=/,.+",
+        );
+        test_escape_case("--aaa=bbb-ccc", "--aaa=bbb-ccc");
+        test_escape_case(
+            "linker=gcc -L/foo -Wl,bar",
+            r#"'linker=gcc -L/foo -Wl,bar'"#,
+        );
+        test_escape_case(r#"--features="default""#, r#"'--features="default"'"#);
+        test_escape_case(r#"'!\$`\\\n "#, r#"''\'''\!'\$`\\\n '"#);
+        test_escape_case("", r#"''"#);
+        test_escape_case(" ", r#"' '"#);
+        test_escape_case("*", r#"'*'"#);
+
+        test_escape_from_bytes(
+            &[0x66, 0x6f, 0x80, 0x6f],
+            &[b'\'', 0x66, 0x6f, 0x80, 0x6f, b'\''],
+        );
+    }
+}

src/lib.rs

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

tests/openssh.rs

@@ -292,6 +292,90 @@ async fn stdout() {
     }
 }
 
+#[tokio::test]
+#[cfg_attr(not(ci), ignore)]
+async fn over_session_ok() {
+    for session in connects().await {
+        let mut command = std::process::Command::new("echo")
+            .arg("foo")
+            .over_ssh(&session)
+            .expect("No env vars or current working dir is set.");
+
+        let child = command.output().await.unwrap();
+        assert_eq!(child.stdout, b"foo\n");
+
+        let child = session
+            .command("echo")
+            .arg("foo")
+            .raw_arg(">")
+            .arg("/dev/stderr")
+            .output()
+            .await
+            .unwrap();
+        assert!(child.stdout.is_empty());
+
+        session.close().await.unwrap();
+    }
+}
+
+#[tokio::test]
+#[cfg_attr(not(ci), ignore)]
+async fn over_session_ok_require_escaping_arguments() {
+    for session in connects().await {
+        let mut command = std::process::Command::new("echo")
+            .arg("\"\'\' foo \'\'\"")
+            .over_ssh(&session)
+            .expect("No env vars or current working dir is set.");
+
+        let child = command.output().await.unwrap();
+        assert_eq!(child.stdout, b"\"\'\' foo \'\'\"\n");
+
+        let child = session
+            .command("echo")
+            .arg("foo")
+            .raw_arg(">")
+            .arg("/dev/stderr")
+            .output()
+            .await
+            .unwrap();
+        assert!(child.stdout.is_empty());
+
+        session.close().await.unwrap();
+    }
+}
+
+/// Test that `over_ssh` errors if the source command has env vars specified.
+#[tokio::test]
+#[cfg_attr(not(ci), ignore)]
+async fn over_session_err_because_env_var() {
+    for session in connects().await {
+        let command_with_env = std::process::Command::new("printenv")
+            .arg("MY_ENV_VAR")
+            .env("MY_ENV_VAR", "foo")
+            .over_ssh(&session);
+        assert!(matches!(
+            command_with_env,
+            Err(openssh::Error::CommandHasEnv)
+        ));
+    }
+}
+
+/// Test that `over_ssh` errors if the source command has a `current_dir` specified.
+#[tokio::test]
+#[cfg_attr(not(ci), ignore)]
+async fn over_session_err_because_cwd() {
+    for session in connects().await {
+        let command_with_current_dir = std::process::Command::new("echo")
+            .arg("foo")
+            .current_dir("/tmp")
+            .over_ssh(&session);
+        assert!(matches!(
+            command_with_current_dir,
+            Err(openssh::Error::CommandHasCwd)
+        ));
+    }
+}
+
 #[tokio::test]
 #[cfg_attr(not(ci), ignore)]
 async fn shell() {