mirror of
https://github.com/rustdesk/hbb_common.git
synced 2026-10-23 08:40:30 +00:00
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:
1 parent
0eb175963c
commit
e999dcea24
3 files changed
+98
No files matched your search
@@ -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");
|
||||
}
|
||||
}
|
||||
Reference in new issue
Block a user