From 5895ca4b5cfacdf0e033333ddbc61b8961fc4268 Mon Sep 17 00:00:00 2001
From: Cursor Agent <cursoragent@cursor.com>
Date: Sun, 12 Jul 2026 02:06:40 +0000
Subject: [PATCH] Fix client onStatus AMF amplification and recv_buffer staging
 cap

- Cap read_onstatus object keys at MAX_OBJECT_KEYS (256), matching
  read_connect and skip_value limits elsewhere in the AMF layer.
- Apply the same 8 MiB recv_buffer staging cap used on the server
  to Client::poll and recv_message so a malicious server cannot
  retain up to 64 MiB of incomplete wire data per connection.

Co-authored-by: Alexander Wagner <info@alexanderwagnerdev.com>
---
 src/client/mod.rs      | 51 ++++++++++++++++++++++++++++++++++++++++++
 src/message/command.rs | 24 ++++++++++++++++++++
 2 files changed, 75 insertions(+)

diff --git a/src/client/mod.rs b/src/client/mod.rs
index 8fb7d07..99ae811 100644
--- a/src/client/mod.rs
+++ b/src/client/mod.rs
@@ -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)]
@@ -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)?;
@@ -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])
@@ -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.
diff --git a/src/message/command.rs b/src/message/command.rs
index 088e3f1..a31edc2 100644
--- a/src/message/command.rs
+++ b/src/message/command.rs
@@ -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("");
@@ -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));
+    }
 }