Merge pull request #86 from openssh-rust/optimize

Optimize and make `request_port_forward` easier to use

Author: jonhoo

.github/workflows/coverage.yml

@@ -49,6 +49,7 @@ jobs:
           # we cannot use 127.0.0.1 (the default here)
           # since we are running from a different container
           TEST_HOST: ssh://test-user@localhost:2222
+          XDG_RUNTIME_DIR: /tmp
       - name: Upload to codecov.io
         uses: codecov/codecov-action@v2
         with:

.github/workflows/minimal.yml

@@ -55,6 +55,7 @@ jobs:
         env:
           # makes all the ignored tests not ignored
           RUSTFLAGS: --cfg=ci
+          XDG_RUNTIME_DIR: /tmp
     services:
       openssh:
         image: linuxserver/openssh-server:amd64-latest

.github/workflows/test.yml

@@ -49,6 +49,7 @@ jobs:
         env:
           # makes all the ignored tests not ignored
           RUSTFLAGS: --cfg=ci
+          XDG_RUNTIME_DIR: /tmp
       - run: docker logs $(docker ps | grep openssh-server | awk '{print $1}')
         name: ssh container log
         if: ${{ failure() }}

src/builder.rs

@@ -1,11 +1,5 @@
 use super::{Error, Session};
 
-#[cfg(feature = "process-mux")]
-use super::process_impl;
-
-#[cfg(feature = "native-mux")]
-use super::native_mux_impl;
-
 use std::borrow::Cow;
 use std::ffi::OsString;
 use std::fs;
