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:
rustdeskandClaude Fable 5.1 committed 2026-09-17 16:29:55 +08:00
1 parent 0eb175963c
commit e999dcea24
3 files changed
+98

No files matched your search

+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");
}
}