Merge pull request #80 from openssh-rust/feature/resume

Feature/resume

Author: jonhoo

.github/workflows/coverage.yml

@@ -30,7 +30,7 @@ jobs:
         uses: taiki-e/install-action@cargo-llvm-cov
       - run: |
           # Wait for startup of openssh-server
-          sleep 30
+          timeout 15 ./wait_for_sshd_start_up.sh
           chmod 600 .test-key
           mkdir /tmp/openssh-rs
           ssh -i .test-key -v -p 2222 -l test-user localhost -o StrictHostKeyChecking=accept-new -o UserKnownHostsFile=/tmp/openssh-rs/known_hosts whoami

.github/workflows/minimal.yml

@@ -29,6 +29,8 @@ jobs:
           toolchain: stable
       - uses: actions/checkout@v2
       - run: |
+          # Wait for startup of openssh-server
+          timeout 15 ./wait_for_sshd_start_up.sh
           chmod 600 .test-key
           mkdir /tmp/openssh-rs
           ssh -i .test-key -v -p 2222 -l test-user 127.0.0.1 -o StrictHostKeyChecking=accept-new -o UserKnownHostsFile=/tmp/openssh-rs/known_hosts whoami

.github/workflows/test.yml

@@ -29,6 +29,8 @@ jobs:
           toolchain: ${{ matrix.toolchain }}
       - uses: actions/checkout@v2
       - run: |
+          # Wait for startup of openssh-server
+          timeout 15 ./wait_for_sshd_start_up.sh
           chmod 600 .test-key
           mkdir /tmp/openssh-rs
           ssh -i .test-key -v -p 2222 -l test-user 127.0.0.1 -o StrictHostKeyChecking=accept-new -o UserKnownHostsFile=/tmp/openssh-rs/known_hosts whoami

run_ci_tests.sh

@@ -41,3 +41,5 @@ cargo test \
     --all-features \
     --no-fail-fast \
     -- --test-threads=3 # Use test-threads=3 so that the output is readable
+
+rm -r .ssh-connection*

src/changelog.rs

@@ -2,6 +2,12 @@
 use crate::*;
 
 /// TODO: RENAME THIS INTO THE NEXT VERSION BEFORE RELEASE
+///
+/// ## Added
+///  - [`Session::resume`]
+///  - [`Session::resume_mux`]
+///  - [`Session::leak`]
+///  - [`Session::force_terminate`]
 #[doc(hidden)]
 pub mod unreleased {}
 

src/native_mux_impl/session.rs

@@ -24,6 +24,10 @@ impl Session {
         }
     }
 
+    pub(crate) fn resume(ctl: Box<Path>, _master_log: Option<Box<Path>>) -> Self {
+        Self { tempdir: None, ctl }
+    }
+
     pub(crate) async fn check(&self) -> Result<(), Error> {
         Connection::connect(&self.ctl)
             .await?
@@ -63,19 +67,37 @@ 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.
-        let tempdir = self.tempdir.take().unwrap();
-
+    async fn close_impl(&self) -> Result<(), Error> {
         Connection::connect(&self.ctl)
             .await?
             .request_stop_listening()
             .await?;
 
-        tempdir.close().map_err(Error::Cleanup)?;
+        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?;
+
+            tempdir.close().map_err(Error::Cleanup)?;
+        } else {
+            self.close_impl().await?;
+        }
 
         Ok(())
     }
+
+    pub(crate) fn detach(mut self) -> (Box<Path>, Option<Box<Path>>) {
+        (
+            self.ctl.clone(),
+            self.tempdir.take().map(TempDir::into_path).map(|mut path| {
+                path.push("log");
+                path.into_boxed_path()
+            }),
+        )
+    }
 }
 
 impl Drop for Session {

src/process_impl/session.rs

@@ -15,7 +15,7 @@ pub(crate) struct Session {
     tempdir: Option<TempDir>,
     ctl: Box<Path>,
     addr: Box<str>,
-    master_log: Box<Path>,
+    master_log: Option<Box<Path>>,
 }
 
 impl Session {
@@ -27,7 +27,19 @@ impl Session {
             tempdir: Some(tempdir),
             ctl,
             addr: addr.into(),
-            master_log: log,
+            master_log: Some(log),
+        }
+    }
+
+    pub(crate) fn resume(ctl: Box<Path>, master_log: Option<Box<Path>>) -> Self {
+        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,
         }
     }
 
