From 3d502c6a518740deb13f987a09585fb499db6d56 Mon Sep 17 00:00:00 2001 From: rustdesk Date: Wed, 22 Jul 2026 17:40:51 +0800 Subject: [PATCH] fix: preserve WebRTC endpoint and send semantics --- src/webrtc.rs | 136 ++++++++++++++++++++++++++++++++++++++++++++------ 1 file changed, 122 insertions(+), 14 deletions(-) diff --git a/src/webrtc.rs b/src/webrtc.rs index e05ae04f7..edcfea33d 100644 --- a/src/webrtc.rs +++ b/src/webrtc.rs @@ -23,8 +23,8 @@ use webrtc::stats::StatsReportType; use base64::engine::general_purpose::STANDARD as BASE64_STANDARD; use base64::Engine; use bytes::{BufMut, Bytes, BytesMut}; -use tokio::sync::{mpsc, watch, Mutex}; -use tokio::time::timeout; +use tokio::sync::{mpsc, watch, Mutex, Semaphore}; +use tokio::time::{timeout, timeout_at, Instant}; use url::Url; use crate::bytes_codec::MAX_FRAME_LENGTH; @@ -52,6 +52,9 @@ pub struct WebRTCStream { // Detached data channel, cached after the first `detach()` so send/recv do not re-lock and // re-fetch it per message. Shared across clones; `detach()` is idempotent. detached: Arc>>>, + // Serialize a complete logical message across clones. Each fragment is a separate SCTP + // message, so serializing only individual writes would allow two large messages to interleave. + send_gate: Arc, // Receive-side reassembly state, guarded by a single mutex so the fragment accumulator // survives `next()` cancellation (e.g. `next_timeout`) instead of losing already-read // fragments mid-message. Assumes a single reader, consistent with the rest of the stream API. @@ -107,6 +110,7 @@ impl Clone for WebRTCStream { send_timeout: self.send_timeout, relay_only: self.relay_only, detached: self.detached.clone(), + send_gate: self.send_gate.clone(), recv_state: self.recv_state.clone(), peer_verified: self.peer_verified.clone(), } @@ -516,6 +520,7 @@ impl WebRTCStream { send_timeout: ms_timeout, relay_only: force_relay, detached: Arc::new(Mutex::new(None)), + send_gate: Arc::new(Semaphore::new(1)), recv_state: Arc::new(Mutex::new(RecvState::default())), peer_verified: Arc::new(AtomicBool::new(false)), }; @@ -543,6 +548,18 @@ impl WebRTCStream { #[inline] pub async fn get_local_endpoint(&self) -> ResultType { + // Preserve the original one-shot endpoint contract: callers that only exchange this SDP + // do not have a separate path for `take_local_ice_rx`, so their endpoint must contain the + // gathered host/srflx/relay candidates. + let mut gather_complete = self.pc.gathering_complete_promise().await; + let _gathering_channel_closed = gather_complete.recv().await; + self.get_local_endpoint_trickle().await + } + + /// Return the current local description immediately for callers that signal candidates via + /// `take_local_ice_rx`. Unlike `get_local_endpoint`, this does not wait for ICE gathering. + #[inline] + pub async fn get_local_endpoint_trickle(&self) -> ResultType { if let Some(local_desc) = self.pc.local_description().await { let sdp = serde_json::to_string(&local_desc)?; let endpoint = Self::sdp_to_endpoint(&sdp); @@ -727,23 +744,46 @@ impl WebRTCStream { pub async fn send_bytes(&mut self, bytes: Bytes) -> ResultType<()> { let send_timeout = self.send_timeout; + let send_gate = self.send_gate.clone(); // Bound the WHOLE data-channel send (wait-for-open + every write) by send_timeout, - // mirroring FramedStream. Without this a write can park indefinitely on SCTP - // pending-queue backpressure and connection.rs's timeout timer never runs. + // including time queued behind another clone. Without this a write can park indefinitely + // on SCTP pending-queue backpressure and connection.rs's timeout timer never runs. if send_timeout > 0 { - match timeout( - Duration::from_millis(send_timeout), - self.send_bytes_inner(bytes), - ) - .await - { + let deadline = Instant::now() + Duration::from_millis(send_timeout); + let _send_permit = match timeout_at(deadline, send_gate.acquire_owned()).await { + Ok(Ok(permit)) => permit, + Ok(Err(err)) => { + return Err(Error::new( + ErrorKind::BrokenPipe, + format!("WebRTC send gate closed: {}", err), + ) + .into()); + } + Err(_) => { + if let Err(err) = self.pc.close().await { + log::warn!("failed to close WebRTC after send timeout: {}", err); + } + return Err(Error::new(ErrorKind::TimedOut, "WebRTC send timeout").into()); + } + }; + match timeout_at(deadline, self.send_bytes_inner(bytes)).await { Ok(res) => res, Err(_) => { - self.pc.close().await.ok(); + // Keep the logical-message permit while closing so no waiting clone can append + // a new message after a partially-written fragment sequence. + if let Err(err) = self.pc.close().await { + log::warn!("failed to close WebRTC after send timeout: {}", err); + } Err(Error::new(ErrorKind::TimedOut, "WebRTC send timeout").into()) } } } else { + let _send_permit = send_gate.acquire_owned().await.map_err(|err| { + Error::new( + ErrorKind::BrokenPipe, + format!("WebRTC send gate closed: {}", err), + ) + })?; self.send_bytes_inner(bytes).await } } @@ -849,7 +889,8 @@ mod tests { use crate::config; use crate::webrtc::WebRTCStream; use crate::webrtc::DEFAULT_ICE_SERVERS; - use std::time::Duration; + use std::{sync::Arc, time::Duration}; + use tokio::sync::Barrier; use tokio::time::timeout; use webrtc::peer_connection::sdp::session_description::RTCSessionDescription; @@ -1112,9 +1153,9 @@ IHR5cCBzcmZseCByYWRkciAwLjAuMC4wIHJwb3J0IDY0MDA4XHJcbmE9ZW5kLW9mLWNhbmRpZGF0ZXNc async fn connect_loopback() -> (WebRTCStream, WebRTCStream) { let mut offerer = WebRTCStream::new("", false, 20000).await.unwrap(); - let offer = offerer.get_local_endpoint().await.unwrap(); + let offer = offerer.get_local_endpoint_trickle().await.unwrap(); let answerer = WebRTCStream::new(&offer, false, 20000).await.unwrap(); - let answer = answerer.get_local_endpoint().await.unwrap(); + let answer = answerer.get_local_endpoint_trickle().await.unwrap(); offerer.set_remote_endpoint(&answer).await.unwrap(); // Bridge trickle candidates directly between the two peers, both directions. @@ -1139,6 +1180,26 @@ IHR5cCBzcmZseCByYWRkciAwLjAuMC4wIHJwb3J0IDY0MDA4XHJcbmE9ZW5kLW9mLWNhbmRpZGF0ZXNc (offerer, answerer) } + // One-shot callers exchange only the endpoints and never consume `take_local_ice_rx`. + #[tokio::test] + async fn test_webrtc_loopback_gathered_endpoints() { + let connect = async { + let mut offerer = WebRTCStream::new("", false, 20000).await.unwrap(); + let offer = offerer.get_local_endpoint().await.unwrap(); + let mut answerer = WebRTCStream::new(&offer, false, 20000).await.unwrap(); + let answer = answerer.get_local_endpoint().await.unwrap(); + offerer.set_remote_endpoint(&answer).await.unwrap(); + + offerer.wait_connected(20000).await.unwrap(); + answerer.wait_connected(20000).await.unwrap(); + offerer.close().await; + answerer.close().await; + }; + timeout(Duration::from_secs(40), connect) + .await + .expect("gathered-endpoint 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). @@ -1186,4 +1247,51 @@ IHR5cCBzcmZseCByYWRkciAwLjAuMC4wIHJwb3J0IDY0MDA4XHJcbmE9ZW5kLW9mLWNhbmRpZGF0ZXNc .expect("webrtc loopback did not complete in time"); } + #[tokio::test] + async fn test_webrtc_concurrent_large_sends_preserve_boundaries() { + let connect = async { + let (offerer, mut answerer) = connect_loopback().await; + let mut sender_a = offerer.clone(); + let mut sender_b = offerer.clone(); + let expected_a = vec![0xAA; 200_000]; + let expected_b = vec![0xBB; 200_000]; + let payload_a = expected_a.clone(); + let payload_b = expected_b.clone(); + let barrier = Arc::new(Barrier::new(3)); + + let barrier_a = barrier.clone(); + let send_a = tokio::spawn(async move { + barrier_a.wait().await; + sender_a.send_raw(payload_a).await + }); + let barrier_b = barrier.clone(); + let send_b = tokio::spawn(async move { + barrier_b.wait().await; + sender_b.send_raw(payload_b).await + }); + + barrier.wait().await; + let receive = async { + let first = answerer.next().await.unwrap().unwrap(); + let second = answerer.next().await.unwrap().unwrap(); + (first, second) + }; + let (send_a, send_b, (first, second)) = tokio::join!(send_a, send_b, receive); + send_a.unwrap().unwrap(); + send_b.unwrap().unwrap(); + + let boundaries_preserved = (first.as_ref() == expected_a.as_slice() + && second.as_ref() == expected_b.as_slice()) + || (first.as_ref() == expected_b.as_slice() + && second.as_ref() == expected_a.as_slice()); + assert!(boundaries_preserved, "concurrent messages were interleaved"); + + offerer.close().await; + answerer.close().await; + }; + timeout(Duration::from_secs(40), connect) + .await + .expect("concurrent WebRTC sends did not complete in time"); + } + }