From e999dcea24462a4291cbada100bbc70ff3e04574 Mon Sep 17 00:00:00 2001 From: rustdesk Date: Thu, 17 Sep 2026 16:27:53 +0800 Subject: [PATCH] cover the message cap's lifecycle on WebSocket and WebRTC The cap's header-only refusal was tested on the TCP codec and its reassembly bound on WebRTC, but not on WebSocket, which is the transport where a frame header buys an allocation, and neither transport had a test that lifting the cap lets a large message through again - the half of the lifecycle a server runs at authorization, and the one a forgotten restore would break silently. Two WebSocket tests build a loopback from a raw TcpStream with no handshake, write server-to-client frames by hand, and check that a header declaring one byte over the cap is refused before any payload exists, and that a message refused under a low cap arrives once the cap is lifted. A WebRTC test does the same lift. Each fails when the corresponding half of set_max_packet_length is made a no-op. The doc on Stream::set_max_packet_length now says the one thing the bound does not do: lowering it applies from the next frame on and does not discard one whose decoding or reassembly has already begun, so it has to be set before the first read of untrusted data. That is how the server uses it; the note is there so nothing later treats it as a switch that can be thrown at any time. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01EZ49AbZJYfm8NTp5yDPMab --- src/stream.rs | 3 +++ src/webrtc.rs | 29 +++++++++++++++++++++ src/websocket.rs | 66 ++++++++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 98 insertions(+) diff --git a/src/stream.rs b/src/stream.rs index d7b3738c5..71132ff6f 100644 --- a/src/stream.rs +++ b/src/stream.rs @@ -37,6 +37,9 @@ 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. + /// + /// Set it before the first read of untrusted data: a lower bound applies from the next frame + /// on and does not discard one whose decoding or reassembly has already begun. #[inline] pub fn set_max_packet_length(&mut self, n: usize) { match self { diff --git a/src/webrtc.rs b/src/webrtc.rs index 147d84293..f33f9141c 100644 --- a/src/webrtc.rs +++ b/src/webrtc.rs @@ -2427,6 +2427,35 @@ IHR5cCBzcmZseCByYWRkciAwLjAuMC4wIHJwb3J0IDY0MDA4XHJcbmE9ZW5kLW9mLWNhbmRpZGF0ZXNc .expect("WebRTC loopback did not complete in time"); } + // The other half of the same lifecycle: what was refused while the bound was low arrives once + // it is lifted, which is what authorization does on the server. + #[tokio::test] + async fn test_webrtc_max_packet_length_restores() { + 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); + + answerer.set_max_packet_length(usize::MAX); + offerer.send_raw(vec![0xCDu8; 200_000]).await.unwrap(); + let got = timeout(Duration::from_secs(10), answerer.next()) + .await + .expect("answerer.next() hung after the bound was lifted") + .unwrap() + .unwrap(); + assert_eq!(got.len(), 200_000, "lifting the bound lets a large message through again"); + + 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 a74e73274..62125fd19 100644 --- a/src/websocket.rs +++ b/src/websocket.rs @@ -420,6 +420,7 @@ pub fn check_ws(endpoint: &str) -> String { mod tests { use super::*; use crate::config::{keys, Config}; + use tokio::{io::AsyncWriteExt, net::TcpListener}; #[test] fn test_check_ws() { @@ -554,4 +555,69 @@ mod tests { assert_eq!(check_ws("127.0.0.1:34567"), "ws://127.0.0.1:34569"); Config::set_options(options); } + // A server-to-client binary frame is unmasked, so it can be put on the wire by hand: the header + // alone, which is all it takes to ask tungstenite for the allocation, or with its payload. + fn ws_binary_frame(len: usize, with_payload: bool) -> Vec { + let mut f = vec![0x82]; + if len < 126 { + f.push(len as u8); + } else if len <= u16::MAX as usize { + f.push(0x7E); + f.extend_from_slice(&(len as u16).to_be_bytes()); + } else { + f.push(0x7F); + f.extend_from_slice(&(len as u64).to_be_bytes()); + } + if with_payload { + f.resize(f.len() + len, 0xCD); + } + f + } + + async fn ws_loopback() -> (WsFramedStream, TcpStream) { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let client = TcpStream::connect(addr).await.unwrap(); + let (server, _) = listener.accept().await.unwrap(); + let ws = WsFramedStream::from_tcp_stream(client, addr).await.unwrap(); + (ws, server) + } + + #[tokio::test] + async fn max_packet_length_refuses_a_frame_on_its_header() { + const CAP: usize = 16 * 1024; + let (mut ws, mut server) = ws_loopback().await; + ws.set_max_packet_length(CAP); + + // One byte over, header only: refused before there is any payload to buffer. + server.write_all(&ws_binary_frame(CAP + 1, false)).await.unwrap(); + match timeout(Duration::from_secs(5), ws.next()).await { + Ok(Some(Err(e))) => assert!(e.to_string().contains("Message too long"), "{}", e), + Ok(other) => panic!( + "expected a refusal, got {:?}", + other.map(|r| r.map(|b| b.len())) + ), + Err(_) => panic!("next() waited for a payload the header should have refused"), + } + } + + #[tokio::test] + async fn max_packet_length_lowered_then_restored() { + const CAP: usize = 16 * 1024; + let (mut ws, mut server) = ws_loopback().await; + + ws.set_max_packet_length(CAP); + server.write_all(&ws_binary_frame(8 * 1024, true)).await.unwrap(); + let got = ws.next().await.unwrap().unwrap(); + assert_eq!(got.len(), 8 * 1024, "a message under the cap still arrives"); + + ws.set_max_packet_length(usize::MAX); + server.write_all(&ws_binary_frame(200_000, true)).await.unwrap(); + let got = timeout(Duration::from_secs(5), ws.next()) + .await + .expect("next() hung after the cap was lifted") + .unwrap() + .unwrap(); + assert_eq!(got.len(), 200_000, "lifting the cap lets a large message through again"); + } }