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 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01EZ49AbZJYfm8NTp5yDPMab
This commit is contained in:
rustdesk
2026-09-17 16:29:55 +08:00
co-authored by Claude Fable 5.1
parent 0eb175963c
commit e999dcea24
3 changed files with 98 additions and 0 deletions
+3
View File
@@ -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 {
+29
View File
@@ -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).
+66
View File
@@ -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<u8> {
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");
}
}