diff --git a/crates/h264/src/decoder.rs b/crates/h264/src/decoder.rs index 5f8941ea..352e3487 100644 --- a/crates/h264/src/decoder.rs +++ b/crates/h264/src/decoder.rs @@ -446,6 +446,9 @@ impl Decoder { } else { let max = sps.max_frame_num(); if sh.frame_num != self.prev_ref_frame_num && sh.frame_num != (self.prev_ref_frame_num + 1) % max { + // A gap is only legal when the SPS allows it (7.4.2.1.1); otherwise the stream is broken + // and filling it would only cost work, so the picture is rejected. + ensure!(sps.gaps_in_frame_num_allowed, "frame_num gap in a stream that does not allow gaps"); // frame_num gap (8.2.5.2): insert "non-existing" frames. Their samples are never used by // conforming streams; copy the latest decoded frame for concealment. let meta = Arc::new(output_meta(sps, pts, false)); diff --git a/crates/h264/src/dpb.rs b/crates/h264/src/dpb.rs index cb064335..5662e299 100644 --- a/crates/h264/src/dpb.rs +++ b/crates/h264/src/dpb.rs @@ -133,7 +133,27 @@ impl Dpb { make_frame: &mut dyn FnMut(u32) -> (FrameRef, Arc), poc_state: &mut dyn FnMut(u32), ) { + if max_frame_num == 0 { + return; + } let mut unused = (prev_ref_frame_num + 1) % max_frame_num; + // Only the last `max_num_ref_frames` "non-existing" frames can survive the sliding window, and + // inserting more than that evicts every earlier short-term reference. Skip straight to that tail + // (keeping the POC state walk for every missing frame_num) so a hostile gap of up to 65535 + // frames costs at most `max_num_ref_frames` frame allocations instead of one per missing frame. + let gap = frame_num.wrapping_add(max_frame_num).wrapping_sub(unused) % max_frame_num; + let keep = u32::try_from(max_num_ref_frames.max(1)).unwrap_or(u32::MAX); + let skip = gap.saturating_sub(keep); + if skip > 0 { + for e in self.entries.iter_mut().filter(|e| e.mark == RefMark::Short) { + e.mark = RefMark::Unused; + } + self.prune(); + for _ in 0..skip { + poc_state(unused); + unused = (unused + 1) % max_frame_num; + } + } let mut guard = 0; while unused != frame_num && guard < max_frame_num { self.sliding_window(unused, max_frame_num, max_num_ref_frames); @@ -418,3 +438,99 @@ impl Default for Dpb { Self::new() } } + +#[cfg(test)] +mod tests { + use super::*; + use crate::picture::Frame; + + fn meta() -> Arc { + Arc::new(OutputMeta { + pts: 0, + key: false, + crop: (0, 0, 0, 0), + full_range: false, + colour_primaries: 2, + transfer_characteristics: 2, + matrix_coefficients: 2, + sar: (1, 1), + draft: false, + }) + } + + fn short(dpb: &mut Dpb, id: u32, frame_num: u32) { + dpb.entries.push(DpbEntry { + frame: Arc::new(Frame::new(id, 0, 1, 1)), + poc: 0, + frame_num, + long_term_frame_idx: 0, + mark: RefMark::Short, + needed_for_output: false, + non_existing: false, + meta: meta(), + }); + } + + /// Fill a gap and return (frames made, POC walk, reference frame_nums left in the DPB). + fn fill(dpb: &mut Dpb, prev: u32, cur: u32, max: u32, refs: usize) -> (u32, Vec, Vec<(u32, bool)>) { + let mut made = 0u32; + let mut walk = Vec::new(); + let mut make = |_f: u32| { + made += 1; + (Arc::new(Frame::new(1000 + made, 0, 1, 1)), meta()) + }; + let mut poc = |f: u32| walk.push(f); + dpb.fill_frame_num_gap(prev, cur, max, refs, &mut make, &mut poc); + let left = dpb.entries.iter().filter(|e| e.mark != RefMark::Unused).map(|e| (e.frame_num, e.non_existing)).collect(); + (made, walk, left) + } + + #[test] + fn huge_gap_makes_only_the_frames_that_stay_referenced() { + let mut dpb = Dpb::new(); + short(&mut dpb, 1, 0); + short(&mut dpb, 2, 1); + let (made, walk, left) = fill(&mut dpb, 1, 65535, 65536, 4); + assert_eq!(made, 4); + // The POC state still walks every missing frame_num, in order. + assert_eq!(walk.len(), 65533); + assert_eq!(walk.first(), Some(&2)); + assert_eq!(walk.last(), Some(&65534)); + assert_eq!(left, vec![(65531, true), (65532, true), (65533, true), (65534, true)]); + } + + #[test] + fn huge_gap_across_wraparound_matches_the_frame_by_frame_result() { + let mut dpb = Dpb::new(); + short(&mut dpb, 1, 65000); + let (made, walk, left) = fill(&mut dpb, 65000, 3, 65536, 2); + assert_eq!(made, 2); + assert_eq!(walk.len(), 538); + assert_eq!(walk.first(), Some(&65001)); + assert_eq!(walk.last(), Some(&2)); + assert_eq!(left, vec![(1, true), (2, true)]); + } + + #[test] + fn small_gap_keeps_earlier_references() { + let mut dpb = Dpb::new(); + short(&mut dpb, 1, 0); + short(&mut dpb, 2, 1); + let (made, walk, left) = fill(&mut dpb, 1, 4, 16, 4); + assert_eq!(made, 2); + assert_eq!(walk, vec![2, 3]); + assert_eq!(left, vec![(0, false), (1, false), (2, true), (3, true)]); + } + + #[test] + fn long_term_references_survive_a_huge_gap() { + let mut dpb = Dpb::new(); + short(&mut dpb, 1, 0); + if let Some(e) = dpb.entries.last_mut() { + e.mark = RefMark::Long; + } + let (made, _, left) = fill(&mut dpb, 0, 1000, 1024, 3); + assert!(made <= 3); + assert_eq!(left, vec![(0, false), (998, true), (999, true)]); + } +} diff --git a/crates/h264/src/synth_tests.rs b/crates/h264/src/synth_tests.rs index b0e0d5d0..8f66cf14 100644 --- a/crates/h264/src/synth_tests.rs +++ b/crates/h264/src/synth_tests.rs @@ -539,3 +539,33 @@ fn poc_type1_reordering() { } } } + +fn gap_stream(gaps: bool) -> Vec { + let mut seed = 7u32; + let a: Vec>> = (0..NMB).map(|_| Some(random_mb(&mut seed))).collect(); + let mut stream = sps(66, 2, gaps, PocType::T2); + stream.extend(pps(false, false)); + stream.extend(slice_nal(&Pic { idr: true, frame_num: 0, ref_idc: 3, mbs: a, ..Default::default() }, false)); + // frame_num 1..=9 are missing. + stream.extend(slice_nal(&Pic { frame_num: 10, ref_idc: 2, mbs: vec![None; NMB], ..Default::default() }, false)); + stream +} + +#[test] +fn frame_num_gap_is_rejected_when_the_sps_forbids_gaps() { + let mut dec = Decoder::with_threads(1); + let res = dec.decode(&gap_stream(false), 0); + let _ = dec.flush(); + assert!(res.is_err() || dec.take_error().is_some(), "a forbidden gap must be reported"); + assert_eq!(dec.stats().frame_num_gaps, 0); +} + +#[test] +fn frame_num_gap_is_filled_when_the_sps_allows_gaps() { + let mut dec = Decoder::with_threads(1); + let mut out = dec.decode(&gap_stream(true), 0).unwrap(); + out.extend(dec.flush()); + assert!(dec.take_error().is_none()); + assert_eq!(out.len(), 2); + assert_eq!(dec.stats().frame_num_gaps, 1); +}