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..30701d956 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(), } } } @@ -338,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(_) @@ -747,12 +760,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 +881,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 +1186,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 +1429,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 +2260,109 @@ 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), + "{} must leave the hint as it found it", + terminal + ); + } + } + + // 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(); + + // 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) }); + + // `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..600 { + match answerer.next_timeout(20).await { + Some(msg) => { + got = Some(msg); + break; + } + None => { + if answerer.rx_progress() > baseline { + moved_early = true; + } + } + } + } + + 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. + 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).