From a32373ff40a4a4419f10740a5110393f889de003 Mon Sep 17 00:00:00 2001 From: Matthew Jackson <1085847+MattJackson@users.noreply.github.com> Date: Wed, 29 Jul 2026 20:47:31 -0700 Subject: [PATCH] Fix fifteen defects across perf, resource, panics and key hygiene MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit All 21 findings held up under verification; 15 fixed here, 6 deferred to files another agent held this round, 0 rejected. **A defect in my own round-2 probe fix.** CHUNK_SECTORS was 1024, and 1024 % 3 == 1 — verified — so every chunk after the first was misaligned against the 6144-byte AACS aligned unit and would be REJECTED by DecryptingSectorSource's alignment gate. On an encrypted disc the forced-subtitle probe I added last round would have read almost nothing past its first chunk. Now 1023 sectors (341 aligned units) with a const assertion that fails the build if it stops dividing, plus set_unit_base per extent so the source's gate is anchored where the extent actually starts. **The same probe skipped sectors on a short read**, advancing by the REQUESTED count rather than the bytes actually returned, so a partial read silently left a gap in the middle of the evidence. It now advances by n/SECTOR_BYTES and clamps n to the buffer. **Its cache key omitted the PGS PID set**, so a playlist declaring an extra subtitle PID got another playlist's verdict for a track that had never been probed. And the key was the whole extent list, so partial clip sharing missed entirely. Both fixed by keying (start_lba, sector_count, pid) — and per-extent keying was shown SOUND rather than assumed: ForcedTracker is two monotone booleans, so per-extent evidence composes by field-wise OR, order- and grouping-independently. Making that honest required per-extent demux state, so an extent's evidence comes only from its own bytes, and memoising only extents whose read reached a designed stop. **A reachable panic in the timeline.** mkvstream::parse_block accepts a TimestampScale up to i64::MAX, so a video frame can set high_ns = i64::MAX and the next passive frame panicked adding the backstep. In release it wrapped negative instead, firing the straggler clamp for essentially every passive frame — audio and subtitles rewritten onto the wrong point of the output timeline. All four sites saturate. **A public constructor divided by zero**: PrefetchedSectorSource::new_with_events with unit_align == 0. Now InvalidInput, matching its batch_sectors sibling. **Two Debug impls printed key material.** DiscInputs (volume_id, mkb, unit_key_ro, samples) and UnitKeyFile both derived Debug. Nothing logs them today — fixed as prevention, because the next tracing::debug! someone adds is the leak. A doc claim that DiscInputs "contains no secrets" was false and is corrected. **An env-var multiply could overflow** in file_sector_source; now bounded at 64 GiB like its writeback sibling, with the parse split out so the bound is testable without touching process env. **The mp4 demuxer allowed one sample per file byte** — ~64x RAM amplification. Now file_len/16, since only vide/soun tracks are indexed and the shortest legal AC-3 frame is 128 bytes. **Two pipeline concurrency defects**: a consumer apply() error was invisible to the producer, and abandon/finalise had a TOCTOU where a caller could report an unfinalised output. Both fixed with compare-exchange state rather than a bool. **Two per-frame copies removed**, both MEASURED rather than reasoned: the AU assembler now hands its allocation to the frame (same pointer, unchanged capacity, proven by asserting the pointer) and tsmux reuses one Annex-B buffer across frames. Both keep capacity deliberately — a naive split_off would have cost more than it saved. **A comment pointed at the wrong file** for a mirrored constant; the mirror is now compiler-enforced with a const assertion converting 90 kHz ticks to ns, so drift fails the build. Deferred to another agent's files, all confirmed: detect_rate's fractional-twin snap, the mp4 reserve's u32 truncation, round_up_grain's overflow, the quadratic base-key gap fill, and MkvStream's frame cap counting frames rather than bytes. Every fix verified red by reverting it. Also noted for later: DecodeSampleSet still derives Debug over multi-MB of on-disc ciphertext. --- src/aacs/inf.rs | 53 ++- src/disc/pgs_forced_probe.rs | 578 ++++++++++++++++++++++++++----- src/io/file_sector_source/mod.rs | 56 ++- src/io/pipeline.rs | 332 ++++++++++++++++-- src/keysource.rs | 68 +++- src/mux/au_assembly.rs | 85 ++++- src/mux/codec/hevc.rs | 15 +- src/mux/mp4/read.rs | 85 ++++- src/mux/timeline.rs | 61 +++- src/mux/tsmux.rs | 104 +++++- src/sector/prefetched.rs | 34 ++ 11 files changed, 1303 insertions(+), 168 deletions(-) diff --git a/src/aacs/inf.rs b/src/aacs/inf.rs index 8136d85..df5e91c 100644 --- a/src/aacs/inf.rs +++ b/src/aacs/inf.rs @@ -5,7 +5,6 @@ use super::mkb::*; /// Parsed Unit_Key_RO.inf file. -#[derive(Debug)] pub struct UnitKeyFile { /// Disc hash (SHA1 of the entire file) — used as KEYDB lookup key pub disc_hash: [u8; 20], @@ -23,6 +22,29 @@ pub struct UnitKeyFile { pub title_cps_unit: Vec, } +/// Redacting `Debug`, per the policy `aacs::types` documents: this struct holds +/// the disc's ENCRYPTED CPS unit keys — exactly the material a keydb entry stores +/// — plus the disc hash they are looked up by. A derived `Debug` printed every key +/// byte verbatim, so any `{:?}` (a downstream crate, an `assert_eq!` failure +/// message, a future `tracing::debug!` in this module) leaked them. Only +/// non-secret shape is printed. Guarded by `unit_key_file_debug_is_redacted`. +impl std::fmt::Debug for UnitKeyFile { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("UnitKeyFile") + // The disc hash is the public keydb lookup key, printed as hex the + // same way `DiscEntry` prints its own — never as raw bytes. + .field("disc_hash", &disc_hash_hex(&self.disc_hash)) + .field("app_type", &self.app_type) + .field("num_bdmv_dir", &self.num_bdmv_dir) + .field("use_skb_mkb", &self.use_skb_mkb) + .field("version", &self.version) + .field("encrypted_keys", &"") + .field("encrypted_keys_len", &self.encrypted_keys.len()) + .field("title_cps_unit", &self.title_cps_unit) + .finish() + } +} + /// Compute disc hash (SHA1 of Unit_Key_RO.inf content). pub fn disc_hash(data: &[u8]) -> [u8; 20] { use sha1::{Digest, Sha1}; @@ -497,4 +519,33 @@ mod vtkf_tests { // Same as applying the shared unwrap directly to the stored enc key. assert_eq!(derived, super::super::derive::decrypt_unit_key(&vuk, &enc)); } + + /// `UnitKeyFile` holds the disc's ENCRYPTED CPS unit keys. A derived `Debug` + /// printed every byte; the hand-written impl must not. Sentinel key byte + /// 0xD5 = decimal 213 (a derived `Debug` renders `[u8; 16]` in decimal), the + /// same probe `aacs::types::redaction_tests` uses. Mutation guard: putting + /// `#[derive(Debug)]` back fails this. + #[test] + fn unit_key_file_debug_is_redacted() { + let f = UnitKeyFile { + disc_hash: [0xD5; 20], + app_type: 1, + num_bdmv_dir: 1, + use_skb_mkb: false, + version: AacsVersion::V20, + encrypted_keys: vec![(0, [0xD5; 16]), (1, [0xD5; 16])], + title_cps_unit: vec![0, 1], + }; + let dbg = format!("{f:?}"); + assert!( + !dbg.contains("213"), + "UnitKeyFile Debug leaked key bytes (decimal 213): {dbg}" + ); + assert!( + dbg.contains("redacted"), + "UnitKeyFile Debug missing redaction marker: {dbg}" + ); + // Non-secret shape is still useful for diagnostics. + assert!(dbg.contains("encrypted_keys_len: 2"), "{dbg}"); + } } diff --git a/src/disc/pgs_forced_probe.rs b/src/disc/pgs_forced_probe.rs index 5bf530e..ebb854c 100644 --- a/src/disc/pgs_forced_probe.rs +++ b/src/disc/pgs_forced_probe.rs @@ -29,8 +29,22 @@ use crate::sector::SectorSource; use std::collections::HashMap; const SECTOR_BYTES: usize = 2048; -/// Read the clip in 2 MiB chunks. -const CHUNK_SECTORS: u16 = 1024; +/// Read the clip in ~2 MiB chunks. +/// +/// A whole number of AACS aligned units (3 sectors / 6144 B), because with a +/// decrypting source — the case this module's doc promises — every read must +/// begin on a unit boundary measured from the extent base or +/// `DecryptingSectorSource` rejects it outright with `DecryptFailed`. At 1024 +/// (`1024 % 3 == 1`) every chunk after the first drifted off the boundary, so +/// content-based forced detection was unreachable past the first chunk of an +/// AACS disc. 1023 = 341 units. +const CHUNK_SECTORS: u16 = 1023; + +// The alignment requirement above is enforced, not just described. +const _: () = assert!( + CHUNK_SECTORS as u32 % crate::aacs::content::ALIGNED_UNIT_SECTORS == 0, + "probe chunks must be a whole number of AACS aligned units" +); /// Hard ceiling on sectors read per probe call (256 MiB). /// @@ -43,12 +57,46 @@ const CHUNK_SECTORS: u16 = 1024; /// non-forced set before accepting the forced verdict. const PROBE_BUDGET_SECTORS: u32 = 131_072; -/// Memoises probe results across titles. Keyed by the title's exact extent -/// list, so a hit returns a result computed from byte-identical input — many -/// playlists on one disc reference the same clips (main feature, play-all, -/// seamless-branch variants), and without this the same physical extents are -/// re-read from the drive once per playlist. -pub(crate) type ForcedProbeCache = HashMap, HashMap>; +/// What one probed extent showed about one PGS track — the two monotone facts a +/// [`ForcedTracker`] accumulates, and nothing else. +/// +/// Keeping the EVIDENCE (rather than a composed forced/not-forced verdict) is +/// what makes per-extent memoisation sound: both fields only ever go from +/// `false` to `true` as more data is seen, so a title's verdict is the +/// field-wise OR over its extents, in any order, with no dependence on how the +/// extents were grouped into playlists. +#[derive(Clone, Copy, Default, PartialEq, Eq, Debug)] +pub(crate) struct TrackEvidence { + /// A PGS display set was actually seen for this track in this extent. + observed: bool, + /// At least one of those display sets was NOT forced. + non_forced: bool, +} + +impl TrackEvidence { + fn merge(&mut self, other: Self) { + self.observed |= other.observed; + self.non_forced |= other.non_forced; + } +} + +/// Memoises probe results across titles, keyed PER PHYSICAL EXTENT and per PGS +/// track — `(start_lba, sector_count, pid)`. +/// +/// Many playlists on one disc reference the same clips (main feature, play-all, +/// seamless-branch variants) but rarely with byte-identical extent LISTS: 00800 +/// = [A, B], 00801 = [A], 00802 = [B] are three different lists over two clips. +/// Keying on the whole list de-duplicated only exactly-identical playlists and +/// re-read every shared clip once per list — up to `PROBE_BUDGET_SECTORS` +/// (256 MiB) of optical-drive time each. Per-extent keying reads each physical +/// extent at most once per disc, and per-track keying means a playlist that +/// declares MORE PGS tracks over the same extents still probes the extra ones +/// instead of silently taking a verdict map that has no entry for them. +/// +/// Only extents whose read reached a DESIGNED stop are memoised (see +/// `probe_and_set_forced`), so one cancellation or read fault is never frozen in +/// as an extent's answer. +pub(crate) type ForcedProbeCache = HashMap<(u32, u32, u16), TrackEvidence>; /// Why the read loop stopped — which decides whether the observations it /// accumulated may be applied as an authoritative verdict. @@ -125,57 +173,109 @@ pub(crate) fn probe_and_set_forced( return; } - // Same extents → same verdicts. Serve from cache rather than re-reading. - let key: Vec<(u32, u32)> = title - .extents + // Same extent → same evidence. Take from the cache what is already known and + // read only the extents that are not (for every declared track). + let mut evidence: HashMap = pg_pids .iter() - .map(|e| (e.start_lba, e.sector_count)) + .map(|&p| (p, TrackEvidence::default())) .collect(); - if let Some(hit) = cache.get(&key) { - apply_verdicts(title, hit); + let mut todo: Vec = Vec::new(); + for ext in &title.extents { + let hits: Option> = pg_pids + .iter() + .map(|&p| cache.get(&(ext.start_lba, ext.sector_count, p)).copied()) + .collect(); + match hits { + Some(known) => { + for (&pid, ev) in pg_pids.iter().zip(known) { + if let Some(slot) = evidence.get_mut(&pid) { + slot.merge(ev); + } + } + } + // At least one declared track has no evidence for this extent — read + // it. (A playlist that declares a PGS PID a previous playlist did not + // lands here, so the extra track is genuinely probed.) + None => todo.push(*ext), + } + } + if todo.is_empty() { + // Every extent's evidence came from a run that reached a designed stop, so + // an absence claim over the composed evidence is as sound as the run that + // produced each part. + apply_verdicts(title, &verdicts(&evidence, true)); return; } - let mut demux = TsDemuxer::new(&pg_pids); - let mut parsers: HashMap = - pg_pids.iter().map(|&p| (p, PgsParser::new())).collect(); - let mut trackers: HashMap = - pg_pids.iter().map(|&p| (p, ForcedTracker::new())).collect(); - - let extents = title.extents.clone(); let mut buf = vec![0u8; CHUNK_SECTORS as usize * SECTOR_BYTES]; let mut sectors_read: u32 = 0; // Record WHY the loop ended rather than leaving it implicit in the control // flow: every exit below names its reason, and the reason decides what may be // asserted from what was observed. - let stop = 'outer: { - for ext in &extents { - let mut lba = ext.start_lba; - let mut remaining = ext.sector_count; - while remaining > 0 { - // Bounded work and a responsive cancel: without these the probe - // reads the entire title whenever a track really is forced. - if halt.is_some_and(|h| h.is_cancelled()) { - break 'outer StopReason::Halted; + let mut stop = StopReason::Exhausted; + 'outer: for ext in &todo { + // Demux/parse state is PER EXTENT, so the evidence an extent yields is + // derived from that extent's own bytes and nothing else — which is what + // makes the per-extent cache entry mean what it claims, and is required + // now that a cache hit can make the read skip an extent in the middle of + // the title (a demuxer carried across a skipped extent would splice two + // non-adjacent byte runs into one PES). Each extent is a clip's own + // contiguous run, so this loses at most a display set that straddles an + // extent boundary of a fragmented file. + let mut demux = TsDemuxer::new(&pg_pids); + let mut parsers: HashMap = + pg_pids.iter().map(|&p| (p, PgsParser::new())).collect(); + let mut trackers: HashMap = + pg_pids.iter().map(|&p| (p, ForcedTracker::new())).collect(); + // AACS aligned units are anchored at THIS extent's start LBA, so tell a + // decrypt-on-read source to gate relative to it rather than absolute disc + // LBA 0 — without this the very first read of a clip whose start_lba is + // not itself 3-aligned is rejected. Mirrors the mux read paths. + reader.set_unit_base(ext.start_lba); + let mut lba = ext.start_lba; + let mut remaining = ext.sector_count; + // `None` = this extent was read to its end, so its evidence is complete + // and may be memoised. `Some(reason)` = the read stopped early. + let mut cut_short: Option = None; + while remaining > 0 { + // Bounded work and a responsive cancel: without these the probe + // reads the entire title whenever a track really is forced. + if halt.is_some_and(|h| h.is_cancelled()) { + cut_short = Some(StopReason::Halted); + break; + } + if sectors_read >= PROBE_BUDGET_SECTORS { + cut_short = Some(StopReason::Budget); + break; + } + let budget_left = PROBE_BUDGET_SECTORS - sectors_read; + let count = remaining.min(CHUNK_SECTORS as u32).min(budget_left) as u16; + let want = count as usize * SECTOR_BYTES; + let n = match reader.read_sectors(lba, count, &mut buf[..want], false) { + Ok(n) => n, + // Best-effort — stop reading, but the data past here was never + // seen, so the observation is a truncated prefix. + Err(_) => { + cut_short = Some(StopReason::ReadFailed); + break; } - if sectors_read >= PROBE_BUDGET_SECTORS { - break 'outer StopReason::Budget; - } - let budget_left = PROBE_BUDGET_SECTORS - sectors_read; - let count = remaining.min(CHUNK_SECTORS as u32).min(budget_left) as u16; - let want = count as usize * SECTOR_BYTES; - let n = match reader.read_sectors(lba, count, &mut buf[..want], false) { - Ok(n) => n, - // Best-effort — stop reading, but the data past here was never - // seen, so the observation is a truncated prefix. - Err(_) => break 'outer StopReason::ReadFailed, - }; - if n == 0 { - // Short read: the extent claimed sectors the source would not - // yield. Same truncated prefix as an error. - break 'outer StopReason::ReadFailed; - } - for pes in demux.feed(&buf[..n]) { + }; + // Advance by what was actually READ, not by what was requested. A + // short-but-nonzero read (a source whose batch is smaller than the + // request — `PrefetchedSectorSource` returns its producer's batch) + // used to advance `lba`/`remaining`/`sectors_read` by the full + // `count`, silently SKIPPING the unread tail of the chunk while + // `stop` stayed `Exhausted` — so an absence-based forced verdict was + // asserted (and memoised) over data that was never seen. The partial + // trailing sector, if any, is left for the next read rather than fed + // twice. + let got = (n.min(want) / SECTOR_BYTES) as u32; + if got == 0 { + // Less than one whole sector: the bytes are real, so feed them, + // but the loop cannot advance (re-reading the same partial sector + // would feed it twice) — so this is the truncated prefix an error + // is. A source that claims sectors and yields none is the same case. + for pes in demux.feed(&buf[..n.min(want)]) { if let (Some(parser), Some(tracker)) = (parsers.get_mut(&pes.pid), trackers.get_mut(&pes.pid)) { @@ -184,64 +284,111 @@ pub(crate) fn probe_and_set_forced( } } } - // Every track has already shown a non-forced set → nothing left to - // learn; stop reading the (huge) clip. - if trackers.values().all(ForcedTracker::settled_not_forced) { - break 'outer StopReason::Exhausted; + cut_short = Some(StopReason::ReadFailed); + break; + } + for pes in demux.feed(&buf[..got as usize * SECTOR_BYTES]) { + if let (Some(parser), Some(tracker)) = + (parsers.get_mut(&pes.pid), trackers.get_mut(&pes.pid)) + { + for frame in parser.parse(&pes) { + tracker.observe(&frame.data); + } } - lba += count as u32; - remaining -= count as u32; - sectors_read += count as u32; + } + lba += got; + remaining -= got; + sectors_read += got; + // Every track has already shown a non-forced set — counting the + // evidence carried in from other extents — so there is nothing left to + // learn; stop reading the (huge) clip. + if pg_pids.iter().all(|p| { + let carried = evidence.get(p).copied().unwrap_or_default().non_forced; + carried + || trackers + .get(p) + .is_some_and(ForcedTracker::settled_not_forced) + }) { + cut_short = Some(StopReason::Exhausted); + break; } } - StopReason::Exhausted - }; - // Drain any buffered final display set. - for (pid, parser) in parsers.iter_mut() { - if let Some(tracker) = trackers.get_mut(pid) { - for frame in parser.flush() { - tracker.observe(&frame.data); + // Drain any buffered final display set of THIS extent. + for (pid, parser) in parsers.iter_mut() { + if let Some(tracker) = trackers.get_mut(pid) { + for frame in parser.flush() { + tracker.observe(&frame.data); + } } } + + // Fold this extent's evidence in, and memoise it if the extent's read + // reached a DESIGNED stop — read to its end, stopped at the sector budget, + // or stopped because every track had already settled. The budget is a + // designed stop for exactly the reason [`StopReason`] documents (a forced + // track's display sets appear throughout, so a bounded prefix is + // representative), and it is the stop that fires on every disc that HAS a + // forced track — excluding it from the cache would mean nothing is ever + // memoised on precisely those discs. + // + // A halt or a read fault is different: the cut-off point is arbitrary, so + // its evidence is real for THIS title (nothing observed is retracted) but + // must not be frozen in as the extent's answer, or one transient fault + // would be replayed onto every other playlist sharing the clip. + let cacheable = cut_short.is_none_or(StopReason::absence_is_conclusive); + for (&pid, t) in trackers.iter() { + let ev = TrackEvidence { + observed: t.observed(), + non_forced: t.settled_not_forced(), + }; + if let Some(slot) = evidence.get_mut(&pid) { + slot.merge(ev); + } + if cacheable { + cache.insert((ext.start_lba, ext.sector_count, pid), ev); + } + } + if let Some(reason) = cut_short { + stop = reason; + break 'outer; + } } - // Collect the verdicts we are entitled to assert and apply them. A track - // absent from the map keeps its vendor-derived flag. - // - // Two gates, both PER TRACK, because the evidence is per track: - // * `observed()` — saw no display set at all, so nothing is known. (Never - // assert "not forced" from having seen nothing.) - // * on a truncated run, `settled_not_forced()` — the track saw an actual - // non-forced display set, which no further reading could retract, so that - // verdict stands even though the run was cut short. A track that merely - // hadn't YET seen a non-forced set is exactly the claim the truncation - // invalidates, so it is dropped and keeps the vendor flag. let conclusive = stop.absence_is_conclusive(); - let verdicts: HashMap = trackers - .iter() - .filter(|(_, t)| t.observed() && (conclusive || t.settled_not_forced())) - .map(|(&pid, t)| (pid, t.is_forced())) - .collect(); - // Only memoise a run that reached a designed stop. The cache key is the - // extent list, so caching a truncated run would replay one read fault (or one - // cancellation) onto every other playlist that shares these clips, and a later - // title would never get the chance to re-read them successfully. - if conclusive { - cache.insert(key, verdicts.clone()); - } else { + let verdicts = verdicts(&evidence, conclusive); + if !conclusive { tracing::debug!( target: "freemkv::scan", stop = ?stop, sectors_read, asserted = verdicts.len(), tracks = pg_pids.len(), - "forced-subtitle probe truncated; verdicts limited and not cached" + "forced-subtitle probe truncated; verdicts limited and truncated extents not cached" ); } apply_verdicts(title, &verdicts); } +/// Compose the per-track verdicts a run is ENTITLED to assert from the evidence +/// it gathered. A track absent from the result keeps its vendor-derived flag. +/// +/// Two gates, both PER TRACK, because the evidence is per track: +/// * `observed` — saw no display set at all, so nothing is known. (Never assert +/// "not forced" from having seen nothing.) +/// * on a truncated run, `non_forced` — the track saw an actual non-forced +/// display set, which no further reading could retract, so that verdict +/// stands even though the run was cut short. A track that merely hadn't YET +/// seen a non-forced set is exactly the claim the truncation invalidates, so +/// it is dropped and keeps the vendor flag. +fn verdicts(evidence: &HashMap, conclusive: bool) -> HashMap { + evidence + .iter() + .filter(|(_, e)| e.observed && (conclusive || e.non_forced)) + .map(|(&pid, e)| (pid, !e.non_forced)) + .collect() +} + /// Set `forced` on every PGS subtitle track named in `verdicts`. A track absent /// from the map was never observed and keeps its vendor-derived flag. fn apply_verdicts(title: &mut DiscTitle, verdicts: &HashMap) { @@ -435,6 +582,11 @@ mod tests { /// A reader that serves a fixed BD-TS byte stream once (across sequential /// `read_sectors` calls), then EOF — so the probe's demux→parse→observe→apply /// path runs on real synthetic PGS content. + /// + /// Sector-granular, like every real [`SectorSource`]: a read that is served + /// from the payload's short tail zero-pads to the sector boundary and reports + /// whole sectors. (The probe accounts in SECTORS, so a source that returned a + /// sub-sector byte count could never advance.) struct TsReader { data: Vec, pos: usize, @@ -453,7 +605,10 @@ mod tests { let n = buf.len().min(self.data.len() - self.pos); buf[..n].copy_from_slice(&self.data[self.pos..self.pos + n]); self.pos += n; - Ok(n) + let padded = n.div_ceil(SECTOR_BYTES) * SECTOR_BYTES; + let out = padded.min(buf.len()); + buf[n..out].fill(0); + Ok(out) } fn capacity_sectors(&self) -> u32 { self.data.len().div_ceil(SECTOR_BYTES) as u32 @@ -532,6 +687,14 @@ mod tests { pos: 0, }; let mut title = pgs_title(pid, false); // vendor label says NOT forced + // One sector, which is exactly what the reader serves: an extent that + // claims more sectors than the source yields is a SHORT read, and a short + // read is (correctly) inconclusive — see + // `read_error_after_partial_content_preserves_vendor_forced`. + title.extents = vec![Extent { + start_lba: 0, + sector_count: 1, + }]; probe_and_set_forced(&mut reader, &mut title, &mut ForcedProbeCache::new(), None); let Stream::Subtitle(s) = &title.streams[0] else { panic!() @@ -705,7 +868,13 @@ mod tests { s.forced, "a forced verdict from a budget-bounded prefix must still be applied" ); - assert_eq!(cache.len(), 1, "a conclusive probe is memoised"); + // One entry per (extent, PGS track): both extents reached a designed stop + // (the first was read to its end, the second stopped at the budget), so + // both are memoised — the budget is the stop that fires on every disc that + // HAS a forced track, so excluding it would memoise nothing there. + assert_eq!(cache.len(), 2, "a conclusive probe is memoised per extent"); + assert!(cache.contains_key(&(0, 4, pid))); + assert!(cache.contains_key(&(100, u32::MAX, pid))); } #[test] @@ -731,9 +900,14 @@ mod tests { !s.forced, "a verdict from a cancelled probe must not overwrite the vendor flag" ); + // The cancel landed inside the SECOND extent, whose read is therefore an + // arbitrary prefix: that extent must not be memoised, or the one + // cancellation would be replayed onto every other playlist sharing the + // clip. (The first extent WAS read to its end before the cancel, so its + // own evidence is sound and keeping it is the point of per-extent keying.) assert!( - cache.is_empty(), - "a cancelled probe must not poison the extent cache" + !cache.contains_key(&(100, u32::MAX, pid)), + "a cancelled probe must not poison the cancelled extent's cache entry" ); } @@ -748,4 +922,228 @@ mod tests { probe_and_set_forced(&mut reader, &mut title, &mut ForcedProbeCache::new(), None); assert_eq!(reader.served, 0, "no PGS PIDs → no reads"); } + + // ── per-extent, per-track memoisation ─────────────────────────────────── + + /// MEASURED: overlapping-but-not-identical extent lists must not re-read the + /// shared clips. A disc's playlists share clips without sharing whole extent + /// LISTS (00800 = [X, Y], 00801 = [X], 00802 = [Y]), and keying the cache on + /// the whole list de-duplicated only exactly-identical playlists: each of the + /// three lists missed, so clip X was read twice and Y twice — up to + /// PROBE_BUDGET_SECTORS (256 MiB) of optical-drive time per miss. + #[test] + fn overlapping_extent_lists_read_each_clip_once() { + let pid = 0x1200u16; + let x = Extent { + start_lba: 0, + sector_count: 600, + }; + let y = Extent { + start_lba: 10_000, + sector_count: 900, + }; + let mut reader = EndlessReader { served: 0 }; + let mut cache = ForcedProbeCache::new(); + + let mut both = pgs_title(pid, true); + both.extents = vec![x, y]; + probe_and_set_forced(&mut reader, &mut both, &mut cache, None); + let after_both = reader.served; + assert_eq!( + after_both, + x.sector_count + y.sector_count, + "the first title reads both clips exactly once" + ); + + // A playlist over X alone, and one over Y alone: every extent is already + // known, so neither costs a single further sector. + let mut only_x = pgs_title(pid, true); + only_x.extents = vec![x]; + probe_and_set_forced(&mut reader, &mut only_x, &mut cache, None); + let mut only_y = pgs_title(pid, true); + only_y.extents = vec![y]; + probe_and_set_forced(&mut reader, &mut only_y, &mut cache, None); + assert_eq!( + reader.served, after_both, + "clips shared with an already-probed playlist must not be re-read" + ); + + // And a list that mixes a known extent with a NEW one reads only the new + // one. + let z = Extent { + start_lba: 50_000, + sector_count: 300, + }; + let mut mixed = pgs_title(pid, true); + mixed.extents = vec![x, z]; + probe_and_set_forced(&mut reader, &mut mixed, &mut cache, None); + assert_eq!( + reader.served, + after_both + z.sector_count, + "a partially-known list reads only the extents it adds" + ); + } + + /// A later playlist that declares MORE PGS tracks over the SAME extents must + /// still probe the extra track. With the cache keyed on the extent list alone, + /// the verdict map it hit had no entry for the new PID, so that track was never + /// probed and silently kept its vendor-label flag — `info` then reported a + /// different forced flag for it depending purely on playlist ordering. + #[test] + fn extra_pgs_track_over_known_extents_is_still_probed() { + let ext = Extent { + start_lba: 0, + sector_count: 600, + }; + let mut reader = EndlessReader { served: 0 }; + let mut cache = ForcedProbeCache::new(); + + let mut one_track = pgs_title(0x1200, true); + one_track.extents = vec![ext]; + probe_and_set_forced(&mut reader, &mut one_track, &mut cache, None); + let after_first = reader.served; + assert_eq!(after_first, ext.sector_count); + + // Same extents, two declared PGS tracks. + let mut two_tracks = pgs_title(0x1200, true); + two_tracks.extents = vec![ext]; + two_tracks.streams.push(Stream::Subtitle(SubtitleStream { + pid: 0x1201, + codec: Codec::Pgs, + language: "fra".into(), + forced: true, + qualifier: LabelQualifier::None, + codec_data: None, + })); + probe_and_set_forced(&mut reader, &mut two_tracks, &mut cache, None); + assert!( + reader.served > after_first, + "a newly declared PGS track must be probed, not served from a verdict \ + map that has no entry for it" + ); + assert!( + cache.contains_key(&(ext.start_lba, ext.sector_count, 0x1201)), + "the new track gets its own per-extent evidence" + ); + } + + /// Records every (lba, count) served and every `set_unit_base` call. + struct AlignSpy { + reads: Vec<(u32, u16)>, + bases: Vec, + } + impl SectorSource for AlignSpy { + fn read_sectors( + &mut self, + lba: u32, + count: u16, + buf: &mut [u8], + _recovery: bool, + ) -> crate::error::Result { + self.reads.push((lba, count)); + let want = count as usize * SECTOR_BYTES; + buf[..want].fill(0); + Ok(want) + } + fn capacity_sectors(&self) -> u32 { + u32::MAX + } + fn set_unit_base(&mut self, lba: u32) { + self.bases.push(lba); + } + } + + /// Every probe read must begin on an AACS aligned-unit boundary measured from + /// the extent's own base, and the probe must declare that base to the source. + /// A `DecryptingSectorSource` holding AACS keys rejects any other read outright + /// (`DecryptFailed`) — and with a 1024-sector chunk (`1024 % 3 == 1`) every + /// read after the first was misaligned, so content-based forced detection was + /// unreachable past the first chunk of an encrypted disc, silently. + #[test] + fn probe_reads_stay_on_aacs_unit_boundaries() { + let pid = 0x1200u16; + // A start_lba that is NOT itself 3-aligned, so absolute `lba % 3` and the + // base-relative gate disagree — the case the gate exists for. + let base = 4_001u32; + let mut reader = AlignSpy { + reads: Vec::new(), + bases: Vec::new(), + }; + let mut title = pgs_title(pid, true); + title.extents = vec![Extent { + start_lba: base, + sector_count: CHUNK_SECTORS as u32 * 3, + }]; + probe_and_set_forced(&mut reader, &mut title, &mut ForcedProbeCache::new(), None); + + assert_eq!( + reader.bases, + vec![base], + "the probe must anchor the source's unit gate at the extent's start_lba" + ); + assert!(reader.reads.len() > 1, "more than one chunk was read"); + for &(lba, _) in &reader.reads { + assert!( + crate::aacs::content::is_unit_aligned(lba, base), + "read at lba {lba} is not on an aligned-unit boundary from base {base}" + ); + } + } + + /// A source that serves only `frac` of the sectors requested, never erroring — + /// what `PrefetchedSectorSource` does (it returns its producer's batch, not + /// `count * 2048`). Records the LBAs it actually served. + struct ShortReader { + frac: u32, + served: Vec<(u32, u32)>, + } + impl SectorSource for ShortReader { + fn read_sectors( + &mut self, + lba: u32, + count: u16, + buf: &mut [u8], + _recovery: bool, + ) -> crate::error::Result { + let give = (count as u32 / self.frac).max(1); + self.served.push((lba, give)); + let n = give as usize * SECTOR_BYTES; + buf[..n].fill(0); + Ok(n) + } + fn capacity_sectors(&self) -> u32 { + u32::MAX + } + } + + /// A short-but-nonzero read must advance by what was READ, not by what was + /// requested. Advancing by the request skipped the unread tail of every chunk + /// — silently, with `StopReason` still `Exhausted`, so the absence-based + /// forced verdict was asserted (and memoised) over sectors nobody read. + #[test] + fn short_reads_do_not_skip_sectors() { + let pid = 0x1200u16; + let count = CHUNK_SECTORS as u32 * 2; + let mut reader = ShortReader { + frac: 4, + served: Vec::new(), + }; + let mut title = pgs_title(pid, true); + title.extents = vec![Extent { + start_lba: 0, + sector_count: count, + }]; + probe_and_set_forced(&mut reader, &mut title, &mut ForcedProbeCache::new(), None); + + // The served ranges must tile the extent exactly: contiguous, no gaps. + let mut next = 0u32; + for &(lba, given) in &reader.served { + assert_eq!(lba, next, "gap: sectors {next}..{lba} were never read"); + next += given; + } + assert_eq!( + next, count, + "every sector of the extent must be read when the source short-reads" + ); + } } diff --git a/src/io/file_sector_source/mod.rs b/src/io/file_sector_source/mod.rs index dac16aa..d1465b6 100644 --- a/src/io/file_sector_source/mod.rs +++ b/src/io/file_sector_source/mod.rs @@ -85,11 +85,27 @@ use crate::consts::{SECTOR_BYTES, SECTOR_BYTES_U64}; /// pressure concurrent writes. Override via `FREEMKV_READ_DROP_CHUNK_MIB`. const READ_DROP_CHUNK_BYTES_DEFAULT: u64 = 32 * 1024 * 1024; +/// Upper bound (in MiB) accepted from `FREEMKV_READ_DROP_CHUNK_MIB`. 64 GiB — +/// generous for any real medium, and small enough that `n * 1024 * 1024` cannot +/// wrap `u64`. Mirrors `WRITEBACK_CHUNK_MIB_MAX`, whose identical multiply is +/// bounded for exactly this reason: without the bound, a value above 2^44 +/// overflows — a panic on the first ISO open in an overflow-checked build, and in +/// release a wrap to a near-zero window that fires `drop_window` on every read. +/// Out-of-range values fall back to the default. +const READ_DROP_CHUNK_MIB_MAX: u64 = 64 * 1024; + fn read_drop_chunk_bytes() -> u64 { - std::env::var("FREEMKV_READ_DROP_CHUNK_MIB") - .ok() - .and_then(|v| v.parse::().ok()) - .filter(|&n| n > 0) + resolve_read_drop_chunk( + std::env::var("FREEMKV_READ_DROP_CHUNK_MIB") + .ok() + .and_then(|v| v.parse::().ok()), + ) +} + +/// The pure part of [`read_drop_chunk_bytes`], split out so the bound is +/// testable without mutating process environment. +fn resolve_read_drop_chunk(mib: Option) -> u64 { + mib.filter(|&n| n > 0 && n <= READ_DROP_CHUNK_MIB_MAX) .map(|n| n * 1024 * 1024) .unwrap_or(READ_DROP_CHUNK_BYTES_DEFAULT) } @@ -518,4 +534,36 @@ mod tests { lba += batch as u32; } } + + /// `FREEMKV_READ_DROP_CHUNK_MIB` must be BOUNDED before the MiB→byte + /// multiply, exactly as its writeback twin bounds the identical multiply. + /// Unbounded, any value above 2^44 overflowed `n * 1024 * 1024`: a panic on + /// the first ISO open in an overflow-checked build, and in release a wrap to + /// a near-zero window that fires `drop_window` on essentially every read. + #[test] + fn read_drop_chunk_env_is_bounded_before_the_multiply() { + // Default when unset / zero / out of range. + assert_eq!(resolve_read_drop_chunk(None), READ_DROP_CHUNK_BYTES_DEFAULT); + assert_eq!( + resolve_read_drop_chunk(Some(0)), + READ_DROP_CHUNK_BYTES_DEFAULT + ); + // The overflow value: `u64::MAX * 1024 * 1024` panicked here. + assert_eq!( + resolve_read_drop_chunk(Some(u64::MAX)), + READ_DROP_CHUNK_BYTES_DEFAULT + ); + assert_eq!( + resolve_read_drop_chunk(Some(READ_DROP_CHUNK_MIB_MAX + 1)), + READ_DROP_CHUNK_BYTES_DEFAULT + ); + // In-range values convert MiB→bytes. Mutation: `* 1024` breaks this. + assert_eq!(resolve_read_drop_chunk(Some(1)), 1024 * 1024); + assert_eq!( + resolve_read_drop_chunk(Some(READ_DROP_CHUNK_MIB_MAX)), + READ_DROP_CHUNK_MIB_MAX * 1024 * 1024 + ); + // And the bound itself keeps the multiply inside u64. + assert!((READ_DROP_CHUNK_MIB_MAX as u128) * 1024 * 1024 <= u64::MAX as u128); + } } diff --git a/src/io/pipeline.rs b/src/io/pipeline.rs index 1039492..1741520 100644 --- a/src/io/pipeline.rs +++ b/src/io/pipeline.rs @@ -33,7 +33,7 @@ //! consumer lag detection). This is critical for diagnosing stalls. use std::sync::Arc; -use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::atomic::{AtomicBool, AtomicU8, Ordering}; use std::thread::{self, JoinHandle}; use std::time::{Duration, Instant}; @@ -117,8 +117,29 @@ fn consumer_panicked(payload: Box) -> Error { Error::PipelineConsumerPanicked } +/// Consumer lifecycle state, shared between the caller and the consumer thread. +/// +/// A plain `AtomicBool` could not make "the caller abandons" and "the consumer +/// commits to finalising" mutually exclusive: the consumer loaded the flag, the +/// caller stored it, and the consumer then finalised the container anyway — the +/// caller reporting the rip as interrupted while a fully finalised MKV (Cues +/// written, Segment size patched) landed on disk, indistinguishable from a +/// complete one. The two transitions are therefore a single compare-exchange each, +/// out of [`ST_RUNNING`]: whoever wins decides, and the loser observes the winner. +mod state { + /// Consumer is running; neither side has committed yet. + pub const RUNNING: u8 = 0; + /// The caller gave up on the consumer and will report failure — the consumer + /// must NOT finalise the output. + pub const ABANDONED: u8 = 1; + /// The consumer has committed to `close()` (finalising the output). The caller + /// can no longer abandon it; it must wait for the result it is about to + /// produce. + pub const CLOSING: u8 = 2; +} + /// After a halt or deadline fires, spin-poll `handle.is_finished()` for -/// [`FINISH_GRACE_SECS`] before accepting the thread leak. This converts +/// `grace` before accepting the thread leak. This converts /// the common "nearly-done" consumer (whose own bounded_syscall just /// returned and is about to drop its output file) into a clean join, /// releasing the file handle without waiting the full grace period. @@ -132,11 +153,12 @@ fn consumer_panicked(payload: Box) -> Error { /// syscall itself; that still returns on its own (or at process exit). fn finish_with_grace( handle: thread::JoinHandle>, - abandoned: &Arc, + state: &Arc, + grace: Duration, leak_err: Error, ) -> Result { - let grace = Instant::now() + Duration::from_secs(FINISH_GRACE_SECS); - while Instant::now() < grace { + let deadline = Instant::now() + grace; + while Instant::now() < deadline { if handle.is_finished() { return match handle.join() { Ok(result) => result, @@ -145,16 +167,50 @@ fn finish_with_grace( } thread::sleep(POLL_INTERVAL); } - // Grace expired. Signal abandonment, then log and leak. Setting the - // flag BEFORE dropping the handle guarantees the leaked consumer - // observes it the moment its wedged syscall returns: it then skips - // any further `apply` and skips `close()`, rather than running on to - // finalise the abandoned output file. - // `Release` here pairs with the `Acquire` loads in the consumer loop so - // the leaked consumer reliably observes the flag the moment its wedged - // syscall returns, even on weak memory models (ARM64/POWER) where - // `Relaxed` gives no cross-thread visibility guarantee. - abandoned.store(true, Ordering::Release); + // Grace expired. CLAIM abandonment, then log and leak. Claiming BEFORE + // dropping the handle guarantees the leaked consumer observes it the moment + // its wedged syscall returns: it then skips any further `apply` and skips + // `close()`, rather than running on to finalise the abandoned output file. + // + // A compare-exchange, not a store, because the consumer may have committed to + // `close()` in the instant between our last `is_finished()` poll and now. It + // then cannot be stopped — the finalise IS happening — so abandoning it would + // report the rip as interrupted while a valid, fully finalised container + // lands on disk. Losing the race means waiting for the result the consumer is + // already producing instead. `AcqRel` pairs with the consumer's own + // compare-exchange and with the `Acquire` loads in its drain loop, so the flag + // is reliably observed even on weak memory models (ARM64/POWER). + if state + .compare_exchange( + state::RUNNING, + state::ABANDONED, + Ordering::AcqRel, + Ordering::Acquire, + ) + .is_err() + { + tracing::warn!( + target: "freemkv::pipeline", + phase = "finish_with_halt_close_in_flight", + "pipeline consumer had already committed to finalising the output; \ + waiting for it rather than reporting an unfinalised output" + ); + let close_deadline = Instant::now() + grace; + while Instant::now() < close_deadline { + if handle.is_finished() { + return match handle.join() { + Ok(result) => result, + Err(payload) => Err(consumer_panicked(payload)), + }; + } + thread::sleep(POLL_INTERVAL); + } + // Still finalising after a second grace window: leak and report the wedge. + // The output may end up finalised by the leaked thread — but that is now a + // wedged-`close()` case, not the check-then-finalise race. + drop(handle); + return Err(leak_err); + } tracing::warn!( target: "freemkv::pipeline", phase = "finish_with_halt_grace_expired", @@ -243,7 +299,21 @@ pub struct Pipeline { /// a syscall the consumer is currently wedged in, but it does bound /// the damage to "whatever write is already in flight" once that /// syscall returns, instead of running on to a clean finalise. - abandoned: Arc, + /// + /// One of [`state::RUNNING`] / [`state::ABANDONED`] / [`state::CLOSING`]; + /// both transitions are compare-exchanges so abandoning and finalising are + /// mutually exclusive rather than racing. + state: Arc, + /// Set by the consumer the moment an `apply` returns `Err`. The consumer keeps + /// draining the channel after that (so the producer never blocks on a dead + /// receiver) — which means a producer watching only `send`'s return value + /// cannot tell the difference between "being consumed" and "being discarded + /// after a fatal write error", and would go on reading the whole remaining + /// disc before `finish()` finally surfaced the error. This flag is that + /// missing edge: [`Pipeline::send_with_halt`] fails fast on it, and + /// [`Pipeline::consumer_failed`] exposes it to producers that use plain + /// [`Pipeline::send`]. + failed: Arc, } impl Pipeline { @@ -277,8 +347,10 @@ impl Pipeline { sink: S, ) -> Result { let (tx, rx) = bounded::(depth); - let abandoned = Arc::new(AtomicBool::new(false)); - let abandoned_consumer = abandoned.clone(); + let state = Arc::new(AtomicU8::new(state::RUNNING)); + let state_consumer = state.clone(); + let failed = Arc::new(AtomicBool::new(false)); + let failed_consumer = failed.clone(); let handle = thread::Builder::new() .name(name.into()) .spawn(move || -> Result { @@ -309,7 +381,7 @@ impl Pipeline { // dead receiver, but we touch the output no further. The // final post-loop abandonment check returns the error // and skips `close()`. - if abandoned_consumer.load(Ordering::Acquire) { + if state_consumer.load(Ordering::Acquire) == state::ABANDONED { continue; } @@ -338,6 +410,12 @@ impl Pipeline { tracing::debug!("Pipeline: apply error, stopping, err={:?}", e); } first_err = Some(e); + // Publish the failure so the producer can stop + // FEEDING a dead write side instead of only learning + // about it at `finish()` — by which time it has read + // the rest of the disc. `Release` pairs with the + // `Acquire` load in `send_with_halt`. + failed_consumer.store(true, Ordering::Release); } } @@ -398,13 +476,38 @@ impl Pipeline { // MKV Cues + patching the segment header) on a file the // caller already reported as failed is exactly the // write race we must not run. - if abandoned_consumer.load(Ordering::Acquire) { - return Err(Error::Halted); - } - match first_err { - Some(e) => Err(e), - None => sink.close(), + // No `close()` on this path, so there is nothing to claim — + // just report, unless the caller has already given up on us. + Some(e) => { + if state_consumer.load(Ordering::Acquire) == state::ABANDONED { + Err(Error::Halted) + } else { + Err(e) + } + } + // CLAIM the finalise. A plain load here left a window in which + // the caller stored `abandoned` AFTER we read it as clear, so + // `close()` ran anyway and finalised (Cues + Segment-size + // patch) an output the caller had already reported as + // interrupted — a truncated rip indistinguishable from a + // complete one. The compare-exchange closes that window: if the + // caller got there first we skip `close()`, and if we get there + // first the caller waits for us instead of abandoning. + None => { + if state_consumer + .compare_exchange( + state::RUNNING, + state::CLOSING, + Ordering::AcqRel, + Ordering::Acquire, + ) + .is_err() + { + return Err(Error::Halted); + } + sink.close() + } } }) .map_err(|e| Error::IoError { source: e })?; @@ -412,10 +515,24 @@ impl Pipeline { Ok(Pipeline { tx, handle, - abandoned, + state, + failed, }) } + /// Whether the consumer's `apply` has already failed fatally. + /// + /// The consumer keeps draining the channel after an `apply` error (so the + /// producer never blocks on a dead receiver), which means `send` keeps + /// succeeding and a producer has no other way to tell that everything it feeds + /// is being discarded. A long-running producer — the mux frame pump reading a + /// 60 GB title off an optical drive — should check this and unwind instead of + /// reading the rest of the disc for a write that has already failed. + /// [`Pipeline::send_with_halt`] checks it automatically. + pub fn consumer_failed(&self) -> bool { + self.failed.load(Ordering::Acquire) + } + /// Push one item. Blocks if the channel is full — that's the /// back-pressure the whole primitive exists to provide. Returns /// the item back if the consumer thread is gone (panicked or @@ -513,6 +630,21 @@ impl Pipeline { let end = Instant::now() + deadline; let mut pending = item; loop { + // The consumer's `apply` has failed fatally: everything sent from here + // is drained and discarded, so hand the item back at once. Without this + // the producer saw every send succeed (the channel is always being + // drained) and went on reading the whole remaining title — an hour of + // drive time on a UHD — for a write that died on the first frame, only + // learning about it at `finish()`. + if self.consumer_failed() { + if debug_enabled() { + tracing::debug!( + "Pipeline send_with_halt: consumer apply failed, returning item={}", + std::any::type_name::() + ); + } + return Err(pending); + } // Pre-check the cheap exit conditions before parking. if halt.is_cancelled() { if debug_enabled() { @@ -567,7 +699,8 @@ impl Pipeline { let Pipeline { tx, handle, - abandoned: _, + state: _, + failed: _, } = self; // Explicit drop, although the destructure already drops `tx` // at end-of-scope. Being explicit keeps the intent obvious. @@ -607,7 +740,8 @@ impl Pipeline { let Pipeline { tx, handle, - abandoned, + state, + failed: _, } = self; drop(tx); let deadline = Instant::now() + Duration::from_secs(JOIN_TIMEOUT_SECS); @@ -620,11 +754,21 @@ impl Pipeline { } if let Some(h) = halt { if h.is_cancelled() { - return finish_with_grace(handle, &abandoned, Error::Halted); + return finish_with_grace( + handle, + &state, + Duration::from_secs(FINISH_GRACE_SECS), + Error::Halted, + ); } } if Instant::now() >= deadline { - return finish_with_grace(handle, &abandoned, Error::PipelineJoinTimeout); + return finish_with_grace( + handle, + &state, + Duration::from_secs(FINISH_GRACE_SECS), + Error::PipelineJoinTimeout, + ); } thread::sleep(POLL_INTERVAL); } @@ -1594,4 +1738,132 @@ mod tests { let res = pipe.finish_with_halt(None); assert!(matches!(res, Ok(190)), "expected Ok(190), got {res:?}"); } + + /// A fatal `apply` error must become visible to the PRODUCER, not only to + /// `finish()`. The consumer keeps draining after the error (so the producer + /// never blocks on a dead receiver), which meant every `send_with_halt` + /// returned `Ok` for the rest of the run: on a 60 GB mkv:// mux that hit + /// ENOSPC on the first frame, the mux driver read the entire remaining title — + /// an hour of optical-drive time — before learning the write had died. + #[test] + fn send_with_halt_fails_fast_once_apply_has_failed() { + struct FailFirst { + failed: Arc, + } + impl Sink for FailFirst { + type Output = (); + fn apply(&mut self, _item: u64) -> Result { + self.failed.fetch_add(1, Ordering::SeqCst); + Err(Error::DecryptFailed) + } + fn close(self) -> Result<(), Error> { + Ok(()) + } + } + let applied = Arc::new(AtomicUsize::new(0)); + let pipe = Pipeline::spawn( + DEFAULT_PIPELINE_DEPTH, + FailFirst { + failed: applied.clone(), + }, + ) + .expect("spawn"); + let halt = crate::halt::Halt::new(); + let deadline = Duration::from_secs(5); + + // Feed one item and wait until the consumer has actually applied (and + // failed on) it, so the check below is deterministic rather than racy. + pipe.send_with_halt(0u64, &halt, deadline) + .expect("the first send lands"); + let until = Instant::now() + Duration::from_secs(2); + while Instant::now() < until && applied.load(Ordering::SeqCst) == 0 { + std::thread::sleep(Duration::from_millis(5)); + } + assert_eq!(applied.load(Ordering::SeqCst), 1, "apply ran and failed"); + + assert!(pipe.consumer_failed(), "the failure must be observable"); + // The very next send must hand the item straight back — the producer's + // signal to stop reading the disc. + assert_eq!( + pipe.send_with_halt(1u64, &halt, deadline), + Err(1u64), + "send_with_halt must fail fast once the consumer's apply has failed" + ); + // The halt was never fired, so this is not a cancellation: the real error + // still comes out of finish(). + assert!(matches!(pipe.finish(), Err(Error::DecryptFailed))); + assert_eq!( + applied.load(Ordering::SeqCst), + 1, + "no further item was applied" + ); + } + + /// The abandon/finalise race. A consumer that has ALREADY committed to + /// `close()` when the grace period expires cannot be stopped — the finalise is + /// happening — so the caller must wait for its result instead of reporting the + /// output as un-finalised. With a plain flag the consumer read it as clear, the + /// caller then stored it, and the caller returned `Err(Halted)` + /// (`completed = false`) while a fully finalised MKV (Cues written, Segment + /// size patched) landed on disk — a truncated rip indistinguishable from a + /// complete one. + #[test] + fn abandon_loses_to_a_close_already_committed() { + let state = Arc::new(AtomicU8::new(state::RUNNING)); + let release = Arc::new(AtomicBool::new(false)); + let in_close = Arc::new(AtomicBool::new(false)); + + let (st, rel, inc) = (state.clone(), release.clone(), in_close.clone()); + let handle = thread::Builder::new() + .name("test-consumer".into()) + .spawn(move || -> Result { + // Exactly what the consumer does before finalising: claim the + // right to close. + assert!( + st.compare_exchange( + state::RUNNING, + state::CLOSING, + Ordering::AcqRel, + Ordering::Acquire + ) + .is_ok(), + "the consumer claims the finalise first" + ); + inc.store(true, Ordering::SeqCst); + // Inside `close()`, finalising the container. + while !rel.load(Ordering::SeqCst) { + thread::sleep(Duration::from_millis(5)); + } + Ok(42) + }) + .expect("spawn"); + + let until = Instant::now() + Duration::from_secs(2); + while Instant::now() < until && !in_close.load(Ordering::SeqCst) { + thread::sleep(Duration::from_millis(5)); + } + assert!(in_close.load(Ordering::SeqCst), "consumer reached close()"); + + // Finish the close only AFTER the first grace window has expired, so the + // caller genuinely reaches the abandon decision with a close in flight. + let rel = release.clone(); + thread::spawn(move || { + // Past the first grace window (and past the 250 ms poll cadence that + // bounds when the window is actually observed), inside the second. + thread::sleep(Duration::from_millis(600)); + rel.store(true, Ordering::SeqCst); + }); + + let grace = Duration::from_millis(300); + let res = finish_with_grace(handle, &state, grace, Error::Halted); + assert!( + matches!(res, Ok(42)), + "a finalise already in flight must be waited for, not abandoned: {res:?}" + ); + assert_eq!( + state.load(Ordering::SeqCst), + state::CLOSING, + "the caller must not have overwritten the consumer's claim" + ); + } } diff --git a/src/keysource.rs b/src/keysource.rs index 885faf5..3922807 100644 --- a/src/keysource.rs +++ b/src/keysource.rs @@ -82,9 +82,12 @@ impl DecodeSampleSet { } /// The public AACS inputs a key source needs to look a disc up. Captured at -/// scan; contains no secrets — only the disc identity and the on-disc AACS -/// structures a source or key server may key on. -#[derive(Debug, Clone)] +/// scan; carries no DERIVED secrets (no media key, VUK or plaintext unit key) — +/// only the disc identity and the on-disc AACS structures a source or key server +/// may key on. The on-disc structures are nonetheless key MATERIAL (the encrypted +/// title keys live in `unit_key_ro`), so [`Debug`] is hand-written and redacting; +/// see the impl below. +#[derive(Clone)] pub struct DiscInputs { /// SHA-1 of `Unit_Key_RO.inf`, `0x`-prefixed hex. The value a keydb keys /// its per-disc entries by, and a key server identifies the disc with. @@ -115,6 +118,30 @@ pub struct DiscInputs { pub volume_label: Option, } +/// Redacting `Debug`, per the policy `aacs::types` documents (and which +/// `aacs::types::Vid` already applies to this very Volume ID). `DiscInputs` is +/// public and returned by [`crate::Disc::inputs`], so a consumer's +/// `tracing::debug!("{inputs:?}")` used to print the Volume ID, the whole +/// `Unit_Key_RO.inf` (the encrypted title keys), the entire MKB and every +/// ciphertext sample verbatim into a log that ends up attached to a bug report. +/// Only non-secret identity and shape (presence, lengths) is printed. +impl std::fmt::Debug for DiscInputs { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("DiscInputs") + .field("disc_hash", &self.disc_hash) + .field("volume_id", &"") + .field("version", &self.version) + .field("mkb", &"") + .field("mkb_len", &self.mkb.len()) + .field("unit_key_ro", &"") + .field("unit_key_ro_len", &self.unit_key_ro.len()) + .field("samples", &"") + .field("samples_len", &self.samples.len()) + .field("volume_label", &self.volume_label) + .finish() + } +} + /// A lazy view of a disc's AACS material, handed to [`KeySource::get_unit_keys`] so a /// source can drive the derivation chain without holding the disc reader. /// @@ -1334,4 +1361,39 @@ mod tests { assert_eq!(k10[1], [0x10; 16], "V10 reads the 2nd key at +48"); assert_ne!(k20[1], k10[1], "the parse stride follows inputs.version"); } + + /// `DiscInputs` is public and returned by `Disc::inputs`, so any consumer's + /// `tracing::debug!("{inputs:?}")` prints it. A derived `Debug` printed the + /// Volume ID (the value `aacs::types::Vid` deliberately renders as + /// `Vid()`), the whole `Unit_Key_RO.inf` (the encrypted title keys), + /// the entire MKB and every ciphertext sample verbatim. Sentinel byte + /// 0xD5 = decimal 213, matching `aacs::types::redaction_tests`. Mutation + /// guard: restoring `#[derive(Debug)]` fails this. + #[test] + fn disc_inputs_debug_is_redacted() { + let inputs = DiscInputs { + disc_hash: "0xAA".into(), + volume_id: [0xD5; 16], + version: 2, + mkb: vec![0xD5; 64], + unit_key_ro: vec![0xD5; 48], + samples: vec![vec![0xD5; 6144]], + volume_label: Some("TITLE_2024".into()), + }; + let dbg = format!("{inputs:?}"); + assert!( + !dbg.contains("213"), + "DiscInputs Debug leaked key material (decimal 213): {dbg}" + ); + assert!( + dbg.contains("redacted"), + "DiscInputs Debug missing redaction marker: {dbg}" + ); + // Non-secret identity and shape stay printable for diagnostics. + assert!(dbg.contains("0xAA"), "{dbg}"); + assert!(dbg.contains("mkb_len: 64"), "{dbg}"); + assert!(dbg.contains("unit_key_ro_len: 48"), "{dbg}"); + assert!(dbg.contains("samples_len: 1"), "{dbg}"); + assert!(dbg.contains("TITLE_2024"), "{dbg}"); + } } diff --git a/src/mux/au_assembly.rs b/src/mux/au_assembly.rs index ecdfe85..03a23dd 100644 --- a/src/mux/au_assembly.rs +++ b/src/mux/au_assembly.rs @@ -321,8 +321,7 @@ impl AuAssembler { self.disc_marks.pop_front(); } - let data = self.buf[..end].to_vec(); - self.buf.drain(..end); + let data = self.take_front(end); self.base += end as u64; self.reset_scan(); out.push(AssembledAu { @@ -336,6 +335,41 @@ impl AuAssembler { out } + /// Detach `buf[..end]` as the emitted AU's own `Vec` and leave `buf` holding + /// the tail. + /// + /// The AU's bytes are HANDED OVER — `buf`'s allocation becomes the returned + /// `Vec` and a fresh buffer (pre-sized to the same capacity, so the next AU + /// accumulates without re-growing) takes its place holding only the short + /// tail. `buf[..end].to_vec()` + `drain(..end)` instead copied every AU out + /// in full: on a UHD HEVC title that is a whole-frame memcpy (hundreds of KB) + /// per coded picture, ~200k times, for bytes that are about to be discarded + /// from `buf` anyway. + /// + /// The allocation COUNT is unchanged (one per AU either way — the frame `Vec` + /// before, the replacement buffer now), so the only difference is the copy + /// that no longer happens. Nothing depends on `buf` keeping its identity: the + /// only state tied to `buf[0]`'s position is `base`/`scan_pos`/`opener_pos`, + /// which the caller updates immediately after. + /// + /// Falls back to a copy when the buffer's capacity is far larger than the AU + /// (a small AU after a multi-MB one): handing over would otherwise attach an + /// oversized idle allocation to a small frame for as long as the frame queues + /// downstream, trading a copy for resident memory. + fn take_front(&mut self, end: usize) -> Vec { + let cap = self.buf.capacity(); + if cap > end.saturating_mul(2) { + let data = self.buf[..end].to_vec(); + self.buf.drain(..end); + return data; + } + let mut tail = Vec::with_capacity(cap.max(self.buf.len() - end)); + tail.extend_from_slice(&self.buf[end..]); + let mut data = std::mem::replace(&mut self.buf, tail); + data.truncate(end); + data + } + /// Reset the incremental boundary-scan cursor. Called whenever `buf[0]` moves /// (an AU drained, or leading bytes discarded) so the next scan starts fresh /// from the new AU opener. @@ -841,4 +875,51 @@ mod tests { "over-cap AU is force-flushed, not buffered forever" ); } + + /// MEASURED: a drained AU must be HANDED the accumulation buffer's + /// allocation, not copied out of it. The emitted `Vec`'s data pointer is the + /// buffer's own pointer — which is only true if no full-frame copy happened. + /// (`buf[..end].to_vec()` allocates fresh, so the pointers differ.) One + /// whole-AU memcpy per coded picture is ~200k memcpys of a few hundred KB + /// each on a UHD feature. + #[test] + fn drained_au_takes_over_the_buffer_allocation_without_copying() { + let mut a = AuAssembler::for_codec(Codec::H264); + // An AU large enough that the buffer's capacity is not >2x its size (the + // small-AU copy path exists so a small frame cannot carry an oversized + // idle allocation downstream). + let au1 = au(0x11, 400 * 1024); + let au2 = au(0x22, 400 * 1024); + let mut stream = au1.clone(); + stream.extend_from_slice(&au2); + + // Push everything except the final byte of AU2's delimiter, so no AU has + // been emitted yet but the buffer holds the whole of AU1. + a.push(&stream[..au1.len() + 3], Some(1), None, None, false); + let before = a.buf.as_ptr(); + let cap_before = a.buf.capacity(); + let out = a.push( + &stream[au1.len() + 3..au1.len() + 4], + None, + None, + None, + false, + ); + assert_eq!(out.len(), 1); + assert_eq!( + out[0].data, au1, + "handover must preserve the AU bytes exactly" + ); + assert_eq!( + out[0].data.as_ptr(), + before, + "the emitted AU must own the buffer's allocation (no whole-frame copy)" + ); + assert_eq!( + a.buf.capacity(), + cap_before, + "the replacement buffer keeps the capacity, so the next AU does not re-grow" + ); + assert_eq!(a.buf.len(), 4, "the buffer holds only AU2's delimiter tail"); + } } diff --git a/src/mux/codec/hevc.rs b/src/mux/codec/hevc.rs index d52bd8b..cbcae38 100644 --- a/src/mux/codec/hevc.rs +++ b/src/mux/codec/hevc.rs @@ -201,12 +201,21 @@ struct ContentLightLevel { // window (HEVC reorder depth tops out ~16 frames, <1 s at 24 fps). 3 s = 270000 // ticks sits well above any legitimate reorder dip and far below any real clip's // duration, so it never false-triggers within a clip. This MIRRORS the mux-side -// `DISCONTINUITY_BACKSTEP_NS` (3 s) in `mux/mkv.rs`, which independently rebases -// the timeline at the same boundaries; here it drives the CRA→BLA rewrite that -// kills the dangling-RASL "Could not find ref with POC N" decode errors a +// `DISCONTINUITY_BACKSTEP_NS` (3 s) in `mux/timeline.rs`, which independently +// rebases the timeline at the same boundaries; here it drives the CRA→BLA rewrite +// that kills the dangling-RASL "Could not find ref with POC N" decode errors a // concatenated multi-clip title otherwise produces. const BACKSTEP_TICKS: i64 = 270_000; +// The mirror above is enforced, not just described: 90 kHz ticks → ns is +// × (1_000_000_000 / 90_000) = × 100_000 / 9, so 270_000 ticks must be exactly +// `DISCONTINUITY_BACKSTEP_NS`. Changing either constant without the other fails +// the build here, which is the drift the comment exists to prevent. +const _: () = assert!( + BACKSTEP_TICKS * 100_000 / 9 == crate::mux::timeline::DISCONTINUITY_BACKSTEP_NS, + "HEVC BACKSTEP_TICKS must mirror mux::timeline::DISCONTINUITY_BACKSTEP_NS" +); + // The 33-bit 90 kHz PES PTS counter wraps at 2^33 ticks (~26.5 h). When the raw // PTS steps backward by approximately a full period — i.e. it landed just past // the wrap — it is a counter wraparound, NOT a clip reset: unwrap it (add 2^33) diff --git a/src/mux/mp4/read.rs b/src/mux/mp4/read.rs index c3b1f3d..4d62cad 100644 --- a/src/mux/mp4/read.rs +++ b/src/mux/mp4/read.rs @@ -35,6 +35,17 @@ const MAX_TRACKS: usize = 512; /// hostile file's allocation without truncating any legitimate track. const MAX_SAMPLE_COUNT: usize = 1 << 24; +/// Smallest number of FILE bytes one indexed sample is assumed to occupy — the +/// divisor that turns `file_len` into a sample-count budget (see `from_reader`). +/// +/// Only `vide` and `soun` tracks are indexed here, and no real coded video or +/// audio access unit is anywhere near this small: the shortest legal AC-3 frame is +/// 128 bytes, an AAC frame is hundreds, and a video sample carries at least a NAL +/// header plus slice data. A 2-hour title runs to thousands of file bytes per +/// sample, so this cannot truncate a genuine track — it only stops a crafted +/// sample table from claiming more samples than the file could possibly hold. +const MIN_FILE_BYTES_PER_SAMPLE: u64 = 16; + /// Absolute ceiling on a single allocation sized from an untrusted MP4 field (a /// per-sample buffer or the `moov` payload). The EOF check alone is not enough: /// `file_len` is cheaply inflatable with a sparse file (`truncate -s 8G`), so a @@ -98,13 +109,16 @@ impl Mp4Reader { // crafted file with a fixed-size `stsz` claiming count=0xFFFFFFFF can't // inflate the `sizes`/`Vec` allocations past the file's own size // (a genuine large title has file_len ≫ sample count, so it is unaffected). - // NOTE on the bound this actually gives: each indexed sample costs about - // 52 bytes of RAM (SampleRef 40 + u32 size 4 + u64 offset 8, plus 4 each - // for the expanded stts/ctts), so the ceiling is ~52x file_len, not 1x — - // capped by MAX_SAMPLE_COUNT. That is still a real bound (a 1 MiB crafted - // file cannot reach the 16M-sample ceiling), just not the "past the file's - // own size" the previous comment implied. - let mut sample_budget = MAX_SAMPLE_COUNT.min(file_len.min(usize::MAX as u64) as usize); + // The bound is in file BYTES PER SAMPLE, not in samples: each indexed + // sample costs ~60 bytes of RAM (SampleRef 40 + u32 size 4 + u64 offset 8, + // plus 4 each for the expanded stts/ctts), so a budget of one sample per + // file byte still let a 16 MiB crafted file (a fixed-size `stsz` declaring + // 16M one-byte samples, a one-entry stsc/stco/stts) force a ~1 GiB eager + // allocation and a 16M-element sort before a single frame was read — a 64x + // amplification. Dividing by MIN_FILE_BYTES_PER_SAMPLE caps the + // amplification at ~4x instead. + let mut sample_budget = MAX_SAMPLE_COUNT + .min((file_len / MIN_FILE_BYTES_PER_SAMPLE).min(usize::MAX as u64) as usize); // Bound the scan at MAX_TRACKS *matches* so a crafted moov packed with tiny // (8-byte) trak headers can't force the scan to materialize a Vec far @@ -1341,10 +1355,10 @@ mod tests { /// NOT the 16M `MAX_SAMPLE_COUNT` ceiling. Mutation check: revert the budget /// to a bare `MAX_SAMPLE_COUNT` and this file yields ~16M samples, failing the /// `<= file_len` (and `< MAX_SAMPLE_COUNT`) assertions below. - #[test] - fn stsz_sample_count_bounded_by_file_len() { - use std::io::Cursor; - // A minimal audio trak, but with a fixed-size stsz lying about its count. + /// A minimal-but-complete audio `trak` whose fixed-size `stsz` LIES about its + /// sample count (`u32::MAX`), with an stsc/stts wide enough to place whatever + /// count survives the budget. The shape the sample-table budget bounds. + fn audio_trak_hostile_count() -> Vec { let mdhd = { let mut p = vec![0u8; 24]; p[12..16].copy_from_slice(&48_000u32.to_be_bytes()); // timescale @@ -1410,6 +1424,13 @@ mod tests { mdia.extend_from_slice(&hdlr); mdia.extend_from_slice(&minf); let trak = mp4_box(b"trak", &mp4_box(b"mdia", &mdia)); + trak + } + + #[test] + fn stsz_sample_count_bounded_by_file_len() { + use std::io::Cursor; + let trak = audio_trak_hostile_count(); let moov = mp4_box(b"moov", &trak); let file_len = moov.len() as u64; @@ -1418,11 +1439,16 @@ mod tests { "fixture stays a few hundred bytes ({file_len})" ); let rd = Mp4Reader::from_reader(Cursor::new(moov), "hostile".into()).unwrap(); - // The index must be bounded by the file's byte length, NOT the 16M ceiling. + // The index must be bounded by the file's byte length in FILE BYTES PER + // SAMPLE, not one sample per byte: each indexed sample costs ~60 bytes of + // RAM, so a one-sample-per-byte budget still let a 16 MiB crafted file + // force a ~1 GiB allocation (64x amplification) before a frame was read. assert!( - (rd.samples.len() as u64) <= file_len, - "sample count {} must be bounded by file_len {file_len}, not the count lie", - rd.samples.len() + (rd.samples.len() as u64) <= file_len / MIN_FILE_BYTES_PER_SAMPLE, + "sample count {} must be bounded by file_len/{MIN_FILE_BYTES_PER_SAMPLE} \ + ({}), not by the count lie", + rd.samples.len(), + file_len / MIN_FILE_BYTES_PER_SAMPLE ); assert!( rd.samples.len() < MAX_SAMPLE_COUNT, @@ -1430,6 +1456,35 @@ mod tests { ); } + /// The RAM amplification the byte-per-sample budget actually bounds: a small + /// crafted file whose `stsz` claims `u32::MAX` samples must not force an eager + /// multi-hundred-MB index. At ~60 bytes of RAM per indexed sample, a budget of + /// one sample per file byte gave ~60x the input size; the assertion below pins + /// the amplification factor rather than a raw count, so it fails if the budget + /// ever goes back to counting samples per byte. + #[test] + fn sample_index_ram_is_bounded_by_a_multiple_of_the_file() { + use std::io::Cursor; + // ~64 KiB of `free` padding so file_len is large enough for the ratio to + // be meaningful, with the same lying stsz as above. + let mut traks = Vec::new(); + traks.extend_from_slice(&audio_trak_hostile_count()); + let mut file = mp4_box(b"moov", &traks); + file.extend_from_slice(&mp4_box(b"free", &vec![0u8; 64 * 1024])); + let file_len = file.len() as u64; + + let rd = Mp4Reader::from_reader(Cursor::new(file), "amp".into()).unwrap(); + // ~60 bytes of RAM per sample; assert the index cannot exceed ~4x the file. + const RAM_PER_SAMPLE: u64 = 60; + let ram = rd.samples.len() as u64 * RAM_PER_SAMPLE; + assert!( + ram <= file_len * 4, + "sample index RAM {ram} B from a {file_len} B file is more than 4x \ + amplification ({} samples)", + rd.samples.len() + ); + } + /// Build an audio `trak` identical to `audio_trak(48_000)` but with the /// named stbl child box omitted. Used to reach the untrusted-input guards /// that drop a track whose `stsz` says samples exist yet whose `stco`/`co64` diff --git a/src/mux/timeline.rs b/src/mux/timeline.rs index cc888c9..1819b5f 100644 --- a/src/mux/timeline.rs +++ b/src/mux/timeline.rs @@ -130,10 +130,22 @@ impl TimelineContinuity { // within a backstep below the frontier keeps the remap to genuine // tail stragglers; a long audio-only tail, a sparse subtitle, or an // EL frame that simply runs ahead is left on the current offset. + // + // Every comparison below saturates. `high` is derived from an + // untrusted container timestamp (an `mkv://` source's + // CLUSTER_TIMESTAMP × TimestampScale is clamped only against + // `i64::MAX`, so a hostile file can put the frontier AT `i64::MAX`), + // and `raw_pts_ns` can be negative (a SimpleBlock's signed relative + // timestamp). Plain `high + BACKSTEP` / `high - BACKSTEP` would then + // overflow: a panic out of the public `Stream::write` path in an + // overflow-checked build, and in release a wrap to the opposite sign + // that fires the straggler clamp on essentially every passive frame. if let Some(high) = self.high_ns { - if mapped > high + DISCONTINUITY_BACKSTEP_NS { + if mapped > high.saturating_add(DISCONTINUITY_BACKSTEP_NS) { let prev_mapped = raw_pts_ns.saturating_add(self.prev_offset_ns); - if prev_mapped <= high && prev_mapped >= high - DISCONTINUITY_BACKSTEP_NS { + if prev_mapped <= high + && prev_mapped >= high.saturating_sub(DISCONTINUITY_BACKSTEP_NS) + { return prev_mapped; } } @@ -147,12 +159,17 @@ impl TimelineContinuity { return adj; }; let adj = raw_pts_ns.saturating_add(self.offset_ns); - if adj < high - DISCONTINUITY_BACKSTEP_NS { + if adj < high.saturating_sub(DISCONTINUITY_BACKSTEP_NS) { // Clip-boundary reset (real multi-clip seam): continue just after the // frontier. Save the previous offset so a lagging non-video tail // frame can be recognised and remapped to the seam (see above). self.prev_offset_ns = self.offset_ns; - let bump = (high - adj).saturating_add(DISCONTINUITY_GAP_NS); + // `high - adj` is a backward step, so positive — but both ends are + // untrusted (`high` up to i64::MAX, `adj` down to i64::MIN), so + // saturate rather than panic in a checked build. + let bump = high + .saturating_sub(adj) + .saturating_add(DISCONTINUITY_GAP_NS); self.offset_ns = self.offset_ns.saturating_add(bump); let adj2 = raw_pts_ns.saturating_add(self.offset_ns); self.high_ns = Some(high.max(adj2)); @@ -464,4 +481,40 @@ mod tests { "frame must stay in the new epoch (> frontier), got {out}" ); } + + /// A saturated frontier must not panic the muxer. An `mkv://` source's + /// tick→ns multiply saturates at `i64::MAX` (mkvstream's `parse_block`), so a + /// hostile TimestampScale/CLUSTER_TIMESTAMP puts `high_ns` AT `i64::MAX`. + /// Every subsequent PASSIVE frame then evaluated `high + BACKSTEP`, which + /// panicked ("attempt to add with overflow") out of the public + /// `Stream::write` path in any overflow-checked build. + #[test] + fn saturated_frontier_does_not_overflow_on_passive_frame() { + let mut tc = TimelineContinuity::new(); + // Video establishes the frontier at the saturation point. + assert_eq!(adj_video(&mut tc, i64::MAX), i64::MAX); + assert_eq!(tc.high_ns, Some(i64::MAX)); + // Passive frame: `high + BACKSTEP` overflowed here. + let out = adj_other(&mut tc, 0); + assert_eq!(out, 0, "a passive frame keeps its own mapping"); + // And a passive frame AT the frontier: `high - BACKSTEP` is the other + // unchecked side of the straggler discriminator. + assert_eq!(adj_other(&mut tc, i64::MAX), i64::MAX); + } + + /// The epoch-decision side of the same arithmetic: `adj < high - BACKSTEP` + /// and the `high - adj` bump both took untrusted ends. A frontier at + /// `i64::MIN`-adjacent values (a negative SimpleBlock-relative timestamp) and + /// a `i64::MAX` frontier are both reachable from container data. + #[test] + fn extreme_video_pts_does_not_overflow_the_epoch_bump() { + let mut tc = TimelineContinuity::new(); + assert_eq!(adj_video(&mut tc, i64::MAX), i64::MAX); + // Hard backward jump to the negative extreme: `high - adj` overflowed. + let out = adj_video(&mut tc, i64::MIN); + // Saturated bump (`i64::MAX`) applied to `i64::MIN` → -1, and the + // frontier never regresses. + assert_eq!(out, -1); + assert_eq!(tc.high_ns, Some(i64::MAX)); + } } diff --git a/src/mux/tsmux.rs b/src/mux/tsmux.rs index 1b32dc4..2debbbe 100644 --- a/src/mux/tsmux.rs +++ b/src/mux/tsmux.rs @@ -70,6 +70,15 @@ pub struct TsMuxer { /// so a header-only `m2ts://` output can't be reported as success — /// mirroring `MkvMuxer.frame_count`. frame_count: u64, + /// Reusable Annex-B conversion buffer for NAL video, kept across frames so + /// the conversion does not allocate (and free) a whole-frame buffer per + /// coded picture — the same reason `MkvMuxer` keeps its `block_group_buf`. + /// A UHD frame is ~310 KB, i.e. an mmap + first-touch page faults + munmap + /// per frame, ~200k times per feature. Cleared (never shrunk) per use, so it + /// settles at the largest frame's size. Taken out of `self` while in use, so + /// the borrow of the converted bytes does not conflict with the `&mut self` + /// the writer needs. + annex_b: Vec, } impl TsMuxer { @@ -84,6 +93,7 @@ impl TsMuxer { video_codec: vec![Codec::Hevc; n], base_pts_ns: None, frame_count: 0, + annex_b: Vec::new(), } } @@ -180,11 +190,20 @@ impl TsMuxer { // start-code ES, never length-prefixed) the bytes pass through // unchanged, so borrow `data` directly rather than copying it; only // NAL video needs an owned Annex-B conversion buffer. - let es_data: std::borrow::Cow<'_, [u8]> = if is_video && self.is_nal_video(track) { - // Size the buffer once for the whole frame. `Vec::new()` re-grew from - // zero capacity on every frame, reallocating repeatedly inside a single - // ~310 KB conversion. The slack covers any prepended parameter sets. - let mut annex_b = Vec::with_capacity(data.len() + 1024); + // Take the reusable conversion buffer out of `self` for the duration of + // this frame: that frees the `&mut self` the writer needs below while the + // converted bytes are still borrowed, and keeps the allocation across + // frames instead of making (and dropping) a whole-frame one per picture. + // It is put back at the end of the function, so a mid-frame error path + // costs only the buffer's capacity, never correctness. + let mut annex_b = std::mem::take(&mut self.annex_b); + let convert = is_video && self.is_nal_video(track); + if convert { + annex_b.clear(); + // Size once for the whole frame; the slack covers any prepended + // parameter sets. `reserve` is a no-op once the buffer has settled at + // the largest frame's size. + annex_b.reserve(data.len() + 1024); if keyframe && !self.params_written[track] { if let Some(ref cp) = self.codec_privates[track] { // avcC and hvcC are DIFFERENT box layouts; parsing one with @@ -221,13 +240,12 @@ impl TsMuxer { // allocations and two full-frame copies. At ~200k frames averaging // ~310 KB of ES on a UHD, that is ~124 GB of pointless memcpy. append_length_prefixed_as_annex_b(&mut annex_b, data); - std::borrow::Cow::Owned(annex_b) - } else { - if is_video { - self.params_written[track] = true; - } - std::borrow::Cow::Borrowed(data) - }; + } else if is_video { + self.params_written[track] = true; + } + // Borrowed from a LOCAL (the taken-out buffer), never from `self`, so the + // `&mut self` writes below are free of it. + let es_data: &[u8] = if convert { &annex_b } else { data }; let pts_90k = if pts_ns >= 0 { (pts_ns as u64).saturating_mul(9) / 100_000 @@ -240,15 +258,32 @@ impl TsMuxer { // access units into multiple PES packets. Each emitted PES carries // the same PTS and starts on its own PUSI packet (only the keyframe // RAI rides the first packet of the first PES). - if is_video || es_data.len() <= MAX_BD_PES_PAYLOAD { - self.write_pes_chain(track, pid, pts_90k, is_video, keyframe, &es_data)?; + // + // The write result is held rather than `?`-propagated so the conversion + // buffer goes back into `self` on every path. + let res = if is_video || es_data.len() <= MAX_BD_PES_PAYLOAD { + self.write_pes_chain(track, pid, pts_90k, is_video, keyframe, es_data) } else { let mut first_pes = true; + let mut res = Ok(()); for chunk in es_data.chunks(MAX_BD_PES_PAYLOAD) { - self.write_pes_chain(track, pid, pts_90k, is_video, keyframe && first_pes, chunk)?; + res = self.write_pes_chain( + track, + pid, + pts_90k, + is_video, + keyframe && first_pes, + chunk, + ); + if res.is_err() { + break; + } first_pes = false; } - } + res + }; + self.annex_b = annex_b; + res?; // A frame that survived the pre-keyframe drop guard above and reached // the writer counts as emitted. `finish()` checks this so a zero-frame // mux fails loudly instead of producing a header-only "success". @@ -1133,4 +1168,41 @@ mod tests { assert_eq!(got.len(), big.len(), "no bytes lost in the PES split"); assert_eq!(got, big, "split audio reassembles byte-for-byte"); } + + /// MEASURED: the Annex-B conversion buffer must be REUSED across video + /// frames, not allocated per frame. Both the allocation's address and its + /// capacity are unchanged after the second and third same-sized frames — if + /// the conversion allocated a fresh `Vec` per frame (the old + /// `Vec::with_capacity(data.len() + 1024)`), the buffer left on the muxer + /// would be empty with zero capacity, and each frame would pay an + /// allocate/first-touch/free cycle over the whole ~310 KB frame. + #[test] + fn annex_b_conversion_buffer_is_reused_across_frames() { + let mut sink: Vec = Vec::new(); + let mut mux = TsMuxer::new(&mut sink, &[VIDEO_PID]); + let idr = fake_hevc_nal(19, 300_000); + mux.write_frame(0, 0, true, &idr).unwrap(); + let cap = mux.annex_b.capacity(); + let ptr = mux.annex_b.as_ptr(); + assert!( + cap >= idr.len(), + "buffer survives the frame with the frame's capacity, got {cap}" + ); + + for i in 1..4 { + let p = fake_hevc_nal(1, 300_000); + mux.write_frame(0, i * 41_000_000, false, &p).unwrap(); + assert_eq!( + mux.annex_b.as_ptr(), + ptr, + "frame {i}: conversion buffer must be the same allocation" + ); + assert_eq!( + mux.annex_b.capacity(), + cap, + "frame {i}: no re-grow once the buffer has settled" + ); + } + mux.finish().unwrap(); + } } diff --git a/src/sector/prefetched.rs b/src/sector/prefetched.rs index d6425bf..6d8ce4e 100644 --- a/src/sector/prefetched.rs +++ b/src/sector/prefetched.rs @@ -142,6 +142,17 @@ impl PrefetchedSectorSource { source: std::io::Error::from(std::io::ErrorKind::InvalidInput), }); } + // A zero alignment is the sibling programming error, and it is worse: the + // producer thread reaches `remaining % unit_align` and panics with a + // divide-by-zero, which `catch_unwind` then reports as + // `DemuxThreadPanicked` — a panic printed through the process hook and a + // misleading error, out of a public constructor that returned `Ok`. Reject + // it here, exactly as `batch_sectors == 0` is rejected. + if unit_align == 0 { + return Err(crate::error::Error::IoError { + source: std::io::Error::from(std::io::ErrorKind::InvalidInput), + }); + } // Accumulate in u64 then clamp: extents can derive from // untrusted nav/MPLS/UDF data, so a naive u32 `sum()` could // panic in debug / wrap in release on a hostile total. The @@ -667,6 +678,29 @@ mod tests { assert!(err.is_err(), "zero batch_sectors must be rejected"); } + /// `unit_align == 0` must be rejected by the constructor, not turned into a + /// divide-by-zero panic on the producer thread. Before the guard, + /// `new_with_events` returned `Ok` and the producer evaluated + /// `remaining % 0`, panicking ("attempt to calculate the remainder with a + /// divisor of zero"); `catch_unwind` then reported the read as + /// `DemuxThreadPanicked` instead of the `InvalidInput` its sibling parameter + /// gets — a panic printed out of a public constructor's own thread. + #[test] + fn zero_unit_align_rejected() { + let res = PrefetchedSectorSource::new_with_events( + EndlessZeroSource, + big_extent(), + 4096, + 0, + None, + None, + ); + let Err(crate::error::Error::IoError { source }) = res else { + panic!("zero unit_align must be rejected with InvalidInput"); + }; + assert_eq!(source.kind(), std::io::ErrorKind::InvalidInput); + } + /// More than 3 sequential direct `read_sectors` calls must succeed. The /// recycle pool seeds PREFETCH_CHANNEL_DEPTH+1 (3) buffers; before /// the fix the direct path dropped each drained buffer, so the 4th