diff --git a/Cargo.toml b/Cargo.toml index 8e1735690..fc53aa354 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -57,7 +57,10 @@ tokio-rustls = { version = "0.26", features = [ "ring", ], default-features = false } tokio-native-tls = "0.3" -tokio-tungstenite = { version = "0.26", features = ["native-tls", "rustls-tls-native-roots", "rustls-tls-webpki-roots"] } +# Fork of v0.26.2: upstream's WebSocketStream exposes get_config but no set_config, so a +# message cap lowered before authorization could never be lifted again. One commit on top of +# the tag; pinned by rev, not branch, because a fork branch can be rewritten under the lockfile. +tokio-tungstenite = { git = "https://github.com/rustdesk-org/tokio-tungstenite", rev = "9c320de8e7c546cb654b017b7f4cbb4a5cc09cd9", version = "0.26", features = ["native-tls", "rustls-tls-native-roots", "rustls-tls-webpki-roots"] } tungstenite = { version = "0.26", features = ["native-tls", "rustls-tls-native-roots", "rustls-tls-webpki-roots"] } rustls-platform-verifier = "0.6" rustls-pki-types = "1.11" diff --git a/src/bytes_codec.rs b/src/bytes_codec.rs index 5bd5590c1..987a01922 100644 --- a/src/bytes_codec.rs +++ b/src/bytes_codec.rs @@ -300,4 +300,21 @@ mod tests { assert!(matches!(codec.decode(&mut buf), Ok(None))); assert!(buf.capacity() <= MAX_REASONABLE_CAPACITY); } + + #[test] + fn max_packet_length_refuses_on_the_header() { + const CAP: usize = 128 * 1024; + let decode_header_for = |n: usize| { + let mut codec = BytesCodec::new(); + codec.set_max_packet_length(CAP); + let mut buf = BytesMut::new(); + buf.put_u32_le((n << 2) as u32 | 0x3); + codec.decode(&mut buf) + }; + + // At the cap the header is accepted and the payload awaited. + assert!(matches!(decode_header_for(CAP), Ok(None))); + // One byte over, and the frame dies on its header - before any of it has been buffered. + assert!(decode_header_for(CAP + 1).is_err()); + } } diff --git a/src/stream.rs b/src/stream.rs index d714fbcef..d7b3738c5 100644 --- a/src/stream.rs +++ b/src/stream.rs @@ -34,6 +34,19 @@ impl Stream { } } + /// The largest message the peer may send. Lowered before authorization, where nothing a peer + /// legitimately sends is large, so an unauthenticated connection cannot make us hold more than + /// that. Each transport clamps `n` to its own ceiling, so `usize::MAX` restores the default. + #[inline] + pub fn set_max_packet_length(&mut self, n: usize) { + match self { + #[cfg(feature = "webrtc")] + Stream::WebRTC(s) => s.set_max_packet_length(n), + Stream::WebSocket(s) => s.set_max_packet_length(n), + Stream::Tcp(s) => s.codec_mut().set_max_packet_length(n), + } + } + #[inline] pub async fn send_bytes(&mut self, bytes: bytes::Bytes) -> ResultType<()> { match self { diff --git a/src/webrtc.rs b/src/webrtc.rs index 5bdbdcae0..147d84293 100644 --- a/src/webrtc.rs +++ b/src/webrtc.rs @@ -99,6 +99,10 @@ pub struct WebRTCStream { // 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, + // The largest message this stream will reassemble, lowered while the peer is unauthenticated. + // Shared by every clone, the `SESSIONS` entry included, because the receive path runs off a + // clone and the bound is about the pc, not about one handle. + max_recv_message: Arc, } /// Whether a pc `new_inner` built is still wanted, shared by every `NewStreamHandoff` handed out @@ -249,6 +253,7 @@ impl Clone for WebRTCStream { local_endpoint: self.local_endpoint.clone(), rx_probe: self.rx_probe.clone(), disconnected: self.disconnected.clone(), + max_recv_message: self.max_recv_message.clone(), } } } @@ -900,6 +905,7 @@ impl WebRTCStream { local_endpoint: Arc::new(local_endpoint), rx_probe: RxProbe::new(), disconnected, + max_recv_message: Arc::new(AtomicUsize::new(MAX_RECV_MESSAGE)), }; // 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, @@ -1193,6 +1199,14 @@ impl WebRTCStream { // not-supported } + /// See `Stream::set_max_packet_length`. Clamped to `MAX_RECV_MESSAGE`: the send path checks + /// against that constant, so a bound above it would only accept what no peer will ever send. + #[inline] + pub fn set_max_packet_length(&mut self, n: usize) { + self.max_recv_message + .store(n.min(MAX_RECV_MESSAGE), Ordering::Relaxed); + } + #[inline] pub fn local_addr(&self) -> SocketAddr { SocketAddr::new(IpAddr::V4(Ipv4Addr::UNSPECIFIED), 0) @@ -1483,7 +1497,7 @@ impl WebRTCStream { // Bound BEFORE growing: this framing carries no length, so an oversize message is only // discovered by accumulating it, and checking after the append would let `BytesMut`'s // reallocate-and-copy hold the old and the new buffer at once. - if acc.len() + (n - 1) > MAX_RECV_MESSAGE { + if acc.len() + (n - 1) > self.max_recv_message.load(Ordering::Relaxed) { // Release the buffer, don't just truncate it: by definition it is near the cap // here, and `recv_state` outlives this call through the `SESSIONS` clone. *acc = BytesMut::new(); @@ -2380,6 +2394,39 @@ IHR5cCBzcmZseCByYWRkciAwLjAuMC4wIHJwb3J0IDY0MDA4XHJcbmE9ZW5kLW9mLWNhbmRpZGF0ZXNc .expect("WebRTC loopback did not complete in time"); } + // The bound `Stream::set_max_packet_length` lowers before authorization: a message above it + // must end the connection as it arrives, not be reassembled. Only the receiver's bound moves, + // so the sender still puts it on the wire, exactly as an unauthenticated peer would. + #[tokio::test] + async fn test_webrtc_max_packet_length_bounds_reassembly() { + let body = async { + let (mut offerer, mut answerer) = connect_loopback().await; + answerer.set_max_packet_length(16 * 1024); + + offerer.send_raw(vec![0xCDu8; 8 * 1024]).await.unwrap(); + let got = answerer.next().await.unwrap().unwrap(); + assert_eq!(got.len(), 8 * 1024, "a message under the bound still arrives"); + + offerer.send_raw(vec![0xCDu8; 200_000]).await.unwrap(); + match timeout(Duration::from_secs(10), answerer.next()).await { + Ok(Some(Err(e))) => assert!( + e.to_string().contains("maximum frame size"), + "unexpected error: {}", + e + ), + Ok(Some(Ok(b))) => panic!("expected a refusal, got {} bytes", b.len()), + Ok(None) => panic!("expected a refusal, got EOF"), + Err(_) => panic!("answerer.next() hung on an oversize message"), + } + + 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). diff --git a/src/websocket.rs b/src/websocket.rs index 7bf210840..29ea5a4e0 100644 --- a/src/websocket.rs +++ b/src/websocket.rs @@ -28,6 +28,11 @@ use tokio_tungstenite::{ use tungstenite::client::IntoClientRequest; use tungstenite::protocol::Role; +/// tungstenite's own defaults. `set_max_packet_length` clamps to these, so `usize::MAX` puts the +/// transport back exactly where it started rather than above it. +const DEFAULT_MAX_MESSAGE_SIZE: usize = 64 << 20; +const DEFAULT_MAX_FRAME_SIZE: usize = 16 << 20; + pub struct WsFramedStream { stream: WebSocketStream>, addr: SocketAddr, @@ -211,6 +216,17 @@ impl WsFramedStream { self.encrypt = None; } + /// Both bounds, not just the message one: tungstenite reserves the whole declared payload of a + /// frame as soon as it passes `max_frame_size` (`protocol/frame/mod.rs`), so that is what keeps + /// a frame header from buying an allocation, while `max_message_size` bounds reassembly. + #[inline] + pub fn set_max_packet_length(&mut self, n: usize) { + self.stream.set_config(|c| { + c.max_message_size = Some(n.min(DEFAULT_MAX_MESSAGE_SIZE)); + c.max_frame_size = Some(n.min(DEFAULT_MAX_FRAME_SIZE)); + }); + } + #[inline] pub async fn from_tcp_stream(stream: TcpStream, addr: SocketAddr) -> ResultType { let ws_stream =