Merge pull request #104 from openssh-rust/fix/cut-off

Fixed stdout/stderr cut off bug when using `RemoteChild::wait_with_output`

Author: jonhoo

Cargo.toml

@@ -49,6 +49,8 @@ dirs = "4.0.0"
 
 openssh-mux-client = { version = "0.15.1", optional = true }
 
+libc = "0.2.137"
+
 [dev-dependencies]
 regex = "1"
 tokio = { version = "1", features = [ "full" ] }

src/changelog.rs

@@ -5,7 +5,13 @@ use crate::*;
 ///
 /// ## Added
 ///  - `impl From<std::os::unix::io::OwnedFd> for Stdio`
+///  -  Add new fn `Stdio::from_raw_fd_owned`
 /// ## Changed
+///  - Mark `FromRawFd` impl for `Stdio` as deprecated
+///  - Mark `From<tokio_pipe::PipeRead>` for `Stdio` as deprecated
+///  - Mark `From<tokio_pipe::PipeWrite>` for `Stdio` as deprecated
+/// ## Fixed
+///  - [`wait_with_output` + `native-mux` cuts off stdout output](https://github.com/openssh-rust/openssh/issues/103)
 /// ## Removed
 #[doc(hidden)]
 pub mod unreleased {}

src/child.rs

