audit: fix AU mark-field loss, VTI tie determinism, and mark/perf issues
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.
This commit is contained in:
@@ -352,3 +352,86 @@ pub fn mkb_type(mkb: &[u8]) -> Option<MkbType> {
|
||||
pub fn mkb_is_uhd(mkb: &[u8]) -> Option<bool> {
|
||||
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<u8> {
|
||||
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<u8> {
|
||||
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"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
+38
-1
@@ -104,7 +104,14 @@ fn parse_vti_clip_order(vti: &[u8]) -> Vec<String> {
|
||||
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<u8>, 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).
|
||||
|
||||
+79
-12
@@ -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<u8>,
|
||||
pts: Option<i64>,
|
||||
dts: Option<i64>,
|
||||
source: Option<SourcePos>,
|
||||
discontinuity: bool,
|
||||
) -> Vec<AssembledAu> {
|
||||
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,
|
||||
|
||||
@@ -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
|
||||
|
||||
+6
-1
@@ -1218,7 +1218,12 @@ impl<W: Write + Seek> MkvMuxer<W> {
|
||||
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 {
|
||||
|
||||
@@ -296,7 +296,7 @@ impl PipelinedPesStream {
|
||||
// → no continuity-gap flag.)
|
||||
let pkts: Vec<PesPacket> = 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,
|
||||
|
||||
Reference in New Issue
Block a user