craft-audio: bound the work of MP4 audio timestamp lookups, not the stts entry count (#910)

crates/audio/src/read/mp4.rs

@@ -4,11 +4,13 @@
 use crate::pcm::{be_u32, be_u64};
 use crate::{AudioError, Result};
 
-/// The most `stts` entries a sound track may have. Symphonia finds each packet's timestamp by
-/// scanning the table from its start, so a table of one-sample runs makes decoding quadratic in
-/// the packet count. Audio encoders write one run per packet length (a handful of entries); this
-/// bound keeps the worst case at a few hundred million steps. Ports soundcraft #1332.
-const MAX_AUDIO_STTS_ENTRIES: u32 = 1 << 14;
+/// The most table steps symphonia may take finding packet timestamps in an audio track, 2^34
+/// (about 15 s of work). It scans the `stts` table from its first entry for every packet, so a
+/// table of one-sample runs makes decoding quadratic in the packet count; [`stts_steps`] counts
+/// the steps. The bound is on that work, not on entries: a long track whose packet durations vary
+/// (a run change every ten AAC packets for three hours is 1.3e10 steps) stays under it, and a table
+/// of many entries with few packets costs next to nothing. Ports soundcraft #1332.
+const MAX_STTS_STEPS: u128 = 1 << 34;
 
 #[derive(Clone, Copy)]
 enum Level {
@@ -17,10 +19,11 @@ enum Level {
     Trak,
     Mdia,
     Minf,
-    /// `audio`: the sample entry is an audio format, so symphonia will decode its packets
-    /// (whatever the handler says: it builds codec parameters from `stsd` alone).
+    /// `packets`: `Some` with the `stsz` sample count when the sample entry is an audio format, so
+    /// symphonia will decode its packets (whatever the handler says: it builds codec parameters
+    /// from `stsd` alone).
     Stbl {
-        audio: bool,
+        packets: Option<u64>,
     },
     Udta,
     Meta,
@@ -65,13 +68,13 @@ fn atoms(b: &[u8], mut pos: usize, end: usize, level: Level) -> Result<bool> {
             (Level::Moov, b"udta") => atoms(b, body, inner, Level::Udta)?,
             (Level::Trak, b"mdia") => atoms(b, body, inner, Level::Mdia)?,
             (Level::Mdia, b"minf") => atoms(b, body, inner, Level::Minf)?,
-            (Level::Minf, b"stbl") => atoms(b, body, inner, Level::Stbl { audio: entry_is_audio(b, body, inner) })?,
+            (Level::Minf, b"stbl") => atoms(b, body, inner, Level::Stbl { packets: audio_packets(b, body, inner) })?,
             (Level::Meta, b"ilst") => atoms(b, body, inner, Level::Ilst)?,
             (Level::Ilst, _) => atoms(b, body, inner, Level::Tag)?,
-            (Level::Stbl { audio }, b"stts") => {
+            (Level::Stbl { packets }, b"stts") => {
                 let fits = table(b, body + 4, 8, inner)?;
-                if audio && be_u32(b, body + 4).is_some_and(|n| n > MAX_AUDIO_STTS_ENTRIES) {
-                    return Err(AudioError::Malformed("MP4 audio timing table has more runs than packets could need".into()));
+                if packets.is_some_and(|packets| stts_steps(b, body + 4, inner, packets) > MAX_STTS_STEPS) {
+                    return Err(AudioError::Malformed("MP4 audio timing table makes finding packet times too slow".into()));
                 }
                 fits
             }
@@ -99,29 +102,54 @@ const AUDIO_FORMATS: [&[u8; 4]; 18] = [
     b"fl32", b"fl64",
 ];
 
-/// Whether the `stbl` atom whose children lie in `pos..end` describes audio: its `stsd` may come
-/// before or after the timing tables, so the children are walked once more for it. Symphonia keeps
-/// the last `stsd` it reads, so the last one decides.
-fn entry_is_audio(b: &[u8], mut pos: usize, end: usize) -> bool {
-    let mut audio = false;
+/// The packet count (`stsz` sample count) of the `stbl` atom whose children lie in `pos..end` when
+/// it describes audio, else `None`: its `stsd` and `stsz` may come before or after the timing
+/// tables, so the children are walked once more. Symphonia keeps the last `stsd` and `stsz` it
+/// reads, so the last of each decides.
+fn audio_packets(b: &[u8], mut pos: usize, end: usize) -> Option<u64> {
+    let (mut audio, mut packets) = (false, 0);
     while pos < end {
-        let (Some(size), Some(kind)) = (be_u32(b, pos), b.get(pos + 4..pos + 8)) else { return audio };
+        let (Some(size), Some(kind)) = (be_u32(b, pos), b.get(pos + 4..pos + 8)) else { break };
         let head = if size == 1 { 16 } else { 8 };
         if kind == b"stsd" {
             // Version and flags, the entry count, then the first entry's size and format.
             audio = b.get(pos + head + 12..pos + head + 16).is_some_and(|f| AUDIO_FORMATS.iter().any(|a| a.as_slice() == f));
+        } else if kind == b"stsz" {
+            // Version and flags, the common sample size, then the count.
+            packets = be_u32(b, pos + head + 8).map_or(0, u64::from);
         }
         let len = match size {
-            0 => return audio,
+            0 => break,
             1 => be_u64(b, pos + 8).unwrap_or(0),
             n => u64::from(n),
         };
         match usize::try_from(len).ok().filter(|&n| n >= head).and_then(|n| pos.checked_add(n)) {
             Some(next) => pos = next,
-            None => return audio,
+            None => break,
         }
     }
-    audio
+    audio.then_some(packets)
+}
+
+/// The steps symphonia takes finding the timestamps of `packets` packets in the `stts` table whose
+/// entry count is at `count`: packet `i` costs the number of entries up to and including its run,
+/// and only the packets that exist (`packets`) are looked up. Entries past the atom's `end` are
+/// never read.
+fn stts_steps(b: &[u8], count: usize, end: usize, packets: u64) -> u128 {
+    let entries = be_u32(b, count).map_or(0, |n| n as usize);
+    let mut left = packets;
+    let mut steps = 0u128;
+    for j in 0..entries {
+        let at = count + 4 + 8 * j;
+        let Some(run) = (at + 8 <= end).then(|| be_u32(b, at)).flatten() else { break };
+        let run = u64::from(run).min(left);
+        steps += u128::from(run) * (j as u128 + 1);
+        left -= run;
+        if left == 0 {
+            break;
+        }
+    }
+    steps
 }
 
 /// A sample table's entries, counted at `count`, must fit in the file: symphonia reserves room for
@@ -153,3 +181,42 @@ fn sample_sizes(b: &[u8], body: usize, end: usize) -> Result<bool> {
     }
     Ok(fits)
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    fn atom(kind: &[u8; 4], body: &[u8]) -> Vec<u8> {
+        [&(8 + body.len() as u32).to_be_bytes()[..], kind, body].concat()
+    }
+
+    /// A `moov` with one `.mp3` track of `packets` packets whose `stts` has `runs` runs of
+    /// `per_run` packets each. Only the atoms the check walks are present.
+    fn track(runs: u32, per_run: u32, packets: u32) -> Vec<u8> {
+        let stsd = atom(b"stsd", &[&[0u8; 4][..], &1u32.to_be_bytes(), &atom(b".mp3", &[0; 28])].concat());
+        let stts_body: Vec<u8> = [0, runs]
+            .iter()
+            .flat_map(|v: &u32| v.to_be_bytes())
+            .chain((0..runs).flat_map(|_| [per_run, 1024].into_iter().flat_map(u32::to_be_bytes)))
+            .collect();
+        let stsz = atom(b"stsz", &[0u32, 1, packets].iter().flat_map(|v| v.to_be_bytes()).collect::<Vec<u8>>());
+        let stbl = atom(b"stbl", &[stsd, atom(b"stts", &stts_body), stsz].concat());
+        atom(b"moov", &atom(b"trak", &atom(b"mdia", &atom(b"minf", &stbl))))
+    }
+
+    #[test]
+    fn the_bound_is_on_work_not_on_entries() {
+        // many entries, few packets: only the packets that exist are looked up
+        assert!(check(&track(100_000, 1, 100), 0).is_ok());
+        // a three-hour track (500 000 AAC packets) whose duration changes every ten packets
+        assert!(check(&track(50_000, 10, 500_000), 0).is_ok());
+        // a run per packet over the same packets: quadratic work, refused
+        assert!(matches!(check(&track(500_000, 1, 500_000), 0), Err(AudioError::Malformed(m)) if m.contains("timing table")));
+        // the same table in a video track is not symphonia's to scan
+        let video = track(500_000, 1, 500_000);
+        let at = video.windows(4).position(|w| w == b".mp3").unwrap();
+        let mut video = video;
+        video[at..at + 4].copy_from_slice(b"avc1");
+        assert!(check(&video, 0).is_ok());
+    }
+}

crates/audio/tests/audio/malformed.rs

@@ -324,7 +324,7 @@ fn mp4(frames: u32, entry: &[u8], swap: impl Fn(&[u8; 4], Vec<u8>) -> Vec<u8>) -
     };
     let ftyp = atom(b"ftyp", b"M4A \0\0\0\0M4A isom");
     let offset = (ftyp.len() + moov(0).len() + 8) as u32;
-    let pcm: Vec<u8> = (0..frames as i16).flat_map(|i| i.wrapping_mul(40).to_le_bytes()).collect();
+    let pcm: Vec<u8> = (0..frames).flat_map(|i| (i as i16).wrapping_mul(40).to_le_bytes()).collect();
     [ftyp, moov(offset), atom(b"mdat", &pcm)].concat()
 }
 
@@ -886,24 +886,23 @@ fn one_sample_runs(n: u32) -> Vec<u8> {
 }
 
 /// soundcraft #1332: symphonia scans an audio track's `stts` table from its start for every
-/// packet, so a table of one-sample runs made decoding quadratic.
+/// packet, so a table of one-sample runs made decoding quadratic. The bound is on that work.
 #[test]
 fn an_audio_timing_table_with_a_run_per_sample_is_bounded() {
-    let frames = 1 << 14;
-    let at_bound = mp3_mp4(frames, |k, a| if k == b"stts" { one_sample_runs(frames) } else { a });
-    assert!(
-        !matches!(decode(&at_bound, None), Err(AudioError::Malformed(m)) if m.contains("timing table")),
-        "a table at the bound is not refused for its size"
-    );
-    let past = mp3_mp4(frames + 1, |k, a| if k == b"stts" { one_sample_runs(frames + 1) } else { a });
-    assert!(matches!(decode(&past, None), Err(AudioError::Malformed(m)) if m.contains("timing table")), "one run past the bound is refused");
+    let frames = 200_000;
+    let hostile = mp3_mp4(frames, |k, a| if k == b"stts" { one_sample_runs(frames) } else { a });
+    let r = decode(&hostile, None);
+    assert!(matches!(&r, Err(AudioError::Malformed(m)) if m.contains("timing table")), "{:?}", r.map(|x| x.0));
+    // a table of many runs over few packets is cheap to scan, so it is not refused for its size
+    let few = mp3_mp4(100, |k, a| if k == b"stts" { one_sample_runs(100_000) } else { a });
+    assert!(!matches!(decode(&few, None), Err(AudioError::Malformed(m)) if m.contains("timing table")));
 }
 
 /// The bound follows the sample entry, as symphonia decodes by it, not the handler.
 #[test]
 fn the_audio_timing_bound_follows_the_sample_entry_not_the_handler() {
     let video_handler = atom(b"hdlr", &[0u32, 0, u32::from_be_bytes(*b"vide"), 0, 0, 0, 0].map(u32::to_be_bytes).concat());
-    let frames = (1 << 14) + 1;
+    let frames = 200_000;
     let f = mp3_mp4(frames, |k, a| match k {
         b"stts" => one_sample_runs(frames),
         b"hdlr" => video_handler.clone(),
@@ -917,7 +916,7 @@ fn the_audio_timing_bound_follows_the_sample_entry_not_the_handler() {
 fn the_last_sample_entry_decides_whether_the_timing_bound_applies() {
     let video = sample_entry(b"avc1", 0, &[], &[]);
     let audio_stsd = atom(b"stsd", &[&[0, 0, 0, 0, 0, 0, 0, 1][..], &sample_entry(b".mp3", 0, &[], &[])].concat());
-    let frames = (1 << 14) + 1;
+    let frames = 200_000;
     let f = mp4(frames, &video, |k, a| if k == b"stts" { [audio_stsd.clone(), one_sample_runs(frames)].concat() } else { a });
     assert!(matches!(decode(&f, None), Err(AudioError::Malformed(m)) if m.contains("timing table")));
 }