@@ -62,6 +62,9 @@ macro_rules! delegate {
 /// available,`Stdio::piped()` should be passed to the corresponding method on
 /// [`Command`](crate::Command).
 ///
+/// NOTE that once `RemoteChild` is dropped, any data written to `stdin` will not be sent to the
+/// remote process and `stdout` and `stderr` will yield EOF immediately.
+///
 /// ```rust,no_run
 /// # async fn foo() {
 /// # let child: openssh::RemoteChild<'static> = unimplemented!();
@@ -169,9 +172,15 @@ impl<'s> RemoteChild<'s> {
 
         // Execute them concurrently to avoid the pipe buffer being filled up
         // and cause the remote process to block forever.
-        let (status, stdout, stderr) = try_join!(self.wait(), stdout_read, stderr_read)?;
+        let (stdout, stderr) = try_join!(stdout_read, stderr_read)?;
         Ok(Output {
-            status,
+            // The self.wait() future terminates the stdout and stderr futures
+            // when it resolves, even if there may still be more data arriving
+            // from the server.
+            //
+            // Therefore, we wait for them first, and only once they're complete
+            // do we wait for the process to have terminated.
+            status: self.wait().await?,
             stdout,
             stderr,
         })

src/native_mux_impl/stdio.rs

@@ -1,12 +1,13 @@
-use crate::stdio::StdioImpl;
-use crate::Error;
-use crate::Stdio;
+use crate::{stdio::StdioImpl, Error, Stdio};
 
-use once_cell::sync::OnceCell;
+use std::{
+    fs::{File, OpenOptions},
+    io,
+    os::unix::io::{AsRawFd, FromRawFd, IntoRawFd, OwnedFd, RawFd},
+};
 
-use std::fs::{File, OpenOptions};
-use std::io;
-use std::os::unix::io::{AsRawFd, RawFd};
+use libc::{c_int, fcntl, F_GETFL, F_SETFL, O_NONBLOCK};
+use once_cell::sync::OnceCell;
 use tokio_pipe::{pipe, PipeRead, PipeWrite};
 
 fn create_pipe() -> Result<(PipeRead, PipeWrite), Error> {
@@ -28,25 +29,73 @@ fn get_null_fd() -> Result<RawFd, Error> {
 }
 
 pub(crate) enum Fd {
-    PipeReadEnd(PipeRead),
-    PipeWriteEnd(PipeWrite),
-
+    Owned(OwnedFd),
     Borrowed(RawFd),
     Null,
 }
 
+fn cvt(ret: c_int) -> io::Result<c_int> {
+    if ret == -1 {
+        Err(io::Error::last_os_error())
+    } else {
+        Ok(ret)
+    }
+}
+
+fn set_blocking_inner(fd: RawFd) -> io::Result<()> {
+    let flags = cvt(unsafe { fcntl(fd, F_GETFL) })?;
+    cvt(unsafe { fcntl(fd, F_SETFL, flags & (!O_NONBLOCK)) })?;
+
+    Ok(())
+}
+
+fn set_blocking(fd: RawFd) -> Result<(), Error> {
+    set_blocking_inner(fd).map_err(Error::ChildIo)
+}
+
 impl Fd {
     pub(crate) fn as_raw_fd_or_null_fd(&self) -> Result<RawFd, Error> {
         use Fd::*;
 
         match self {
-            PipeReadEnd(fd) => Ok(AsRawFd::as_raw_fd(fd)),
-            PipeWriteEnd(fd) => Ok(AsRawFd::as_raw_fd(fd)),
-
+            Owned(owned_fd) => Ok(owned_fd.as_raw_fd()),
             Borrowed(rawfd) => Ok(*rawfd),
             Null => get_null_fd(),
         }
     }
+
+    /// # Safety
+    ///
+    /// `T::into_raw_fd` must return a valid fd and transfers
+    /// the ownershipt of it.
+    unsafe fn new_owned<T: IntoRawFd>(fd: T) -> Result<Self, Error> {
+        let raw_fd = fd.into_raw_fd();
+        Ok(Fd::Owned(OwnedFd::from_raw_fd(raw_fd)))
+    }
+}
+
+impl TryFrom<PipeRead> for Fd {
+    type Error = Error;
+
+    fn try_from(pipe_read: PipeRead) -> Result<Self, Error> {
+        // Safety:
+        //
+        // PipeRead::into_raw_fd returns a valid fd and transfers the
+        // ownership of it.
+        unsafe { Self::new_owned(pipe_read) }
+    }
+}
+
+impl TryFrom<PipeWrite> for Fd {
+    type Error = Error;
+
+    fn try_from(pipe_write: PipeWrite) -> Result<Self, Error> {
+        // Safety:
+        //
+        // PipeWrite::into_raw_fd returns a valid fd and transfers the
+        // ownership of it.
+        unsafe { Self::new_owned(pipe_write) }
+    }
 }
 
 impl Stdio {
@@ -56,9 +105,19 @@ impl Stdio {
             StdioImpl::Null => Ok((Fd::Null, None)),
             StdioImpl::Pipe => {
                 let (read, write) = create_pipe()?;
-                Ok((Fd::PipeReadEnd(read), Some(write)))
+
+                // read end will be sent to ssh multiplex server
+                // and it expects blocking fd.
+                set_blocking(read.as_raw_fd())?;
+                Ok((read.try_into()?, Some(write)))
+            }
+            StdioImpl::Fd(fd, owned) => {
+                let raw_fd = fd.as_raw_fd();
+                if *owned {
+                    set_blocking(raw_fd)?;
+                }
+                Ok((Fd::Borrowed(raw_fd), None))
             }
-            StdioImpl::Fd(fd) => Ok((Fd::Borrowed(fd.as_raw_fd()), None)),
         }
     }
 
@@ -68,9 +127,19 @@ impl Stdio {
             StdioImpl::Null => Ok((Fd::Null, None)),
             StdioImpl::Pipe => {
                 let (read, write) = create_pipe()?;
-                Ok((Fd::PipeWriteEnd(write), Some(read)))
+
+                // write end will be sent to ssh multiplex server
+                // and it expects blocking fd.
+                set_blocking(write.as_raw_fd())?;
+                Ok((write.try_into()?, Some(read)))
+            }
+            StdioImpl::Fd(fd, owned) => {
+                let raw_fd = fd.as_raw_fd();
+                if *owned {
+                    set_blocking(raw_fd)?;
+                }
+                Ok((Fd::Borrowed(raw_fd), None))
             }
-            StdioImpl::Fd(fd) => Ok((Fd::Borrowed(fd.as_raw_fd()), None)),
         }
     }
 

src/stdio.rs

@@ -18,7 +18,7 @@ pub(crate) enum StdioImpl {
     /// Read/Write to a newly created pipe
     Pipe,
     /// Read/Write to custom fd
-    Fd(OwnedFd),
+    Fd(OwnedFd, bool),
     /// Inherit stdin/stdout/stderr
     Inherit,
 }
@@ -40,13 +40,40 @@ impl Stdio {
     }
 
     /// The child inherits from the corresponding parent descriptor.
+    ///
+    /// NOTE that the stdio fd must be in blocking mode, otherwise
+    /// ssh might not flush all output since it considers
+    /// (`EAGAIN`/`EWOULDBLOCK`) as an error
     pub const fn inherit() -> Self {
         Self(StdioImpl::Inherit)
     }
+
+    /// `Stdio::from_raw_fd_owned` takes ownership of the fd passed in
+    /// and closes the fd on drop.
+    ///
+    /// NOTE that the fd will be put into blocking mode, then it will be
+    /// closed when `Stdio` is dropped.
+    ///
+    /// # Safety
+    ///
+    /// * `fd` - must be a valid fd and must give its ownership to `Stdio`.
+    pub unsafe fn from_raw_fd_owned(fd: RawFd) -> Self {
+        Self(StdioImpl::Fd(OwnedFd::from_raw_fd(fd), true))
+    }
 }
+/// **Deprecated, use [`Stdio::from_raw_fd_owned`] instead.**
+///
+/// FromRawFd takes ownership of the fd passed in
+/// and closes the fd on drop.
+///
+/// NOTE that the fd must be in blocking mode, otherwise
+/// ssh might not flush all output since it considers
+/// (`EAGAIN`/`EWOULDBLOCK`) as an error
+#[allow(useless_deprecated)]
+#[deprecated(since = "0.9.8", note = "Use Stdio::from_raw_fd_owned instead")]
 impl FromRawFd for Stdio {
     unsafe fn from_raw_fd(fd: RawFd) -> Self {
-        Self(StdioImpl::Fd(OwnedFd::from_raw_fd(fd)))
+        Self(StdioImpl::Fd(OwnedFd::from_raw_fd(fd), false))
     }
 }
 impl From<Stdio> for process::Stdio {
@@ -59,14 +86,16 @@ impl From<Stdio> for process::Stdio {
             // safety: StdioImpl(fd) is only constructed from known-valid and
             // owned file descriptors by virtue of the safety requirement
             // for invoking from_raw_fd.
-            StdioImpl::Fd(fd) => unsafe { process::Stdio::from_raw_fd(IntoRawFd::into_raw_fd(fd)) },
+            StdioImpl::Fd(fd, _) => unsafe {
+                process::Stdio::from_raw_fd(IntoRawFd::into_raw_fd(fd))
+            },
         }
     }
 }
 
 impl From<OwnedFd> for Stdio {
     fn from(fd: OwnedFd) -> Self {
-        Self(StdioImpl::Fd(fd))
+        Self(StdioImpl::Fd(fd, true))
     }
 }
 
