diff --git a/crates/muninn-daemon/src/media_packetizer.rs b/crates/muninn-daemon/src/media_packetizer.rs index 653de16..d1400f2 100644 --- a/crates/muninn-daemon/src/media_packetizer.rs +++ b/crates/muninn-daemon/src/media_packetizer.rs @@ -1,5 +1,6 @@ use anyhow::{Result, anyhow}; use odin_core::{MuninnMediaAudioPacketRecord, MuninnMediaVideoAccessUnitRecord}; +use std::collections::BTreeMap; #[derive(Clone, Debug, PartialEq, Eq)] pub struct VideoAccessUnit { @@ -184,6 +185,99 @@ pub fn packetize_audio_packet( }) } +pub fn reassemble_video_access_unit( + chunks: &[MuninnMediaVideoAccessUnitRecord], +) -> Result { + let first = chunks + .first() + .ok_or_else(|| anyhow!("video access unit requires at least one chunk"))?; + if first.chunk_count == 0 { + return Err(anyhow!("video access unit chunk_count must be non-zero")); + } + + let mut by_index = BTreeMap::new(); + for chunk in chunks { + if chunk.stream_id != first.stream_id { + return Err(anyhow!( + "video access unit chunks have mixed stream_id values" + )); + } + if chunk.session_id != first.session_id { + return Err(anyhow!( + "video access unit chunks have mixed session_id values" + )); + } + if chunk.frame_id != first.frame_id { + return Err(anyhow!( + "video access unit chunks have mixed frame_id values" + )); + } + if chunk.codec != first.codec { + return Err(anyhow!("video access unit chunks have mixed codec values")); + } + if chunk.pts_ticks != first.pts_ticks + || chunk.duration_ticks != first.duration_ticks + || chunk.timebase_num != first.timebase_num + || chunk.timebase_den != first.timebase_den + || chunk.deadline_ticks != first.deadline_ticks + { + return Err(anyhow!("video access unit chunks have mixed timing values")); + } + if chunk.keyframe != first.keyframe + || chunk.dependency_frame_id != first.dependency_frame_id + { + return Err(anyhow!( + "video access unit chunks have mixed dependency values" + )); + } + if chunk.chunk_count != first.chunk_count { + return Err(anyhow!( + "video access unit chunks have mixed chunk_count values" + )); + } + if chunk.chunk_index >= first.chunk_count { + return Err(anyhow!( + "video access unit chunk_index {} is outside chunk_count {}", + chunk.chunk_index, + first.chunk_count + )); + } + if chunk.payload.is_empty() { + return Err(anyhow!("video access unit chunk payload must be non-empty")); + } + if by_index + .insert(chunk.chunk_index, chunk.payload.as_slice()) + .is_some() + { + return Err(anyhow!( + "video access unit has duplicate chunk_index {}", + chunk.chunk_index + )); + } + } + + if by_index.len() != first.chunk_count as usize { + return Err(anyhow!( + "video access unit is missing chunks: expected {}, received {}", + first.chunk_count, + by_index.len() + )); + } + + let mut bytes = Vec::new(); + for index in 0..first.chunk_count { + let payload = by_index + .get(&index) + .ok_or_else(|| anyhow!("video access unit is missing chunk_index {index}"))?; + bytes.extend_from_slice(payload); + } + + Ok(VideoAccessUnit { + bytes, + keyframe: first.keyframe, + }) +} + fn h264_annex_b_nal_units(input: &[u8]) -> Result>> { let mut starts = Vec::new(); let mut offset = 0_usize; @@ -428,4 +522,91 @@ mod tests { assert!(error.to_string().contains("payload must be non-empty")); } + + #[test] + fn reassembles_video_chunks_in_chunk_index_order() -> Result<()> { + let access_unit = VideoAccessUnit { + bytes: vec![1, 2, 3, 4, 5], + keyframe: true, + }; + let records = packetize_video_access_unit( + VideoFramePacketizeOptions { + stream_id: "muninn.raven.av.rudp", + session_id: "session-1", + codec: "h264", + frame_id: 9, + pts_ticks: 27_000, + duration_ticks: 3_000, + timebase_num: 1, + timebase_den: 90_000, + deadline_ticks: 28_800, + max_payload_bytes: 2, + }, + &access_unit, + )?; + let shuffled = vec![records[2].clone(), records[0].clone(), records[1].clone()]; + + let reassembled = reassemble_video_access_unit(&shuffled)?; + + assert_eq!(reassembled, access_unit); + Ok(()) + } + + #[test] + fn rejects_video_chunks_from_mixed_frames() -> Result<()> { + let access_unit = VideoAccessUnit { + bytes: vec![1, 2, 3, 4], + keyframe: false, + }; + let mut records = packetize_video_access_unit( + VideoFramePacketizeOptions { + stream_id: "muninn.raven.av.rudp", + session_id: "session-1", + codec: "h264", + frame_id: 9, + pts_ticks: 27_000, + duration_ticks: 3_000, + timebase_num: 1, + timebase_den: 90_000, + deadline_ticks: 28_800, + max_payload_bytes: 2, + }, + &access_unit, + )?; + records[1].frame_id = 10; + + let error = reassemble_video_access_unit(&records).unwrap_err(); + + assert!(error.to_string().contains("mixed frame_id")); + Ok(()) + } + + #[test] + fn rejects_incomplete_video_chunk_sets() -> Result<()> { + let access_unit = VideoAccessUnit { + bytes: vec![1, 2, 3, 4, 5], + keyframe: false, + }; + let mut records = packetize_video_access_unit( + VideoFramePacketizeOptions { + stream_id: "muninn.raven.av.rudp", + session_id: "session-1", + codec: "h264", + frame_id: 9, + pts_ticks: 27_000, + duration_ticks: 3_000, + timebase_num: 1, + timebase_den: 90_000, + deadline_ticks: 28_800, + max_payload_bytes: 2, + }, + &access_unit, + )?; + records.pop(); + + let error = reassemble_video_access_unit(&records).unwrap_err(); + + assert!(error.to_string().contains("missing chunks")); + Ok(()) + } }