Skip to content
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
51 changes: 51 additions & 0 deletions src/client/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -39,6 +39,11 @@ const MAX_RECV_BYTES_PER_POLL: usize = 256 * 1024;
const MAX_RECV_BYTES_PER_COMMAND_WAIT: usize = 256 * 1024;
/// Maximum time to wait for the initial TCP connect before failing.
const TCP_CONNECT_TIMEOUT_SECS: u64 = 10;
/// Cap incomplete wire data staged in `recv_buffer` between chunk reads.
/// Mirrors the server-side staging limit in `session::conn` so a malicious
/// peer cannot retain up to `BUFFER_MAX_SIZE` (64 MiB) per client connection
/// when message budgets defer draining.
const MAX_RECV_BUFFER_BYTES: usize = 2 * DEFAULT_MAX_MSG_LENGTH as usize;

/// Client connection states.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
Expand Down Expand Up @@ -324,6 +329,14 @@ impl Client {
};
if n > 0 {
let chunk_len = n as usize;
if self
.recv_buffer
.available()
.saturating_add(chunk_len)
> MAX_RECV_BUFFER_BYTES
{
return Err(ErrorCode::Protocol);
}
self.recv_buffer
.write(&buf[..chunk_len])
.map_err(|_| ErrorCode::Internal)?;
Expand Down Expand Up @@ -562,6 +575,14 @@ impl Client {
};
if n > 0 {
let chunk_len = n as usize;
if self
.recv_buffer
.available()
.saturating_add(chunk_len)
> MAX_RECV_BUFFER_BYTES
{
return Err(ErrorCode::Protocol);
}
*recv_budget -= chunk_len;
self.recv_buffer
.write(&tmp[..chunk_len])
Expand Down Expand Up @@ -687,6 +708,36 @@ mod tests {
assert!(MAX_RECV_BYTES_PER_POLL >= 65536);
}

#[test]
fn recv_buffer_staging_cap_matches_server_limit() {
assert_eq!(
MAX_RECV_BUFFER_BYTES,
2 * DEFAULT_MAX_MSG_LENGTH as usize
);
}

#[test]
fn poll_rejects_recv_buffer_growth_past_staging_cap() {
use std::io::Write;
use std::os::unix::io::IntoRawFd;
use std::os::unix::net::UnixStream;

let (client_end, mut peer) = UnixStream::pair().unwrap();
client_end.set_nonblocking(true).unwrap();
peer.set_nonblocking(true).unwrap();
peer.write_all(&[0x01, 0x02, 0x03]).unwrap();

let mut client = Client::new();
client.state = ClientState::Playing;
client.transport = Some(Transport::new_plain(client_end.into_raw_fd()));
client
.recv_buffer
.write(&vec![0u8; MAX_RECV_BUFFER_BYTES])
.unwrap();

assert_eq!(client.poll(0), Err(ErrorCode::Protocol));
}

#[test]
fn command_wait_recv_budget_bounds_connect_handshake_amplification() {
// 64 max-size AMF commands would be 256 MiB without a byte cap.
Expand Down
24 changes: 24 additions & 0 deletions src/message/command.rs
Original file line number Diff line number Diff line change
Expand Up @@ -300,7 +300,12 @@ pub fn read_onstatus(buf: &mut Buffer) -> Result<()> {
amf0::read_object_begin(buf)?;
let mut level = [0u8; 32];
let mut level_len = 0usize;
let mut keys = 0usize;
while !amf0::is_object_end(buf) {
keys += 1;
if keys > amf0::MAX_OBJECT_KEYS {
return Err(ErrorCode::Amf);
}
let mut key = [0u8; 256];
let key_len = amf0::read_object_key(buf, &mut key)?;
let key_str = std::str::from_utf8(&key[..key_len]).unwrap_or("");
Expand Down Expand Up @@ -477,4 +482,23 @@ mod tests {

assert!(read_onstatus(&mut buf).is_ok());
}

#[test]
fn read_onstatus_rejects_excessive_object_keys() {
let mut buf = Buffer::new();
amf0::write_string(&mut buf, "onStatus").unwrap();
amf0::write_number(&mut buf, 1.0).unwrap();
amf0::write_null(&mut buf).unwrap();
amf0::write_object_begin(&mut buf).unwrap();
for i in 0..=amf0::MAX_OBJECT_KEYS {
let key = format!("k{i}");
amf0::write_object_key(&mut buf, &key).unwrap();
amf0::write_boolean(&mut buf, false).unwrap();
}
amf0::write_object_key(&mut buf, "level").unwrap();
amf0::write_string(&mut buf, "status").unwrap();
amf0::write_object_end(&mut buf).unwrap();

assert_eq!(read_onstatus(&mut buf), Err(ErrorCode::Amf));
}
}