@@ -77,14 +106,31 @@ macro_rules! impl_from_for_stdio {
                 let fd = arg.into_raw_fd();
                 // safety: $type must have a valid into_raw_fd implementation
                 // and must not be RawFd.
-                unsafe { Self::from_raw_fd(fd) }
+                Self(StdioImpl::Fd(unsafe { OwnedFd::from_raw_fd(fd) }, true))
+            }
+        }
+    };
+    (deprecated $type:ty) => {
+        #[allow(useless_deprecated)]
+        #[deprecated(
+            since = "0.9.8",
+            note = "Use From<OwnedFd> for Stdio or Stdio::from_raw_fd_owned instead"
+        )]
+        /// **Deprecated, use `From<OwnedFd> for Stdio` or
+        /// [`Stdio::from_raw_fd_owned`] instead.**
+        impl From<$type> for Stdio {
+            fn from(arg: $type) -> Self {
+                let fd = arg.into_raw_fd();
+                // safety: $type must have a valid into_raw_fd implementation
+                // and must not be RawFd.
+                Self(StdioImpl::Fd(unsafe { OwnedFd::from_raw_fd(fd) }, true))
             }
         }
     };
 }
 
-impl_from_for_stdio!(tokio_pipe::PipeWrite);
-impl_from_for_stdio!(tokio_pipe::PipeRead);
+impl_from_for_stdio!(deprecated tokio_pipe::PipeWrite);
+impl_from_for_stdio!(deprecated tokio_pipe::PipeRead);
 
 impl_from_for_stdio!(process::ChildStdin);
 impl_from_for_stdio!(process::ChildStdout);

tests/openssh.rs

@@ -1,15 +1,16 @@
 use once_cell::sync::Lazy;
 use regex::Regex;
-use std::env;
-use std::io;
-use std::io::Write;
-use std::net::IpAddr;
-use std::path::PathBuf;
-use std::time::Duration;
+use std::{
+    env,
+    io::{self, Write},
+    net::IpAddr,
+    path::PathBuf,
+    process,
+    time::Duration,
+};
 use tempfile::tempdir;
-
-use tokio::io::{AsyncReadExt, AsyncWriteExt};
 use tokio::{
+    io::{AsyncReadExt, AsyncWriteExt},
     net::{UnixListener, UnixStream},
     time::sleep,
 };
@@ -52,7 +53,7 @@ async fn session_builder_connect(mut builder: SessionBuilder, addr: &str) -> Vec
     sessions
 }
 
-async fn connects() -> Vec<Session> {
+async fn connects_with_name() -> Vec<(Session, &'static str)> {
     let mut sessions = Vec::with_capacity(2);
 
     let mut builder = SessionBuilder::default();
@@ -63,17 +64,25 @@ async fn connects() -> Vec<Session> {
 
     #[cfg(feature = "process-mux")]
     {
-        sessions.push(builder.connect(&addr()).await.unwrap());
+        sessions.push((builder.connect(&addr()).await.unwrap(), "process-mux"));
     }
 
     #[cfg(feature = "native-mux")]
     {
-        sessions.push(builder.connect_mux(&addr()).await.unwrap());
+        sessions.push((builder.connect_mux(&addr()).await.unwrap(), "native-mux"));
     }
 
     sessions
 }
 
+async fn connects() -> Vec<Session> {
+    connects_with_name()
+        .await
+        .into_iter()
+        .map(|(session, _name)| session)
+        .collect()
+}
+
 async fn connects_err(host: &str) -> Vec<Error> {
     session_builder_connects_err(host, SessionBuilder::default()).await
 }
@@ -942,3 +951,25 @@ async fn test_sftp_subsystem() {
         sftp.close().await.unwrap();
     }
 }
+
+#[tokio::test]
+#[cfg_attr(not(ci), ignore)]
+async fn test_read_large_file_bug() {
+    for (session, name) in connects_with_name().await {
+        eprintln!("Testing {name} implementation");
+
+        let bs = 1024;
+        let count = 20480;
+
+        let process::Output { status, stdout, .. } = session
+            .shell(format!("dd if=/dev/zero bs={bs} count={count}"))
+            .output()
+            .await
+            .unwrap();
+
+        assert!(status.success());
+
+        stdout.iter().copied().for_each(|byte| assert_eq!(byte, 0));
+        assert_eq!(stdout.len(), bs * count);
+    }
+}