Merge branch 'main' into worktree-agent-a91cd2cf29779b84e

src/aacs/decrypt.rs

@@ -91,7 +91,11 @@ pub fn is_aacs_scrambled(unit: &[u8]) -> bool {
 /// unit looks like clear MPEG-TS. Syncs sit at offset 4 and every 192 bytes
 /// after (4-byte TP_extra_header + 188-byte TS packet). An encrypted body
 /// scrambles all but the first (which lives in the clear 16-byte seed).
-fn ts_syncs_intact(unit: &[u8]) -> bool {
+/// Count the MPEG-TS sync bytes (`0x47`) present at the BD-TS packet stride
+/// (offset 4 and every 192 bytes after). A clear or correctly-decrypted m2ts
+/// unit shows ~one per packet; an encrypted unit, or a non-content unit
+/// decrypted under a key that doesn't apply, shows ~none.
+pub fn ts_sync_count(unit: &[u8]) -> usize {
     let mut count = 0;
     let mut offset = 4;
     while offset < unit.len() {
@@ -100,13 +104,19 @@ fn ts_syncs_intact(unit: &[u8]) -> bool {
         }
         offset += TS_PACKET_LEN;
     }
-    // One sync byte is checked per 192-byte BD-TS packet (at offset 4 of
-    // each). `total` is exactly that packet count; the old
-    // `(len - 4) / TS_PACKET_LEN + 1` over-counted by one for lengths of
-    // the form `4 + k·192` (harmless for the always-6144 aligned unit, but
-    // wrong in general and it biased the majority threshold).
-    let total = unit.len() / TS_PACKET_LEN;
-    count > total / 2
+    count
+}
+
+/// Number of BD-TS packets in the unit — the maximum possible sync count.
+pub fn ts_packet_total(unit: &[u8]) -> usize {
+    // One sync byte per 192-byte BD-TS packet (at offset 4 of each). The old
+    // `(len - 4) / TS_PACKET_LEN + 1` over-counted by one for lengths of the
+    // form `4 + k·192`.
+    unit.len() / TS_PACKET_LEN
+}
+
+fn ts_syncs_intact(unit: &[u8]) -> bool {
+    ts_sync_count(unit) > ts_packet_total(unit) / 2
 }
 
 /// Verify a decrypted unit looks like clear MPEG-TS (sync bytes intact).
@@ -152,6 +162,64 @@ pub fn decrypt_unit(unit: &mut [u8], unit_key: &[u8; 16]) -> bool {
     verify_ts(unit)
 }
 
+/// Fast, NON-MUTATING unit-key validation for the brute-force key search.
+///
+/// `decrypt_unit` pays a full 6128-byte (383-block) CBC decrypt before
+/// `verify_ts` can reject a wrong key — but in a brute scan ~every candidate is
+/// wrong. In CBC the plaintext of block *i* is `AES_dec(C_i) XOR C_{i-1}`, so
+/// the FIRST restored TS sync byte (payload offset 196, which lands in CBC
+/// block 11 of the `unit[16..]` region) can be recovered with a SINGLE block
+/// decrypt instead of 383. A wrong key fails this 1-byte gate ~255/256 of the
+/// time for the cost of one AES block; the rare survivor is then confirmed with
+/// the full [`decrypt_unit`], so the set of accepted keys is bit-for-bit
+/// identical to the slow path.
+///
+/// The caller MUST pass an aligned, already-[`is_aacs_scrambled`] unit
+/// (`unit.len() >= ALIGNED_UNIT_LEN`). The brute pre-filters its units, so the
+/// per-candidate scramble re-scan is intentionally skipped here.
+///
+/// NOTE: this is a search accelerator — it never writes the input and never
+/// participates in the content decrypt path. Aggregate correctness (does a key
+/// validate against *any* of the disc's units) is preserved because a true key
+/// restores offset-196 on every standard BD-TS unit.
+pub fn unit_key_validates(unit: &[u8], unit_key: &[u8; 16]) -> bool {
+    if unit.len() < ALIGNED_UNIT_LEN {
+        return false;
+    }
+    // Per-unit decrypt key: AES-ECB-encrypt the 16-byte plaintext header with
+    // the unit key, XOR with the header (same derivation as `decrypt_unit`).
+    let mut header = [0u8; 16];
+    header.copy_from_slice(&unit[..16]);
+    let derived = aes_ecb_encrypt(unit_key, &header);
+    let mut decrypt_key = [0u8; 16];
+    for i in 0..16 {
+        decrypt_key[i] = derived[i] ^ header[i];
+    }
+
+    // Cheap gate: recover ONLY payload byte 196 (the 2nd BD-TS packet's sync).
+    // The CBC region is `unit[16..]`; payload offset 196 → region offset 180 =
+    // block 11, byte 4. P[11] = AES_dec(C[11]) XOR C[10]; C[10] is raw
+    // ciphertext (no decrypt needed). Constant offsets for the fixed 6144 unit.
+    const SYNC_PAYLOAD_OFF: usize = 196;
+    let region_off = SYNC_PAYLOAD_OFF - 16; // 180
+    let blk = region_off / 16; // 11
+    let byte = region_off % 16; // 4
+    let c11 = 16 + blk * 16; // absolute offset of C[11] in `unit` (=192)
+    let cipher = Aes128::new(GenericArray::from_slice(&decrypt_key));
+    let mut b = GenericArray::clone_from_slice(&unit[c11..c11 + 16]);
+    cipher.decrypt_block(&mut b);
+    let prev = unit[c11 - 16 + byte]; // C[10] byte (region block 10)
+    if b[byte] ^ prev != TS_SYNC {
+        return false;
+    }
+
+    // Survivor (~1/256 of candidates): confirm with the authoritative full
+    // decrypt + verify, so the verdict matches `decrypt_unit` exactly.
+    let mut full = [0u8; ALIGNED_UNIT_LEN];
+    full.copy_from_slice(&unit[..ALIGNED_UNIT_LEN]);
+    decrypt_unit(&mut full, unit_key)
+}
+
 /// Decrypt one aligned unit trying multiple unit keys. Returns the key index that worked.
 pub fn decrypt_unit_try_keys(unit: &mut [u8], unit_keys: &[[u8; 16]]) -> Option<usize> {
     if !is_aacs_scrambled(unit) {

src/aacs/mod.rs

@@ -24,7 +24,7 @@ pub mod variants;
 // AES primitives (aes_ecb_encrypt, aes_ecb_decrypt, aes_cbc_decrypt) are pub(crate) in decrypt.rs.
 pub use decrypt::{
     ALIGNED_UNIT_LEN, decrypt_bus, decrypt_unit, decrypt_unit_full, decrypt_unit_try_keys,
-    is_aacs_scrambled,
+    is_aacs_scrambled, unit_key_validates,
 };
 pub use keydb::{DeviceKey, DiscEntry, HostCert, KeyDb};
 pub use keys::probe;

src/mux/codec/hevc.rs

@@ -21,6 +21,15 @@ const NAL_BLA_W_LP: u8 = 16;
 const NAL_RSV_IRAP_VCL23: u8 = 23;
 
 pub struct HevcParser {
+    // First-seen parameter set of each type → seeds the MKV codecPrivate (hvcC).
+    // This is the ONLY copy the player gets out-of-band, and a player re-applies
+    // it at every keyframe (ffmpeg's hvcC→Annex-B insertion). A stream may
+    // redefine a parameter set mid-title under the SAME id with a different body
+    // (Fight Club redefines PPS id 0 partway through). Any occurrence whose body
+    // DIFFERS from this codecPrivate copy must therefore be emitted IN-BAND at
+    // each point it appears (i.e. at every keyframe of the redefined segment) so
+    // it overrides the re-applied codecPrivate set; otherwise those frames decode
+    // against the wrong parameter set → CABAC/cu_qp_delta desync.
     vps: Option<Vec<u8>>,
     sps: Option<Vec<u8>>,
     pps: Option<Vec<u8>>,
@@ -42,6 +51,32 @@ impl HevcParser {
     }
 }
 
+/// Handle a VPS/SPS/PPS NAL.
+///
+/// - First of its type → seeds codecPrivate (`first`); stripped from frame data
+///   (the player gets it from hvcC).
+/// - Identical to the codecPrivate copy → stripped (the player already re-applies
+///   it from hvcC at each keyframe; BD streams repeat param sets at every IRAP).
+/// - DIFFERENT body from the codecPrivate copy (a mid-title redefinition of the
+///   same id) → emitted IN-BAND (length-prefixed) at EVERY occurrence, so it
+///   overrides the hvcC copy the player re-applies at each keyframe. Emitting it
+///   only once is not enough — the next keyframe's hvcC re-insertion would revert
+///   it. This matches what a conforming muxer produces and fixes the Fight Club
+///   PPS-id-0 redefinition.
+fn handle_param_set(first: &mut Option<Vec<u8>>, nal: &[u8], frame_data: &mut Vec<u8>) {
+    match first {
+        None => {
+            first.replace(nal.to_vec()); // seeds codecPrivate; stripped here
+        }
+        Some(f) if f.as_slice() == nal => {} // == codecPrivate → player has it
+        Some(_) => {
+            // Differs from codecPrivate → emit in-band so it wins at this AU.
+            frame_data.extend_from_slice(&(nal.len() as u32).to_be_bytes());
+            frame_data.extend_from_slice(nal);
+        }
+    }
+}
+
 impl CodecParser for HevcParser {
     fn parse(&mut self, pes: &PesPacket) -> Vec<Frame> {
         if pes.data.is_empty() {
@@ -83,13 +118,13 @@ impl CodecParser for HevcParser {
 
                     match nal_type {
                         NAL_VPS => {
-                            self.vps = Some(data[nal_start..end].to_vec());
+                            handle_param_set(&mut self.vps, &data[nal_start..end], &mut frame_data)
                         }
                         NAL_SPS => {
-                            self.sps = Some(data[nal_start..end].to_vec());
+                            handle_param_set(&mut self.sps, &data[nal_start..end], &mut frame_data)
                         }
                         NAL_PPS => {
-                            self.pps = Some(data[nal_start..end].to_vec());
+                            handle_param_set(&mut self.pps, &data[nal_start..end], &mut frame_data)
                         }
                         NAL_AUD => {} // Skip access unit delimiters
                         t if (NAL_BLA_W_LP..=NAL_RSV_IRAP_VCL23).contains(&t) => {
@@ -561,6 +596,70 @@ mod tests {
         );
     }
 
+    // --- parameter-set redefinition (Fight Club bug) ---
+
+    /// A parameter set REDEFINED mid-stream (same id, different body) must be
+    /// emitted INLINE so the decoder re-activates it. Fight Club redefines PPS
+    /// id 0 partway through the title; the old parser kept only the first PPS,
+    /// so the second segment decoded against the wrong PPS (CABAC desync).
+    #[test]
+    fn redefined_pps_emitted_inline() {
+        let mut parser = HevcParser::new();
+        let pps = |body: u8| {
+            let mut v = vec![0x00, 0x00, 0x01];
+            v.extend_from_slice(&hevc_nal_header(34)); // PPS
+            v.extend_from_slice(&[body, body]);
+            v
+        };
+        let slice = || {
+            let mut v = vec![0x00, 0x00, 0x01];
+            v.extend_from_slice(&hevc_nal_header(1)); // TRAIL_R
+            v.extend_from_slice(&[0x10, 0x20]);
+            v
+        };
+        // count PPS (type 34) NALs in length-prefixed frame data
+        let count_pps = |fd: &[u8]| {
+            let (mut n, mut o) = (0usize, 0usize);
+            while o + 4 <= fd.len() {
+                let len =
+                    u32::from_be_bytes([fd[o], fd[o + 1], fd[o + 2], fd[o + 3]]) as usize;
+                o += 4;
+                if o < fd.len() && (fd[o] >> 1) & 0x3F == 34 {
+                    n += 1;
+                }
+                o += len;
+            }
+            n
+        };
+
+        // PES1: first PPS-A → seeds codecPrivate, stripped from frame.
+        let mut d = pps(0xAA);
+        d.extend(slice());
+        let f = parser.parse(&make_pes(d, Some(0)));
+        assert_eq!(count_pps(&f[0].data), 0, "first PPS goes to codecPrivate");
+
+        // PES2: PPS-B (redefinition, different body) → emitted INLINE.
+        let mut d = pps(0xBB);
+        d.extend(slice());
+        let f = parser.parse(&make_pes(d, Some(1)));
+        assert_eq!(count_pps(&f[0].data), 1, "redefined PPS must be inline");
+
+        // PES3: PPS-B repeated — still differs from codecPrivate(A), so emitted
+        // AGAIN. Every keyframe of the redefined segment must carry it, because
+        // the player re-applies the hvcC (codecPrivate) copy at each keyframe;
+        // emitting once would be reverted at the next keyframe.
+        let mut d = pps(0xBB);
+        d.extend(slice());
+        let f = parser.parse(&make_pes(d, Some(2)));
+        assert_eq!(count_pps(&f[0].data), 1, "redefined PPS re-emitted every occurrence");
+
+        // PES4: back to PPS-A (== codecPrivate) → stripped (hvcC supplies it).
+        let mut d = pps(0xAA);
+        d.extend(slice());
+        let f = parser.parse(&make_pes(d, Some(3)));
+        assert_eq!(count_pps(&f[0].data), 0, "occurrence equal to codecPrivate stripped");
+    }
+
     // --- empty PES ---
 
     #[test]

src/mux/mkv.rs

@@ -212,6 +212,11 @@ pub struct MkvMuxer<W: Write + Seek> {
     cluster_size_pos: u64,
     cluster_ts_ms: i64,
     base_pts_ms: Option<i64>,
+    /// Last block timecode (ms, relative to base_pts) written PER TRACK, to
+    /// enforce strictly-monotonic per-track timestamps — players/ffmpeg reject
+    /// non-monotonic DTS, and some audio PES PTS land on the same millisecond
+    /// (or tick back 1ms from rounding).
+    last_pts_ms: std::collections::HashMap<usize, i64>,
     cues: Vec<CuePoint>,
     frame_count: u64,
     seek_fixups: Vec<SeekPositionFixup>,
@@ -229,6 +234,19 @@ const CLUSTER_DURATION_MS: i64 = 5000;
 /// the `as i16` cast can never wrap.
 const MAX_BLOCK_REL_MS: i64 = i16::MAX as i64;
 
+/// Force a per-track block timestamp to be strictly later than the previous one
+/// written for that track. `prev` is the last timestamp for the track (`None`
+/// for the first frame). Fixes non-monotonic DTS: some audio PES PTS truncate to
+/// the same millisecond as the prior frame (or tick back 1ms from rounding),
+/// which ffmpeg/strict players reject. The nudge is at most a few ms — sub-frame
+/// and inaudible — and never moves a timestamp earlier.
+fn monotonic_ts(prev: Option<i64>, pts_ms: i64) -> i64 {
+    match prev {
+        Some(p) => pts_ms.max(p.saturating_add(1)),
+        None => pts_ms,
+    }
+}
+
 impl<W: Write + Seek> MkvMuxer<W> {
     /// Create a new MKV muxer: writes EBML header, Segment start, Info, Tracks, Chapters.
     pub fn new(
@@ -425,6 +443,7 @@ impl<W: Write + Seek> MkvMuxer<W> {
             cluster_size_pos: 0,
             cluster_ts_ms: 0,
             base_pts_ms: None,
+            last_pts_ms: std::collections::HashMap::new(),
             cues: Vec::new(),
             frame_count: 0,
             seek_fixups,
@@ -453,6 +472,13 @@ impl<W: Write + Seek> MkvMuxer<W> {
         let base = *self.base_pts_ms.get_or_insert(raw_ms);
         let pts_ms = raw_ms - base;
 
+        // Enforce strictly-monotonic per-track block timestamps. Some audio PES
+        // PTS truncate to the same millisecond as the previous frame (or, rarely,
+        // tick back 1ms), which surfaces as "non-monotonic DTS" and is rejected
+        // by ffmpeg/strict players. Nudge to prev+1ms — sub-frame, inaudible,
+        // and A/V sync is unaffected at millisecond granularity.
+        let pts_ms = monotonic_ts(self.last_pts_ms.get(&track_idx).copied(), pts_ms);
+
         // Cluster boundaries normally coincide with a video keyframe so every
         // Cues entry resolves to a seekable IDR at the cluster start.
         let is_video_key = keyframe && track_idx == 0;
@@ -482,6 +508,10 @@ impl<W: Write + Seek> MkvMuxer<W> {
             self.start_cluster(pts_ms)?;
         }
 
+        // Committed to writing this frame — record its (monotonic) timestamp so
+        // the next block on this track is forced strictly later.
+        self.last_pts_ms.insert(track_idx, pts_ms);
+
         let relative_ts = (pts_ms - self.cluster_ts_ms) as i16;
         match duration_ns {
             Some(dur_ns) => {
@@ -829,6 +859,33 @@ mod tests {
         );
     }
 
+    #[test]
+    fn monotonic_ts_forces_strictly_increasing() {
+        // First frame passes through unchanged.
+        assert_eq!(monotonic_ts(None, 1000), 1000);
+        // A repeated millisecond is nudged to prev+1.
+        assert_eq!(monotonic_ts(Some(1000), 1000), 1001);
+        // A backwards tick is nudged forward, never earlier.
+        assert_eq!(monotonic_ts(Some(1001), 1000), 1002);
+        // A genuine advance is left alone.
+        assert_eq!(monotonic_ts(Some(1000), 1040), 1040);
+        // Simulate a stream of audio PTS that round to dup/back-tick ms and
+        // confirm the emitted sequence is strictly increasing.
+        let raw = [1000i64, 1000, 1000, 999, 1032, 1032, 1064];
+        let mut prev: Option<i64> = None;
+        let mut out = Vec::new();
+        for &p in &raw {
+            let t = monotonic_ts(prev, p);
+            out.push(t);
+            prev = Some(t);
+        }
+        assert!(
+            out.windows(2).all(|w| w[1] > w[0]),
+            "not strictly monotonic: {out:?}"
+        );
+        assert_eq!(out, [1000, 1001, 1002, 1003, 1032, 1033, 1064]);
+    }
+
     #[test]
     fn mkv_multiple_tracks() {
         let buf = Cursor::new(Vec::new());