Only set_blocking only on OwnedFd

Do not `set_blocking` on `StdioImpl::inherit` as we do not own them
and also do not `set_blocking` on `Stdio` constructed using
`FromRawFd::from_raw_fd`

Signed-off-by: Jiahao XU <Jiahao_XU@outlook.com>

Author: NobodyXu

src/native_mux_impl/command.rs

@@ -1,6 +1,6 @@
 use super::Error;
 use super::RemoteChild;
-use super::{stdio::set_blocking, ChildStderr, ChildStdin, ChildStdout, Stdio};
+use super::{ChildStderr, ChildStdin, ChildStdout, Stdio};
 
 use std::borrow::Cow;
 use std::ffi::OsStr;
@@ -71,27 +71,6 @@ impl Command {
             stderr.as_raw_fd_or_null_fd()?,
         ];
 
-        let is_inherited = [
-            self.stdin_v.is_inherited(),
-            self.stdout_v.is_inherited(),
-            self.stderr_v.is_inherited(),
-        ];
-
-        stdios
-            .into_iter()
-            .zip(is_inherited)
-            .filter_map(|(stdio, is_inherited)| if !is_inherited { Some(stdio) } else { None })
-            // Note that once we do this, these file descriptors
-            // (and the Fds we got them from above) should no longer be used in
-            // any async context, as they'd start blocking.
-            //
-            // We give away the descriptors in stdios when we pass them to
-            // open_new_session (which doesn't use them in an async context),
-            // and the Fds are dropped when this function returns, meaning no
-            // owned file descriptors we set to be blocking here can be used in
-            // an async context in the future.
-            .try_for_each(|stdio| set_blocking(stdio).map_err(Error::ChildIo))?;
-
         let cmd = NonZeroByteSlice::new(&self.cmd).ok_or(Error::InvalidCommand)?;
 
         let session = Session::builder()

src/native_mux_impl/stdio.rs

@@ -42,7 +42,7 @@ fn cvt(ret: c_int) -> io::Result<c_int> {
     }
 }
 
-pub(super) fn set_blocking(fd: RawFd) -> io::Result<()> {
+fn set_blocking(fd: RawFd) -> io::Result<()> {
     let flags = cvt(unsafe { fcntl(fd, F_GETFL) })?;
     cvt(unsafe { fcntl(fd, F_SETFL, flags & (!O_NONBLOCK)) })?;
 
@@ -64,14 +64,21 @@ impl Fd {
     ///
     /// `T::into_raw_fd` must return a valid fd and transfers
     /// the ownershipt of it.
-    unsafe fn new_owned<T: IntoRawFd>(fd: T) -> Self {
+    unsafe fn new_owned<T: IntoRawFd>(fd: T) -> Result<Self, Error> {
         let raw_fd = fd.into_raw_fd();
-        Fd::Owned(OwnedFd::from_raw_fd(raw_fd))
+        // Create owned_fd so that the fd will be closed on error
+        let owned_fd = OwnedFd::from_raw_fd(raw_fd);
+
+        set_blocking(raw_fd).map_err(Error::ChildIo)?;
+
+        Ok(Fd::Owned(owned_fd))
     }
 }
 
-impl From<PipeRead> for Fd {
-    fn from(pipe_read: PipeRead) -> Self {
+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
@@ -80,8 +87,10 @@ impl From<PipeRead> for Fd {
     }
 }
 
-impl From<PipeWrite> for Fd {
-    fn from(pipe_write: PipeWrite) -> Self {
+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
@@ -97,9 +106,15 @@ impl Stdio {
             StdioImpl::Null => Ok((Fd::Null, None)),
             StdioImpl::Pipe => {
                 let (read, write) = create_pipe()?;
-                Ok((read.into(), Some(write)))
+                Ok((read.try_into()?, Some(write)))
+            }
+            StdioImpl::Fd(fd, owned) => {
+                let raw_fd = fd.as_raw_fd();
+                if *owned {
+                    set_blocking(raw_fd).map_err(Error::ChildIo)?;
+                }
+                Ok((Fd::Borrowed(raw_fd), None))
             }
-            StdioImpl::Fd(fd) => Ok((Fd::Borrowed(fd.as_raw_fd()), None)),
         }
     }
 
@@ -109,9 +124,15 @@ impl Stdio {
             StdioImpl::Null => Ok((Fd::Null, None)),
             StdioImpl::Pipe => {
                 let (read, write) = create_pipe()?;
-                Ok((write.into(), Some(read)))
+                Ok((write.try_into()?, Some(read)))
+            }
+            StdioImpl::Fd(fd, owned) => {
+                let raw_fd = fd.as_raw_fd();
+                if *owned {
+                    set_blocking(raw_fd).map_err(Error::ChildIo)?;
+                }
+                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,18 +40,23 @@ 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)
     }
-
-    #[cfg(feature = "native-mux")]
-    pub(super) fn is_inherited(&self) -> bool {
-        matches!(self.0, StdioImpl::Inherit)
-    }
 }
+/// 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
 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 {
@@ -64,14 +69,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))
     }
 }
 
@@ -82,7 +89,7 @@ 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))
             }
         }
     };