diff --git a/src/aacs/derive.rs b/src/aacs/derive.rs index 8c98573..624d225 100644 --- a/src/aacs/derive.rs +++ b/src/aacs/derive.rs @@ -367,6 +367,7 @@ pub(crate) fn resolve_dk_node( /// independent reproduction harnesses (e.g. `examples/prove_hkd_aacs.rs`) can /// exercise the exact same parser + verify primitives the production walk uses. /// These are thin wrappers — no new logic. +#[doc(hidden)] pub mod probe { use super::super::crypto::aes_ecb_decrypt; diff --git a/src/disc/hddvd.rs b/src/disc/hddvd.rs index e83caf3..065ec1c 100644 --- a/src/disc/hddvd.rs +++ b/src/disc/hddvd.rs @@ -215,12 +215,13 @@ fn probe_evo_streams(reader: &mut dyn SectorSource, extents: &[Extent]) -> Vec, data: Vec) -> crate::mux::ps::PsPacket { + crate::mux::ps::PsPacket { + stream_id, + sub_stream_id: sub, + pts: None, + dts: None, + data, + source: None, + } + } + + #[test] + fn collect_es_routes_only_vc1_0xfd_to_video() { + use crate::mux::ps::hddvd_extended_pid; + // The 0xFD guard: only the VC-1 extension (0x55) is video. An HD-audio + // 0xFD sub-stream (e.g. 0x72) that arrives FIRST must NOT stamp video_pid + // with its PID or pollute the video sample — else the real video track is + // lost. (Routing 0xFD audio to its own track is deferred.) + let mut video = Vec::new(); + let mut video_pid: Option = None; + let mut audio = BTreeMap::new(); + // Audio-on-0xFD (ext 0x72) first — must be ignored by the video path. + collect_es( + &ps_pkt(0xFD, Some(0x72), vec![0xAA; 32]), + &mut video, + &mut video_pid, + &mut audio, + ); + assert!( + video.is_empty(), + "0xFD audio sub-stream not routed to video" + ); + assert_eq!(video_pid, None, "0xFD audio did not stamp the video PID"); + // Then the real VC-1 video (ext 0x55). + collect_es( + &ps_pkt(0xFD, Some(0x55), vec![0xBB; 32]), + &mut video, + &mut video_pid, + &mut audio, + ); + assert_eq!( + video_pid, + Some(hddvd_extended_pid(0x55)), + "video PID stamped from the VC-1 0xFD sub-stream (0xFD55)" + ); + assert_eq!(video.len(), 32, "VC-1 0xFD payload accumulated as video"); + } + /// End-to-end: an `.evo` whose video rides the extended-stream-id 0xFD yields /// a VC-1 video track routed to `0xFD00 | ext` (0xFD55) — the PID the demuxer /// derives from the same stream_id_extension, so mux-time routing lines up. diff --git a/src/disc/mod.rs b/src/disc/mod.rs index feaa1b0..e781552 100644 --- a/src/disc/mod.rs +++ b/src/disc/mod.rs @@ -4854,6 +4854,25 @@ mod tests { ); } + #[test] + fn inject_unit_keys_labels_fmts_as_uhd_family() { + // FMTS is AACS 2.1 — a UHD-family, bus-encrypted format. Injecting a UK + // on an FMTS disc must synthesize the UHD version + bus encryption, not + // mislabel it AACS 1.0 / bus-off (which would break FMTS decryption on + // the mapfile-recovered-UK path). + let mut disc = make_test_disc(1000, "FMTS"); + disc.format = DiscFormat::Fmts; + disc.encrypted = true; + disc.inject_unit_keys(vec![(0u32, [0x22u8; 16])]); + let aacs = disc.aacs.as_ref().expect("aacs state synthesized"); + assert_eq!( + aacs.version, + crate::aacs::mkb::AACS_MAJOR_UHD, + "FMTS is AACS 2.x (UHD major), not BD" + ); + assert!(aacs.bus_encryption, "FMTS is bus-encrypted like UHD"); + } + /// Build an AacsState carrying the given unit keys (other fields are inert /// defaults — these tests only exercise the unit-key/decrypt-keys plumbing). fn aacs_with(unit_keys: Vec<(u32, [u8; 16])>) -> AacsState { diff --git a/src/mux/au_assembly.rs b/src/mux/au_assembly.rs index bef4b16..2c5022e 100644 --- a/src/mux/au_assembly.rs +++ b/src/mux/au_assembly.rs @@ -121,6 +121,11 @@ pub(crate) struct AuAssembler { /// boundary is "the next opener after a frame is already seen"). Meaningless /// for `Mode::StartCode`. Reset with `scan_pos`. seen_unit: bool, + /// Pre-sync opener-search cursor: the offset up to which the buffer has been + /// searched for the FIRST AU opener with none found. Resumes the opener scan + /// so a long run of junk with no start code (hostile/corrupt input) costs + /// O(bytes) total, not O(buffer) per push. Reset when `buf[0]` moves. + opener_pos: usize, } impl AuAssembler { @@ -143,6 +148,7 @@ impl AuAssembler { disc_marks: VecDeque::new(), scan_pos: 0, seen_unit: false, + opener_pos: 0, } } @@ -159,6 +165,7 @@ impl AuAssembler { disc_marks: VecDeque::new(), scan_pos: 0, seen_unit: false, + opener_pos: 0, } } @@ -245,11 +252,11 @@ impl AuAssembler { if matches!(self.mode, Mode::Passthrough) { return Vec::new(); } - let mode = self.mode; let mut out = Vec::new(); loop { - // Locate the AU start code that opens the buffered run. - let Some(a0) = au_opener(mode, &self.buf) else { + // Locate the AU start code that opens the buffered run (resumes from + // opener_pos so an unsynced junk run is scanned once, not per push). + let Some(a0) = self.au_opener_resumable() else { // No AU boundary buffered. Bound memory: drop all but a 3-byte // tail (enough to catch a start-code prefix straddling the cut) // once over the cap; otherwise wait for more data. @@ -329,6 +336,22 @@ impl AuAssembler { fn reset_scan(&mut self) { self.scan_pos = 0; self.seen_unit = false; + self.opener_pos = 0; + } + + /// Locate the first AU opener in `buf`, resuming the search from `opener_pos` + /// (bytes already searched with no opener) so a long unsynced run costs + /// O(bytes) total, not O(buffer) per push. Advances `opener_pos` on a miss. + fn au_opener_resumable(&mut self) -> Option { + match au_opener_from(self.mode, &self.buf, self.opener_pos) { + Some(o) => Some(o), + None => { + // Nothing yet; next call resumes here (back up 3 for a straddling + // start-code prefix). Never advance past what is searchable. + self.opener_pos = self.buf.len().saturating_sub(3).max(self.opener_pos); + None + } + } } /// Find the end of the AU that opens at `buf[0]`, resuming from `scan_pos` @@ -401,13 +424,13 @@ impl AuAssembler { /// Offset of the start code that opens the next AU in `buf` (at or after 0), or /// `None` if no AU-opening start code is buffered yet. -fn au_opener(mode: Mode, buf: &[u8]) -> Option { +fn au_opener_from(mode: Mode, buf: &[u8], from: usize) -> Option { match mode { - Mode::StartCode(marker) => find_start_code(buf, 0, marker), + Mode::StartCode(marker) => find_start_code(buf, from, marker), // Any of the three AU-opening BDU types opens a VC-1 access unit. - Mode::Vc1 => find_vc1_start(buf, 0), + Mode::Vc1 => find_vc1_start(buf, from), // A sequence header, GOP header, or picture opens an MPEG-2 access unit. - Mode::Mpeg2 => find_mpeg2_start(buf, 0), + Mode::Mpeg2 => find_mpeg2_start(buf, from), Mode::Passthrough => None, } } diff --git a/src/mux/codec/mpeg2.rs b/src/mux/codec/mpeg2.rs index e3fb718..1c6379c 100644 --- a/src/mux/codec/mpeg2.rs +++ b/src/mux/codec/mpeg2.rs @@ -46,10 +46,10 @@ const PICTURE_CODE: u8 = 0x00; /// Picture coding type: I-frame. const PICTURE_TYPE_I: u8 = 1; -/// Hard cap on the access-unit reassembly buffer. A real MPEG-2 frame is well -/// under 1 MiB (DVD I-frames ~100 KB); past this cap a corrupt stream that -/// never produces a second access-unit boundary is force-flushed as a single -/// frame rather than driving unbounded allocation. +/// The access-unit reassembly cap now lives in [`crate::mux::au_assembly`] (the +/// `AuAssembler` owns cross-PES buffering); this mirror exists only so the +/// force-flush test below can size an over-cap fixture against the same bound. +#[cfg(test)] const MAX_AU_BUFFER: usize = 8 * 1024 * 1024; /// Cap on frames held awaiting the first PES PTS anchor. A DVD stamps a PTS in @@ -259,8 +259,11 @@ impl Mpeg2Parser { }, }); // Safety cap: a stream with no GOP/sequence boundaries would buffer - // unbounded. Force-flush a pathologically long run as its own GOP. - if self.gop_buf.len() >= MAX_PENDING_FRAMES { + // unbounded. Force-flush a pathologically long run as its own GOP — + // bounded by BOTH the frame count and the total buffered bytes, so a + // crafted stream of few-but-huge pictures cannot over-allocate either. + let gop_bytes: usize = self.gop_buf.iter().map(|p| p.frame.data.len()).sum(); + if self.gop_buf.len() >= MAX_PENDING_FRAMES || gop_bytes >= MAX_PENDING_BYTES { self.flush_gop(out); } } @@ -1402,13 +1405,12 @@ mod tests { let mut data = make_picture_header(PICTURE_TYPE_I); // > MAX_AU_BUFFER of slice bytes with no following picture/seq/GOP. data.extend(std::iter::repeat_n(0xAA, MAX_AU_BUFFER + 1024)); - let frames = parser.parse(&make_pes(data, Some(0))); - assert!( - frames.is_empty(), - "over-cap AU is force-COMPLETED (bounded) but buffered in its GOP" - ); - let frames = parser.flush(); - assert_eq!(frames.len(), 1, "force-flushed at EOF, not dropped"); + // The AU assembler force-completes the ~8 MiB AU (no boundary), and the + // GOP byte cap (MAX_PENDING_BYTES) then force-flushes that oversized GOP + // during parse rather than buffering it unbounded. + let mut frames = parser.parse(&make_pes(data, Some(0))); + frames.extend(parser.flush()); + assert_eq!(frames.len(), 1, "over-cap AU force-flushed, not dropped"); assert!(frames[0].keyframe); } diff --git a/src/mux/codec/reorder.rs b/src/mux/codec/reorder.rs index daccadd..34e7a56 100644 --- a/src/mux/codec/reorder.rs +++ b/src/mux/codec/reorder.rs @@ -44,6 +44,13 @@ const FALLBACK_FRAME_DUR_NS: i64 = 1_001_000_000 / 24; /// reassembly buffer grows unbounded on disc-controlled input. const MAX_GOP_FRAMES: usize = 600; +/// Byte cap on the buffered GOP, complementing [`MAX_GOP_FRAMES`]. A GOP holds a +/// couple hundred MB at most in practice; this force-completes a run of +/// few-but-huge access units so a crafted/corrupt stream cannot over-allocate +/// (the AU assembler caps each frame at 8 MiB, so 600 frames alone could reach +/// ~5 GiB without this). +const MAX_GOP_BYTES: usize = 64 * 1024 * 1024; + /// One buffered coded picture awaiting its GOP's completion. struct Pending { /// Explicit PES PTS (ns) for this AU, or `None` when the source omitted it. @@ -72,6 +79,9 @@ struct Gop { pub(crate) struct SparsePtsReorder { /// Frames of the GOP currently accumulating, in decode order. cur: Vec, + /// Total `data` bytes buffered in `cur` — the byte-cap counter, reset each + /// time `cur` is drained into a completed GOP. + cur_bytes: usize, /// The previously-completed GOP, held one step so its duration can be /// calibrated from the next GOP's anchor before it is emitted. held: Option, @@ -86,6 +96,7 @@ impl SparsePtsReorder { pub(crate) fn new() -> Self { Self { cur: Vec::new(), + cur_bytes: 0, held: None, dur_ns: 0, next_start_ns: 0, @@ -102,11 +113,15 @@ impl SparsePtsReorder { // A keyframe opens a new GOP: the picture already accumulated in `cur` is // a complete GOP. Complete it (this frame belongs to the NEW GOP). Also // force-complete a pathologically long run that never signalled a - // keyframe, so a crafted/corrupt stream cannot buffer without bound. + // keyframe — bounded by BOTH frame count and total buffered bytes, so a + // crafted/corrupt stream of few-but-huge access units cannot buffer + // without bound. let mut out = Vec::new(); - if (frame.keyframe || self.cur.len() >= MAX_GOP_FRAMES) && !self.cur.is_empty() { + let over_cap = self.cur.len() >= MAX_GOP_FRAMES || self.cur_bytes >= MAX_GOP_BYTES; + if (frame.keyframe || over_cap) && !self.cur.is_empty() { out = self.complete_current_gop(); } + self.cur_bytes += frame.data.len(); self.cur.push(Pending { explicit, ctype, @@ -131,6 +146,7 @@ impl SparsePtsReorder { return Vec::new(); } let pend = std::mem::take(&mut self.cur); + self.cur_bytes = 0; let dispidx = display_indices(pend.iter().map(|p| p.ctype)); let count = pend.len() as i64; let anchor = pend @@ -310,6 +326,29 @@ mod tests { ); } + #[test] + fn force_flushes_a_gop_that_exceeds_the_byte_cap() { + use CodingType::*; + // Few-but-huge access units with no keyframe must not accumulate past the + // byte cap: a handful of ~MAX_GOP_BYTES/4-sized frames force-completes the + // GOP well before the frame-count cap, bounding memory. + let big = MAX_GOP_BYTES / 4 + 1; + let mut r = SparsePtsReorder::new(); + let mut emitted = 0usize; + // Enough huge frames to trigger several byte-cap completions (a GOP is + // held one step for duration calibration, so the first emit lands after + // the second cap fires) — well under the 600-frame count cap. + for i in 0..16 { + let mut f = frame(P, false); + f.data = vec![0u8; big]; + emitted += r.push((i == 0).then_some(0), f).len(); + } + assert!( + emitted >= 1, + "byte cap force-flushed (emitted {emitted}) before the frame-count cap" + ); + } + #[test] fn force_flushes_a_gop_that_never_signals_a_keyframe() { use CodingType::*;