From c81a6e05cdd480d2377979be05573c6736adc8df Mon Sep 17 00:00:00 2001 From: Matthew Jackson <1085847+MattJackson@users.noreply.github.com> Date: Thu, 9 Jul 2026 17:30:41 -0700 Subject: [PATCH] audit: fix AU mark-field loss, VTI tie determinism, and mark/perf issues MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Round-4 findings from the 10-phase release audit (the first fully clean round; it dug into the new #22/#18 refactor code): - AuAssembler closed each AU from only the FRONT mark's fields, so when one PES fragment carried the source and a later fragment of the same AU carried the PTS, the second field was dropped — a regression vs the old separate pts/source mark deques. Now merge the first Some of each field across all in-range marks. - parse_vti_clip_order picked the largest residue bucket with HashMap::into_values().max_by_key(), nondeterministic on a size tie (randomized HashMap iteration) — could select a different clip table run-to-run. Break ties by smallest offset. - Bound the marks/disc_marks deques (MAX_MARKS): the buf-size cap prunes marks only when bytes accumulate, so a run of zero-length timed fragments could grow them without bound on hostile input. - Add push_owned so the PS path moves the PES payload into a passthrough AU with no copy (MPEG-2 video + all audio), removing a per-PES malloc+memcpy the refactor had introduced on the DVD path. - Back-patch the MKV duration from the block END (start + its own duration) so it covers the final frame instead of understating by one. - Add direct tests for the MKB record-framing walker; drop a stale drain_complete_aus doc comment left on process_au. --- src/aacs/mkb.rs | 83 +++++++++++++++++++++++++++++++++ src/disc/hddvd.rs | 39 +++++++++++++++- src/mux/au_assembly.rs | 91 ++++++++++++++++++++++++++++++++----- src/mux/codec/mpeg2.rs | 3 -- src/mux/mkv.rs | 7 ++- src/mux/pipelined_stream.rs | 2 +- 6 files changed, 207 insertions(+), 18 deletions(-) diff --git a/src/aacs/mkb.rs b/src/aacs/mkb.rs index 19f7af2..2300c48 100644 --- a/src/aacs/mkb.rs +++ b/src/aacs/mkb.rs @@ -352,3 +352,86 @@ pub fn mkb_type(mkb: &[u8]) -> Option { pub fn mkb_is_uhd(mkb: &[u8]) -> Option { mkb_type(mkb).map(MkbType::is_uhd) } + +#[cfg(test)] +mod tests { + use super::*; + + /// One MKB record: 1 type byte + big-endian 24-bit total length + body. + fn rec(rec_type: u8, body: &[u8]) -> Vec { + let len = 4 + body.len(); + let mut v = vec![rec_type, (len >> 16) as u8, (len >> 8) as u8, len as u8]; + v.extend_from_slice(body); + v + } + + /// Type-and-Version record (0x10): body = 4-byte MKBType + 4-byte version. + fn type_and_version(mkb_type: u32, version: u32) -> Vec { + let mut body = mkb_type.to_be_bytes().to_vec(); + body.extend_from_slice(&version.to_be_bytes()); + rec(REC_TYPE_AND_VERSION, &body) + } + + #[test] + fn walker_frames_records_and_stops_at_end_marker() { + let mut mkb = type_and_version(MKB_20_CATEGORY_C, 77); + mkb.extend(rec(REC_VKD_TABLE, &[0xAA; 16])); + mkb.extend([0x00, 0x00, 0x00, 0x00]); // end marker + mkb.extend(rec(0x99, &[0xFF; 8])); // must NOT be walked (past the marker) + + let recs = walk_mkb(&mkb); + assert_eq!(recs.len(), 2, "walk stops at the 00 000000 end marker"); + assert_eq!(recs[0].rec_type, REC_TYPE_AND_VERSION); + assert_eq!(recs[1].rec_type, REC_VKD_TABLE); + assert_eq!(recs[1].body, vec![0xAA; 16]); + } + + #[test] + fn walker_stops_on_malformed_or_out_of_bounds_length() { + // A record whose declared length runs past the buffer end must terminate + // the walk rather than panic or read OOB. + let mkb = vec![REC_VKD_TABLE, 0x00, 0xFF, 0xFF, 0x01, 0x02]; // len=0xFFFF, only 6 bytes + assert!( + walk_mkb(&mkb).is_empty(), + "over-long record yields no records" + ); + // A sub-4 length (shorter than the header itself) is also rejected. + let short = vec![REC_VKD_TABLE, 0x00, 0x00, 0x02]; + assert!(walk_mkb(&short).is_empty(), "sub-4 length is rejected"); + // A truncated header (< 4 bytes) yields nothing. + assert!(walk_mkb(&[0x10, 0x00]).is_empty()); + } + + #[test] + fn mkb_type_and_version_decode_from_the_type_record() { + let mut mkb = type_and_version(MKB_21_CATEGORY_C, 100); + mkb.extend([0x00, 0x00, 0x00, 0x00]); + assert_eq!(mkb_type_raw(&mkb), Some(MKB_21_CATEGORY_C)); + assert_eq!(mkb_version(&mkb), Some(100)); + assert_eq!(mkb_is_uhd(&mkb), Some(true), "2.1 Category C is UHD"); + + let bd = type_and_version(MKB_TYPE_4_PRERECORDED, 68); + assert_eq!( + mkb_is_uhd(&bd), + Some(false), + "AACS 1.0 prerecorded is not UHD" + ); + // No Type record → None (not a panic, not a fabricated value). + assert_eq!(mkb_version(&rec(REC_VKD_TABLE, &[0; 16])), None); + assert_eq!(mkb_type_raw(&[]), None); + } + + #[test] + fn trim_mkb_keeps_only_the_framed_records() { + let mut mkb = type_and_version(MKB_20_CATEGORY_C, 1); + let content_len = mkb.len(); // the single framed record, no end marker + mkb.extend([0x00, 0x00, 0x00, 0x00]); // end marker + mkb.extend([0xDE; 4096]); // trailing padding past the end marker + let trimmed = trim_mkb(mkb); + assert_eq!( + trimmed.len(), + content_len, + "trim keeps the framed records, dropping the end marker and padding" + ); + } +} diff --git a/src/disc/hddvd.rs b/src/disc/hddvd.rs index 4921e63..66a1f18 100644 --- a/src/disc/hddvd.rs +++ b/src/disc/hddvd.rs @@ -104,7 +104,14 @@ fn parse_vti_clip_order(vti: &[u8]) -> Vec { count += 1; } } - let Some(mut best) = buckets.into_values().max_by_key(|g| g.len()) else { + // Pick the largest residue bucket (the clip table). On a size tie, break + // deterministically by the bucket's smallest offset — `HashMap` iteration + // order is randomized, so `max_by_key` alone could pick a different bucket + // run-to-run on identical bytes. + let Some(mut best) = buckets + .into_values() + .max_by_key(|g| (g.len(), std::cmp::Reverse(g.iter().map(|(o, _)| *o).min()))) + else { return Vec::new(); }; best.sort_by_key(|(o, _)| *o); @@ -539,6 +546,36 @@ mod tests { assert!(parse_vti_clip_order(b"not a vti").is_empty()); } + #[test] + fn parse_vti_clip_order_is_deterministic_on_a_bucket_size_tie() { + // Two residue buckets of EQUAL size must resolve to the SAME winner every + // call — `HashMap` iteration is randomized, so a `max_by_key` without a + // deterministic tie-break could pick a different bucket run-to-run on + // identical bytes. Build a VTI whose stray `.EVO` names tie the real + // table's bucket count, then assert the result is stable across calls. + let mut vti = vec![0u8; 0x600]; + vti[..HDDVD_VTI_MAGIC.len()].copy_from_slice(HDDVD_VTI_MAGIC); + let put = |v: &mut Vec, off: usize, name: &str| { + v[off..off + name.len()].copy_from_slice(name.as_bytes()); + }; + // Bucket A (residue 0x42): two names at stride 0x140. + put(&mut vti, 0x142, "A1.EVO"); + put(&mut vti, 0x282, "A2.EVO"); + // Bucket B (residue 0x50): two names — same count, different residue. + put(&mut vti, 0x150, "B1.EVO"); + put(&mut vti, 0x290, "B2.EVO"); + + let first = parse_vti_clip_order(&vti); + for _ in 0..20 { + assert_eq!( + parse_vti_clip_order(&vti), + first, + "tie-break must be deterministic across repeated calls" + ); + } + assert!(!first.is_empty()); + } + #[test] fn is_feature_clip_matches_the_feature_naming_variants() { // Layer-break split (Shaun / Anchorman) and the divide form (Harry Potter). diff --git a/src/mux/au_assembly.rs b/src/mux/au_assembly.rs index 74c881d..58a704c 100644 --- a/src/mux/au_assembly.rs +++ b/src/mux/au_assembly.rs @@ -32,6 +32,13 @@ use std::collections::VecDeque; /// at the cap rather than buffering without bound on hostile/corrupt input. const MAX_AU_BUFFER: usize = 8 * 1024 * 1024; +/// Cap on buffered timing/discontinuity marks. A real access unit spans a few +/// hundred PES fragments at most; this bounds the mark deques so a run of +/// zero-length (or start-code-free) timed fragments — which grow no buffer bytes +/// and so never trip the `MAX_AU_BUFFER` mark-prune — cannot accumulate marks +/// without bound on hostile/corrupt disc input. +const MAX_MARKS: usize = 64 * 1024; + /// One AU-complete unit drained from the buffer: its elementary-stream bytes plus /// the timing/source/discontinuity of the fragment that opened the AU. pub(crate) struct AssembledAu { @@ -155,7 +162,31 @@ impl AuAssembler { } } - /// Feed one PES fragment; return every AU that is now complete. + /// Feed one PES fragment the caller OWNS; return every AU now complete. For + /// a self-framing (`Passthrough`) stream the payload is MOVED straight into + /// the emitted unit with no copy — the common DVD/HD-DVD case (MPEG-2 video, + /// all audio). A buffering mode copies into `buf` exactly as [`Self::push`]. + pub(crate) fn push_owned( + &mut self, + data: Vec, + pts: Option, + dts: Option, + source: Option, + discontinuity: bool, + ) -> Vec { + if matches!(self.mode, Mode::Passthrough) { + return vec![AssembledAu { + data, + pts, + dts, + source, + discontinuity, + }]; + } + self.push(&data, pts, dts, source, discontinuity) + } + + /// Feed one PES fragment (borrowed); return every AU that is now complete. pub(crate) fn push( &mut self, data: &[u8], @@ -183,9 +214,20 @@ impl AuAssembler { dts, source, }); + // Backstop: the `buf`-size cap prunes marks only when bytes accumulate. + // A run of zero-length (or start-code-free) timed fragments grows no + // bytes, so bound the deque directly — drop the oldest (stalest) mark, + // which belongs to an already-emitted or lost AU. A real AU spans far + // fewer fragments than this cap. + if self.marks.len() > MAX_MARKS { + self.marks.pop_front(); + } } if discontinuity { self.disc_marks.push_back(off); + if self.disc_marks.len() > MAX_MARKS { + self.disc_marks.pop_front(); + } } self.buf.extend_from_slice(data); self.drain(false) @@ -245,19 +287,18 @@ impl AuAssembler { } let end_abs = self.base + end as u64; - // The AU's own timing/source/discontinuity: by the mark-drain - // invariant (stale marks below `base` were already dropped) the front - // mark, if it sits before this AU's end, belongs to this AU. + // The AU's own timing/source: take the FIRST Some of each field + // across every mark in this AU's range [base, end_abs), independently + // — one PES fragment may carry the source while a later fragment of + // the same AU carries the PTS (and vice versa), so reading only the + // front mark would drop the other field. This restores the semantics + // of the pre-consolidation separate pts/source mark deques. let (mut pts, mut dts, mut source) = (None, None, None); - if let Some(m) = self.marks.front() { - if m.off < end_abs { - pts = m.pts; - dts = m.dts; - source = m.source; - } - } while self.marks.front().is_some_and(|m| m.off < end_abs) { - self.marks.pop_front(); + let m = self.marks.pop_front().unwrap(); + pts = pts.or(m.pts); + dts = dts.or(m.dts); + source = source.or(m.source); } let mut discontinuity = false; if self.disc_marks.front().is_some_and(|&o| o < end_abs) { @@ -499,6 +540,32 @@ mod tests { assert_eq!(out2[0].data, au2); } + #[test] + fn au_merges_pts_and_source_from_different_fragments() { + // One fragment of an AU may carry the source stamp while a later fragment + // of the SAME AU carries the PTS (each PES gets a source; only the anchor + // gets a PTS). The AU must keep BOTH — reading only the front mark would + // drop whichever field the first fragment lacked. + let src = crate::pes::SourcePos::at_byte(4242); + let mut a = AuAssembler::for_codec(Codec::H264); + let full = au(0xAB, 80); + // Fragment 1: source only, no PTS. + assert!(a.push(&full[..30], None, None, Some(src), false).is_empty()); + // Fragment 2 (same AU): PTS only, no source. + assert!( + a.push(&full[30..], Some(9000), None, None, false) + .is_empty() + ); + let out = a.flush(); + assert_eq!(out.len(), 1); + assert_eq!(out[0].pts, Some(9000), "PTS from the 2nd fragment retained"); + assert_eq!( + out[0].source.map(|s| s.byte), + Some(4242), + "source from the 1st fragment retained" + ); + } + #[test] fn discontinuity_flag_attaches_to_the_au_it_opens() { // A discontinuity-flagged fragment opens AU2; that flag must land on AU2, diff --git a/src/mux/codec/mpeg2.rs b/src/mux/codec/mpeg2.rs index ba4bf8c..e3fb718 100644 --- a/src/mux/codec/mpeg2.rs +++ b/src/mux/codec/mpeg2.rs @@ -174,9 +174,6 @@ impl Mpeg2Parser { parse_aspect_ratio(hdr) } - /// Drain every complete access unit from `buf`, returning one Frame each. - /// When `force` is true (EOF flush, or buffer-cap backstop) the trailing - /// in-progress access unit is emitted even without a following boundary. /// Process one reassembled access unit (from [`AuAssembler`]): decode its /// per-picture coding info, capture a new sequence header, and buffer the /// picture into the current GOP for display-order timestamping. The AU's diff --git a/src/mux/mkv.rs b/src/mux/mkv.rs index acba74d..5130c81 100644 --- a/src/mux/mkv.rs +++ b/src/mux/mkv.rs @@ -1218,7 +1218,12 @@ impl MkvMuxer { self.last_pts_ticks.insert(track_idx, pts_ticks); // Track the highest block timestamp so a missing source duration can be // back-patched from the real muxed runtime at finish(). - self.max_block_ticks = self.max_block_ticks.max(pts_ticks); + // Track the block END (start + its own duration when known), not just + // the start, so a back-patched Segment Duration covers the final frame's + // full presentation instead of understating the runtime by one frame. + let block_end_ticks = + pts_ticks + duration_ns.map_or(0, |d| (d as i64 / TIMESTAMP_SCALE_NS).max(1)); + self.max_block_ticks = self.max_block_ticks.max(block_end_ticks); let relative_ts = (pts_ticks - self.cluster_ts_ticks) as i16; match duration_ns { diff --git a/src/mux/pipelined_stream.rs b/src/mux/pipelined_stream.rs index 3ada30e..5c069cf 100644 --- a/src/mux/pipelined_stream.rs +++ b/src/mux/pipelined_stream.rs @@ -296,7 +296,7 @@ impl PipelinedPesStream { // → no continuity-gap flag.) let pkts: Vec = match self.au_asm.get_mut(track) { Some(asm) => asm - .push(&ps.data, pts_i64, dts_i64, src, false) + .push_owned(ps.data, pts_i64, dts_i64, src, false) .into_iter() .map(|au| PesPacket { source: au.source,