From 84acebf007e0e1a3fc3e449f93cddf62e6fbbeb9 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Fri, 18 Sep 2026 02:26:39 +0000 Subject: [PATCH] fix(conn): account aggregate sub-tags in per-recv message budget Decrement the per-recv message budget for every aggregate sub-tag, including unknown types, and reject zero-length sub-tags (mirroring the multitrack guard). Closes the CPU amplification hole left when issue #233 budgeted only audio/video/script aggregate fan-out. Fixes #302 Co-authored-by: Alexander Wagner --- src/client/mod.rs | 59 ++++++++++++++++++++++++++++++++++++--------- src/session/conn.rs | 46 ++++++++++++++++++++++++++++++----- 2 files changed, 87 insertions(+), 18 deletions(-) diff --git a/src/client/mod.rs b/src/client/mod.rs index abba8fc9..13299a26 100644 --- a/src/client/mod.rs +++ b/src/client/mod.rs @@ -804,6 +804,13 @@ impl Client { if body + data_size > payload.len() { return Err(ErrorCode::Protocol); } + if data_size == 0 { + return Err(ErrorCode::Protocol); + } + if *messages_processed >= MAX_MESSAGES_PER_POLL { + break; + } + *messages_processed += 1; if !have_base { sub_base_ts = ts; have_base = true; @@ -814,10 +821,6 @@ impl Client { if let Some(cb) = self.on_frame_cb { match tag_type { msg_dispatch::RTMP_MSG_AUDIO => { - if *messages_processed >= MAX_MESSAGES_PER_POLL { - break; - } - *messages_processed += 1; self.deliver_av_frame_cb( cb, FrameType::Audio, @@ -827,10 +830,6 @@ impl Client { )?; } msg_dispatch::RTMP_MSG_VIDEO => { - if *messages_processed >= MAX_MESSAGES_PER_POLL { - break; - } - *messages_processed += 1; self.deliver_av_frame_cb( cb, FrameType::Video, @@ -840,10 +839,6 @@ impl Client { )?; } msg_dispatch::RTMP_MSG_AMF0_DATA => { - if *messages_processed >= MAX_MESSAGES_PER_POLL { - break; - } - *messages_processed += 1; self.deliver_script_frame_cb(cb, out_ts, tag_payload); } _ => { @@ -1929,6 +1924,46 @@ mod tests { assert_eq!(*CALLBACKS.lock().unwrap(), MAX_MESSAGES_PER_POLL); } + #[test] + fn aggregate_unknown_subtags_consume_message_budget() { + let filler = [0x00]; + let mut aggregate = Vec::new(); + for i in 0..(MAX_MESSAGES_PER_POLL + 8) { + aggregate.push(0x01); + aggregate.push(0x00); + aggregate.push(0x00); + aggregate.push(0x01); + aggregate.extend_from_slice(&[ + (i >> 16) as u8, + (i >> 8) as u8, + i as u8, + (i >> 24) as u8, + ]); + aggregate.extend_from_slice(&[0, 0, 0]); + aggregate.push(filler[0]); + let prev_tag_size = 12u32; + aggregate.extend_from_slice(&prev_tag_size.to_be_bytes()); + } + + let mut wire = Buffer::new(); + let mut cmsg = ChunkMessage::default(); + cmsg.csid = 6; + cmsg.fmt = 0; + cmsg.msg_length = aggregate.len() as u32; + cmsg.msg_type_id = msg_dispatch::RTMP_MSG_AGGREGATE; + cmsg.msg_stream_id = 1; + chunk_write(&mut wire, &cmsg, &aggregate, aggregate.len(), 128).unwrap(); + + let mut client = Client::new(); + client.recv_buffer.write(wire.peek()).unwrap(); + let mut messages_processed = 0; + client + .drain_ready_messages(&mut messages_processed) + .unwrap(); + + assert_eq!(messages_processed, MAX_MESSAGES_PER_POLL); + } + #[test] fn drain_ready_messages_splits_multitrack_video() { use std::sync::{LazyLock, Mutex}; diff --git a/src/session/conn.rs b/src/session/conn.rs index 9ed28097..459c9d43 100644 --- a/src/session/conn.rs +++ b/src/session/conn.rs @@ -1026,6 +1026,16 @@ impl Conn { if body + data_size > payload.len() { return Err(ErrorCode::Protocol); } + // Mirror multitrack: zero-length sub-tags pack the maximum sub-tag + // count into the smallest wire footprint and bypass per-sub-tag + // processing budgets when their type is not audio/video/script. + if data_size == 0 { + return Err(ErrorCode::Protocol); + } + if *messages_budget == 0 { + return Ok(()); + } + *messages_budget = messages_budget.saturating_sub(1); if !have_base { sub_base_ts = ts; @@ -1040,12 +1050,7 @@ impl Conn { match tag_type { msg_dispatch::RTMP_MSG_AUDIO | msg_dispatch::RTMP_MSG_VIDEO - | msg_dispatch::RTMP_MSG_AMF0_DATA => { - if *messages_budget == 0 { - return Ok(()); - } - *messages_budget = messages_budget.saturating_sub(1); - } + | msg_dispatch::RTMP_MSG_AMF0_DATA => {} _ => { pos = body + data_size + 4; continue; @@ -4993,6 +4998,35 @@ mod tests { assert_eq!(conn.pending_relay.len(), 3); } + #[test] + fn aggregate_unknown_subtags_consume_message_budget() { + let mut conn = Conn::new(); + let filler = [0x00]; + let mut aggregate = Vec::new(); + for i in 0..10 { + aggregate.extend(flv_subtag(0x01, i, &filler)); + } + + let mut messages_budget = 3; + conn.handle_aggregate(1, 0, &aggregate, &mut messages_budget) + .unwrap(); + + assert_eq!(messages_budget, 0); + assert!(conn.pending_relay.is_empty()); + } + + #[test] + fn aggregate_rejects_zero_size_subtags() { + let mut conn = Conn::new(); + let aggregate = flv_subtag(0x08, 0, &[]); + let mut messages_budget = MAX_MESSAGES_PER_RECV; + + assert_eq!( + conn.handle_aggregate(1, 0, &aggregate, &mut messages_budget), + Err(ErrorCode::Protocol) + ); + } + #[test] fn aggregate_message_rejects_subtag_size_overrunning_payload() { let mut conn = Conn::new();