@@ -130,13 +142,12 @@ impl Session {
         }
     }
 
-    pub(crate) async fn close(mut self) -> Result<(), Error> {
-        let mut exit_cmd = self.new_cmd(&["-O", "exit"]);
-
-        // Take self.tempdir so that drop would do nothing
-        let tempdir = self.tempdir.take().unwrap();
-
-        let exit = exit_cmd.output().await.map_err(Error::Ssh)?;
+    async fn close_impl(&self) -> Result<(), Error> {
+        let exit = self
+            .new_cmd(&["-O", "exit"])
+            .output()
+            .await
+            .map_err(Error::Ssh)?;
 
         if let Some(master_error) = self.discover_master_error() {
             return Err(master_error);
@@ -164,13 +175,29 @@ impl Session {
             )));
         }
 
-        tempdir.close().map_err(Error::Cleanup)?;
+        Ok(())
+    }
+
+    pub(crate) async fn close(mut self) -> Result<(), Error> {
+        // Take self.tempdir so that drop would do nothing
+        if let Some(tempdir) = self.tempdir.take() {
+            self.close_impl().await?;
+
+            tempdir.close().map_err(Error::Cleanup)?;
+        } else {
+            self.close_impl().await?;
+        }
 
         Ok(())
     }
 
+    pub(crate) fn detach(mut self) -> (Box<Path>, Option<Box<Path>>) {
+        self.tempdir.take().map(TempDir::into_path);
+        (self.ctl.clone(), self.master_log.take())
+    }
+
     fn discover_master_error(&self) -> Option<Error> {
-        let err = match fs::read_to_string(&self.master_log) {
+        let err = match fs::read_to_string(self.master_log.as_ref()?) {
             Ok(err) => err,
             Err(e) => return Some(Error::Master(e)),
         };

src/session.rs

@@ -65,6 +65,31 @@ impl From<native_mux_impl::Session> for Session {
 // TODO: UserKnownHostsFile for custom known host fingerprint.
 
 impl Session {
+    /// Resume the connection using path to control socket and
+    /// path to ssh multiplex output log.
+    ///
+    /// If you do not use `-E` option (or redirection) to write
+    /// the log of the ssh multiplex master to the disk, you can
+    /// simply pass `None` to `master_log`.
+    ///
+    /// [`Session`] created this way will not be terminated on drop,
+    /// but can be forced terminated by [`Session::close`].
+    ///
+    /// This connects to the ssh multiplex master using process mux impl.
+    #[cfg(feature = "process-mux")]
+    #[cfg_attr(docsrs, doc(cfg(feature = "process-mux")))]
+    pub fn resume(ctl: Box<Path>, master_log: Option<Box<Path>>) -> Self {
+        process_impl::Session::resume(ctl, master_log).into()
+    }
+
+    /// Same as [`Session::resume`] except that it connects to
+    /// the ssh multiplex master using native mux impl.
+    #[cfg(feature = "native-mux")]
+    #[cfg_attr(docsrs, doc(cfg(feature = "native-mux")))]
+    pub fn resume_mux(ctl: Box<Path>, master_log: Option<Box<Path>>) -> Self {
+        native_mux_impl::Session::resume(ctl, master_log).into()
+    }
+
     /// Connect to the host at the given `host` over SSH using process impl, which will
     /// spawn a new ssh process for each `Child` created.
     ///
@@ -292,7 +317,18 @@ impl Session {
     }
 
     /// Terminate the remote connection.
+    ///
+    /// This destructor terminates the ssh multiplex server
+    /// regardless of how it was created.
     pub async fn close(self) -> Result<(), Error> {
         delegate!(self.0, imp, { imp.close().await })
     }
+
+    /// Detach the lifetime of underlying ssh multiplex master
+    /// from this `Session`.
+    ///
+    /// Return (path to control socket, path to ssh multiplex output log)
+    pub fn detach(self) -> (Box<Path>, Option<Box<Path>>) {
+        delegate!(self.0, imp, { imp.detach() })
+    }
 }

start_sshd.sh

@@ -28,9 +28,10 @@ chmod 600 .test-key
 mkdir -p "$RUNTIME_DIR/openssh-rs/"
 
 echo Waiting for sshd to be up
-while ! ssh -i .test-key -v -p 2222 -l test-user $HOSTNAME -o StrictHostKeyChecking=no -o UserKnownHostsFile="$RUNTIME_DIR/openssh-rs/known_hosts" whoami; do
-    sleep 3
-done
+timeout 30 ./wait_for_sshd_start_up.sh
+
+# Add the ip to  known_hosts file
+ssh -i .test-key -v -p 2222 -l test-user $HOSTNAME -o StrictHostKeyChecking=no -o UserKnownHostsFile="$RUNTIME_DIR/openssh-rs/known_hosts" whoami
 
 # Create sshd_started in runtime directory so that it is auto removed on restart.
 touch "$RUNTIME_DIR/openssh-rs/sshd_started"

tests/openssh.rs

@@ -829,6 +829,96 @@ async fn local_socket_forward() {
     }
 }
 
+#[tokio::test]
+#[cfg_attr(not(ci), ignore)]
+#[cfg(feature = "process-mux")]
+async fn test_detach_and_resume_process_mux() {
+    for session1 in connects().await {
+        session1.check().await.unwrap();
+
+        // First detach
+        let (ctl1, master_log1) = session1.detach();
+
+        // First resume
+        let session2 = Session::resume(ctl1, master_log1);
+        session2.check().await.unwrap();
+
+        // Second detach to ensure detach handles tempdir
+        // set to None correctly.
+        let (ctl2, master_log2) = session2.detach();
+
+        // Second resume to ensure close handles tempdir set to None correctly
+        let session3 = Session::resume(ctl2, master_log2);
+        session3.check().await.unwrap();
+
+        session3.close().await.unwrap();
+    }
+
+    // test close
+    for session1 in connects().await {
+        session1.check().await.unwrap();
+
+        let (ctl1, master_log1) = session1.detach();
+
+        let ctl = ctl1.clone();
+
+        let session2 = Session::resume(ctl1, master_log1);
+        session2.check().await.unwrap();
+
+        session2.close().await.unwrap();
+
+        // Wait for ssh multiplex master to clean up and exit.
+        sleep(Duration::from_secs(3)).await;
+
+        assert!(!ctl.exists());
+    }
+}
+
+#[tokio::test]
+#[cfg_attr(not(ci), ignore)]
+#[cfg(feature = "native-mux")]
+async fn test_detach_and_resume_native_mux() {
+    for session1 in connects().await {
+        session1.check().await.unwrap();
+
+        // First detach
+        let (ctl1, master_log1) = session1.detach();
+
+        // First resume_mux
+        let session2 = Session::resume_mux(ctl1, master_log1);
+        session2.check().await.unwrap();
+
+        // Second detach to ensure detach handles tempdir
+        // set to None correctly.
+        let (ctl2, master_log2) = session2.detach();
+
+        // Second resume_mux to ensure close handles tempdir set to None correctly
+        let session3 = Session::resume_mux(ctl2, master_log2);
+        session3.check().await.unwrap();
+
+        session3.close().await.unwrap();
+    }
+
+    // test close
+    for session1 in connects().await {
+        session1.check().await.unwrap();
+
+        let (ctl1, master_log1) = session1.detach();
+
+        let ctl = ctl1.clone();
+
+        let session2 = Session::resume_mux(ctl1, master_log1);
+        session2.check().await.unwrap();
+
+        session2.close().await.unwrap();
+
+        // Wait for ssh multiplex master to clean up and exit.
+        sleep(Duration::from_secs(3)).await;
+
+        assert!(!ctl.exists());
+    }
+}
+
 #[tokio::test]
 #[cfg_attr(not(ci), ignore)]
 async fn test_sftp_subsystem() {