diff --git a/libs/hbb_common b/libs/hbb_common index 6918d3593..26da70418 160000 --- a/libs/hbb_common +++ b/libs/hbb_common @@ -1 +1 @@ -Subproject commit 6918d35935b48b03ab24f11dfda48b6379807cbf +Subproject commit 26da70418c1ad560385386da3b0112cdfd63f400 diff --git a/src/client.rs b/src/client.rs index 5b059a7cc..d17f62587 100644 --- a/src/client.rs +++ b/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 => { diff --git a/src/common.rs b/src/common.rs index 648bc6b5c..57e91db85 100644 --- a/src/common.rs +++ b/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 { diff --git a/src/port_forward.rs b/src/port_forward.rs index 0cc640663..2c240eba8 100644 --- a/src/port_forward.rs +++ b/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>, remote_host: String, remote_port: i32, + tunnel: Arc, ) -> 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() => { - log::info!("new connection from {:?}", addr); - let id = id.clone(); - let password = password.clone(); - 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, &remote_host, remote_port).await { - Ok(Some(stream)) => { - let interface = interface.clone(); - tokio::spawn(async move { - if let Err(err) = run_forward(forward, stream).await { - interface.msgbox("error", "Error", &err.to_string(), ""); + // `addr` above the loop is the listener's own address; `run_rdp` needs + // that port. The accepted peer's address gets its own name so it can + // never shadow it. + Ok((forward, peer_addr)) = listener.accept() => { + log::debug!("new connection from {:?}", peer_addr); + let claim = tunnel.try_claim(); + let claim = match claim { + Claim::Wait => { + // Keep servicing the UI while the establishing accept holds the + // prompt. `None` means establishment failed: this accept is dropped. + let resolved: Option = loop { + 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 => { - break; + other => other, + }; + 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) => { - interface.on_establish_connection_error(err.to_string()); + Claim::Legacy => { + 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, + 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, + 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, mut prebuf: Vec) -> (TcpStream, Vec) { + 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, - remote_port: i32, -) -> ResultType> { +) -> ResultType> { 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(); - if !buffer.is_empty() { - allow_err!(stream.send_bytes(buffer.into()).await); - } - Ok(Some(stream)) + Ok(Some(LoginOutcome { + stream, + mux, + prebuf: buffer, + local_eof, + })) } @@ -334,170 +445,50 @@ async fn run_forward(forward: Framed, 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 - /// before building the login, where the real one looks passwords up. - #[derive(Clone)] - struct Ui { - lc: Arc>, - pause: Duration, - } - - #[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) {} - 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> { - 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 { - 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 peer_supports_mux_reads_the_features_bit() { + let mut pi = PeerInfo::new(); + assert!(!peer_supports_mux(&pi)); + pi.features = Some(Features { port_forward_mux: false, ..Default::default() }).into(); + assert!(!peer_supports_mux(&pi)); + pi.features = Some(Features { port_forward_mux: true, ..Default::default() }).into(); + assert!(peer_supports_mux(&pi)); } #[test] - fn mappings_logging_in_at_once_each_carry_their_own_target() { - let rt = tokio::runtime::Builder::new_current_thread() - .enable_all() - .build() - .unwrap(); - rt.block_on(async { - let mut ui = window(); - 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)); - }); + fn port_forward_mux_defaults_to_on() { + use hbb_common::config::{keys, option2bool}; + // The `enable-` prefix is what makes an unset value mean on. Getting the + // key's name wrong silently flips the default, so pin it here. + assert!(option2bool(keys::OPTION_ENABLE_PORT_FORWARD_MUX, "")); + assert!(option2bool(keys::OPTION_ENABLE_PORT_FORWARD_MUX, "Y")); + assert!(!option2bool(keys::OPTION_ENABLE_PORT_FORWARD_MUX, "N")); } #[test] - fn a_mapping_answers_the_prompt_with_its_own_challenge() { - let rt = tokio::runtime::Builder::new_current_thread() - .enable_all() - .build() - .unwrap(); + fn take_socket_hands_back_a_working_socket_and_the_prebuf() { + use hbb_common::tokio::io::{AsyncReadExt, AsyncWriteExt}; + let rt = tokio::runtime::Builder::new_current_thread().enable_all().build().unwrap(); rt.block_on(async { - let ui = window(); - let (mut a, mut a_peer) = loopback().await; - let (mut b, mut b_peer) = loopback().await; - // A's hash arrived last, so it is the one the handler holds. - assert!(login_with_hash(&ui, "pw", hash("a"), "a", 1, &mut a).await); - login_at(&mut a_peer).await; - let typed = (String::new(), String::new(), "pw".to_owned(), false); - login_from_ui(&ui, &hash("b"), 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)); - }); - } - - #[test] - fn a_password_typed_before_this_connections_hash_answers_it_when_it_comes() { - 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)); + let l = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = l.local_addr().unwrap(); + let mut client = TcpStream::connect(addr).await.unwrap(); + let (server, _) = l.accept().await.unwrap(); + let mut framed = Framed::new(server, BytesCodec::new()); + client.write_all(b"abc").await.unwrap(); + // Read through the codec, as connect_and_login does during login. + let pulled = framed.next().await.unwrap().unwrap(); + assert_eq!(&pulled[..], b"abc"); + let (mut sock, prebuf) = take_socket(framed, pulled.to_vec()); + 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]; + sock.read_exact(&mut buf).await.unwrap(); + assert_eq!(&buf, b"def"); }); } } diff --git a/src/ui_session_interface.rs b/src/ui_session_interface.rs index 03b59a497..e6946d4b2 100644 --- a/src/ui_session_interface.rs +++ b/src/ui_session_interface.rs @@ -1944,6 +1944,7 @@ pub async fn io_loop(handler: Session, 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(handler: Session, 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::>::new(); @@ -1976,6 +1977,7 @@ pub async fn io_loop(handler: Session, 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(handler: Session, round: u32) { receiver, &key, &token, + tunnel, ) .await; }); @@ -2023,6 +2026,7 @@ pub async fn io_loop(handler: Session, round: u32) { receiver, &key, &token, + tunnel, ) .await; } @@ -2042,6 +2046,7 @@ async fn start_one_port_forward( receiver: mpsc::UnboundedReceiver, key: &str, token: &str, + tunnel: std::sync::Arc, ) { if let Err(err) = crate::port_forward::listen( handler.get_id(), @@ -2054,6 +2059,7 @@ async fn start_one_port_forward( handler.lc.clone(), remote_host, remote_port, + tunnel, ) .await {