@@ -185,10 +179,8 @@ impl SessionBuilder {
     #[cfg(feature = "process-mux")]
     #[cfg_attr(docsrs, doc(cfg(feature = "process-mux")))]
     pub async fn connect<S: AsRef<str>>(&self, destination: S) -> Result<Session, Error> {
-        let destination = destination.as_ref();
-        let (builder, destination) = self.resolve(destination);
-        let tempdir = builder.launch_master(destination).await?;
-        Ok(process_impl::Session::new(tempdir, destination).into())
+        self.connect_impl(destination.as_ref(), Session::new_process_mux)
+            .await
     }
 
     /// Connect to the host at the given `host` over SSH using native mux, which will
@@ -207,10 +199,18 @@ impl SessionBuilder {
     #[cfg(feature = "native-mux")]
     #[cfg_attr(docsrs, doc(cfg(feature = "native-mux")))]
     pub async fn connect_mux<S: AsRef<str>>(&self, destination: S) -> Result<Session, Error> {
-        let destination = destination.as_ref();
+        self.connect_impl(destination.as_ref(), Session::new_native_mux)
+            .await
+    }
+
+    async fn connect_impl(
+        &self,
+        destination: &str,
+        f: fn(TempDir) -> Session,
+    ) -> Result<Session, Error> {
         let (builder, destination) = self.resolve(destination);
         let tempdir = builder.launch_master(destination).await?;
-        Ok(native_mux_impl::Session::new(tempdir).into())
+        Ok(f(tempdir))
     }
 
     fn resolve<'a, 'b>(&'a self, mut destination: &'b str) -> (Cow<'a, Self>, &'b str) {

src/changelog.rs

@@ -2,6 +2,21 @@
 use crate::*;
 
 /// TODO: RENAME THIS INTO THE NEXT VERSION BEFORE RELEASE
+///
+/// ## Added
+///  - `From<SocketAddr> for Socket<'static>`
+///  - `From<Cow<'a, Path>> for Socket<'a>`
+///  - `From<&'a Path> for Socket<'a>`
+///  - `From<PathBuf> for Socket<'static>`
+///  - `From<Box<Path>> for Socket<'static>`
+///  - `From<(IpAddr, u16)> for Socket<'static>`
+///  - `From<(Ipv4Addr, u16)> for Socket<'static>`
+///  - `From<(Ipv6Addr, u16)> for Socket<'static>`
+///
+/// ## Changed
+///  - [`Session::request_port_forward`] now takes `impl Into<...>`
+///    to make it much easier to use.
+///  - [`Socket::new`] now returns `Socket<'static>`
 #[doc(hidden)]
 pub mod unreleased {}
 

src/command.rs

@@ -23,7 +23,7 @@ impl From<super::process_impl::Command> for CommandImp {
 }
 
 #[cfg(feature = "native-mux")]
-impl<'s> From<super::native_mux_impl::Command> for CommandImp {
+impl From<super::native_mux_impl::Command> for CommandImp {
     fn from(imp: super::native_mux_impl::Command) -> Self {
         CommandImp::NativeMuxImpl(imp)
     }
@@ -139,7 +139,7 @@ impl<'s> Command<'s> {
     /// To pass multiple unescaped arguments see [`raw_args`](Command::raw_args).
     pub fn raw_arg<S: AsRef<OsStr>>(&mut self, arg: S) -> &mut Self {
         delegate!(&mut self.imp, imp, {
-            imp.raw_arg(arg);
+            imp.raw_arg(arg.as_ref());
         });
         self
     }

src/native_mux_impl/session.rs

@@ -1,4 +1,4 @@
-use super::{Command, Error, ForwardType, Socket};
+use super::{Command, Error};
 
 use std::ffi::OsStr;
 use std::os::unix::ffi::OsStrExt;
@@ -51,9 +51,9 @@ impl Session {
 
     pub(crate) async fn request_port_forward(
         &self,
-        forward_type: impl Into<ForwardType>,
-        listen_socket: impl Into<Socket<'_>>,
-        connect_socket: impl Into<Socket<'_>>,
+        forward_type: crate::ForwardType,
+        listen_socket: crate::Socket<'_>,
+        connect_socket: crate::Socket<'_>,
     ) -> Result<(), Error> {
         Connection::connect(&self.ctl)
             .await?
@@ -76,17 +76,13 @@ impl Session {
         Ok(())
     }
 
-    pub(crate) async fn close(mut self) -> Result<(), Error> {
-        // This also set self.tempdir to None so that Drop::drop would do nothing.
-        if let Some(tempdir) = self.tempdir.take() {
-            self.close_impl().await?;
+    pub(crate) async fn close(mut self) -> Result<Option<TempDir>, Error> {
+        // Take self.tempdir so that drop would do nothing
+        let tempdir = self.tempdir.take();
 
-            tempdir.close().map_err(Error::Cleanup)?;
-        } else {
-            self.close_impl().await?;
-        }
+        self.close_impl().await?;
 
-        Ok(())
+        Ok(tempdir)
     }
 
     pub(crate) fn detach(mut self) -> (Box<Path>, Option<Box<Path>>) {

src/native_mux_impl/stdio.rs

@@ -62,10 +62,7 @@ impl Stdio {
         }
     }
 
-    fn to_output(
-        &self,
-        get_inherit_rawfd: impl FnOnce() -> RawFd,
-    ) -> Result<(Fd, Option<PipeRead>), Error> {
+    fn to_output(&self, get_inherit_rawfd: fn() -> RawFd) -> Result<(Fd, Option<PipeRead>), Error> {
         match &self.0 {
             StdioImpl::Inherit => Ok((Fd::Borrowed(get_inherit_rawfd()), None)),
             StdioImpl::Null => Ok((Fd::Null, None)),

src/port_forwarding.rs

@@ -7,8 +7,8 @@ use std::ffi::OsStr;
 use std::borrow::Cow;
 use std::fmt;
 use std::io;
-use std::net::{SocketAddr, ToSocketAddrs};
-use std::path::Path;
+use std::net::{self, SocketAddr, ToSocketAddrs};
+use std::path::{Path, PathBuf};
 
 /// Type of forwarding
 #[derive(Copy, Clone, Debug, Eq, PartialEq)]
@@ -46,16 +46,64 @@ pub enum Socket<'a> {
     /// Tcp socket.
     TcpSocket(SocketAddr),
 }
-impl Socket<'_> {
-    /// Create a new TcpSocket
-    pub fn new<T: ToSocketAddrs>(addr: &T) -> Result<Self, io::Error> {
-        let mut it = addr.to_socket_addrs()?;
 
-        let addr = it.next().ok_or_else(|| {
-            io::Error::new(io::ErrorKind::Other, "no more socket addresses to try")
-        })?;
+impl From<SocketAddr> for Socket<'static> {
+    fn from(addr: SocketAddr) -> Self {
+        Socket::TcpSocket(addr)
+    }
+}
+
+macro_rules! impl_from_addr {
+    ($ip:ty) => {
+        impl From<($ip, u16)> for Socket<'static> {
+            fn from((ip, port): ($ip, u16)) -> Self {
+                SocketAddr::new(ip.into(), port).into()
+            }
+        }
+    };
+}
+
+impl_from_addr!(net::IpAddr);
+impl_from_addr!(net::Ipv4Addr);
+impl_from_addr!(net::Ipv6Addr);
+
+impl<'a> From<Cow<'a, Path>> for Socket<'a> {
+    fn from(path: Cow<'a, Path>) -> Self {
+        Socket::UnixSocket { path }
+    }
+}
+
+impl<'a> From<&'a Path> for Socket<'a> {
+    fn from(path: &'a Path) -> Self {
+        Socket::UnixSocket {
+            path: Cow::Borrowed(path),
+        }
+    }
+}
+
+impl From<PathBuf> for Socket<'static> {
+    fn from(path: PathBuf) -> Self {
+        Socket::UnixSocket {
+            path: Cow::Owned(path),
+        }
+    }
+}
 
-        Ok(Socket::TcpSocket(addr))
+impl From<Box<Path>> for Socket<'static> {
+    fn from(path: Box<Path>) -> Self {
+        Socket::UnixSocket {
+            path: Cow::Owned(path.into()),
+        }
+    }
+}
+
+impl Socket<'_> {
+    /// Create a new TcpSocket
+    pub fn new<T: ToSocketAddrs>(addr: &T) -> Result<Socket<'static>, io::Error> {
+        addr.to_socket_addrs()?
+            .next()
+            .ok_or_else(|| io::Error::new(io::ErrorKind::Other, "no more socket addresses to try"))
+            .map(Socket::TcpSocket)
     }
 
     #[cfg(feature = "process-mux")]

src/process_impl/session.rs

@@ -14,19 +14,17 @@ use tempfile::TempDir;
 pub(crate) struct Session {
     tempdir: Option<TempDir>,
     ctl: Box<Path>,
-    addr: Box<str>,
     master_log: Option<Box<Path>>,
 }
 
 impl Session {
-    pub(crate) fn new(tempdir: TempDir, addr: &str) -> Self {
+    pub(crate) fn new(tempdir: TempDir) -> Self {
         let log = tempdir.path().join("log").into_boxed_path();
         let ctl = tempdir.path().join("master").into_boxed_path();
 
         Self {
             tempdir: Some(tempdir),
             ctl,
-            addr: addr.into(),
             master_log: Some(log),
         }
     }
@@ -35,10 +33,6 @@ impl Session {
         Self {
             tempdir: None,
             ctl,
-            // ssh don't care about the addr as long as we have the ctl.
-            //
-            // I tested this behavior on OpenSSH_8.9p1 Ubuntu-3, OpenSSL 3.0.2 15 Mar 2022
-            addr: "none".into(),
             master_log,
         }
     }
@@ -51,7 +45,10 @@ impl Session {
             .arg("-o")
             .arg("BatchMode=yes")
             .args(args)
-            .arg(&*self.addr);
+            // ssh does not care about the addr as long as we have passed
+            // `-S &*self.ctl`.
+            // It is tested on OpenSSH 8.2p1, 8.9p1, 9.0p1
+            .arg("none");
         cmd
     }
 
@@ -107,18 +104,18 @@ impl Session {
 
     pub(crate) async fn request_port_forward(
         &self,
-        forward_type: impl Into<ForwardType>,
-        listen_socket: impl Into<Socket<'_>>,
-        connect_socket: impl Into<Socket<'_>>,
+        forward_type: ForwardType,
+        listen_socket: Socket<'_>,
+        connect_socket: Socket<'_>,
     ) -> Result<(), Error> {
-        let flag = match forward_type.into() {
+        let flag = match forward_type {
             ForwardType::Local => OsStr::new("-L"),
             ForwardType::Remote => OsStr::new("-R"),
         };
 
-        let mut forwarding = listen_socket.into().as_os_str().into_owned();
+        let mut forwarding = listen_socket.as_os_str().into_owned();
         forwarding.push(":");
-        forwarding.push(connect_socket.into().as_os_str());
+        forwarding.push(connect_socket.as_os_str());
 
         let port_forwarding = self
             .new_cmd(&[OsStr::new("-fNT"), flag, &*forwarding])
@@ -178,17 +175,13 @@ impl Session {
         Ok(())
     }
 
-    pub(crate) async fn close(mut self) -> Result<(), Error> {
+    pub(crate) async fn close(mut self) -> Result<Option<TempDir>, Error> {
         // Take self.tempdir so that drop would do nothing
-        if let Some(tempdir) = self.tempdir.take() {
-            self.close_impl().await?;
+        let tempdir = self.tempdir.take();
 
-            tempdir.close().map_err(Error::Cleanup)?;
-        } else {
-            self.close_impl().await?;
-        }
+        self.close_impl().await?;
 
-        Ok(())
+        Ok(tempdir)
     }
 
     pub(crate) fn detach(mut self) -> (Box<Path>, Option<Box<Path>>) {