fix(conn): close aggregate unknown sub-tag budget bypass (#302)

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};

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();