From c50c1fef326de134ad2b15e7cfbf92133c9cefc8 Mon Sep 17 00:00:00 2001 From: rustdesk Date: Tue, 8 Sep 2026 14:28:32 +0800 Subject: [PATCH 1/3] webrtc: report receive progress and ICE's disconnected hint A peer that is killed, logged out or rebooted sends no reset over UDP, so a session only learns of it by timing out. Nothing here acts on that; these are the two observations a caller needs to. `rx_progress` counts arriving fragments. A message reaches the caller only once it is whole, so a large one - a clipboard image is the case that occurs - leaves the receive loop with nothing to observe for as long as the transfer takes, and a peer sending steadily looks exactly like one that died. A counter, not a timestamp: `Instant::now()` costs an order of magnitude more per read, and callers only compare successive samples. `webrtc_disconnected` exposes the ICE state webrtc-rs already raises ~5s after it stops hearing from the peer. It is deliberately not a fourth `WebRTCConnectionState`: that enum gates send and receive readiness and session cache admissibility, and a transient state must change neither. TCP and WebSocket answer for neither, on purpose: this change leaves their existing close, error and timeout behaviour exactly as it was. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_019aokqJuhjvB3kijXtAg5Ns --- src/lib.rs | 1 + src/rx_probe.rs | 51 +++++++++++++++++ src/stream.rs | 36 ++++++++++++ src/webrtc.rs | 145 ++++++++++++++++++++++++++++++++++++++++++++++++ 4 files changed, 233 insertions(+) create mode 100644 src/rx_probe.rs diff --git a/src/lib.rs b/src/lib.rs index 2974f2d8f..20dfa4312 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -56,6 +56,7 @@ pub use uuid; pub mod fingerprint; pub use flexi_logger; pub mod log_throttle; +pub mod rx_probe; pub mod stream; pub mod websocket; #[cfg(feature = "webrtc")] diff --git a/src/rx_probe.rs b/src/rx_probe.rs new file mode 100644 index 000000000..520fa3867 --- /dev/null +++ b/src/rx_probe.rs @@ -0,0 +1,51 @@ +use std::sync::{ + atomic::{AtomicU64, Ordering}, + Arc, +}; + +/// Bytes received from the peer, counted before framing. +/// +/// A message reaches the application only once it is complete, so a large one — a clipboard image +/// is the case that occurs in practice — leaves the receive loop with nothing to observe for as +/// long as the transfer takes, and a peer that is sending steadily looks identical to one that +/// died. This counter moves on every read, so liveness stays observable across such a message. +/// +/// A counter rather than a timestamp on purpose: `Instant::now()` costs an order of magnitude more +/// than a relaxed add, and the reader only ever compares successive samples, so the clock is +/// needed once per sampling tick instead of once per read. +#[derive(Clone, Default)] +pub struct RxProbe(Arc); + +impl RxProbe { + pub fn new() -> Self { + Self::default() + } + + #[inline] + pub fn add(&self, n: usize) { + self.0.fetch_add(n as u64, Ordering::Relaxed); + } + + /// Changes whenever bytes arrived. Compare successive samples; the absolute value and its + /// wrap-around carry no meaning. + #[inline] + pub fn get(&self) -> u64 { + self.0.load(Ordering::Relaxed) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn clones_share_one_counter() { + let probe = RxProbe::new(); + let cloned = probe.clone(); + assert_eq!(probe.get(), 0); + cloned.add(7); + assert_eq!(probe.get(), 7); + probe.add(3); + assert_eq!(cloned.get(), 10); + } +} diff --git a/src/stream.rs b/src/stream.rs index 2df826915..d714fbcef 100644 --- a/src/stream.rs +++ b/src/stream.rs @@ -64,6 +64,42 @@ impl Stream { } } + /// An opaque token that changes whenever bytes have arrived from the peer, a fragment of a + /// message that is still incomplete included. Compare successive samples; the value itself + /// means nothing. `None` where the transport reports no such thing. + /// + /// A message reaches the caller only once it is whole, so `next()` alone cannot distinguish a + /// peer sending a large one — a clipboard image is the case that occurs — from one that died + /// mid-message. Sampling this alongside the last received message tells them apart. + /// + /// TCP and WebSocket answer `None` on purpose: this change leaves their existing close, error + /// and 30s-timeout behaviour exactly as it was, and giving them a value would alter it. + #[inline] + pub fn rx_progress(&self) -> Option { + match self { + #[cfg(feature = "webrtc")] + Stream::WebRTC(s) => Some(s.rx_progress()), + #[allow(unreachable_patterns)] + _ => None, + } + } + + /// Whether ICE has stopped hearing from the peer on a WebRTC transport. Transient and + /// recoverable — it says the peer is worth suspecting, never that the session is over — and + /// always false for transports this change leaves alone. + /// + /// Maintained by ICE's own task, so unlike anything derived from the session's read loop it + /// stays accurate while that loop is busy sending. + #[inline] + pub fn webrtc_disconnected(&self) -> bool { + match self { + #[cfg(feature = "webrtc")] + Stream::WebRTC(s) => s.is_disconnected(), + #[allow(unreachable_patterns)] + _ => false, + } + } + #[inline] pub fn is_secured(&self) -> bool { match self { diff --git a/src/webrtc.rs b/src/webrtc.rs index 14422c6a2..39cf6ae9a 100644 --- a/src/webrtc.rs +++ b/src/webrtc.rs @@ -48,6 +48,7 @@ use url::Url; use crate::config; use crate::protobuf::Message; +use crate::rx_probe::RxProbe; use crate::sodiumoxide::crypto::secretbox::Key; use crate::ResultType; @@ -92,6 +93,12 @@ pub struct WebRTCStream { // Taken once at construction: `pc.local_description()` grows with every gathered candidate // and this must not (why: `UDP_ENDPOINT_BUDGET`). Shared so cloning the struct stays cheap. local_endpoint: Arc, + // Counts arriving fragments, which `recv_state` cannot answer for: that lock is deliberately + // held across the whole reassembly (see the field), so a sampler would block on it. + rx_probe: RxProbe, + // ICE has stopped hearing from the peer but the connection is not terminal yet. Kept out of + // `state_notify` on purpose - see where it is set. + disconnected: Arc, } /// Whether a pc `new_inner` built is still wanted, shared by every `NewStreamHandoff` handed out @@ -240,6 +247,8 @@ impl Clone for WebRTCStream { peer_verified: self.peer_verified.clone(), handoff: self.handoff.clone(), local_endpoint: self.local_endpoint.clone(), + rx_probe: self.rx_probe.clone(), + disconnected: self.disconnected.clone(), } } } @@ -747,12 +756,20 @@ impl WebRTCStream { // pc and the ICE-candidate sender's forwarding task. Upgrade inside the handler; if the pc // is already gone there is nothing left in SESSIONS to evict. let pc_for_close = Arc::downgrade(&pc); + // Deliberately NOT a `WebRTCConnectionState` variant: that enum gates send and receive + // readiness (`wait_for_connect_result`) and session-cache admissibility, so a fourth state + // would have to be given a meaning at each of those. `Disconnected` is transient and must + // change neither, so it is reported beside them instead of through them. + let disconnected = Arc::new(AtomicBool::new(false)); + let disconnected_for_state = disconnected.clone(); pc.on_peer_connection_state_change(Box::new(move |s: RTCPeerConnectionState| { let stream_for_close2 = stream_for_close.clone(); let on_connection_notify = notify_tx.clone(); let pc_for_close2 = pc_for_close.clone(); + let disconnected_for_state = disconnected_for_state.clone(); Box::pin(async move { log::debug!("WebRTC session peer connection state: {}", s); + Self::note_health(&disconnected_for_state, s); match s { // `Disconnected` is a transient, recoverable ICE state (webrtc-ice fires it // after ~5s without consent and returns to `Connected` when traffic resumes). @@ -860,6 +877,8 @@ impl WebRTCStream { peer_verified: Arc::new(AtomicBool::new(false)), handoff: Default::default(), local_endpoint: Arc::new(local_endpoint), + rx_probe: RxProbe::new(), + disconnected, }; // Insert into the session cache, but never `await pc.close()` while holding this lock: // `close()` fires the peer-connection-state handler inline, which itself locks SESSIONS, @@ -1163,6 +1182,30 @@ impl WebRTCStream { self.send_timeout = ms; } + /// See `Stream::rx_progress`. + #[inline] + pub fn rx_progress(&self) -> u64 { + self.rx_probe.get() + } + + /// Track whether ICE still hears from the peer. `Disconnected` is raised about 5s after it + /// stops and is cleared again when traffic resumes, so it is a hint to be acted on by whoever + /// asks, never a terminal state; the terminal ones are handled by the caller of this and + /// deliberately leave the flag alone, since a closed session has no use for a hint. + fn note_health(disconnected: &AtomicBool, s: RTCPeerConnectionState) { + match s { + RTCPeerConnectionState::Disconnected => disconnected.store(true, Ordering::SeqCst), + RTCPeerConnectionState::Connected => disconnected.store(false, Ordering::SeqCst), + _ => {} + } + } + + /// See `Stream::webrtc_disconnected`. + #[inline] + pub fn is_disconnected(&self) -> bool { + self.disconnected.load(Ordering::SeqCst) + } + #[inline] pub fn set_key(&mut self, _key: Key) { // WebRTC traffic is DTLS-encrypted regardless; the secretbox key is unused. @@ -1382,6 +1425,11 @@ impl WebRTCStream { ))); } }; + // Before the framing checks below: a fragment that turns out to be malformed still + // proves the peer sent it, which is all this counter claims. + if n > 0 { + self.rx_probe.add(n); + } if n == 0 { // Not exclusively a stream reset: webrtc-data maps the empty-message PPIDs to // n == 0 as well and `read` drops the flag that would separate them. Both mean @@ -2208,6 +2256,103 @@ IHR5cCBzcmZseCByYWRkciAwLjAuMC4wIHJwb3J0IDY0MDA4XHJcbmE9ZW5kLW9mLWNhbmRpZGF0ZXNc .expect("gathered-endpoint WebRTC loopback did not complete in time"); } + // The state mapping the health flag is built from. ICE raises `Disconnected` only on a real + // consent failure, which an in-process loopback cannot produce, so this drives the mapping + // directly; the loopback test below covers the wiring from the flag out to callers. + #[test] + fn test_note_health_maps_only_the_transient_states() { + use crate::webrtc::RTCPeerConnectionState; + use std::sync::atomic::{AtomicBool, Ordering}; + let flag = AtomicBool::new(false); + + WebRTCStream::note_health(&flag, RTCPeerConnectionState::Disconnected); + assert!(flag.load(Ordering::SeqCst), "Disconnected must raise the hint"); + + WebRTCStream::note_health(&flag, RTCPeerConnectionState::Connected); + assert!(!flag.load(Ordering::SeqCst), "Connected must clear it"); + + // The terminal states have their own path and must not touch the hint: a session that + // ends has no use for one, and clearing it here would hide a peer that went silent first. + for terminal in [ + RTCPeerConnectionState::Failed, + RTCPeerConnectionState::Closed, + ] { + WebRTCStream::note_health(&flag, RTCPeerConnectionState::Disconnected); + WebRTCStream::note_health(&flag, terminal); + assert!( + flag.load(Ordering::SeqCst), + "{terminal} must leave the hint as it found it" + ); + } + } + + // Pins the wiring and the polarity of the health flag: a pair that just connected must not + // look suspect. + #[tokio::test] + async fn test_webrtc_disconnected_is_false_on_a_live_pair() { + let body = async { + let (offerer, answerer) = connect_loopback().await; + assert!(!offerer.is_disconnected()); + assert!(!answerer.is_disconnected()); + offerer.close().await; + answerer.close().await; + }; + timeout(Duration::from_secs(30), body) + .await + .expect("WebRTC loopback did not complete in time"); + } + + // A message far above the fragment size must show as activity while it is still arriving: + // `next()` yields nothing until the last fragment, so a receiver watching only completed + // messages cannot tell a peer mid-transfer from one that died. + #[tokio::test] + async fn test_webrtc_rx_progress_moves_before_a_large_message_completes() { + let body = async { + let (mut offerer, mut answerer) = connect_loopback().await; + + offerer.send_raw(b"warmup".to_vec()).await.unwrap(); + assert_eq!(&answerer.next().await.unwrap().unwrap()[..], b"warmup"); + let baseline = answerer.rx_progress(); + + let big = vec![0x5Au8; 512 * 1024]; + let expected = big.len(); + let sender = tokio::spawn(async move { offerer.send_raw(big).await.map(|_| offerer) }); + + // The receive loop's shape: sampling happens in a sibling arm, so the probe has to be + // readable while `next()` is in flight and a cancelled `next()` has to keep the + // partial message. + let mut moved_early = false; + let mut got = None; + for _ in 0..2_000 { + tokio::select! { + msg = answerer.next() => { + got = msg; + break; + } + _ = tokio::time::sleep(Duration::from_millis(1)) => { + if answerer.rx_progress() > baseline { + moved_early = true; + } + } + } + } + + assert_eq!(got.expect("message").unwrap().len(), expected); + assert!( + moved_early, + "rx_progress must advance while the message is still incomplete" + ); + let offerer = sender.await.unwrap().expect("send"); + // `WebRTCStream` has no `Drop`, so both pcs stay in `SESSIONS` unless closed here; + // `test_cancelled_new_does_not_leak_the_pc` scans that map and would report them. + offerer.close().await; + answerer.close().await; + }; + timeout(Duration::from_secs(60), body) + .await + .expect("WebRTC loopback did not complete in time"); + } + // In-process offerer<->answerer loopback exercising the send/next data plane that the framing, // empty-message, and EOF fixes live in. Connects over host candidates (works offline; any // configured/default STUN just fails in the background without blocking the host pair). From 166bb3d2ac3a59eeaafab6ac0aac7b98c8d93778 Mon Sep 17 00:00:00 2001 From: rustdesk Date: Tue, 8 Sep 2026 16:52:28 +0800 Subject: [PATCH 2/3] webrtc: do not hand out a cached session ICE has lost the peer on `is_reusable_for` rejected a closed session but not one ICE currently reports Disconnected. Nothing else does either - that state is deliberately not part of the connection-state enum - so a caller could be given a session that already carries the hint, and start off suspecting a peer it never used. Also drops a race from the receive-progress test: it raced `next()` against a sleeping sibling in a select, so a transfer that landed inside one poll left the assertion looking false, and a slow enough one exhausted the loop and failed with a misleading message. It now reads the counter after a `next_timeout` that returned nothing, which is a point where the message is known to be incomplete. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_019aokqJuhjvB3kijXtAg5Ns --- src/webrtc.rs | 29 +++++++++++++++++++---------- 1 file changed, 19 insertions(+), 10 deletions(-) diff --git a/src/webrtc.rs b/src/webrtc.rs index 39cf6ae9a..52960df06 100644 --- a/src/webrtc.rs +++ b/src/webrtc.rs @@ -347,6 +347,10 @@ impl WebRTCStream { fn is_reusable_for(&self, force_relay: bool) -> bool { self.relay_only == force_relay && !self.handoff.closing.load(Ordering::Acquire) + // ICE has stopped hearing from the peer on this one. Not terminal, so nothing else + // rejects it, but handing it to a caller that is about to judge liveness starts them + // off already suspecting a pc they never used. + && !self.is_disconnected() && !matches!( *self.state_notify.borrow(), WebRTCConnectionState::Closed(_) @@ -2314,22 +2318,24 @@ IHR5cCBzcmZseCByYWRkciAwLjAuMC4wIHJwb3J0IDY0MDA4XHJcbmE9ZW5kLW9mLWNhbmRpZGF0ZXNc assert_eq!(&answerer.next().await.unwrap().unwrap()[..], b"warmup"); let baseline = answerer.rx_progress(); - let big = vec![0x5Au8; 512 * 1024]; + // Large enough that it cannot land inside one poll of the loop below, so a probe + // that only counted whole messages would leave the counter untouched throughout. + let big = vec![0x5Au8; 4 * 1024 * 1024]; let expected = big.len(); let sender = tokio::spawn(async move { offerer.send_raw(big).await.map(|_| offerer) }); - // The receive loop's shape: sampling happens in a sibling arm, so the probe has to be - // readable while `next()` is in flight and a cancelled `next()` has to keep the - // partial message. + // `next_timeout` returning None leaves the partial message in place, so the counter + // is read at a point where nothing has been delivered. No select and no sleeping + // sibling: the assertion cannot turn on which branch a poll happened to take. let mut moved_early = false; let mut got = None; - for _ in 0..2_000 { - tokio::select! { - msg = answerer.next() => { - got = msg; + for _ in 0..600 { + match answerer.next_timeout(20).await { + Some(msg) => { + got = Some(msg); break; } - _ = tokio::time::sleep(Duration::from_millis(1)) => { + None => { if answerer.rx_progress() > baseline { moved_early = true; } @@ -2337,11 +2343,14 @@ IHR5cCBzcmZseCByYWRkciAwLjAuMC4wIHJwb3J0IDY0MDA4XHJcbmE9ZW5kLW9mLWNhbmRpZGF0ZXNc } } - assert_eq!(got.expect("message").unwrap().len(), expected); assert!( moved_early, "rx_progress must advance while the message is still incomplete" ); + assert_eq!( + got.expect("message never completed").unwrap().len(), + expected + ); let offerer = sender.await.unwrap().expect("send"); // `WebRTCStream` has no `Drop`, so both pcs stay in `SESSIONS` unless closed here; // `test_cancelled_new_does_not_leak_the_pc` scans that map and would report them. From 6db01ec2f60b51de10ee3179161192f748e5a66c Mon Sep 17 00:00:00 2001 From: rustdesk Date: Wed, 9 Sep 2026 16:38:02 +0800 Subject: [PATCH 3/3] review: name the state in the health-hint assert hbb_common is edition 2018, where a lone string literal is not treated as a format string, so the capture printed itself and the loop over the two terminal states could not say which one had failed. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_019aokqJuhjvB3kijXtAg5Ns --- src/webrtc.rs | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/webrtc.rs b/src/webrtc.rs index 52960df06..30701d956 100644 --- a/src/webrtc.rs +++ b/src/webrtc.rs @@ -2285,7 +2285,8 @@ IHR5cCBzcmZseCByYWRkciAwLjAuMC4wIHJwb3J0IDY0MDA4XHJcbmE9ZW5kLW9mLWNhbmRpZGF0ZXNc WebRTCStream::note_health(&flag, terminal); assert!( flag.load(Ordering::SeqCst), - "{terminal} must leave the hint as it found it" + "{} must leave the hint as it found it", + terminal ); } }