From b798c9c209f9d7e23f804992e136fd7f6d2b3aa1 Mon Sep 17 00:00:00 2001
From: Jae-Won Chung <jwnchung@umich.edu>
Date: Tue, 10 Sep 2024 09:54:43 -0400
Subject: [PATCH] Rename API and to `read` 0 test in remote forward test

---
 src/native_mux_impl/session.rs |  8 ++++++--
 src/process_impl/session.rs    |  2 +-
 src/session.rs                 |  8 +++-----
 tests/openssh.rs               | 35 +++++++++++++++++-----------------
 4 files changed, 27 insertions(+), 26 deletions(-)

diff --git a/src/native_mux_impl/session.rs b/src/native_mux_impl/session.rs
index a712967f0..06cdecf88 100644
--- a/src/native_mux_impl/session.rs
+++ b/src/native_mux_impl/session.rs
@@ -67,13 +67,17 @@ impl Session {
         Ok(())
     }
 
-    pub(crate) async fn cancel_port_forward(
+    pub(crate) async fn close_port_forward(
         &self,
         forward_type: crate::ForwardType,
         listen_socket: crate::Socket<'_>,
         connect_socket: crate::Socket<'_>,
     ) -> Result<(), Error> {
-        unimplemented!("Port forwarding cancellation is not implemented yet")
+        Connection::connect(&self.ctl).await?.close_port_forward(
+            forward_type.into(),
+            &listen_socket.into(),
+            &connect_socket.into(),
+        )
     }
 
     async fn close_impl(&self) -> Result<(), Error> {
diff --git a/src/process_impl/session.rs b/src/process_impl/session.rs
index 81bd79f6e..612ba3132 100644
--- a/src/process_impl/session.rs
+++ b/src/process_impl/session.rs
@@ -139,7 +139,7 @@ impl Session {
         }
     }
 
-    pub(crate) async fn cancel_port_forward(
+    pub(crate) async fn close_port_forward(
         &self,
         forward_type: ForwardType,
         listen_socket: Socket<'_>,
diff --git a/src/session.rs b/src/session.rs
index 1e86edce3..121874c57 100644
--- a/src/session.rs
+++ b/src/session.rs
@@ -487,19 +487,17 @@ impl Session {
         })
     }
 
-    /// Cancel a previously established local/remote port forwarding.
+    /// Close a previously established local/remote port forwarding.
     ///
     /// The same set of arguments should be passed as when the port forwarding was requested.
-    ///
-    /// Currently, cancelling port forwarding is only supported for the process mux impl.
-    pub async fn cancel_port_forward(
+    pub async fn close_port_forward(
         &self,
         forward_type: impl Into<ForwardType>,
         listen_socket: impl Into<Socket<'_>>,
         connect_socket: impl Into<Socket<'_>>,
     ) -> Result<(), Error> {
         delegate!(&self.0, imp, {
-            imp.cancel_port_forward(
+            imp.close_port_forward(
                 forward_type.into(),
                 listen_socket.into(),
                 connect_socket.into(),
diff --git a/tests/openssh.rs b/tests/openssh.rs
index f627485d2..5dd812282 100644
--- a/tests/openssh.rs
+++ b/tests/openssh.rs
@@ -851,17 +851,18 @@ async fn remote_socket_forward() {
 
         assert_eq!(DATA, &buffer);
 
+        eprintln!("Canceling port forward");
+        session
+            .close_port_forward(ForwardType::Remote, (loopback(), *port), &*unix_socket)
+            .await
+            .unwrap();
+
+        eprintln!("Trying to connect again");
+        assert_eq!(output.read(&mut buffer).await.unwrap(), 0);
+
         drop(output);
         drop(output_listener);
 
-        if !cfg!(feature = "native-mux") {
-            eprintln!("Canceling port forward");
-            session
-                .cancel_port_forward(ForwardType::Remote, (loopback(), *port), &*unix_socket)
-                .await
-                .unwrap();
-        }
-
         eprintln!("Waiting for session to end");
         let output = child.wait_with_output().await.unwrap();
         eprintln!("remote_socket_forward: {:#?}", output);
@@ -910,17 +911,15 @@ async fn local_socket_forward() {
 
         drop(output);
 
-        if !cfg!(feature = "native-mux") {
-            eprintln!("Canceling port forward");
-            session
-                .cancel_port_forward(ForwardType::Local, &*unix_socket, (loopback(), port))
-                .await
-                .unwrap();
+        eprintln!("Closing port forward");
+        session
+            .close_port_forward(ForwardType::Local, &*unix_socket, (loopback(), port))
+            .await
+            .unwrap();
 
-            eprintln!("Trying to connect again");
-            let e = UnixStream::connect(&unix_socket).await.unwrap_err();
-            assert_eq!(e.kind(), io::ErrorKind::ConnectionRefused);
-        }
+        eprintln!("Trying to connect again");
+        let e = UnixStream::connect(&unix_socket).await.unwrap_err();
+        assert_eq!(e.kind(), io::ErrorKind::ConnectionRefused);
 
         eprintln!("Waiting for session to end");
         let output = child.wait_with_output().await.unwrap();