From 6be5198886bb637abfb58b97936456d0b14f8fd2 Mon Sep 17 00:00:00 2001 From: MattJackson <1085847+MattJackson@users.noreply.github.com> Date: Fri, 5 Jun 2026 16:10:54 -0700 Subject: [PATCH] =?UTF-8?q?libfreemkv:=20v1.0=20hardening=20=E2=80=94=20co?= =?UTF-8?q?dec/EBML/TS=20robustness=20+=20DTS=20parser=20fixes?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Audit-driven fixes (rounds 1–3): - hevc: correct hvcC profile/level SPS offsets (HEVC has a 2-byte NAL header) - mkv: map all DTS variants to the registered A_DTS codec id; force a new cluster before the i16 cluster-relative timestamp can overflow - ebml/mkvstream: bound untrusted EBML sizes (no multi-GB allocs); reject uint>8 (was an OOB panic) and non-{0,4,8} float widths (were a desync) - ts: skip PES-header bytes that span a TS packet boundary; add the PMT section_len/prog_info_len bounds the PAT parser already had - ac3: preserve a 0x0B77 syncword split across a PES boundary; cap buffer - dts: validate each next-core boundary by decoded core size (a 0x7FFE8001 pattern inside XLL payload no longer false-splits/drops the lossless extension); reject sub-minimum core frames; fix forced-emit PTS base - lpcm: DVD program-stream PCM no longer double-strips the BD LPCM header - vc1/mpeg2: do not emit a parameter-set-only PES as a standalone frame - pgs/truehd: cap the pending reassembly buffer (parity with ac3/dts) - aacs: ts_syncs_intact uses the exact packet count - prefetched: capacity-guard the recycled-buffer set_len - Cargo.toml: exclude project docs from the published crate Convergence: a third independent audit pass found no remaining material (CRITICAL/HIGH/MEDIUM) issues. Full precommit (fmt + clippy -D + tests, Rust 1.86) green. --- Cargo.toml | 2 + src/aacs/decrypt.rs | 7 +- src/mux/codec/ac3.rs | 126 ++++++++++++++++++++-- src/mux/codec/dts.rs | 217 +++++++++++++++++++++++++++++++++++--- src/mux/codec/hevc.rs | 155 +++++++++++++++++++++++---- src/mux/codec/lpcm.rs | 98 ++++++++++++++++-- src/mux/codec/mod.rs | 12 ++- src/mux/codec/mpeg2.rs | 77 +++++++++++++- src/mux/codec/pgs.rs | 40 ++++++- src/mux/codec/truehd.rs | 12 +++ src/mux/codec/vc1.rs | 56 +++++++++- src/mux/disc.rs | 3 +- src/mux/ebml.rs | 57 +++++++--- src/mux/mkv.rs | 155 +++++++++++++++++++++++++-- src/mux/mkvstream.rs | 219 ++++++++++++++++++++++++++++++++++++--- src/mux/resolve.rs | 3 +- src/mux/ts.rs | 66 ++++++++++-- src/sector/prefetched.rs | 15 +-- 18 files changed, 1201 insertions(+), 119 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index 77b5213..bf9ae49 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -8,6 +8,8 @@ description = "Open source raw disc access library for optical drives" repository = "https://github.com/freemkv/libfreemkv" keywords = ["bluray", "uhd", "optical", "scsi", "disc"] categories = ["hardware-support", "multimedia"] +# Keep internal AI-instruction / private notes out of the published crate. +exclude = [] [dependencies] serde = { version = "1", features = ["derive"] } diff --git a/src/aacs/decrypt.rs b/src/aacs/decrypt.rs index 23bbf19..f248155 100644 --- a/src/aacs/decrypt.rs +++ b/src/aacs/decrypt.rs @@ -100,7 +100,12 @@ fn ts_syncs_intact(unit: &[u8]) -> bool { } offset += TS_PACKET_LEN; } - let total = (unit.len() - 4) / TS_PACKET_LEN + 1; + // One sync byte is checked per 192-byte BD-TS packet (at offset 4 of + // each). `total` is exactly that packet count; the old + // `(len - 4) / TS_PACKET_LEN + 1` over-counted by one for lengths of + // the form `4 + k·192` (harmless for the always-6144 aligned unit, but + // wrong in general and it biased the majority threshold). + let total = unit.len() / TS_PACKET_LEN; count > total / 2 } diff --git a/src/mux/codec/ac3.rs b/src/mux/codec/ac3.rs index 311a019..7319e9c 100644 --- a/src/mux/codec/ac3.rs +++ b/src/mux/codec/ac3.rs @@ -6,6 +6,14 @@ use super::{CodecParser, Frame, PesPacket, pts_to_ns}; +/// Hard cap on the carry-over buffer. An AC-3/E-AC-3 frame is at most 8192 +/// bytes (the `frame_size > 8192` reject below), so a single straddling frame +/// plus a little slack never needs more than this. If the buffer grows past +/// the cap without yielding a frame (pathological / never-syncing input) we +/// drop it and resync rather than accumulate one PES worth of data per call +/// for the whole title. +const MAX_AC3_BUF: usize = 64 * 1024; + pub struct Ac3Parser { /// Leftover bytes from previous PES (incomplete frame at end). buf: Vec, @@ -81,19 +89,38 @@ impl CodecParser for Ac3Parser { pos = start + frame_size; } - // Keep unconsumed data for next call - // `pos` points to the start of unconsumed data (either a partial sync or leftover) + // Keep unconsumed data for the next call. `pos` is the start of the + // unconsumed region: either a partial frame that straddles this PES + // boundary — which, by construction, begins at a syncword (every byte + // before `pos` was emitted as a frame or skipped as pre-sync junk) — or + // trailing bytes too short to size/complete a frame. Carry from `pos`, + // NOT from the next syncword: discarding bytes between `pos` and the + // next sync would drop the partial frame we are deliberately keeping + // across the boundary. let keep_from = if pos < data.len() { - // Find the last sync word position in the unconsumed region - find_ac3_sync(&data[pos..]) - .map(|o| pos + o) - .unwrap_or(data.len()) + // A syncword at/after `pos` marks the carry-over start (anything + // before it is junk with no sync). With no full sync, retain the + // whole tail — including a lone trailing 0x0B that may be the first + // half of a syncword split across the PES boundary. + match find_ac3_sync(&data[pos..]) { + Some(o) => pos + o, + None if data.last() == Some(&0x0B) => data.len() - 1, + None => data.len(), + } } else { data.len() }; if keep_from < data.len() { - self.buf = data[keep_from..].to_vec(); + let tail = &data[keep_from..]; + if tail.len() > MAX_AC3_BUF { + // No frame could be parsed out of a buffer this large — this is + // not valid AC-3 here. Drop it and resync on the next PES rather + // than grow without bound on pathological input. + self.buf.clear(); + } else { + self.buf = tail.to_vec(); + } } else { self.buf.clear(); } @@ -276,6 +303,91 @@ mod tests { assert_eq!(frames[0].data.len(), 160); } + #[test] + fn sync_word_split_across_pes_is_preserved() { + // A frame whose 0x0B77 syncword straddles the PES boundary (0x0B at the + // tail of PES 1, 0x77 at the head of PES 2) must still be emitted whole. + // Previously the lone trailing 0x0B was dropped and the frame lost. + let mut parser = Ac3Parser::new(); + let frame_data = make_ac3_frame(0, 2); // 160 bytes, starts with 0x0B 0x77 + + // PES 1: a complete frame, then a single 0x0B (first half of next sync). + let mut pes1_data = frame_data.clone(); + pes1_data.push(0x0B); + let pes1 = PesPacket { + pid: 0, + pts: Some(90000), + dts: None, + data: pes1_data, + }; + let frames1 = parser.parse(&pes1); + assert_eq!(frames1.len(), 1, "first complete frame emitted"); + + // PES 2: 0x77 (second half of sync) + rest of the second frame. + let mut pes2_data = vec![0x77]; + pes2_data.extend_from_slice(&frame_data[2..]); + let pes2 = PesPacket { + pid: 0, + pts: Some(93000), + dts: None, + data: pes2_data, + }; + let frames2 = parser.parse(&pes2); + assert_eq!(frames2.len(), 1, "split-sync frame must be recovered"); + assert_eq!(frames2[0].data.len(), 160); + } + + #[test] + fn buffer_stays_bounded_across_many_garbage_pes() { + // Finding 14: the carry-over buffer must never grow without bound. Feed + // many large PES packets that contain no usable frame and assert the + // retained buffer stays tiny — carry-from-`pos` drops all pre-sync junk, + // and a never-completing frame is bounded by the 8192-byte frame cap and + // the MAX_AC3_BUF resync guard. + let mut parser = Ac3Parser::new(); + for i in 0..256 { + // Vary the trailing byte so we also exercise the lone-0x0B retain. + let mut data = vec![0x55u8; 8192]; + if i % 3 == 0 { + *data.last_mut().unwrap() = 0x0B; + } + let pes = PesPacket { + pid: 0, + pts: None, + dts: None, + data, + }; + let frames = parser.parse(&pes); + assert!(frames.is_empty()); + assert!( + parser.buf.len() <= MAX_AC3_BUF, + "buffer grew to {} (cap {})", + parser.buf.len(), + MAX_AC3_BUF + ); + } + // After all that garbage the retained tail is at most a single partial + // syncword byte — never an accumulation of whole PES packets. + assert!(parser.buf.len() <= 1, "retained {} bytes", parser.buf.len()); + } + + #[test] + fn split_sync_below_cap_is_still_retained() { + // The cap must not break the normal split-sync straddle: a short tail + // ending in 0x0B (well under the cap) is retained so the next PES can + // complete the syncword. + let mut parser = Ac3Parser::new(); + let data = vec![0x00, 0x00, 0x0B]; + let pes = PesPacket { + pid: 0, + pts: None, + dts: None, + data, + }; + assert!(parser.parse(&pes).is_empty()); + assert_eq!(parser.buf, vec![0x0B], "lone trailing 0x0B retained"); + } + #[test] fn ac3_frame_size_table() { // fscod=0 (48kHz), frmsizecod=0: 64 words = 128 bytes diff --git a/src/mux/codec/dts.rs b/src/mux/codec/dts.rs index db7cf0f..302917d 100644 --- a/src/mux/codec/dts.rs +++ b/src/mux/codec/dts.rs @@ -44,6 +44,19 @@ impl DtsParser { /// this without a clean boundary we resync rather than stall or balloon. const MAX_AU_BYTES: usize = 65536; +/// Minimum plausible DTS core frame size. The core header alone is ~10-14 +/// bytes; a decoded `core_size` below this means we matched a false/corrupt +/// core sync (the 14-bit `fsize` field decoded to a tiny value) rather than a +/// real frame, so we resync instead of emitting a junk access unit. +const MIN_CORE_FRAME_BYTES: usize = 10; + +/// Sentinel for "no valid PTS base captured yet". Real PTS-in-ns values are +/// non-negative (derived from the unsigned 90 kHz PES timestamp), so a negative +/// value can never collide with a genuine timestamp. Used to mark the PTS base +/// invalid after a forced flush so the next PES sets it regardless of buffer +/// state. +const PTS_UNSET: i64 = -1; + impl CodecParser for DtsParser { fn parse(&mut self, pes: &PesPacket) -> Vec { if pes.data.is_empty() { @@ -65,7 +78,11 @@ impl CodecParser for DtsParser { // old per-PES emit that dropped the extension PES packets and // downgraded the track to lossy DTS core (the Dunkirk / Fight Club // bug). The PTS is the core frame's PTS, captured when the unit began. - if self.buf.is_empty() { + // Capture the access unit's PTS base on a fresh buffer, or whenever a + // prior forced (safety-valve) flush left it invalidated — in the + // forced case the bytes still in `buf` are not a real core frame, so + // the first PES to arrive after the flush carries the correct base. + if self.buf.is_empty() || self.pending_pts == PTS_UNSET { self.pending_pts = pts_ns; } self.buf.extend_from_slice(&pes.data); @@ -96,7 +113,15 @@ impl CodecParser for DtsParser { break; } let core_size = dts_core_frame_size(&self.buf); - if core_size == 0 || core_size > MAX_AU_BYTES { + // `dts_core_frame_size` returns a 14-bit `fsize + 1`, so it is + // always in [1, 16384]; the bare `== 0` / `> MAX_AU_BYTES` checks + // can never fire. A real DTS core header is at least ~10-14 bytes, + // so any decoded size below that came from a false/corrupt sync. + // Reject it (drain the 4 syncword bytes and resync) instead of + // letting a tiny bogus size close the current access unit at a junk + // boundary and drop the trailing extension substreams. The + // `> MAX_AU_BYTES` upper bound is kept as a harmless guard. + if !(MIN_CORE_FRAME_BYTES..=MAX_AU_BYTES).contains(&core_size) { // Bogus core sync — skip past it and resync. self.buf.drain(..4); continue; @@ -105,12 +130,25 @@ impl CodecParser for DtsParser { break; // core frame not fully buffered yet — wait } - // The access unit ends at the next core sync. Search begins after - // this core's syncword so we don't re-match it. Anything between - // the core and that next sync is this unit's extension substream(s). - let au_end = match find_sync(&self.buf[core_size..], &DTS_CORE_SYNC) { - Some(rel) => core_size + rel, - None => { + // The access unit ends at the next *valid* core sync. The search + // begins after this core's syncword so we don't re-match it. + // Anything between the core and that next sync is this unit's + // extension substream(s) — which can themselves contain byte + // sequences matching the core syncword, so a raw `find_sync` match + // is not enough: a candidate is only a real boundary if its decoded + // core size is plausible. `next_core_boundary` skips bogus matches. + // + // `forced` distinguishes a real next-core boundary from a forced + // safety-valve flush. On a forced flush the access unit was NOT + // closed by a new core sync, so the bytes following it are not a + // fresh core frame and the current PES's PTS (which on a forced + // flush is an extension-substream PES, carrying its own later + // timestamp) must NOT become the next unit's PTS base. + let mut forced = false; + let au_end = match next_core_boundary(&self.buf, core_size) { + NextCore::Found(end) => end, + NextCore::NeedMore => break, // candidate sync needs more header + NextCore::None => { // No next core sync buffered yet. The trailing extension // substream PES packets may still be arriving, so WAIT for // them rather than emit a core-only (lossy) frame — unless @@ -119,6 +157,7 @@ impl CodecParser for DtsParser { if self.buf.len() <= MAX_AU_BYTES { break; } + forced = true; self.buf.len() } }; @@ -131,11 +170,18 @@ impl CodecParser for DtsParser { duration_ns: None, }); self.buf.drain(..au_end); - // The next access unit (now at buf start) belongs to a later PTS. - // We can't know it exactly until its core PES arrives, but the - // current PES's PTS is the best available approximation when the - // boundary fell inside this PES; refine on the next call's start. - self.pending_pts = pts_ns; + if forced { + // Safety-valve flush: the next access unit's real core PES has + // not arrived. Invalidate the PTS so the next PES sets it + // regardless of buffer state, rather than inheriting this + // (non-core) PES's timestamp. + self.pending_pts = PTS_UNSET; + } else { + // Real boundary: the next access unit (now at buf start) begins + // at a core sync that arrived inside this PES, so this PES's PTS + // is the correct base for it. + self.pending_pts = pts_ns; + } } frames @@ -151,12 +197,21 @@ impl CodecParser for DtsParser { return Vec::new(); } let core_size = dts_core_frame_size(&self.buf); - if core_size == 0 || self.buf.len() < core_size { + // `dts_core_frame_size` returns a 14-bit `fsize + 1` (never 0), so the + // old `== 0` check was dead; reject a sub-minimum core like `parse()`. + if core_size < MIN_CORE_FRAME_BYTES || self.buf.len() < core_size { self.buf.clear(); return Vec::new(); } let au = std::mem::take(&mut self.buf); - let pts_ns = self.pending_pts; + // A non-empty buffer here means a PES arrived after any prior forced + // flush (which fully drains `buf`), so `pending_pts` was reset to that + // PES's real PTS. Clamp the sentinel to 0 defensively all the same. + let pts_ns = if self.pending_pts == PTS_UNSET { + 0 + } else { + self.pending_pts + }; vec![Frame { pts_ns, keyframe: true, @@ -177,6 +232,40 @@ fn find_sync(data: &[u8], pattern: &[u8; 4]) -> Option { (0..=data.len() - 4).find(|&i| data[i..i + 4] == *pattern) } +/// Result of scanning for the next valid core sync that closes an access unit. +enum NextCore { + /// A valid next core sync was found; the access unit ends at this offset. + Found(usize), + /// A candidate core sync was found but its header isn't fully buffered yet, + /// so its validity can't be decided — wait for more data. + NeedMore, + /// No (further) core sync found in the buffer. + None, +} + +/// Find the next *valid* core sync after the current core frame, to delimit the +/// access unit. Extension-substream payload can contain byte sequences that +/// match the core syncword, so each candidate is validated by decoding its +/// core size: a match whose decoded size is implausible (< MIN_CORE_FRAME_BYTES +/// or > MAX_AU_BYTES) is a false sync and is skipped, continuing the search. +fn next_core_boundary(buf: &[u8], core_size: usize) -> NextCore { + let mut from = core_size; + while let Some(rel) = find_sync(&buf[from..], &DTS_CORE_SYNC) { + let pos = from + rel; + // Need the candidate's core header to judge it. + if buf.len() - pos < 10 { + return NextCore::NeedMore; + } + let sz = dts_core_frame_size(&buf[pos..]); + if (MIN_CORE_FRAME_BYTES..=MAX_AU_BYTES).contains(&sz) { + return NextCore::Found(pos); + } + // False sync inside extension payload — skip it and keep searching. + from = pos + 4; + } + NextCore::None +} + /// DTS core frame size from header bits. /// fsize is at bits 46-59 (14 bits) of the header: bytes 5-7. fn dts_core_frame_size(data: &[u8]) -> usize { @@ -349,6 +438,104 @@ mod tests { assert_eq!(tail[0].data.len(), 512 + 300); } + /// Build 4 bytes that look like a DTS core sync but whose `fsize` field + /// decodes to a tiny `core_size` (< MIN_CORE_FRAME_BYTES). With the + /// dead-code guards this passed validation and could close an access unit + /// at a junk boundary; with the fix it must be drained and resynced past. + fn bogus_tiny_core_sync() -> Vec { + // Core sync + zero header bytes. fsize = 0 → core_size = 1 (< 10). + let mut v = vec![0u8; 10]; + v[0..4].copy_from_slice(&DTS_CORE_SYNC); + // bytes 5,6,7 left zero → fsize = 0 → dts_core_frame_size = 1. + assert_eq!(dts_core_frame_size(&v), 1); + v + } + + #[test] + fn bogus_tiny_core_sync_does_not_split_or_drop_real_au() { + // A real core frame followed by an extension substream that happens to + // contain a false core sync whose fsize decodes tiny. The bogus sync + // must NOT close the real access unit early (dropping the rest of the + // extension) nor emit a junk few-byte frame — it must be skipped, and + // the whole core + extension preserved as one access unit. + let mut parser = DtsParser::new(); + + // Frame 1: core(512) + an extension whose body embeds a bogus tiny + // core sync midway through. + let mut ext = make_dts_ext(256); + // Embed the bogus core sync inside the extension body (offset 64). + let bogus = bogus_tiny_core_sync(); + ext[64..64 + bogus.len()].copy_from_slice(&bogus); + + let mut frame1 = make_dts_core(512); + frame1.extend_from_slice(&ext); + + // No next REAL core yet → frame 1 held. + assert!( + parser.parse(&make_pes(frame1, Some(90000))).is_empty(), + "bogus tiny core sync must not close the AU; wait for a real core" + ); + + // Frame 2's real core arrives — closes frame 1. + let f = parser.parse(&make_pes(make_dts_core(640), Some(93000))); + assert_eq!(f.len(), 1, "exactly one real access unit emitted"); + assert_eq!( + f[0].data.len(), + 512 + 256, + "AU must be the full core + extension, not split at the bogus sync" + ); + assert_eq!(f[0].pts_ns, pts_to_ns(90000), "AU keeps the core's PTS"); + + let tail = parser.flush(); + assert_eq!(tail.len(), 1); + assert_eq!(tail[0].data.len(), 640); + } + + #[test] + fn forced_emit_does_not_corrupt_next_au_pts() { + // When the buffer exceeds MAX_AU_BYTES with no next core sync, the + // parser force-emits for forward progress. The current PES at that + // point is an extension-substream PES (later PTS). The forced path must + // NOT make that extension PTS the base of the NEXT access unit. + let mut parser = DtsParser::new(); + + // Core PES at the real PTS, then a giant extension (no next core) that + // pushes the buffer past MAX_AU_BYTES, forcing an emit. + let core_pts = 90000i64; + assert!( + parser + .parse(&make_pes(make_dts_core(512), Some(core_pts))) + .is_empty() + ); + let ext_pts = 120000i64; // later extension-PES timestamp + let big_ext = make_dts_ext(MAX_AU_BYTES + 1024); + let f = parser.parse(&make_pes(big_ext, Some(ext_pts))); + assert_eq!(f.len(), 1, "oversized buffer force-emits one AU"); + assert_eq!( + f[0].pts_ns, + pts_to_ns(core_pts), + "forced AU keeps the core PTS" + ); + + // The next REAL core PES arrives with its own PTS. Its AU must inherit + // THIS core's PTS, not the prior extension PES timestamp. + let next_core_pts = 150000i64; + assert!( + parser + .parse(&make_pes(make_dts_core(512), Some(next_core_pts))) + .is_empty() + ); + let next_next_pts = 180000i64; + let f2 = parser.parse(&make_pes(make_dts_core(512), Some(next_next_pts))); + assert_eq!(f2.len(), 1); + assert_eq!( + f2[0].pts_ns, + pts_to_ns(next_core_pts), + "AU after a forced emit must use the next core's PTS, not the \ + stale extension PTS" + ); + } + #[test] fn codec_private_none() { let parser = DtsParser::new(); diff --git a/src/mux/codec/hevc.rs b/src/mux/codec/hevc.rs index 16a6372..6b1f76a 100644 --- a/src/mux/codec/hevc.rs +++ b/src/mux/codec/hevc.rs @@ -62,6 +62,11 @@ impl CodecParser for HevcParser { while let Some(sc_pos) = find_start_code(data, pos) { if let Some(nal_start) = skip_start_code(data, sc_pos) { let next = find_start_code(data, nal_start).unwrap_or(data.len()); + // Strip the leading zeros of the following start code. For a + // conforming bitstream this is lossless: rbsp_trailing_bits() + // sets a stop-one bit, so the final byte of any RBSP is never + // 0x00 — the only trailing zeros here belong to the next + // 00 00 (00) 01 prefix. let mut end = next; while end > nal_start && data[end - 1] == 0x00 { end -= 1; @@ -124,34 +129,41 @@ impl CodecParser for HevcParser { // Full HEVCDecoderConfigurationRecord is complex — for now, concatenate let mut record = Vec::new(); - // Minimal HEVCDecoderConfigurationRecord header + // Minimal HEVCDecoderConfigurationRecord header. + // + // The stored SPS NAL is [2-byte HEVC NAL header][SPS RBSP...]. + // The RBSP begins at sps[2]; profile_tier_level() begins one byte + // later, after sps_video_parameter_set_id u(4) + + // sps_max_sub_layers_minus1 u(3) + sps_temporal_id_nesting_flag u(1) + // (= sps[2], a full byte). So the profile_tier_level fields are: + // sps[3] general_profile_space u(2)+tier u(1)+profile_idc u(5) + // sps[4..8] general_profile_compatibility_flags u(32) + // sps[8..14] general_constraint_indicator_flags 48 bits + // sps[14] general_level_idc u(8) + // (Byte-aligned read; emulation-prevention bytes within the first + // 15 SPS bytes are not handled — extremely rare and matches the + // pre-existing simplification.) record.push(1); // configurationVersion - // General profile space, tier flag, profile IDC from SPS - if sps.len() > 3 { - record.push(sps[1]); // general_profile_space + general_tier_flag + general_profile_idc + // general_profile_space + general_tier_flag + general_profile_idc + record.push(if sps.len() > 3 { sps[3] } else { 0 }); + // general_profile_compatibility_flags (4 bytes) — SPS bytes 4..8 + if sps.len() > 7 { + record.extend_from_slice(&sps[4..8]); } else { - record.push(0); + let avail = sps.len().saturating_sub(4).min(4); + record.extend_from_slice(&sps[sps.len().min(4)..sps.len().min(8)]); + record.extend_from_slice(&vec![0u8; 4 - avail]); } - // general_profile_compatibility_flags (4 bytes) — from SPS bytes 2..6 - if sps.len() > 5 { - record.extend_from_slice(&sps[2..6]); + // general_constraint_indicator_flags (6 bytes) — SPS bytes 8..14 + if sps.len() > 13 { + record.extend_from_slice(&sps[8..14]); } else { - record.extend_from_slice(&[0, 0, 0, 0]); + let avail = sps.len().saturating_sub(8).min(6); + record.extend_from_slice(&sps[sps.len().min(8)..sps.len().min(14)]); + record.extend_from_slice(&vec![0u8; 6 - avail]); } - // general_constraint_indicator_flags (6 bytes) — from SPS bytes 6..12 - if sps.len() > 11 { - record.extend_from_slice(&sps[6..12]); - } else { - let avail = sps.len().saturating_sub(6).min(6); - if avail > 0 { - record.extend_from_slice(&sps[6..6 + avail]); - record.extend_from_slice(&vec![0u8; 6 - avail]); - } else { - record.extend_from_slice(&[0, 0, 0, 0, 0, 0]); - } - } - // general_level_idc - record.push(if sps.len() > 12 { sps[12] } else { 0 }); + // general_level_idc — SPS byte 14 + record.push(if sps.len() > 14 { sps[14] } else { 0 }); // min_spatial_segmentation_idc (4 + 12 bits) record.extend_from_slice(&[0xF0, 0x00]); // parallelismType (6 + 2 bits) @@ -268,6 +280,103 @@ mod tests { ); } + #[test] + fn hvcc_profile_tier_level_offsets() { + // The hvcC fixed header must read profile_tier_level from the SPS + // RBSP, not from the NAL header. Stored SPS = [2-byte NAL header][RBSP]. + // RBSP layout (byte-aligned): + // sps[2] sps_vps_id/max_sub_layers/temporal_nesting + // sps[3] general_profile_space+tier+profile_idc + // sps[4..8] general_profile_compatibility_flags + // sps[8..14] general_constraint_indicator_flags + // sps[14] general_level_idc + let mut parser = HevcParser::new(); + + // Distinct, recognizable values for each field. + let sps_rbsp: [u8; 13] = [ + 0xAB, // sps[2] (vps_id etc.) — must NOT leak into profile fields + 0x21, // sps[3] profile byte: space=0, tier=0, profile_idc=1 + 0x60, 0x00, 0x00, 0x00, // sps[4..8] compat flags + 0x90, 0x00, 0x00, 0x00, 0x00, 0x00, // sps[8..14] constraint flags + 0x7B, // sps[14] level_idc = 123 + ]; + + let mut data = Vec::new(); + // VPS + data.extend_from_slice(&[0x00, 0x00, 0x01]); + data.extend_from_slice(&hevc_nal_header(32)); + data.extend_from_slice(&[0xAA, 0xBB, 0xCC]); + // SPS — 2-byte header + the structured RBSP above + data.extend_from_slice(&[0x00, 0x00, 0x01]); + data.extend_from_slice(&hevc_nal_header(33)); + data.extend_from_slice(&sps_rbsp); + // PPS + data.extend_from_slice(&[0x00, 0x00, 0x01]); + data.extend_from_slice(&hevc_nal_header(34)); + data.extend_from_slice(&[0xDD, 0xEE]); + + let pes = make_pes(data, Some(0)); + parser.parse(&pes); + + let cp = parser + .codec_private() + .expect("codec_private should be Some"); + + // record[0] = configurationVersion + assert_eq!(cp[0], 1, "configurationVersion"); + // record[1] = general_profile_space+tier+profile_idc <- sps[3] + assert_eq!( + cp[1], 0x21, + "profile byte must come from SPS RBSP, not NAL hdr" + ); + // record[2..6] = general_profile_compatibility_flags <- sps[4..8] + assert_eq!(&cp[2..6], &[0x60, 0x00, 0x00, 0x00], "compatibility flags"); + // record[6..12] = general_constraint_indicator_flags <- sps[8..14] + assert_eq!( + &cp[6..12], + &[0x90, 0x00, 0x00, 0x00, 0x00, 0x00], + "constraint flags" + ); + // record[12] = general_level_idc <- sps[14] + assert_eq!(cp[12], 0x7B, "level_idc must come from sps[14]"); + } + + #[test] + fn hvcc_short_sps_does_not_panic() { + // A truncated SPS must still produce a fixed header without panicking + // and zero-pad the missing profile/level bytes. + let mut parser = HevcParser::new(); + + let mut data = Vec::new(); + data.extend_from_slice(&[0x00, 0x00, 0x01]); + data.extend_from_slice(&hevc_nal_header(32)); + data.extend_from_slice(&[0xAA]); + // SPS with only 3 RBSP bytes (stored len = 5): forces every guard path + data.extend_from_slice(&[0x00, 0x00, 0x01]); + data.extend_from_slice(&hevc_nal_header(33)); + data.extend_from_slice(&[0x11, 0x22, 0x33]); + data.extend_from_slice(&[0x00, 0x00, 0x01]); + data.extend_from_slice(&hevc_nal_header(34)); + data.extend_from_slice(&[0xDD]); + + let pes = make_pes(data, Some(0)); + parser.parse(&pes); + + let cp = parser + .codec_private() + .expect("codec_private should be Some"); + // sps stored = [hdr0, hdr1, 0x11, 0x22, 0x33], len 5. + // profile byte = sps[3] = 0x22; everything past sps[4]=0x33 is absent. + assert_eq!(cp[0], 1); + assert_eq!(cp[1], 0x22, "profile byte = sps[3]"); + // compat flags: only sps[4]=0x33 present, rest zero-padded. + assert_eq!(&cp[2..6], &[0x33, 0x00, 0x00, 0x00]); + // constraint flags: none present, all zero. + assert_eq!(&cp[6..12], &[0x00, 0x00, 0x00, 0x00, 0x00, 0x00]); + // level_idc: absent, zero. + assert_eq!(cp[12], 0x00); + } + #[test] fn codec_private_none_before_params() { let parser = HevcParser::new(); diff --git a/src/mux/codec/lpcm.rs b/src/mux/codec/lpcm.rs index 4e61ea6..7954b07 100644 --- a/src/mux/codec/lpcm.rs +++ b/src/mux/codec/lpcm.rs @@ -1,24 +1,39 @@ //! BD/DVD LPCM (Linear PCM) audio parser. //! -//! BD LPCM PES packets have a 4-byte header: +//! BD LPCM PES packets (TS stream type 0x80) carry a 4-byte header on the +//! elementary-stream payload: //! Bytes 0-1: audio frame number //! Byte 2: reserved //! Byte 3: quantization (bits 7-6), sample rate (bits 5-4), channel assignment (bits 3-0) +//! This header is part of the ES payload, so the BD parser must strip it. //! -//! DVD LPCM (private stream 1, sub-stream 0xA0-0xA7) has a 3-byte header. +//! DVD LPCM lives in private stream 1 (sub-stream 0xA0-0xA7). Its 7-byte +//! private sub-header (sub_id + frames + first-access-unit-ptr(2) + emphasis + +//! quant/freq + channels) is stripped by `PsDemuxer` while demuxing the +//! Program Stream. By the time a DVD LPCM `PesPacket` reaches this parser its +//! `data` is already raw PCM, so the parser must NOT strip any further bytes — +//! doing so drops one sample pair per PES and drifts the audio. //! -//! The raw PCM data follows the header. No framing is needed — each PES -//! payload minus its header is one complete audio frame. +//! The two origins are distinguished by the `strip_header` flag: BD = strip, +//! DVD = leave intact. The raw PCM data is otherwise one complete audio frame +//! per PES; no framing is needed. //! //! For MKV: codec ID "A_PCM/INT/BIG" (BD) or "A_PCM/INT/LIT" (DVD). //! All frames are keyframes (uncompressed audio). use super::{CodecParser, Frame, PesPacket, pts_to_ns}; -/// BD LPCM header size in bytes. +/// BD LPCM header size in bytes (present on BD-TS LPCM, absent on DVD-PS LPCM +/// because `PsDemuxer` already stripped the private sub-header). const BD_LPCM_HEADER_SIZE: usize = 4; -pub struct LpcmParser; +pub struct LpcmParser { + /// Whether to strip the 4-byte BD LPCM header from each PES payload. + /// + /// `true` for BD-TS LPCM (header still present), `false` for DVD-PS LPCM + /// (header already removed by `PsDemuxer`). + strip_header: bool, +} impl Default for LpcmParser { fn default() -> Self { @@ -27,23 +42,36 @@ impl Default for LpcmParser { } impl LpcmParser { + /// BD-TS LPCM parser: strips the 4-byte BD LPCM header from each PES. pub fn new() -> Self { - Self + Self { strip_header: true } + } + + /// DVD-PS LPCM parser: `PsDemuxer` already stripped the private sub-header, + /// so the payload is raw PCM and no further bytes are removed. + pub fn new_dvd() -> Self { + Self { + strip_header: false, + } } } impl CodecParser for LpcmParser { fn parse(&mut self, pes: &PesPacket) -> Vec { - // Skip the BD LPCM header (4 bytes). + let offset = if self.strip_header { + BD_LPCM_HEADER_SIZE + } else { + 0 + }; // If the PES is too short to contain header + data, return nothing. - if pes.data.len() <= BD_LPCM_HEADER_SIZE { + if pes.data.len() <= offset { return Vec::new(); } let pts_ns = pes.pts.map(pts_to_ns).unwrap_or(0); vec![Frame { pts_ns, keyframe: true, - data: pes.data[BD_LPCM_HEADER_SIZE..].to_vec(), + data: pes.data[offset..].to_vec(), duration_ns: None, }] } @@ -84,6 +112,56 @@ mod tests { assert_eq!(frames[0].pts_ns, 1_000_000_000); // 90000 ticks = 1 second } + #[test] + fn bd_lpcm_strips_4_byte_header() { + // BD-TS LPCM: the 4-byte BD header is part of the ES payload and must + // be stripped, leaving exactly the PCM bytes. + let mut parser = LpcmParser::new(); + let header = vec![0x00, 0x01, 0x00, 0b1001_0001]; + let pcm = vec![0x11, 0x22, 0x33, 0x44, 0x55, 0x66, 0x77, 0x88]; + let mut data = header; + data.extend_from_slice(&pcm); + + let frames = parser.parse(&make_pes(data, Some(0))); + assert_eq!(frames.len(), 1); + assert_eq!(frames[0].data, pcm, "BD must strip exactly 4 header bytes"); + } + + #[test] + fn dvd_lpcm_preserves_all_pcm_bytes() { + // DVD-PS LPCM: PsDemuxer already removed the 7-byte private sub-header, + // so the payload handed to this parser is raw PCM. The DVD parser must + // NOT strip any further bytes (the round-2 audit Finding 3 bug: the BD + // 4-byte strip dropped one sample pair per PES, drifting the audio). + let mut parser = LpcmParser::new_dvd(); + let pcm = vec![0xDE, 0xAD, 0xBE, 0xEF, 0xCA, 0xFE, 0x01, 0x02]; + let frames = parser.parse(&make_pes(pcm.clone(), Some(90000))); + + assert_eq!(frames.len(), 1); + assert_eq!( + frames[0].data, pcm, + "DVD must preserve every PCM byte — no second strip" + ); + assert_eq!(frames[0].pts_ns, 1_000_000_000); + } + + #[test] + fn dvd_lpcm_emits_short_payload_bd_would_drop() { + // A 4-byte raw-PCM DVD payload (1 sample pair at 16-bit stereo). The BD + // parser drops <= 4 bytes as "header only"; the DVD parser must emit it. + let mut bd = LpcmParser::new(); + let mut dvd = LpcmParser::new_dvd(); + let pcm = vec![0xAA, 0xBB, 0xCC, 0xDD]; + + assert!( + bd.parse(&make_pes(pcm.clone(), Some(0))).is_empty(), + "BD treats 4 bytes as header-only" + ); + let frames = dvd.parse(&make_pes(pcm.clone(), Some(0))); + assert_eq!(frames.len(), 1); + assert_eq!(frames[0].data, pcm); + } + #[test] fn always_keyframe() { let mut parser = LpcmParser::new(); diff --git a/src/mux/codec/mod.rs b/src/mux/codec/mod.rs index 3bd9d1a..8d77aaa 100644 --- a/src/mux/codec/mod.rs +++ b/src/mux/codec/mod.rs @@ -102,7 +102,16 @@ impl CodecParser for PassthroughParser { /// Create the appropriate parser for a codec, with optional codec private data. /// /// For DvdSub, `codec_data` should be the pre-formatted VobSub .idx palette header. -pub fn parser_for_codec(codec: Codec, codec_data: Option>) -> Box { +/// +/// `is_dvd_ps` selects the DVD program-stream variant where it matters: DVD +/// LPCM arrives with its private sub-header already stripped by the +/// `PsDemuxer`, so the LPCM parser must NOT strip the 4-byte BD LPCM header +/// again (that would drop one PCM sample pair per PES → progressive drift). +pub fn parser_for_codec( + codec: Codec, + codec_data: Option>, + is_dvd_ps: bool, +) -> Box { match codec { Codec::H264 => Box::new(h264::H264Parser::new()), Codec::Hevc => Box::new(hevc::HevcParser::new()), @@ -112,6 +121,7 @@ pub fn parser_for_codec(codec: Codec, codec_data: Option>) -> Box Box::new(dts::DtsParser::new()), Codec::TrueHd => Box::new(truehd::TrueHdParser::new()), Codec::Pgs => Box::new(pgs::PgsParser::new()), + Codec::Lpcm if is_dvd_ps => Box::new(lpcm::LpcmParser::new_dvd()), Codec::Lpcm => Box::new(lpcm::LpcmParser::new()), Codec::DvdSub => Box::new(dvdsub::DvdSubParser::new(codec_data)), _ => Box::new(PassthroughParser::new(true)), diff --git a/src/mux/codec/mpeg2.rs b/src/mux/codec/mpeg2.rs index ff72559..b10e194 100644 --- a/src/mux/codec/mpeg2.rs +++ b/src/mux/codec/mpeg2.rs @@ -99,6 +99,7 @@ impl CodecParser for Mpeg2Parser { let pts_ns = pes.dts.or(pes.pts).map(pts_to_ns).unwrap_or(0); let data = &pes.data; let mut keyframe = false; + let mut has_picture = false; // Scan for start codes in the elementary stream data. let mut pos = 0; @@ -176,6 +177,7 @@ impl CodecParser for Mpeg2Parser { PICTURE_CODE => { // Picture header: bytes after start code contain temporal_reference // (10 bits) + picture_coding_type (3 bits). + has_picture = true; if sc + 5 < data.len() { let picture_coding_type = (data[sc + 5] >> 3) & 0x07; if picture_coding_type == PICTURE_TYPE_I { @@ -190,6 +192,23 @@ impl CodecParser for Mpeg2Parser { } } + // A PES that carried a sequence header but no picture start code is a + // parameter-set-only access unit: it has no coded picture to emit. + // Emitting it as a standalone keyframe would put bare sequence-header + // bytes into frame data with no picture. The sequence header is + // captured into codec_private above and is re-emitted in-band on the + // next real picture's PES, so dropping the empty access unit loses + // nothing. Mirrors how the H.264/HEVC parsers skip parameter-set-only + // access units. + // + // Conservative: only drop when this PES actually contained a sequence + // header and no picture. A PES with neither (e.g. a slice + // continuation) still passes through unchanged, preserving real + // keyframe detection. + if !has_picture && contains_seq_header(data) { + return Vec::new(); + } + vec![Frame { pts_ns, keyframe, @@ -241,6 +260,21 @@ fn parse_aspect_ratio(hdr: &[u8]) -> Option<(u8, u8)> { Some(ASPECT_RATIOS[ar_code]) } +/// Returns true if `data` contains a sequence-header start code (00 00 01 B3). +fn contains_seq_header(data: &[u8]) -> bool { + let mut pos = 0; + while let Some(sc) = find_start_code(data, pos) { + if sc + 3 >= data.len() { + break; + } + if data[sc + 3] == SEQ_HEADER_CODE { + return true; + } + pos = sc + 4; + } + false +} + /// Find the position of the next start code (00 00 01) at or after `from`. fn find_start_code(data: &[u8], from: usize) -> Option { if data.len() < from + 3 { @@ -429,15 +463,15 @@ mod tests { assert!(has_ext, "codec_private should include sequence extension"); } - // --- I-frame with sequence header = keyframe --- + // --- sequence header + picture = keyframe --- #[test] - fn sequence_header_implies_keyframe() { + fn sequence_header_with_picture_is_keyframe() { let mut parser = Mpeg2Parser::new(); let mut data = Vec::new(); data.extend_from_slice(&make_seq_header(720, 480, 3, 4)); - // Even without an explicit picture header, a sequence header implies I-frame. + data.extend_from_slice(&make_picture_header(PICTURE_TYPE_I)); data.extend_from_slice(&[0xFF; 16]); let pes = make_pes(data, Some(0)); @@ -445,6 +479,43 @@ mod tests { assert_eq!(frames.len(), 1); assert!(frames[0].keyframe); + // codecPrivate is still captured. + assert!(parser.codec_private().is_some()); + } + + // --- parameter-set-only PES (seq header, no picture) emits no frame --- + + #[test] + fn sequence_header_only_pes_emits_no_frame() { + let mut parser = Mpeg2Parser::new(); + + // A PES carrying only a sequence header (+ extension), no picture. + let mut data = Vec::new(); + data.extend_from_slice(&make_seq_header(1920, 1080, 3, 4)); + data.extend_from_slice(&[0x00, 0x00, 0x01, SEQ_EXT_CODE]); + data.extend_from_slice(&[0x14, 0x8A, 0x00, 0x01, 0x00, 0x00]); + + let pes = make_pes(data, Some(0)); + let frames = parser.parse(&pes); + + // No coded picture → no frame emitted, but the sequence header is + // still captured for codecPrivate. + assert!( + frames.is_empty(), + "parameter-set-only PES should not emit a frame" + ); + assert!( + parser.codec_private().is_some(), + "sequence header should still be captured into codec_private" + ); + + // A following picture-bearing PES emits the real keyframe. + let mut data2 = Vec::new(); + data2.extend_from_slice(&make_picture_header(PICTURE_TYPE_I)); + data2.extend_from_slice(&[0xFF; 16]); + let frames2 = parser.parse(&make_pes(data2, Some(3600))); + assert_eq!(frames2.len(), 1); + assert!(frames2[0].keyframe); } // --- PTS conversion --- diff --git a/src/mux/codec/pgs.rs b/src/mux/codec/pgs.rs index 17196a2..e9b2af5 100644 --- a/src/mux/codec/pgs.rs +++ b/src/mux/codec/pgs.rs @@ -18,6 +18,13 @@ use super::{CodecParser, Frame, PesPacket, pts_to_ns}; const SEGMENT_PCS: u8 = 0x16; +// Upper bound on a pending display set's accumulated bytes. Real PGS +// display sets are small (a 1080p RLE bitmap plus palette is well under +// 1 MB); a stream that keeps appending non-PCS segments without ever +// emitting a PCS is malformed. Cap accumulation to bound memory and +// drop further appends until the next PCS resyncs the parser. Mirrors +// the MAX_*_BYTES / MAX_*_BUF caps in the DTS and AC-3 parsers. +const MAX_PGS_PENDING_BYTES: usize = 4 * 1024 * 1024; // Offset within the PES payload at which number_of_composition_objects // lives in a PCS: 3-byte segment header + 10 bytes of PCS fields // (video_w/h, frame_rate, comp_num, comp_state, palette_update, @@ -90,7 +97,12 @@ impl CodecParser for PgsParser { // a pending display, append; otherwise emit as-is. None => { if let Some((_, ref mut buf)) = self.pending { - buf.extend_from_slice(&pes.data); + // Bound accumulation: a well-formed display set is small. + // Past the cap, drop further appends (malformed stream); + // the next PCS will take/replace `pending` and resync. + if buf.len() + pes.data.len() <= MAX_PGS_PENDING_BYTES { + buf.extend_from_slice(&pes.data); + } } else { out.push(Frame { pts_ns, @@ -180,6 +192,32 @@ mod tests { assert!(data.windows(5).any(|w| w == [0x15, 0x00, 0x02, 0xAA, 0xBB])); } + #[test] + fn pending_buffer_is_capped() { + let mut parser = PgsParser::new(); + // Open a display set. + let _ = parser.parse(&make_pes(pcs_bytes(1), Some(90000))); + + // Flood with non-PCS segments far exceeding the cap. + let chunk = vec![0x15u8; 256 * 1024]; // 256 KB ODS-like segment + let floods = (MAX_PGS_PENDING_BYTES / chunk.len()) + 32; + for _ in 0..floods { + let frames = parser.parse(&make_pes(chunk.clone(), Some(90000))); + assert!(frames.is_empty(), "non-PCS appends should not emit"); + } + + // The pending buffer must not have grown without bound. + let pending_len = parser.pending.as_ref().map(|(_, b)| b.len()).unwrap_or(0); + assert!( + pending_len <= MAX_PGS_PENDING_BYTES, + "pending buffer {pending_len} exceeded cap {MAX_PGS_PENDING_BYTES}" + ); + + // A following PCS still resyncs and emits the (capped) pending set. + let frames = parser.parse(&make_pes(pcs_bytes(0), Some(180000))); + assert_eq!(frames.len(), 1); + } + #[test] fn codec_private_none() { let parser = PgsParser::new(); diff --git a/src/mux/codec/truehd.rs b/src/mux/codec/truehd.rs index cb332c8..162bfc2 100644 --- a/src/mux/codec/truehd.rs +++ b/src/mux/codec/truehd.rs @@ -16,6 +16,12 @@ use super::{CodecParser, Frame, PesPacket, pts_to_ns}; /// Duration of one TrueHD access unit in nanoseconds (1/1200 second). const AU_DURATION_NS: i64 = 833_333; +/// Hard cap on the reassembly buffer. A valid TrueHD/MAT access unit is +/// well under 32 KiB; if the buffer grows far past that without yielding a +/// frame the stream is malformed, so we drop it and resync rather than grow +/// without bound. Parity with the AC-3 / DTS / PGS caps. +const MAX_TRUEHD_BUF: usize = 256 * 1024; + pub struct TrueHdParser { buf: Vec, next_pts_ns: i64, @@ -145,6 +151,12 @@ impl CodecParser for TrueHdParser { self.next_pts_ns += AU_DURATION_NS; } + // Bound memory on malformed input: a stream that never yields a + // complete frame must not grow the buffer without limit. + if self.buf.len() > MAX_TRUEHD_BUF { + self.buf.clear(); + } + frames } diff --git a/src/mux/codec/vc1.rs b/src/mux/codec/vc1.rs index 959b6b1..24dbde9 100644 --- a/src/mux/codec/vc1.rs +++ b/src/mux/codec/vc1.rs @@ -44,6 +44,7 @@ impl CodecParser for Vc1Parser { // Use DTS when available (monotonic for B-frame content), fall back to PTS let ts_ns = pes.dts.or(pes.pts).map(pts_to_ns).unwrap_or(0); let mut has_seq_header = false; + let mut has_entry_point = false; let mut frame_start: Option = None; // Scan for start codes (00 00 01 XX) @@ -67,6 +68,7 @@ impl CodecParser for Vc1Parser { SC_ENTRY_POINT => { let end = find_next_sc(data, i + 4).unwrap_or(data.len()); self.entry_point = Some(data[i..end].to_vec()); + has_entry_point = true; } SC_FRAME => { // Frame data starts at this start code @@ -85,11 +87,23 @@ impl CodecParser for Vc1Parser { // Keyframe = this PES contains a sequence header (I-frame indicator in BD) let keyframe = has_seq_header; - // Strip sequence header + entry point from frame data — those are in codecPrivate. - // Only include data from the frame start code onwards. + // Strip sequence header + entry point from frame data — those are in + // codecPrivate, not coded-picture data. Only include data from the + // frame start code onwards. let frame_data = match frame_start { Some(start) => &data[start..], - None => data, // no frame start code found, pass through entire PES + None => { + // No frame start code. If this PES carried only parameter sets + // (sequence header / entry point, captured above into + // codecPrivate), there is no coded picture to emit — drop it + // rather than passing parameter bytes through as a bogus + // keyframe. Mirrors how the H.264/HEVC parsers skip + // parameter-set-only access units. + if has_seq_header || has_entry_point { + return Vec::new(); + } + data // genuine picture payload with no leading 0x0D — pass through + } }; vec![Frame { @@ -343,6 +357,42 @@ mod tests { assert_eq!(&frames[0].data[0..4], &[0x00, 0x00, 0x01, SC_FRAME]); } + // --- parameter-set-only PES (seq header + entry point, no frame SC) --- + + #[test] + fn param_set_only_pes_emits_no_frame() { + let mut parser = Vc1Parser::new(); + + // Sequence header + entry point, but NO frame start code (0x0D). + let mut data = Vec::new(); + data.extend_from_slice(&[0x00, 0x00, 0x01, SC_SEQUENCE_HEADER]); + data.extend_from_slice(&[0xAA, 0xBB, 0xCC, 0xDD, 0xEE, 0xFF]); + data.extend_from_slice(&[0x00, 0x00, 0x01, SC_ENTRY_POINT]); + data.extend_from_slice(&[0x11, 0x22, 0x33, 0x44]); + + let pes = make_pes(data, Some(90000)); + let frames = parser.parse(&pes); + + // No coded picture → no frame emitted (parameter bytes must not be + // passed through as a bogus keyframe). + assert!( + frames.is_empty(), + "parameter-set-only PES should not emit a frame" + ); + // But codecPrivate is still captured. + assert!(parser.seq_header.is_some()); + assert!(parser.entry_point.is_some()); + assert!(parser.codec_private().is_some()); + + // A following frame-bearing PES still emits its picture. + let mut data2 = Vec::new(); + data2.extend_from_slice(&[0x00, 0x00, 0x01, SC_FRAME]); + data2.extend_from_slice(&[0x55, 0x66, 0x77]); + let frames2 = parser.parse(&make_pes(data2, Some(180000))); + assert_eq!(frames2.len(), 1); + assert_eq!(&frames2[0].data[0..4], &[0x00, 0x00, 0x01, SC_FRAME]); + } + // --- empty PES --- #[test] diff --git a/src/mux/disc.rs b/src/mux/disc.rs index 40419ce..e783f3b 100644 --- a/src/mux/disc.rs +++ b/src/mux/disc.rs @@ -195,7 +195,8 @@ impl DiscStream { }; pids.push(pid); pid_to_track.push((pid, idx)); - parsers.push((pid, super::codec::parser_for_codec(codec, None))); + let is_dvd_ps = matches!(content_format, crate::disc::ContentFormat::MpegPs); + parsers.push((pid, super::codec::parser_for_codec(codec, None, is_dvd_ps))); } let mut ts_demuxer = None; diff --git a/src/mux/ebml.rs b/src/mux/ebml.rs index be1a093..ce240a7 100644 --- a/src/mux/ebml.rs +++ b/src/mux/ebml.rs @@ -264,6 +264,13 @@ pub fn read_element_header(r: &mut impl Read) -> io::Result<(u32, u64, usize)> { /// Read an unsigned integer value of `len` bytes. pub fn read_uint_val(r: &mut impl Read, len: usize) -> io::Result { + // An EBML unsigned integer is at most 8 bytes. A malformed element + // claiming `len > 8` would index past this stack buffer and panic + // (DoS on untrusted input) — reject it at the source so every caller + // is safe, not just the ones that pre-check. + if len > 8 { + return Err(crate::error::Error::MkvInvalid.into()); + } let mut buf = [0u8; 8]; r.read_exact(&mut buf[..len])?; let mut val = 0u64; @@ -273,23 +280,33 @@ pub fn read_uint_val(r: &mut impl Read, len: usize) -> io::Result { Ok(val) } -/// Read a float value (4 or 8 bytes). +/// Read a float value. EBML floats are exactly 0, 4, or 8 bytes. +/// +/// The previous `else` branch read a fixed 8 bytes for ANY non-4 length, +/// so a malformed element with `len > 8` left `len - 8` unconsumed bytes +/// (mis-read as the next EBML header → desync of the rest of the parent +/// element) and `len < 4` over-read. Consume exactly `len` bytes and +/// reject anything that isn't a valid float width. pub fn read_float_val(r: &mut impl Read, len: usize) -> io::Result { - if len == 4 { - let mut buf = [0u8; 4]; - r.read_exact(&mut buf)?; - Ok(f32::from_be_bytes(buf) as f64) - } else { - let mut buf = [0u8; 8]; - r.read_exact(&mut buf)?; - Ok(f64::from_be_bytes(buf)) + match len { + 0 => Ok(0.0), + 4 => { + let mut buf = [0u8; 4]; + r.read_exact(&mut buf)?; + Ok(f32::from_be_bytes(buf) as f64) + } + 8 => { + let mut buf = [0u8; 8]; + r.read_exact(&mut buf)?; + Ok(f64::from_be_bytes(buf)) + } + _ => Err(crate::error::Error::MkvInvalid.into()), } } /// Read a UTF-8 string value of `len` bytes. pub fn read_string_val(r: &mut impl Read, len: usize) -> io::Result { - let mut buf = vec![0u8; len]; - r.read_exact(&mut buf)?; + let mut buf = read_exact_bounded(r, len)?; // Strip trailing nulls while buf.last() == Some(&0) { buf.pop(); @@ -299,8 +316,22 @@ pub fn read_string_val(r: &mut impl Read, len: usize) -> io::Result { /// Read binary data of `len` bytes. pub fn read_binary_val(r: &mut impl Read, len: usize) -> io::Result> { - let mut buf = vec![0u8; len]; - r.read_exact(&mut buf)?; + read_exact_bounded(r, len) +} + +/// Read exactly `len` bytes WITHOUT trusting `len` to size the allocation. +/// +/// `vec![0u8; len]` on an attacker-controlled EBML size would allocate +/// gigabytes before the read fails. Instead we cap the reader to `len` +/// and grow the buffer as bytes actually arrive: a malformed element that +/// claims a huge length but supplies few bytes allocates only what it +/// delivers, then errors on the short read. +fn read_exact_bounded(r: &mut impl Read, len: usize) -> io::Result> { + let mut buf = Vec::new(); + let got = r.take(len as u64).read_to_end(&mut buf)?; + if got != len { + return Err(io::ErrorKind::UnexpectedEof.into()); + } Ok(buf) } diff --git a/src/mux/mkv.rs b/src/mux/mkv.rs index bd06ff6..e067d7d 100644 --- a/src/mux/mkv.rs +++ b/src/mux/mkv.rs @@ -112,19 +112,22 @@ impl MkvTrack { } pub fn audio(a: &AudioStream) -> Self { - // Codec ID strings must distinguish the DTS family — strict - // players (Plex transcoder, some hardware decoders, some AV - // receivers) reject lossless DTS-HD MA payload when the - // track advertises plain `A_DTS` because it implies the - // bitstream is the 1.5 Mbps "core" only. Fix is just to - // emit the right ID per BD-STN codec field. + // The Matroska codec-ID registry defines `A_DTS` for the entire + // DTS family — the spec text for `A_DTS` explicitly states it + // "Supports DTS, DTS-ES, DTS-96/26, DTS-HD High Resolution Audio + // and DTS-HD Master Audio." Players distinguish core vs HD-HRA vs + // HD-MA by parsing the DTS bitstream extension substreams, not by + // the container codec ID. The previously-emitted `A_DTS/MA` and + // `A_DTS/HR` suffixes are NOT registered codec IDs; strict parsers + // (libmatroska) and some hardware renderers fail to recognise the + // track at all. Emit plain `A_DTS` for every DTS variant — the + // lossless MA / HRA payload bytes are unchanged, only the + // container codec-ID string differs. let codec_id = match a.codec { Codec::Ac3 => "A_AC3", Codec::Ac3Plus => "A_EAC3", Codec::TrueHd => "A_TRUEHD", - Codec::DtsHdMa => "A_DTS/MA", - Codec::DtsHdHr => "A_DTS/HR", - Codec::Dts => "A_DTS", + Codec::DtsHdMa | Codec::DtsHdHr | Codec::Dts => "A_DTS", Codec::Lpcm => "A_PCM/INT/BIG", _ => "A_AC3", }; @@ -220,6 +223,12 @@ pub struct MkvMuxer { /// New cluster every 5 seconds. const CLUSTER_DURATION_MS: i64 = 5000; +/// Maximum block-relative timestamp expressible in the signed 16-bit +/// SimpleBlock/Block field (`i16::MAX` ms). A frame further than this from +/// the open cluster's timestamp forces a new cluster (see `write_frame`) so +/// the `as i16` cast can never wrap. +const MAX_BLOCK_REL_MS: i64 = i16::MAX as i64; + impl MkvMuxer { /// Create a new MKV muxer: writes EBML header, Segment start, Info, Tracks, Chapters. pub fn new( @@ -444,7 +453,7 @@ impl MkvMuxer { let base = *self.base_pts_ms.get_or_insert(raw_ms); let pts_ms = raw_ms - base; - // Cluster boundaries must coincide with a video keyframe so every + // Cluster boundaries normally coincide with a video keyframe so every // Cues entry resolves to a seekable IDR at the cluster start. let is_video_key = keyframe && track_idx == 0; let needs_new_cluster = !self.cluster_open @@ -460,6 +469,17 @@ impl MkvMuxer { track: track_idx + 1, cluster_pos: self.cluster_pos - self.segment_start, }); + } else if (pts_ms - self.cluster_ts_ms) > MAX_BLOCK_REL_MS { + // The block-relative timestamp is a signed 16-bit value, so a + // frame more than i16::MAX ms (~32.767 s) past the current + // cluster's timestamp would silently wrap on the `as i16` cast, + // corrupting A/V sync. The keyframe-driven boundary above only + // fires on a video keyframe — a long audio-only stretch, or a + // very long GOP with no intervening keyframe, can drift past the + // i16 range. Force a fresh cluster here even without a keyframe + // to keep the cast in range. This cluster is not keyframe-aligned + // so it gets no Cues entry (Cues stay IDR-only for seekability). + self.start_cluster(pts_ms)?; } let relative_ts = (pts_ms - self.cluster_ts_ms) as i16; @@ -679,6 +699,42 @@ mod tests { } } + fn audio_stream(codec: Codec) -> AudioStream { + use crate::disc::{AudioChannels, LabelPurpose, SampleRate}; + AudioStream { + pid: 0x1100, + codec, + channels: AudioChannels::Surround51, + language: "eng".into(), + sample_rate: SampleRate::S48, + secondary: false, + purpose: LabelPurpose::Normal, + label: String::new(), + } + } + + #[test] + fn dts_variants_map_to_registered_a_dts_codec_id() { + // The Matroska codec-ID registry defines `A_DTS` for the whole DTS + // family (core, DTS-HD HRA, DTS-HD MA). The `/MA` and `/HR` suffixes + // are not registered and break strict parsers, so every DTS variant + // must emit plain `A_DTS`. + for codec in [Codec::Dts, Codec::DtsHdMa, Codec::DtsHdHr] { + let track = MkvTrack::audio(&audio_stream(codec)); + assert_eq!( + track.codec_id, "A_DTS", + "{codec:?} must map to registered codec ID A_DTS, got {}", + track.codec_id + ); + } + // Sanity: the non-DTS variants keep their distinct IDs. + assert_eq!(MkvTrack::audio(&audio_stream(Codec::Ac3)).codec_id, "A_AC3"); + assert_eq!( + MkvTrack::audio(&audio_stream(Codec::TrueHd)).codec_id, + "A_TRUEHD" + ); + } + #[test] fn dolby_vision_config_profile7() { // dvcC for disc Profile 7 dual-layer: version 1.0, profile 7, all of @@ -1377,6 +1433,85 @@ mod tests { ); } + /// Collect every (cluster_ts_ms, block_relative_ts_i16, absolute_ms) for + /// all SimpleBlocks across all clusters, so a test can assert that the + /// reconstructed absolute timestamp (cluster_ts + relative_ts) is correct + /// and that no relative_ts ever wrapped the i16 range. + fn all_block_timestamps(data: &[u8]) -> Vec<(i64, i16, i64)> { + let mut out = Vec::new(); + for (body_start, body_size, cluster_ts) in find_clusters(data) { + let body = &data[body_start..body_start + body_size as usize]; + let mut cursor = Cursor::new(body); + // Skip CLUSTER_TIMESTAMP. + let (tid, tsize, _) = ebml::read_element_header(&mut cursor).unwrap(); + assert_eq!(tid, ebml::CLUSTER_TIMESTAMP); + cursor.seek(io::SeekFrom::Current(tsize as i64)).unwrap(); + while (cursor.position() as usize) < body.len() { + let (id, sz, _) = ebml::read_element_header(&mut cursor).unwrap(); + if id == ebml::SIMPLE_BLOCK { + let bstart = cursor.position() as usize; + let b0 = body[bstart]; + let vint_len = if b0 & 0x80 != 0 { 1 } else { 2 }; + let ts_pos = bstart + vint_len; + let rel = i16::from_be_bytes([body[ts_pos], body[ts_pos + 1]]); + out.push((cluster_ts as i64, rel, cluster_ts as i64 + rel as i64)); + } + cursor.seek(io::SeekFrom::Current(sz as i64)).unwrap(); + } + } + out + } + + #[test] + fn long_audio_gap_forces_cluster_no_i16_overflow() { + // Regression for the `(pts_ms - cluster_ts_ms) as i16` truncation: + // a single video keyframe at t=0 opens one cluster, then a long + // audio-only stretch (no further video keyframe) drifts well past + // i16::MAX ms (~32.767 s). Without the overflow guard the audio + // blocks past 32.767 s would write a wrapped (negative) relative + // timestamp into the SimpleBlock. With the guard a fresh cluster is + // forced so every relative_ts stays in range and reconstructs to the + // true absolute timestamp. + let tracks = [make_video_track(), make_audio_track()]; + let mut frames: Vec<(usize, i64, bool, Vec)> = Vec::new(); + // One video keyframe at t=0 (opens the first cluster). + frames.push((0, 0, true, vec![0xAB; 16])); + // Audio frames every 100 ms out to 60 s — past the 32.767 s i16 limit + // and past two i16 spans, with NO further video keyframe. + let mut t_ms = 0i64; + while t_ms <= 60_000 { + frames.push((1, t_ms * 1_000_000, true, vec![0xCD; 16])); + t_ms += 100; + } + + let (data, _) = mux_to_bytes(&tracks, &[], &frames); + + let blocks = all_block_timestamps(&data); + assert!(!blocks.is_empty()); + // Every block's relative timestamp must be within i16 range (it is by + // type), AND must reconstruct to a non-negative, monotonic-ish + // absolute timestamp matching the source — i.e. no silent wrap. + for (cluster_ts, rel, abs) in &blocks { + assert!( + *rel as i64 >= 0 && (*rel as i64) <= MAX_BLOCK_REL_MS, + "block relative_ts {rel} out of [0, i16::MAX] range \ + (cluster_ts={cluster_ts}, abs={abs}) — i16 overflow" + ); + } + // The latest audio frame is at 60_000 ms; its reconstructed absolute + // timestamp must equal that, proving no truncation occurred. + let max_abs = blocks.iter().map(|(_, _, abs)| *abs).max().unwrap(); + assert_eq!(max_abs, 60_000, "last block must reconstruct to 60_000 ms"); + // The overflow guard must have opened more than one cluster (the + // single keyframe alone would otherwise yield exactly one). + let clusters = find_clusters(&data); + assert!( + clusters.len() >= 2, + "expected the i16 guard to force extra clusters, got {}", + clusters.len() + ); + } + #[test] fn pre_first_keyframe_frames_dropped() { let tracks = [make_video_track()]; diff --git a/src/mux/mkvstream.rs b/src/mux/mkvstream.rs index 06fc5d0..98c2337 100644 --- a/src/mux/mkvstream.rs +++ b/src/mux/mkvstream.rs @@ -14,6 +14,49 @@ fn skip_bytes(r: &mut impl Read, n: u64) -> io::Result<()> { Ok(()) } +// ── Sanity caps for untrusted EBML element sizes ────────────── +// +// Sizes come straight from the EBML stream (file or network) and are +// otherwise cast to `usize` and used to allocate/read. An adversarial +// or corrupt container can claim a multi-GB element and trigger an OOM +// allocation, or claim an integer element wider than 8 bytes and panic +// the fixed 8-byte reader. Every untrusted size is validated against +// one of these caps before allocation. + +/// Largest accepted SIMPLE_BLOCK payload. A block is a small vint track +/// header + 2-byte rel-ts + 1-byte flags + one frame of elementary data. +/// UHD HEVC keyframes run a few MB; 64 MiB is generously above any real +/// single-frame block while still bounding a hostile allocation. +const MAX_BLOCK_SIZE: u64 = 64 * 1024 * 1024; +/// Largest accepted CODEC_PRIVATE payload. hvcC/avcC/setup blobs are a +/// few KB in practice; 16 MiB is far above any legitimate value. +const MAX_CODEC_PRIVATE: u64 = 16 * 1024 * 1024; +/// Largest accepted string element (TITLE/CODEC_ID/LANGUAGE/TRACK_NAME). +const MAX_STRING_LEN: u64 = 64 * 1024; +/// EBML unsigned-int elements are at most 8 bytes wide. +const MAX_UINT_LEN: u64 = 8; + +/// Reject an untrusted element size that exceeds `cap` before it is used +/// to allocate or read. Returns the size as `usize` when within bounds. +fn checked_size(size: u64, cap: u64) -> io::Result { + if size > cap { + return Err(crate::error::Error::MkvInvalid.into()); + } + Ok(size as usize) +} + +/// Read a bounded unsigned int. Guards against `size > 8` (which would +/// otherwise index out of the fixed 8-byte buffer in `read_uint_val`) +/// before delegating. +fn read_uint_bounded(r: &mut impl Read, size: u64) -> io::Result { + ebml::read_uint_val(r, checked_size(size, MAX_UINT_LEN)?) +} + +/// Read a bounded UTF-8 string element. +fn read_string_bounded(r: &mut impl Read, size: u64) -> io::Result { + ebml::read_string_val(r, checked_size(size, MAX_STRING_LEN)?) +} + use crate::disc::*; use std::io::{self, Read}; @@ -108,11 +151,12 @@ impl crate::pes::Stream for MkvStream { match id { ebml::CLUSTER => continue, ebml::CLUSTER_TIMESTAMP => { - rs.cluster_ts_ms = ebml::read_uint_val(&mut rs.reader, size as usize)? as i64; + rs.cluster_ts_ms = read_uint_bounded(&mut rs.reader, size)? as i64; continue; } ebml::SIMPLE_BLOCK => { - let block = ebml::read_binary_val(&mut rs.reader, size as usize)?; + let block = + ebml::read_binary_val(&mut rs.reader, checked_size(size, MAX_BLOCK_SIZE)?)?; if block.len() < 4 { continue; } @@ -198,7 +242,10 @@ impl crate::pes::Stream for MkvStream { /// Returns (DiscTitle, codec_privates: Vec<(track_number, codec_private_bytes)>) fn parse_mkv_header(r: &mut impl Read) -> MkvHeaderResult { let mut title = String::new(); - let mut duration_ms = 0.0f64; + // EBML `DURATION` is a float expressed in TimestampScale ticks, not + // milliseconds (Matroska spec). Named accordingly; converted to + // seconds below as ticks * ts_scale_ns / 1e9. + let mut duration_ticks = 0.0f64; let mut ts_scale: u64 = 1_000_000; let mut streams: Vec = Vec::new(); let mut codec_privates: Vec<(u16, Vec)> = Vec::new(); @@ -235,9 +282,9 @@ fn parse_mkv_header(r: &mut impl Read) -> MkvHeaderResult { let (cid, cs, hlen) = ebml::read_element_header(r)?; remaining = remaining.saturating_sub(hlen as u64 + cs); match cid { - ebml::TIMESTAMP_SCALE => ts_scale = ebml::read_uint_val(r, cs as usize)?, - ebml::DURATION => duration_ms = ebml::read_float_val(r, cs as usize)?, - ebml::TITLE => title = ebml::read_string_val(r, cs as usize)?, + ebml::TIMESTAMP_SCALE => ts_scale = read_uint_bounded(r, cs)?, + ebml::DURATION => duration_ticks = ebml::read_float_val(r, cs as usize)?, + ebml::TITLE => title = read_string_bounded(r, cs)?, _ => { skip_bytes(r, cs)?; } @@ -274,7 +321,7 @@ fn parse_mkv_header(r: &mut impl Read) -> MkvHeaderResult { let disc_title = DiscTitle { playlist: title, - duration_secs: duration_ms * (ts_scale as f64) / 1_000_000_000.0, + duration_secs: duration_ticks * (ts_scale as f64) / 1_000_000_000.0, streams, ..DiscTitle::empty() }; @@ -296,20 +343,25 @@ fn parse_track( let (cid, cs, hlen) = ebml::read_element_header(r)?; remaining = remaining.saturating_sub(hlen as u64 + cs); match cid { - ebml::TRACK_NUMBER => tnum = ebml::read_uint_val(r, cs as usize)? as u16, - ebml::TRACK_TYPE => ttype = ebml::read_uint_val(r, cs as usize)?, - ebml::CODEC_ID => codec_id = ebml::read_string_val(r, cs as usize)?, - ebml::CODEC_PRIVATE => codec_priv = Some(ebml::read_binary_val(r, cs as usize)?), - ebml::LANGUAGE => lang = ebml::read_string_val(r, cs as usize)?, - ebml::TRACK_NAME => name = ebml::read_string_val(r, cs as usize)?, - ebml::FLAG_FORCED => forced = ebml::read_uint_val(r, cs as usize)? != 0, + ebml::TRACK_NUMBER => tnum = read_uint_bounded(r, cs)? as u16, + ebml::TRACK_TYPE => ttype = read_uint_bounded(r, cs)?, + ebml::CODEC_ID => codec_id = read_string_bounded(r, cs)?, + ebml::CODEC_PRIVATE => { + codec_priv = Some(ebml::read_binary_val( + r, + checked_size(cs, MAX_CODEC_PRIVATE)?, + )?) + } + ebml::LANGUAGE => lang = read_string_bounded(r, cs)?, + ebml::TRACK_NAME => name = read_string_bounded(r, cs)?, + ebml::FLAG_FORCED => forced = read_uint_bounded(r, cs)? != 0, ebml::VIDEO => { let mut vrem = cs; while vrem > 0 { let (vid, vs, vhlen) = ebml::read_element_header(r)?; vrem = vrem.saturating_sub(vhlen as u64 + vs); if vid == ebml::PIXEL_HEIGHT { - ph = ebml::read_uint_val(r, vs as usize)? as u32; + ph = read_uint_bounded(r, vs)? as u32; } else { skip_bytes(r, vs)?; } @@ -322,7 +374,7 @@ fn parse_track( arem = arem.saturating_sub(ahlen as u64 + as_); match aid { ebml::SAMPLING_FREQUENCY => sr = ebml::read_float_val(r, as_ as usize)?, - ebml::CHANNELS => ch = ebml::read_uint_val(r, as_ as usize)? as u8, + ebml::CHANNELS => ch = read_uint_bounded(r, as_)? as u8, _ => { skip_bytes(r, as_)?; } @@ -428,3 +480,138 @@ fn block_vint(d: &[u8]) -> (u64, usize) { } (0, 1) // Unsupported 5+ byte VINT — treat as track 0 } + +#[cfg(test)] +mod tests { + use super::*; + use crate::pes::Stream as _; + use std::io::Cursor; + + // `From for io::Error` encodes the numeric code into the + // Display string as "E{code}: ...". Check the prefix. + fn is_mkv_invalid(e: &io::Error) -> bool { + e.kind() == io::ErrorKind::InvalidData + && e.to_string() + .starts_with(&format!("E{}", crate::error::E_MKV_INVALID)) + } + + #[test] + fn checked_size_rejects_over_cap() { + // Within cap → Ok with usize value. + assert_eq!(checked_size(100, 256).unwrap(), 100); + assert_eq!(checked_size(256, 256).unwrap(), 256); + // Over cap → MkvInvalid, never a giant allocation. + let e = checked_size(257, 256).unwrap_err(); + assert!(is_mkv_invalid(&e)); + // A hostile multi-GB block size is rejected as MkvInvalid. + let e = checked_size(4 * 1024 * 1024 * 1024, MAX_BLOCK_SIZE).unwrap_err(); + assert!(is_mkv_invalid(&e)); + } + + #[test] + fn read_uint_bounded_rejects_oversized_int() { + // size > 8 would index out of the fixed 8-byte buffer in + // read_uint_val (panic / OOB). The guard turns it into a clean + // MkvInvalid error instead. + let mut data = Cursor::new(vec![0u8; 16]); + let e = read_uint_bounded(&mut data, 9).unwrap_err(); + assert!(is_mkv_invalid(&e)); + } + + #[test] + fn read_uint_bounded_accepts_valid_width() { + // 8 bytes is the max legal EBML uint width and must still work. + let mut data = Cursor::new(vec![0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x01, 0x02]); + assert_eq!(read_uint_bounded(&mut data, 8).unwrap(), 0x0102); + } + + #[test] + fn read_string_bounded_rejects_huge_string() { + // Claimed string length far above the cap must not allocate. + let mut data = Cursor::new(vec![0u8; 16]); + let e = read_string_bounded(&mut data, MAX_STRING_LEN + 1).unwrap_err(); + assert!(is_mkv_invalid(&e)); + } + + /// Build a minimal MKV (EBML header + Segment + Info + Tracks) so the + /// reader is positioned in the cluster body, then append the given + /// cluster bytes. Returns the full byte stream. + fn minimal_mkv_with_cluster(cluster_body: &[u8]) -> Vec { + let mut out = Vec::new(); + // EBML header (empty body). + ebml::write_id(&mut out, ebml::EBML).unwrap(); + ebml::write_size(&mut out, 0).unwrap(); + // Segment (unknown size so the reader streams children). + ebml::write_id(&mut out, ebml::SEGMENT).unwrap(); + ebml::write_unknown_size(&mut out).unwrap(); + // Empty Info. + ebml::write_id(&mut out, ebml::INFO).unwrap(); + ebml::write_size(&mut out, 0).unwrap(); + // Empty Tracks. + ebml::write_id(&mut out, ebml::TRACKS).unwrap(); + ebml::write_size(&mut out, 0).unwrap(); + out.extend_from_slice(cluster_body); + out + } + + #[test] + fn simple_block_oversized_size_is_rejected() { + // Cluster containing a SIMPLE_BLOCK that claims a 2 GiB payload. + // The reader must reject it (MkvInvalid) rather than attempt a + // multi-GB allocation. Header parse stops at CLUSTER, so the + // SIMPLE_BLOCK is hit on the first read(). + let mut cluster = Vec::new(); + ebml::write_id(&mut cluster, ebml::CLUSTER).unwrap(); + ebml::write_unknown_size(&mut cluster).unwrap(); + ebml::write_id(&mut cluster, ebml::SIMPLE_BLOCK).unwrap(); + ebml::write_size(&mut cluster, 2 * 1024 * 1024 * 1024).unwrap(); + // No payload follows — but we must fail on the size check, before + // any read of the body. + let bytes = minimal_mkv_with_cluster(&cluster); + + let mut stream = MkvStream::open(Cursor::new(bytes)).unwrap(); + let e = stream.read().unwrap_err(); + assert!(is_mkv_invalid(&e)); + } + + #[test] + fn well_formed_simple_block_round_trips() { + // A small, well-formed SIMPLE_BLOCK must still parse into a frame. + // We need at least one stream so the track index is in range, so + // give Tracks one video TRACK_ENTRY (track number 1). + let mut out = Vec::new(); + ebml::write_id(&mut out, ebml::EBML).unwrap(); + ebml::write_size(&mut out, 0).unwrap(); + ebml::write_id(&mut out, ebml::SEGMENT).unwrap(); + ebml::write_unknown_size(&mut out).unwrap(); + ebml::write_id(&mut out, ebml::INFO).unwrap(); + ebml::write_size(&mut out, 0).unwrap(); + + // Tracks → one TRACK_ENTRY (track number 1, type 1 = video). + let mut entry = Vec::new(); + ebml::write_uint(&mut entry, ebml::TRACK_NUMBER, 1).unwrap(); + ebml::write_uint(&mut entry, ebml::TRACK_TYPE, 1).unwrap(); + let mut track_entry = Vec::new(); + ebml::write_id(&mut track_entry, ebml::TRACK_ENTRY).unwrap(); + ebml::write_size(&mut track_entry, entry.len() as u64).unwrap(); + track_entry.extend_from_slice(&entry); + ebml::write_id(&mut out, ebml::TRACKS).unwrap(); + ebml::write_size(&mut out, track_entry.len() as u64).unwrap(); + out.extend_from_slice(&track_entry); + + // Cluster with a SIMPLE_BLOCK: track vint=0x81 (track 1), + // rel_ts=0x0000, flags=0x80 (keyframe), then 4 bytes of data. + ebml::write_id(&mut out, ebml::CLUSTER).unwrap(); + ebml::write_unknown_size(&mut out).unwrap(); + let block = [0x81u8, 0x00, 0x00, 0x80, 0xAA, 0xBB, 0xCC, 0xDD]; + ebml::write_id(&mut out, ebml::SIMPLE_BLOCK).unwrap(); + ebml::write_size(&mut out, block.len() as u64).unwrap(); + out.extend_from_slice(&block); + + let mut stream = MkvStream::open(Cursor::new(out)).unwrap(); + let frame = stream.read().unwrap().expect("expected a frame"); + assert_eq!(frame.track, 0); + assert!(frame.keyframe); + assert_eq!(frame.data, vec![0xAA, 0xBB, 0xCC, 0xDD]); + } +} diff --git a/src/mux/resolve.rs b/src/mux/resolve.rs index 6633f8d..1979916 100644 --- a/src/mux/resolve.rs +++ b/src/mux/resolve.rs @@ -384,7 +384,8 @@ fn build_demux_state(title: &DiscTitle, format: ContentFormat) -> DemuxState { }; pids.push(pid); pid_to_track.push((pid, idx)); - parsers.push((pid, super::codec::parser_for_codec(codec, None))); + let is_dvd_ps = matches!(format, ContentFormat::MpegPs); + parsers.push((pid, super::codec::parser_for_codec(codec, None, is_dvd_ps))); } let (ts, ps) = match format { ContentFormat::MpegPs => (None, Some(super::ps::PsDemuxer::new())), diff --git a/src/mux/ts.rs b/src/mux/ts.rs index 3f4b864..c8494d2 100644 --- a/src/mux/ts.rs +++ b/src/mux/ts.rs @@ -35,6 +35,15 @@ struct PesAssembler { pts: Option, dts: Option, active: bool, + /// PES-header bytes still to be skipped on the next continuation + /// packet(s). A PES header (9 + PES_header_data_length, up to 264 + /// bytes) can exceed a single 184-byte TS payload, spilling into the + /// following continuation packet. Those spillover bytes are NOT + /// elementary-stream data and must be skipped, or the PES start code + /// (`00 00 01 …`) and timestamp bytes get injected into the ES — for + /// HEVC/H264 that reads as a spurious start code / corrupt slice + /// payload. Tracks how many header bytes remain across packets. + header_remaining: usize, } /// Initial capacity for a fresh PES buffer. Sized to cover the @@ -54,6 +63,7 @@ impl PesAssembler { pts: None, dts: None, active: false, + header_remaining: 0, } } @@ -235,12 +245,35 @@ impl TsDemuxer { let payload = &ts[payload_start..]; if pusi { - let (pts, dts, pes_data_start) = parse_pes_header(payload); + // `header_len` is the FULL (uncapped) PES-header length: + // 0 = malformed (payload is not a PES start), else 6/9+N. + let (pts, dts, header_len) = parse_pes_header(payload); if let Some(prev) = asm.start(pts, dts) { completed.push(prev); } - if pes_data_start < payload.len() { - asm.push(&payload[pes_data_start..]); + if header_len == 0 { + // PUSI packet whose payload is not a valid PES start. Do + // NOT push it — those bytes are not elementary-stream data + // and would inject a spurious start code / garbage. + asm.header_remaining = 0; + } else if header_len <= payload.len() { + // Header fits in this packet (the common case). + asm.header_remaining = 0; + if header_len < payload.len() { + asm.push(&payload[header_len..]); + } + } else { + // Header spills past this packet — skip the remainder on + // the following continuation packet(s). + asm.header_remaining = header_len - payload.len(); + } + } else if asm.header_remaining > 0 { + // Continuation packet still inside a PES header that spanned + // the boundary — consume header bytes before any ES data. + let skip = asm.header_remaining.min(payload.len()); + asm.header_remaining -= skip; + if skip < payload.len() { + asm.push(&payload[skip..]); } } else { asm.push(payload); @@ -260,7 +293,14 @@ impl TsDemuxer { } /// Parse a PES packet header, extracting PTS and DTS. -/// Returns (pts, dts, offset_to_elementary_stream_data). +/// +/// Returns `(pts, dts, header_len)` where `header_len` is the FULL, +/// UNCAPPED PES-header length in bytes (`9 + PES_header_data_length`, or +/// 6 for stream IDs without the standard extension). `0` signals the +/// payload is not a valid PES start (malformed / too short). The caller +/// must treat `header_len` as bytes-to-skip and carry any remainder past +/// this packet's payload into the next continuation packet — the header +/// can exceed one TS payload, and the spillover is header, not ES data. fn parse_pes_header(data: &[u8]) -> (Option, Option, usize) { // PES packet: 00 00 01 [stream_id] [length:2] [flags...] if data.len() < 9 || data[0] != 0x00 || data[1] != 0x00 || data[2] != 0x01 { @@ -288,7 +328,10 @@ fn parse_pes_header(data: &[u8]) -> (Option, Option, usize) { let pts_dts_flags = (data[7] >> 6) & 0x03; let header_data_len = data[8] as usize; - let data_start = (9 + header_data_len).min(data.len()); + // Full, uncapped header length. PTS/DTS (if present) live in the + // first ~19 bytes, always within this packet's payload, so they parse + // here; only the *skip* length may extend into the next packet. + let header_len = 9 + header_data_len; let mut pts = None; let mut dts = None; @@ -300,7 +343,7 @@ fn parse_pes_header(data: &[u8]) -> (Option, Option, usize) { dts = parse_timestamp(&data[14..19]); } - (pts, dts, data_start) + (pts, dts, header_len) } /// Parse a 5-byte PTS/DTS timestamp (33 bits in 90kHz). @@ -404,10 +447,19 @@ pub fn scan_streams(data: &[u8]) -> Option> { let section_len = (((data[pmt_start + 1] & 0x0F) as usize) << 8) | data[pmt_start + 2] as usize; + // section_length counts the bytes after this field, including the + // trailing 4-byte CRC; `< 4` would underflow `end` below. Guard it + // exactly like the PAT parser above. + if section_len < 4 { + offset += BD_TS_PACKET_SIZE; + continue; + } let prog_info_len = (((data[pmt_start + 10] & 0x0F) as usize) << 8) | data[pmt_start + 11] as usize; let mut pos = pmt_start + 12 + prog_info_len; - let end = pmt_start + 3 + section_len - 4; + // Clamp the section end to the buffer; a malformed section_len or + // prog_info_len must never drive reads past `data`. + let end = (pmt_start + 3 + section_len - 4).min(data.len()); while pos + 5 <= data.len() && pos < end { let stream_type = data[pos]; diff --git a/src/sector/prefetched.rs b/src/sector/prefetched.rs index 5bceb24..ea0c1d0 100644 --- a/src/sector/prefetched.rs +++ b/src/sector/prefetched.rs @@ -157,14 +157,15 @@ impl PrefetchedSectorSource { Ok(b) => b, Err(_) => return, // consumer dropped both channels }; - if buf.len() < bytes { - buf.resize(bytes, 0); - } else { - // Re-expose the full extent; previous truncate - // shrank the visible len without freeing pages. - // SAFETY: capacity is at least `bytes` after - // construction with `vec![0u8; batch_bytes]`. + if bytes <= buf.capacity() { + // Re-expose `bytes` without zero-filling pages that + // `read_sectors` is about to overwrite. The capacity + // guard makes the `set_len` provably sound even if a + // recycled buffer ever comes back smaller than the + // `vec![0u8; batch_bytes]` it was born with. unsafe { buf.set_len(bytes) }; + } else { + buf.resize(bytes, 0); } let lba = extent.start_lba + offset; match reader.read_sectors(lba, sectors, &mut buf[..bytes], false) {