mux: restore three drive_mux regressions from the mux hoist

The mux drive loop hoisted out of autorip/CLI into libfreemkv's
drive_mux dropped three guarantees the hand-rolled loops enforced:

A. Header-buffer OOM cap. The header pump buffered every PES frame
   until codec_private resolved with no byte accounting, so a title
   whose video header never resolves (damaged/undecryptable HEVC on
   marginal UHD) buffered the whole 30-90 GB title into RAM until
   OOM-kill. Restore the 512 MiB HEADER_BUFFER_CAP_BYTES cap autorip's
   run_mux enforced; once exceeded, fail fast with the same
   Error::MkvInvalid the headers-never-resolved gate already returns.

B. Watchdog feed during the header-frame drain. After headers resolve
   the buffered header frames are flushed to the sink, but
   on_write_progress (the sole feed for autorip's hard watchdog) fired
   only in the steady-state loop, leaving the watchdog unfed during a
   long/slow drain and able to false-escalate. Feed on_write_progress
   per buffered frame during the drain, matching the steady-state path.

C. Per-title CSS key on MuxInput::Session. The Session arm set keys via
   disc.decrypt_keys() unconditionally, which returns the largest
   title's VTS CSS key; muxing a bonus title in a different VTS
   descrambled with the wrong key (DiscStream's per-title crack only
   fires on DecryptKeys::None). Special-case DVD to None so the pipeline
   cracks the correct per-title key, matching resolve.rs::input(). AACS
   and clear discs still resolve from decrypt_keys().

Adds mutation-verified synthetic tests for each (no live drive).

src/mux/driver.rs

@@ -46,6 +46,17 @@ use super::resolve::{InputOptions, StreamUrl, build_iso_pipeline, input, output,
 /// Mirrors autorip's `MUX_SEND_DEADLINE_SECS`.
 const MUX_SEND_DEADLINE_SECS: u64 = 60;
 
+/// Ceiling on bytes buffered while waiting for `headers_ready()` to resolve.
+/// Normally that resolves after a handful of frames; a damaged/undecryptable
+/// title whose video `codec_private` never resolves would otherwise grow the
+/// pre-headers buffer without bound (the whole 30-90 GB title, one PES frame at
+/// a time) until the process is OOM-killed. 512 MiB is far more than any real
+/// codec-private resolution needs but small enough to fail fast rather than
+/// swap the box to death. Once exceeded the mux is refused exactly as the
+/// headers-never-resolved gate refuses it (`Error::MkvInvalid`). Mirrors
+/// autorip's pre-refactor `HEADER_BUFFER_CAP_BYTES`.
+const HEADER_BUFFER_CAP_BYTES: usize = 512 * 1024 * 1024;
+
 /// Where [`mux_stream`] reads its PES frames from. The driver owns the
 /// construction of the underlying [`Stream`] so consumers stop hand-rolling
 /// `DiscStream::new` / `build_iso_pipeline` / `input()`.
@@ -268,7 +279,9 @@ pub fn mux_stream(
                     .meta_title
                     .clone()
                     .unwrap_or_else(|| disc.volume_id.clone());
-                (title, disc.content_format, disc.decrypt_keys(), playlist)
+                // DVD CSS is per-VTS: resolve the per-title key via the pipeline
+                // (see `session_mux_keys`), never the whole-disc `decrypt_keys()`.
+                (title, disc.content_format, session_mux_keys(disc), playlist)
             };
             // A missing staged reader ("already consumed" / never staged) is a
             // clean error, not a panic (contract Q2).
@@ -361,6 +374,25 @@ pub fn mux_stream(
 ///
 /// Sector numbers are the library's `u64`; the `MuxEvents` LBA hooks take `u32`
 /// (the disc's LBA space), so they are narrowed with `as u32`.
+/// Decrypt keys for the live `Session` mux of `disc`.
+///
+/// A DVD is handed [`DecryptKeys::None`] so [`DiscStream::new`](crate::mux::DiscStream)
+/// cracks the CORRECT per-title CSS key: CSS title keys are per-VTS, and
+/// [`Disc::decrypt_keys`](crate::disc::Disc::decrypt_keys) returns the LARGEST
+/// title's VTS key — muxing a bonus title in a DIFFERENT VTS with it descrambles
+/// to corrupt MPEG-PS. `resolve_dvd_title_key`'s crack-guard only fires when
+/// `keys` is `None`, so passing the whole-disc key would skip the per-title
+/// crack entirely. This mirrors the sibling `resolve.rs::input()` DVD
+/// special-case (and the pre-refactor CLI `pipe.rs`). AACS / genuinely-clear
+/// discs resolve from `decrypt_keys()` with no read.
+fn session_mux_keys(disc: &crate::disc::Disc) -> DecryptKeys {
+    if matches!(disc.format, crate::disc::DiscFormat::Dvd) {
+        DecryptKeys::None
+    } else {
+        disc.decrypt_keys()
+    }
+}
+
 fn reader_event_fn(events: Arc<dyn MuxEvents>) -> crate::sector::prefetched::EventFn {
     Box::new(move |e: Event| match e.kind {
         EventKind::BytesRead { bytes, total } => events.on_read_progress(bytes, total),
@@ -420,6 +452,7 @@ fn drive_mux(
     // can't write a track header without codec init data. The loop breaks on
     // EOF/None too, so the gate below re-checks.
     let mut buffered: Vec<PesFrame> = Vec::new();
+    let mut buffered_bytes: usize = 0;
     while !stream.headers_ready() {
         if halt.is_cancelled() {
             return Ok(MuxOutcome {
@@ -432,7 +465,19 @@ fn drive_mux(
             });
         }
         match stream.read()? {
-            Some(frame) => buffered.push(frame),
+            Some(frame) => {
+                buffered_bytes = buffered_bytes.saturating_add(frame.data.len());
+                buffered.push(frame);
+                // Bounded header buffer: a title whose codec_private never
+                // resolves would otherwise buffer the entire (tens-of-GB)
+                // stream into RAM until OOM. Treat cap-exceeded identically to
+                // "headers never resolved" — fail fast with the SAME
+                // `Error::MkvInvalid` the header gate below returns, instead of
+                // swapping the box to death.
+                if buffered_bytes > HEADER_BUFFER_CAP_BYTES {
+                    return Err(Error::MkvInvalid.into());
+                }
+            }
             None => break,
         }
     }
@@ -482,6 +527,12 @@ fn drive_mux(
             interrupted = true;
             break;
         }
+        // Feed write-side progress during the drain exactly as the steady-state
+        // loop below does. The watchdog is fed only from `on_write_progress`;
+        // no `stream.read()` runs here (reads finished in the header pump), so
+        // omitting this leaves the watchdog unfed for the whole header-drain
+        // phase and it can false-escalate on a long/slow drain.
+        events.on_write_progress(bytes.load(Ordering::Relaxed), total_bytes);
     }
 
     // Then the remainder of the stream.
@@ -593,6 +644,10 @@ mod tests {
         codec_private_ready: bool,
         /// If set, `read()` cancels this halt once `reads` reaches the value.
         cancel_halt: Option<(Halt, usize)>,
+        /// If set, every successful `read()` bumps this shared counter so a test
+        /// can observe how many frames the pump consumed after the stream was
+        /// moved into `drive_mux`.
+        read_observer: Option<Arc<std::sync::atomic::AtomicUsize>>,
     }
 
     fn audio_stream() -> crate::disc::Stream {
@@ -621,6 +676,7 @@ mod tests {
                 reads: 0,
                 codec_private_ready: true,
                 cancel_halt: None,
+                read_observer: None,
             }
         }
         fn with_frames(mut self, n: usize) -> Self {
@@ -658,6 +714,9 @@ mod tests {
             let f = self.frames.pop_front();
             if f.is_some() {
                 self.reads += 1;
+                if let Some(obs) = &self.read_observer {
+                    obs.fetch_add(1, Ordering::SeqCst);
+                }
             }
             Ok(f)
         }
@@ -1109,4 +1168,150 @@ mod tests {
             "our-phase forensic units must still be read (kept LBAs 1030/1054)"
         );
     }
+
+    // ── Regression A: header-buffer cap fails fast instead of OOM ───────────
+    //
+    // A stream whose headers never resolve but that keeps yielding frames must
+    // be refused once the pre-headers buffer passes `HEADER_BUFFER_CAP_BYTES`,
+    // BEFORE draining the whole (unbounded, OOM-inducing) stream. Frames are
+    // zero-filled `Vec`s (`alloc_zeroed`, lazily paged) so the nominal size is
+    // cheap in RSS. We hand the pump FAR more frames than the cap needs and
+    // assert both that it fails with the same `MkvInvalid` outcome AND that it
+    // stopped EARLY (bounded read count) — the latter is what distinguishes the
+    // capped path from the un-capped one.
+    //
+    // Mutation: removing the cap makes the pump drain all `MANY` frames (reads
+    // == MANY) before failing at the header gate → the `reads_seen` bound fails.
+    #[test]
+    fn header_buffer_cap_fails_fast_instead_of_oom() {
+        use std::sync::atomic::AtomicUsize;
+        const FRAME: usize = 64 * 1024 * 1024; // 64 MiB per frame
+        let cap_frames = HEADER_BUFFER_CAP_BYTES / FRAME; // 8 frames == cap
+        const MANY: usize = 200; // vastly more than the cap needs
+
+        let reads_seen = Arc::new(AtomicUsize::new(0));
+        let mut fs = FakeStream::new(1).never_ready();
+        fs.read_observer = Some(reads_seen.clone());
+        for i in 0..MANY {
+            fs.frames.push_back(PesFrame {
+                track: 0,
+                pts: i as i64,
+                keyframe: true,
+                data: vec![0u8; FRAME],
+                duration_ns: None,
+                source: None,
+                coding: None,
+            });
+        }
+        let halt = Halt::new();
+        let err = drive_mux(Box::new(fs), "null://", &halt, &NoopEvents, None)
+            .expect_err("over-cap header buffer must fail fast, not OOM");
+        assert_eq!(
+            err.to_string(),
+            format!("E{}", crate::error::E_MKV_INVALID),
+            "cap-exceeded returns the same MkvInvalid as headers-never-resolved"
+        );
+        let reads = reads_seen.load(Ordering::SeqCst);
+        assert!(
+            reads <= cap_frames + 1,
+            "must fail after ~{cap_frames} frames (cap), not drain all {MANY} (read {reads})"
+        );
+    }
+
+    // ── Regression B: watchdog fed during the buffered header drain ─────────
+    //
+    // Headers resolve only after K frames are buffered; those K frames are then
+    // flushed to the sink in the drain loop. Every flushed frame must fire
+    // `on_write_progress` (the sole watchdog feed) — with exactly K frames and
+    // no remainder, ALL `on_write_progress` calls come from the drain.
+    //
+    // Mutation: dropping the `events.on_write_progress(...)` call in the drain
+    // loop leaves the count at 0 → this fails.
+    #[test]
+    fn write_progress_fed_during_header_drain() {
+        const K: usize = 6;
+        let mut fs = FakeStream::new(1).with_frames(K);
+        // Headers stay unresolved until all K frames have been read+buffered.
+        fs.headers_ready_after = K;
+
+        struct WriteProgressCounter {
+            writes: AtomicU64,
+        }
+        impl MuxEvents for WriteProgressCounter {
+            fn on_write_progress(&self, _bytes_written: u64, _bytes_total: u64) {
+                self.writes.fetch_add(1, Ordering::SeqCst);
+            }
+        }
+        let events = WriteProgressCounter {
+            writes: AtomicU64::new(0),
+        };
+        let halt = Halt::new();
+        let out = drive_mux(Box::new(fs), "null://", &halt, &events, None)
+            .expect("K-frame stream muxes cleanly");
+        assert!(out.completed);
+        assert_eq!(
+            events.writes.load(Ordering::SeqCst),
+            K as u64,
+            "each of the K buffered header frames must feed on_write_progress on drain"
+        );
+    }
+
+    // ── Regression C: Session key selection special-cases DVD ───────────────
+    //
+    // `session_mux_keys` is the seam the `MuxInput::Session` arm uses to pick
+    // decrypt keys. A DVD must get `None` (so DiscStream cracks the correct
+    // per-title/per-VTS CSS key); a non-DVD passes `decrypt_keys()` through.
+    // Building a full multi-VTS DiscSession fixture is infeasible in a unit
+    // test, so this asserts the is_dvd→None DECISION at the tightest seam.
+    //
+    // Mutation: reverting to unconditional `disc.decrypt_keys()` returns
+    // `Css{..}` for the DVD case → the `None` assertion fails.
+    fn disc_with_css(format: crate::disc::DiscFormat) -> crate::disc::Disc {
+        crate::disc::Disc {
+            volume_id: "TEST".into(),
+            meta_title: None,
+            format,
+            capacity_sectors: 0,
+            capacity_bytes: 0,
+            layers: 1,
+            titles: Vec::new(),
+            region: crate::disc::DiscRegion::Free,
+            aacs: None,
+            css: Some(crate::css::CssState {
+                title_key: [1, 2, 3, 4, 5],
+                crack_span: None,
+            }),
+            encrypted: true,
+            aacs_error: None,
+            css_error: None,
+            content_format: crate::disc::ContentFormat::MpegPs,
+        }
+    }
+
+    #[test]
+    fn session_mux_keys_uses_none_for_dvd() {
+        // DVD: must be None so the pipeline cracks the correct per-title key,
+        // even though decrypt_keys() would hand back a whole-disc Css key.
+        let dvd = disc_with_css(crate::disc::DiscFormat::Dvd);
+        assert!(
+            matches!(dvd.decrypt_keys(), DecryptKeys::Css { .. }),
+            "precondition: a CSS disc's decrypt_keys() is Css{{..}}"
+        );
+        assert!(
+            matches!(session_mux_keys(&dvd), DecryptKeys::None),
+            "a DVD must be handed None so DiscStream cracks the per-title key"
+        );
+    }
+
+    #[test]
+    fn session_mux_keys_passes_decrypt_keys_for_non_dvd() {
+        // Non-DVD (e.g. BluRay): the whole-disc key IS the per-title key, so it
+        // passes through — proving the special-case keys on FORMAT, not on the
+        // mere presence of a key.
+        let bd = disc_with_css(crate::disc::DiscFormat::BluRay);
+        assert!(
+            matches!(session_mux_keys(&bd), DecryptKeys::Css { .. }),
+            "a non-DVD passes decrypt_keys() through unchanged"
+        );
+    }
 }