Feat: Port forwarding closure (#57)

* Implement port forwarding closure

* Fix fmt

* Format nits and doc comments

* Fix fmt

* Test socket accept after closure

* Try read instead

* `read` returning 0 seems to mean a closed socket

* Check `read` 0 for local forwarding as well

* Fix fmt

Author: jaywonchung

crates/mux-client/src/connection.rs

@@ -375,11 +375,60 @@ impl Connection {
         Ok(())
     }
 
+    async fn send_close_fwd_request(&mut self, request_id: u32, fwd: &Fwd<'_>) -> Result<()> {
+        let (fwd_mode, listen_socket, connect_socket) = fwd.as_serializable();
+        let (listen_addr, listen_port) = listen_socket.as_serializable();
+        let (connect_addr, connect_port) = connect_socket.as_ref().as_serializable();
+
+        let serialized_listen_port = serialize_u32(listen_port);
+        let serialized_connect_port = serialize_u32(connect_port);
+
+        let listen_addr_len: u32 = listen_addr.get_len_as_u32()?;
+        let connect_addr_len: u32 = connect_addr.get_len_as_u32()?;
+
+        let request = Request::CloseFwd {
+            request_id,
+            fwd_mode,
+        };
+
+        // Serialize
+        self.reset_serializer();
+
+        request.serialize(&mut self.serializer)?;
+        let serialized_header = self.serializer.create_header(
+            // len
+            4 +
+            listen_addr_len +
+            // port
+            4 +
+            // len
+            4 +
+            connect_addr_len
+            // port
+            + 4,
+        )?;
+
+        let serialized_listen_addr_len = serialize_u32(listen_addr_len);
+        let serialized_connect_addr_len = serialize_u32(connect_addr_len);
+
+        // Write them to self.raw_conn
+        let mut io_slices = [
+            IoSlice::new(&serialized_header),
+            IoSlice::new(&self.serializer.output),
+            IoSlice::new(&serialized_listen_addr_len),
+            IoSlice::new(listen_addr.into_inner()),
+            IoSlice::new(&serialized_listen_port),
+            IoSlice::new(&serialized_connect_addr_len),
+            IoSlice::new(connect_addr.into_inner()),
+            IoSlice::new(&serialized_connect_port),
+        ];
+
+        write_vectored_all(&mut self.raw_conn, &mut io_slices).await?;
+
+        Ok(())
+    }
+
     /// Request for local/remote port forwarding.
-    ///
-    /// # Warning
-    ///
-    /// Local port forwarding hasn't been tested yet.
     pub async fn request_port_forward(
         &mut self,
         forward_type: ForwardType,
@@ -426,6 +475,53 @@ impl Connection {
         }
     }
 
+    /// Request for local/remote port forwarding closure.
+    pub async fn close_port_forward(
+        &mut self,
+        forward_type: ForwardType,
+        listen_socket: &Socket<'_>,
+        connect_socket: &Socket<'_>,
+    ) -> Result<()> {
+        use ForwardType::*;
+        use Response::*;
+
+        let fwd = match forward_type {
+            Local => Fwd::Local {
+                listen_socket,
+                connect_socket,
+            },
+            Remote => Fwd::Remote {
+                listen_socket,
+                connect_socket,
+            },
+        };
+
+        let request_id = self.get_request_id();
+        self.send_close_fwd_request(request_id, &fwd).await?;
+
+        match self.read_response().await? {
+            Ok { response_id } => Self::check_response_id(request_id, response_id),
+            PermissionDenied {
+                response_id,
+                reason,
+            } => {
+                Self::check_response_id(request_id, response_id)?;
+                Err(Error::PermissionDenied(reason))
+            }
+            Failure {
+                response_id,
+                reason,
+            } => {
+                Self::check_response_id(request_id, response_id)?;
+                Err(Error::RequestFailure(reason))
+            }
+            response => Err(Error::invalid_server_response(
+                &"Ok, PermissionDenied or Failure",
+                &response,
+            )),
+        }
+    }
+
     /// **UNTESTED** Return remote port opened for dynamic forwarding.
     pub async fn request_dynamic_forward(
         &mut self,
@@ -622,26 +718,27 @@ mod tests {
     }
     run_test!(test_unordered_open_new_session, test_open_new_session_impl);
 
-    async fn test_remote_socket_forward_impl(mut conn: Connection) {
+    async fn test_remote_socket_forward_impl(mut conn0: Connection, mut conn1: Connection) {
         let path = Path::new("/tmp/openssh-remote-forward.socket");
 
         let output_listener = TcpListener::bind(("127.0.0.1", 1234)).await.unwrap();
 
         eprintln!("Requesting port forward");
-        conn.request_port_forward(
-            ForwardType::Remote,
-            &Socket::UnixSocket { path: path.into() },
-            &Socket::TcpSocket {
-                port: 1234,
-                host: "127.0.0.1".into(),
-            },
-        )
-        .await
-        .unwrap();
+        conn0
+            .request_port_forward(
+                ForwardType::Remote,
+                &Socket::UnixSocket { path: path.into() },
+                &Socket::TcpSocket {
+                    port: 1234,
+                    host: "127.0.0.1".into(),
+                },
+            )
+            .await
+            .unwrap();
 
         eprintln!("Creating remote process");
         let cmd = format!("/usr/bin/socat OPEN:/data,rdonly UNIX-CONNECT:{:#?}", path);
-        let (established_session, stdios) = create_remote_process(conn, &cmd).await;
+        let (established_session, stdios) = create_remote_process(conn0, &cmd).await;
 
         eprintln!("Waiting for connection");
         let (mut output, _addr) = output_listener.accept().await.unwrap();
@@ -655,8 +752,24 @@ mod tests {
 
         assert_eq!(DATA, &buffer);
 
-        drop(output);
+        eprintln!("Closing port forward");
+        conn1
+            .close_port_forward(
+                ForwardType::Remote,
+                &Socket::UnixSocket { path: path.into() },
+                &Socket::TcpSocket {
+                    port: 1234,
+                    host: "127.0.0.1".into(),
+                },
+            )
+            .await
+            .unwrap();
+
+        eprintln!("Checking whether the forwarded socket is closed");
+        assert_eq!(output.read(&mut buffer).await.unwrap(), 0);
+
         drop(output_listener);
+        drop(output);
         drop(stdios);
 
         eprintln!("Waiting for session to end");
@@ -667,13 +780,13 @@ mod tests {
                 if exit_value.unwrap() == 0
         );
     }
-    run_test!(
+    run_test2!(
         test_unordered_remote_socket_forward,
         test_remote_socket_forward_impl
     );
 
     async fn test_local_socket_forward_impl(conn0: Connection, mut conn1: Connection) {
-        let path = Path::new("/tmp/openssh-local-forward.socket").into();
+        let path: Cow<'_, _> = Path::new("/tmp/openssh-local-forward.socket").into();
 
         eprintln!("Creating remote process");
         let cmd = format!("socat -u OPEN:/data UNIX-LISTEN:{:#?} >/dev/stderr", path);
@@ -689,7 +802,7 @@ mod tests {
                     port: 1235,
                     host: "127.0.0.1".into(),
                 },
-                &Socket::UnixSocket { path },
+                &Socket::UnixSocket { path: path.clone() },
             )
             .await
             .unwrap();
@@ -705,6 +818,22 @@ mod tests {
 
         assert_eq!(DATA, buffer);
 
+        eprintln!("Closing port forward");
+        conn1
+            .close_port_forward(
+                ForwardType::Local,
+                &Socket::TcpSocket {
+                    port: 1235,
+                    host: "127.0.0.1".into(),
+                },
+                &Socket::UnixSocket { path },
+            )
+            .await
+            .unwrap();
+
+        eprintln!("Checking whether the forwarded socket is closed");
+        assert_eq!(output.read(&mut buffer).await.unwrap(), 0);
+
         drop(output);
         drop(stdios);
 

crates/mux-client/src/constants.rs

@@ -10,6 +10,7 @@ def_constants!(MUX_MSG_HELLO, 0x00000001);
 def_constants!(MUX_C_NEW_SESSION, 0x10000002);
 def_constants!(MUX_C_ALIVE_CHECK, 0x10000004);
 def_constants!(MUX_C_OPEN_FWD, 0x10000006);
+def_constants!(MUX_C_CLOSE_FWD, 0x10000007);
 def_constants!(MUX_C_STOP_LISTENING, 0x10000009);
 def_constants!(MUX_S_OK, 0x80000001);
 def_constants!(MUX_S_PERMISSION_DENIED, 0x80000002);

crates/mux-client/src/request.rs

@@ -47,6 +47,10 @@ pub(crate) enum Request {
     /// `Request::RemotePort`.
     OpenFwd { request_id: u32, fwd_mode: u32 },
 
+    /// A server may reply with `Response::Ok`, `Response::PermissionDenied`,
+    /// or `Response::Failure`.
+    CloseFwd { request_id: u32, fwd_mode: u32 },
+
     /// A client may request the master to stop accepting new multiplexing requests
     /// and remove its listener socket.
     ///
@@ -87,6 +91,15 @@ impl Serialize for Request {
                 "OpenFwd",
                 &(*request_id, fwd_mode),
             ),
+            CloseFwd {
+                request_id,
+                fwd_mode,
+            } => serializer.serialize_newtype_variant(
+                "Request",
+                MUX_C_CLOSE_FWD,
+                "CloseFwd",
+                &(*request_id, fwd_mode),
+            ),
             StopListening { request_id } => serializer.serialize_newtype_variant(
                 "Request",
                 MUX_C_STOP_LISTENING,