Merge pull request #994 from SusieeTheLinuxUser/fix/928-h264-idct-overflow

Video: clamp H.264 IDCT inputs and temporal-direct mvs to the spec ranges (#928)

crates/video/src/h264/slicedec.rs

@@ -570,6 +570,19 @@ impl<'a> SliceDecoder<'a> {
         (refs, mvs, false)
     }
 
+    /// Temporal direct motion vectors (8.4.1.2.1): mvL0 = (mvCol * DistScaleFactor + 128) >> 8
+    /// and mvL1 = mvL0 − mvCol. Motion vectors are always within the spec's −2^15 to 2^15 − 1
+    /// range, which crafted streams leave (stored mvs wrap by design), so the derivation runs in
+    /// i32 and clamps to that range instead of overflowing.
+    fn temporal_direct_mvs(dsf: i32, mv_col: [i16; 2]) -> ([i16; 2], [i16; 2]) {
+        const MV_MIN: i32 = -(1 << 15);
+        const MV_MAX: i32 = (1 << 15) - 1;
+        // mvL1 comes from the unclamped scaled value, as in the spec; only the results are clamped.
+        let m0 = mv_col.map(|c| (dsf * i32::from(c) + 128) >> 8);
+        let clamp = |v: i32| v.clamp(MV_MIN, MV_MAX) as i16;
+        (m0.map(clamp), [clamp(m0[0] - i32::from(mv_col[0])), clamp(m0[1] - i32::from(mv_col[1]))])
+    }
+
     /// Derive direct-mode motion for the 8x8 blocks in `b8_mask` of the current MB.
     fn derive_direct(&mut self, b8_mask: u8) -> Result<()> {
         let spatial = self.sh.direct_spatial_mv_pred;
@@ -604,13 +617,7 @@ impl<'a> SliceDecoder<'a> {
                         }
                     };
                     let (dsf, ltr) = self.tdirect.get(ref0 as usize).copied().unwrap_or((256, true));
-                    let (mv0, mv1) = if ltr {
+                    let (mv0, mv1) = if ltr { (mv_col, [0, 0]) } else { Self::temporal_direct_mvs(dsf, mv_col) };
-                        (mv_col, [0, 0])
-                    } else {
-                        let s = |c: i16| ((dsf * c as i32 + 128) >> 8) as i16;
-                        let m0 = [s(mv_col[0]), s(mv_col[1])];
-                        (m0, [m0[0] - mv_col[0], m0[1] - mv_col[1]])
-                    };
                     ([ref0, 0], [mv0, mv1])
                 };
                 self.set_motion(x, y, unit, unit, refs, mvs);
