port forward: share one multiplexed tunnel across a window's listeners

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01EZ49AbZJYfm8NTp5yDPMab

src/client.rs

@@ -2769,6 +2769,7 @@ impl LoginConfigHandler {
             ConnType::PORT_FORWARD | ConnType::RDP => lr.set_port_forward(PortForward {
                 host: self.port_forward.0.clone(),
                 port: self.port_forward.1,
+                multiplex: crate::common::get_port_forward_mux_enabled(),
                 ..Default::default()
             }),
             ConnType::TERMINAL => {

src/common.rs

@@ -1167,6 +1167,13 @@ pub fn get_ipv6_punch_enabled() -> bool {
     )
 }
 
+pub fn get_port_forward_mux_enabled() -> bool {
+    config::option2bool(
+        keys::OPTION_ENABLE_PORT_FORWARD_MUX,
+        &get_local_option(keys::OPTION_ENABLE_PORT_FORWARD_MUX),
+    )
+}
+
 pub fn get_local_option(key: &str) -> String {
     let v = LocalConfig::get_option(key);
     if key == keys::OPTION_ENABLE_UDP_PUNCH || key == keys::OPTION_ENABLE_IPV6_PUNCH {

src/port_forward.rs

@@ -1,6 +1,7 @@
 use std::sync::{Arc, RwLock};
 
 use crate::client::*;
+use crate::port_forward_mux::{Claim, Tunnel, CHANNEL_WINDOW};
 use hbb_common::{
     allow_err, bail,
     config::READ_TIMEOUT,
@@ -79,6 +80,7 @@ pub async fn listen(
     lc: Arc<RwLock<LoginConfigHandler>>,
     remote_host: String,
     remote_port: i32,
+    tunnel: Arc<Tunnel>,
 ) -> ResultType<()> {
     let listener = tcp::new_listener(format!("127.0.0.1:{}", port), true).await?;
     let addr = listener.local_addr()?;
@@ -90,29 +92,94 @@ pub async fn listen(
     let mut ui_receiver = ui_receiver;
     loop {
         tokio::select! {
-            Ok((forward, addr)) = listener.accept() => {
+            // `addr` above the loop is the listener's own address; `run_rdp` needs
-                log::info!("new connection from {:?}", addr);
+            // that port. The accepted peer's address gets its own name so it can
-                let id = id.clone();
+            // never shadow it.
-                let password = password.clone();
+            Ok((forward, peer_addr)) = listener.accept() => {
-                let mut forward = Framed::new(forward, BytesCodec::new());
+                log::debug!("new connection from {:?}", peer_addr);
-                let mut close_port_forward = false;
+                let claim = tunnel.try_claim();
-                match connect_and_login(&id, &password, &mut ui_receiver, interface.clone(), &mut forward, key, token, is_rdp, &mut close_port_forward, &remote_host, remote_port).await {
+                let claim = match claim {
-                    Ok(Some(stream)) => {
+                    Claim::Wait => {
-                        let interface = interface.clone();
+                        // Keep servicing the UI while the establishing accept holds the
-                        tokio::spawn(async move {
+                        // prompt. `None` means establishment failed: this accept is dropped.
-                            if let Err(err) = run_forward(forward, stream).await {
+                        let resolved: Option<Claim> = loop {
-                                interface.msgbox("error", "Error", &err.to_string(), "");
+                            tokio::select! {
+                                // `tokio::select!` brings its own `Poll::{Ready, Pending}` into
+                                // scope, so `Ready` must be qualified here or it resolves to that.
+                                r = tunnel.wait_ready() => break match r {
+                                    crate::port_forward_mux::Ready::Muxed(h) => Some(Claim::Muxed(h)),
+                                    crate::port_forward_mux::Ready::Legacy => Some(Claim::Legacy),
+                                    crate::port_forward_mux::Ready::Failed => None,
+                                },
+                                d = ui_receiver.recv() => match d {
+                                    Some(Data::Close) => return Ok(()),
+                                    Some(Data::NewRDP) => run_rdp(addr.port(), &rdp_display_name(&lc, &id)),
+                                    _ => {}
+                                },
                             }
-                            log::info!("connection from {:?} closed", addr);
+                        };
-                       });
+                        match resolved {
+                            Some(c) => c,
+                            None => {
+                                log::debug!("tunnel failed while {:?} waited; dropping it", peer_addr);
+                                continue;
+                            }
+                        }
                     }
-                    _ if close_port_forward => {
+                    other => other,
-                        break;
+                };
+                match claim {
+                    Claim::Muxed(handle) => {
+                        if let Err(e) = handle.open(&remote_host, remote_port, forward, Vec::new()) {
+                            log::debug!("cannot open channel for {:?}: {}", peer_addr, e);
+                        }
                     }
-                    Err(err) => {
+                    Claim::Legacy => {
-                        interface.on_establish_connection_error(err.to_string());
+                        lc.write().unwrap().port_forward = (remote_host.clone(), remote_port);
+                        let mut forward = Framed::new(forward, BytesCodec::new());
+                        let mut close_port_forward = false;
+                        match connect_and_login(&id, &password, &mut ui_receiver, interface.clone(), &mut forward, key, token, is_rdp, &mut close_port_forward).await {
+                            Ok(Some(outcome)) => run_legacy(outcome, forward, peer_addr, interface.clone()),
+                            _ if close_port_forward => break,
+                            Err(err) => interface.on_establish_connection_error(err.to_string()),
+                            _ => {}
+                        }
                     }
-                    _ => {}
+                    Claim::Claimed => {
+                        lc.write().unwrap().port_forward = (remote_host.clone(), remote_port);
+                        let mut forward = Framed::new(forward, BytesCodec::new());
+                        let mut close_port_forward = false;
+                        match connect_and_login(&id, &password, &mut ui_receiver, interface.clone(), &mut forward, key, token, is_rdp, &mut close_port_forward).await {
+                            Ok(Some(outcome)) if outcome.mux => {
+                                let handle = tunnel.set_muxed(outcome.stream, interface.clone());
+                                if !outcome.local_eof {
+                                    let (socket, prebuf) = take_socket(forward, outcome.prebuf);
+                                    if let Err(e) = handle.open(&remote_host, remote_port, socket, prebuf) {
+                                        log::debug!("cannot open channel for {:?}: {}", peer_addr, e);
+                                    }
+                                }
+                            }
+                            Ok(Some(outcome)) => {
+                                tunnel.set_legacy();
+                                if outcome.local_eof {
+                                    log::debug!("legacy peer and local {:?} already gone", peer_addr);
+                                } else {
+                                    run_legacy(outcome, forward, peer_addr, interface.clone());
+                                }
+                            }
+                            _ if close_port_forward => {
+                                tunnel.set_failed();
+                                break;
+                            }
+                            Err(err) => {
+                                tunnel.set_failed();
+                                interface.on_establish_connection_error(err.to_string());
+                            }
+                            _ => tunnel.set_failed(),
+                        }
+                    }
+                    // Resolved above; a stray `Wait` just drops this accept.
+                    Claim::Wait => continue,
                 }
             }
             d = ui_receiver.recv() => {
@@ -132,6 +199,45 @@ pub async fn listen(
     Ok(())
 }
 
+/// Today's raw pipe, for peers without multiplexing.
+fn run_legacy(
+    outcome: LoginOutcome,
+    forward: Framed<TcpStream, BytesCodec>,
+    addr: std::net::SocketAddr,
+    interface: impl Interface,
+) {
+    let mut stream = outcome.stream;
+    let prebuf = outcome.prebuf;
+    tokio::spawn(async move {
+        stream.set_raw();
+        if !prebuf.is_empty() {
+            allow_err!(stream.send_bytes(prebuf.into()).await);
+        }
+        if let Err(err) = run_forward(forward, stream).await {
+            interface.msgbox("error", "Error", &err.to_string(), "");
+        }
+        log::info!("connection from {:?} closed", addr);
+    });
+}
+
+pub(crate) struct LoginOutcome {
+    stream: Stream,
+    mux: bool,
+    prebuf: Vec<u8>,
+    local_eof: bool,
+}
+
+fn peer_supports_mux(pi: &PeerInfo) -> bool {
+    pi.features.as_ref().map(|f| f.port_forward_mux).unwrap_or(false)
+}
+
+/// `into_inner()` would drop bytes the codec pulled but never yielded.
+fn take_socket(forward: Framed<TcpStream, BytesCodec>, mut prebuf: Vec<u8>) -> (TcpStream, Vec<u8>) {
+    let parts = forward.into_parts();
+    prebuf.extend_from_slice(&parts.read_buf);
+    (parts.io, prebuf)
+}
+
 async fn connect_and_login(
     id: &str,
     password: &str,
@@ -142,9 +248,7 @@ async fn connect_and_login(
     token: &str,
     is_rdp: bool,
     close_port_forward: &mut bool,
-    remote_host: &str,
+) -> ResultType<Option<LoginOutcome>> {
-    remote_port: i32,
-) -> ResultType<Option<Stream>> {
     let conn_type = if is_rdp {
         ConnType::RDP
     } else {
@@ -160,6 +264,8 @@ async fn connect_and_login(
         }
     }
     let mut buffer = Vec::new();
+    let mut local_eof = false;
+    let mut mux = false;
     let mut received = false;
     let mut challenge = None;
     let mut pending_login = None;
@@ -192,6 +298,7 @@ async fn connect_and_login(
                                 }
                             }
                             Some(login_response::Union::PeerInfo(pi)) => {
+                                mux = peer_supports_mux(&pi);
                                 interface.handle_peer_info(pi);
                                 break;
                             }
@@ -222,20 +329,24 @@ async fn connect_and_login(
                     _ => {}
                 }
             },
-            res = forward.next() => {
+            // Stop pulling once the pre-read buffer is a window deep; the
+            // rest waits in the kernel until the channel opens. A local EOF
+            // no longer aborts the login: the tunnel may still be wanted.
+            res = forward.next(), if !local_eof && buffer.len() < CHANNEL_WINDOW as usize => {
                 if let Some(Ok(bytes)) = res {
                     buffer.extend(bytes);
                 } else {
-                    return Ok(None);
+                    local_eof = true;
                 }
             },
         }
     }
-    stream.set_raw();
+    Ok(Some(LoginOutcome {
-    if !buffer.is_empty() {
+        stream,
-        allow_err!(stream.send_bytes(buffer.into()).await);
+        mux,
-    }
+        prebuf: buffer,
-    Ok(Some(stream))
+        local_eof,
+    }))
 }
 
 
@@ -334,170 +445,50 @@ async fn run_forward(forward: Framed<TcpStream, BytesCodec>, stream: Stream) ->
 }
 
 #[cfg(test)]
-mod login_tests {
+mod tests {
     use super::*;
-    use async_trait::async_trait;
-    use hbb_common::{
-        tcp::FramedStream,
-        tokio::time::{sleep, Duration},
-    };
-    use sha2::{Digest, Sha256};
 
-    /// A window's interface over its shared handler. `handle_hash` can pause
+    #[test]
-    /// before building the login, where the real one looks passwords up.
+    fn peer_supports_mux_reads_the_features_bit() {
-    #[derive(Clone)]
+        let mut pi = PeerInfo::new();
-    struct Ui {
+        assert!(!peer_supports_mux(&pi));
-        lc: Arc<RwLock<LoginConfigHandler>>,
+        pi.features = Some(Features { port_forward_mux: false, ..Default::default() }).into();
-        pause: Duration,
+        assert!(!peer_supports_mux(&pi));
-    }
+        pi.features = Some(Features { port_forward_mux: true, ..Default::default() }).into();
-
+        assert!(peer_supports_mux(&pi));
-    #[async_trait]
-    impl Interface for Ui {
-        fn send(&self, _data: Data) {}
-        fn msgbox(&self, _msgtype: &str, _title: &str, _text: &str, _link: &str) {}
-        fn handle_login_error(&self, _err: &str) -> bool {
-            false
-        }
-        fn handle_peer_info(&self, _pi: PeerInfo) {}
-        fn set_multiple_windows_session(&self, _sessions: Vec<WindowsSession>) {}
-        async fn handle_hash(&self, pass: &str, hash: Hash, peer: &mut Stream) -> bool {
-            sleep(self.pause).await;
-            crate::client::handle_hash(self.lc.clone(), pass, hash, self, peer).await
-        }
-        async fn handle_login_from_ui(
-            &self,
-            os_username: String,
-            os_password: String,
-            password: String,
-            remember: bool,
-            peer: &mut Stream,
-        ) {
-            crate::client::handle_login_from_ui(
-                self.lc.clone(),
-                os_username,
-                os_password,
-                password,
-                remember,
-                peer,
-            )
-            .await
-        }
-        async fn handle_test_delay(&self, _t: TestDelay, _peer: &mut Stream) {}
-        fn get_lch(&self) -> Arc<RwLock<LoginConfigHandler>> {
-            self.lc.clone()
-        }
-    }
-
-    fn window() -> Ui {
-        let mut lc = LoginConfigHandler::default();
-        lc.conn_type = ConnType::PORT_FORWARD;
-        Ui {
-            lc: Arc::new(RwLock::new(lc)),
-            pause: Duration::ZERO,
-        }
-    }
-
-    /// (our end, the peer's end) of one connection.
-    async fn loopback() -> (Stream, Stream) {
-        let l = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
-        let addr = l.local_addr().unwrap();
-        let client = tokio::net::TcpStream::connect(addr).await.unwrap();
-        let (server, _) = l.accept().await.unwrap();
-        (
-            Stream::Tcp(FramedStream::from(client, addr)),
-            Stream::Tcp(FramedStream::from(server, addr)),
-        )
-    }
-
-    async fn login_at(peer: &mut Stream) -> LoginRequest {
-        let bytes = peer.next().await.unwrap().unwrap();
-        Message::parse_from_bytes(&bytes)
-            .unwrap()
-            .login_request()
-            .clone()
-    }
-
-    fn target(lr: &LoginRequest) -> (String, i32) {
-        (lr.port_forward().host.clone(), lr.port_forward().port)
-    }
-
-    fn hash(challenge: &str) -> Hash {
-        Hash {
-            salt: "salt".to_owned(),
-            challenge: challenge.to_owned(),
-            ..Default::default()
-        }
-    }
-
-    /// What the peer expects for password `pw` under `hash(challenge)`.
-    fn digest(challenge: &str) -> Vec<u8> {
-        let mut h = Sha256::new();
-        h.update("pw");
-        h.update("salt");
-        let salted = h.finalize();
-        let mut h2 = Sha256::new();
-        h2.update(&salted[..]);
-        h2.update(challenge);
-        h2.finalize()[..].to_vec()
     }
 
     #[test]
-    fn mappings_logging_in_at_once_each_carry_their_own_target() {
+    fn port_forward_mux_defaults_to_on() {
-        let rt = tokio::runtime::Builder::new_current_thread()
+        use hbb_common::config::{keys, option2bool};
-            .enable_all()
+        // The `enable-` prefix is what makes an unset value mean on. Getting the
-            .build()
+        // key's name wrong silently flips the default, so pin it here.
-            .unwrap();
+        assert!(option2bool(keys::OPTION_ENABLE_PORT_FORWARD_MUX, ""));
-        rt.block_on(async {
+        assert!(option2bool(keys::OPTION_ENABLE_PORT_FORWARD_MUX, "Y"));
-            let mut ui = window();
+        assert!(!option2bool(keys::OPTION_ENABLE_PORT_FORWARD_MUX, "N"));
-            ui.pause = Duration::from_millis(50);
-            let (mut a, mut a_peer) = loopback().await;
-            let (mut b, mut b_peer) = loopback().await;
-            tokio::join!(
-                login_with_hash(&ui, "pw", hash("a"), "a", 1, &mut a),
-                login_with_hash(&ui, "pw", hash("b"), "b", 2, &mut b),
-            );
-            assert_eq!(target(&login_at(&mut a_peer).await), ("a".to_owned(), 1));
-            assert_eq!(target(&login_at(&mut b_peer).await), ("b".to_owned(), 2));
-        });
     }
 
     #[test]
-    fn a_mapping_answers_the_prompt_with_its_own_challenge() {
+    fn take_socket_hands_back_a_working_socket_and_the_prebuf() {
-        let rt = tokio::runtime::Builder::new_current_thread()
+        use hbb_common::tokio::io::{AsyncReadExt, AsyncWriteExt};
-            .enable_all()
+        let rt = tokio::runtime::Builder::new_current_thread().enable_all().build().unwrap();
-            .build()
-            .unwrap();
         rt.block_on(async {
-            let ui = window();
+            let l = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
-            let (mut a, mut a_peer) = loopback().await;
+            let addr = l.local_addr().unwrap();
-            let (mut b, mut b_peer) = loopback().await;
+            let mut client = TcpStream::connect(addr).await.unwrap();
-            // A's hash arrived last, so it is the one the handler holds.
+            let (server, _) = l.accept().await.unwrap();
-            assert!(login_with_hash(&ui, "pw", hash("a"), "a", 1, &mut a).await);
+            let mut framed = Framed::new(server, BytesCodec::new());
-            login_at(&mut a_peer).await;
+            client.write_all(b"abc").await.unwrap();
-            let typed = (String::new(), String::new(), "pw".to_owned(), false);
+            // Read through the codec, as connect_and_login does during login.
-            login_from_ui(&ui, &hash("b"), typed, "b", 2, &mut b).await;
+            let pulled = framed.next().await.unwrap().unwrap();
-            let lr = login_at(&mut b_peer).await;
+            assert_eq!(&pulled[..], b"abc");
-            assert_eq!(lr.password, digest("b"));
+            let (mut sock, prebuf) = take_socket(framed, pulled.to_vec());
-            assert_eq!(target(&lr), ("b".to_owned(), 2));
+            assert_eq!(prebuf, b"abc".to_vec());
-        });
+            // Bytes written after the handoff arrive on the bare socket.
-    }
+            client.write_all(b"def").await.unwrap();
-
+            let mut buf = [0u8; 3];
-    #[test]
+            sock.read_exact(&mut buf).await.unwrap();
-    fn a_password_typed_before_this_connections_hash_answers_it_when_it_comes() {
+            assert_eq!(&buf, b"def");
-        let rt = tokio::runtime::Builder::new_current_thread()
-            .enable_all()
-            .build()
-            .unwrap();
-        rt.block_on(async {
-            let ui = window();
-            let (mut b, mut b_peer) = loopback().await;
-            // The prompt's password reached B before its hash, and no other
-            // mapping has stored it in the handler yet.
-            let typed = (String::new(), String::new(), "pw".to_owned(), false);
-            assert!(hash_arrived(&ui, "", hash("b"), Some(typed), "b", 2, &mut b).await);
-            let lr = login_at(&mut b_peer).await;
-            assert_eq!(lr.password, digest("b"));
-            assert_eq!(target(&lr), ("b".to_owned(), 2));
         });
     }
 }

src/ui_session_interface.rs

@@ -1944,6 +1944,7 @@ pub async fn io_loop<T: InvokeUiSession>(handler: Session<T>, round: u32) {
     let key = crate::get_key(false).await;
     #[cfg(not(any(target_os = "android", target_os = "ios")))]
     if handler.is_port_forward() {
+        let tunnel = std::sync::Arc::new(crate::port_forward_mux::Tunnel::new());
         if handler.is_rdp() {
             let port = handler
                 .get_option("rdp_port".to_owned())
@@ -1958,7 +1959,7 @@ pub async fn io_loop<T: InvokeUiSession>(handler: Session<T>, round: u32) {
                 handler.get_option("rdp_password".to_owned()),
             );
             log::info!("Remote rdp port: {}", port);
-            start_one_port_forward(handler, 0, "".to_owned(), port, receiver, &key, &token).await;
+            start_one_port_forward(handler, 0, "".to_owned(), port, receiver, &key, &token, tunnel).await;
         } else if handler.args.len() == 0 {
             let pfs = handler.lc.read().unwrap().port_forwards.clone();
             let mut queues = HashMap::<i32, mpsc::UnboundedSender<Data>>::new();
@@ -1976,6 +1977,7 @@ pub async fn io_loop<T: InvokeUiSession>(handler: Session<T>, round: u32) {
                         let handler = handler.clone();
                         let key = key.clone();
                         let token = token.clone();
+                        let tunnel = tunnel.clone();
                         tokio::spawn(async move {
                             start_one_port_forward(
                                 handler,
@@ -1985,6 +1987,7 @@ pub async fn io_loop<T: InvokeUiSession>(handler: Session<T>, round: u32) {
                                 receiver,
                                 &key,
                                 &token,
+                                tunnel,
                             )
                             .await;
                         });
@@ -2023,6 +2026,7 @@ pub async fn io_loop<T: InvokeUiSession>(handler: Session<T>, round: u32) {
                 receiver,
                 &key,
                 &token,
+                tunnel,
             )
             .await;
         }
@@ -2042,6 +2046,7 @@ async fn start_one_port_forward<T: InvokeUiSession>(
     receiver: mpsc::UnboundedReceiver<Data>,
     key: &str,
     token: &str,
+    tunnel: std::sync::Arc<crate::port_forward_mux::Tunnel>,
 ) {
     if let Err(err) = crate::port_forward::listen(
         handler.get_id(),
@@ -2054,6 +2059,7 @@ async fn start_one_port_forward<T: InvokeUiSession>(
         handler.lc.clone(),
         remote_host,
         remote_port,
+        tunnel,
     )
     .await
     {