@@ -1474,7 +1481,38 @@ impl<'a> SliceDecoder<'a> {
 
 #[cfg(test)]
 mod tests {
-    use super::skip_run_fits;
+    use super::*;
+
+    /// Temporal direct (8.4.1.2.1) with co-located mvs at the i16 extremes: crafted streams get
+    /// there because stored mvs wrap by design (mvd is truncated, the mvp add wraps), so the
+    /// derivation must stay within the spec's −2^15 to 2^15 − 1 motion-vector range instead of
+    /// overflowing (it panicked in debug/overflow-checks builds).
+    #[test]
+    fn temporal_direct_mvs_at_i16_extremes_stay_in_the_spec_range() {
+        // DistScaleFactor extremes from the clamp in SliceDecoder::new (±1024/1023).
+        for dsf in [-1024, -1, 0, 1, 1023] {
+            for c0 in [i16::MAX, i16::MIN, 0, 40, -4] {
+                for c1 in [i16::MAX, i16::MIN, 0, 40, -4] {
+                    let (m0, m1) = SliceDecoder::temporal_direct_mvs(dsf, [c0, c1]);
+                    for mv in m0.into_iter().chain(m1) {
+                        assert!((-(1 << 15)..=(1 << 15) - 1).contains(&i32::from(mv)));
+                    }
+                }
+            }
+        }
+        // mvL0 = (mvCol * DistScaleFactor + 128) >> 8, mvL1 = mvL0 − mvCol (8.4.1.2.1).
+        let (m0, m1) = SliceDecoder::temporal_direct_mvs(-1024, [-32768, 32767]);
+        assert_eq!(m0, [32767, -32768]);
+        assert_eq!(m1, [32767, -32768]);
+        // mvL1 is derived from the unclamped mvL0: (512 * 32767 + 128) >> 8 = 65534, minus 32767.
+        let (m0, m1) = SliceDecoder::temporal_direct_mvs(512, [32767, 0]);
+        assert_eq!(m0, [32767, 0]);
+        assert_eq!(m1, [32767, 0]);
+        // DistScaleFactor 256 (tb == td) reproduces the co-located mv: mvL1 = 0.
+        let (m0, m1) = SliceDecoder::temporal_direct_mvs(256, [40, -4]);
+        assert_eq!(m0, [40, -4]);
+        assert_eq!(m1, [0, 0]);
+    }
 
     #[test]
     fn skip_run_must_fit_remaining_macroblocks_without_overflow() {

crates/video/src/h264/transform.rs

@@ -91,6 +91,14 @@ fn clip_u8(v: i32) -> u8 {
     v.clamp(0, 255) as u8
 }
 
+/// Clamp a dequantised coefficient to the spec's r range: the inverse-transform output is
+/// unspecified when r falls outside −2^15 to 2^15 − 1 (8.5.12.1, likewise 8.5.13.1), and crafted
+/// streams get there because scale4/scale8 wrap. Clamping keeps the butterflies below overflow.
+#[inline(always)]
+fn r_clamp(v: i32) -> i32 {
+    v.clamp(-(1 << 15), (1 << 15) - 1)
+}
+
 /// Inverse 4x4 transform of raster-order `d` and add to `dst` (8.5.12.2, 8.5.14).
 ///
 /// Rows first, then columns (the spec's order: the `>> 1` steps make it matter). The column
@@ -99,10 +107,11 @@ pub fn idct4_add(d: &[i32; 16], dst: &mut [u8], stride: usize) {
     let mut t = [[0i32; 4]; 4];
     for (i, row) in t.iter_mut().enumerate() {
         let r = &d[i * 4..i * 4 + 4];
-        let e = r[0] + r[2];
+        let (r0, r1, r2, r3) = (r_clamp(r[0]), r_clamp(r[1]), r_clamp(r[2]), r_clamp(r[3]));
-        let f = r[0] - r[2];
+        let e = r0 + r2;
-        let g = (r[1] >> 1) - r[3];
+        let f = r0 - r2;
-        let h = r[1] + (r[3] >> 1);
+        let g = (r1 >> 1) - r3;
+        let h = r1 + (r3 >> 1);
         *row = [e + h, f + g, f - g, e - h];
     }
     let mut o = [[0i32; 4]; 4];
@@ -163,7 +172,7 @@ fn idct8_1d(d: [i32; 8]) -> [i32; 8] {
 pub fn idct8_add(d: &[i32; 64], dst: &mut [u8], stride: usize) {
     let mut t = [[0i32; 8]; 8];
     for (row, src) in t.iter_mut().zip(d.as_chunks::<8>().0) {
-        *row = idct8_1d(*src);
+        *row = idct8_1d(std::array::from_fn(|k| r_clamp(src[k])));
     }
     let mut o = [[0i32; 8]; 8];
     for j in 0..8 {
@@ -292,6 +301,35 @@ mod tests {
         assert_eq!(scale8(1, 320, 30), (320 + 1) >> 1);
     }
 
+    /// Crafted streams scale levels far outside the spec's r range (scale4/scale8 wrap by
+    /// design, so r can be any i32): the spec leaves the inverse-transform output unspecified
+    /// when r falls outside −2^15 to 2^15 − 1 (8.5.12.1), so the IDCT must clamp its input into
+    /// the defined range instead of overflowing the butterfly arithmetic (it panicked in
+    /// debug/overflow-checks builds).
+    #[test]
+    fn idct4_huge_coefficients_are_clamped_to_the_spec_r_range() {
+        let mut d = [0i32; 16];
+        for (i, v) in d.iter_mut().enumerate() {
+            *v = if i % 2 == 0 { i32::MAX } else { i32::MIN };
+        }
+        let (mut a, mut b) = ([0u8; 16], [0u8; 16]);
+        idct4_add(&[i32::MAX; 16], &mut a, 4);
+        idct4_add(&[32767; 16], &mut b, 4);
+        assert_eq!(a, b);
+        let (mut a, mut b) = ([0u8; 16], [0u8; 16]);
+        idct4_add(&d, &mut a, 4);
+        idct4_add(&std::array::from_fn(|i| if i % 2 == 0 { 32767 } else { -32768 }), &mut b, 4);
+        assert_eq!(a, b);
+    }
+
+    #[test]
+    fn idct8_huge_coefficients_are_clamped_to_the_spec_r_range() {
+        let (mut a, mut b) = ([0u8; 64], [0u8; 64]);
+        idct8_add(&[i32::MIN; 64], &mut a, 8);
+        idct8_add(&[-32768; 64], &mut b, 8);
+        assert_eq!(a, b);
+    }
+
     /// Corrupt streams can carry huge coefficient levels: dequantisation wraps (as release
     /// builds always did) instead of panicking with "attempt to multiply with overflow" in
     /// debug builds (found by mutation fuzzing an H.264 transport stream).