Merge pull request #238 from dylanmtaylor/security/media-allocation-bounds

security: bound craft-media metadata, sample reservation, header stripping and cue bookkeeping (CodeRabbit)

crates/media/src/isobmff/demux.rs

@@ -10,6 +10,11 @@ use std::collections::HashMap;
 const MAX_SAMPLES: usize = 1 << 26;
 /// Hard cap on the size of an in-memory `moov`/`moof` box.
 const MAX_HEADER_BOX: u64 = 1 << 30;
+/// Most text all metadata entries may hold together; entries past it are dropped. One `keys`
+/// name can be named by many `ilst` items, each a copy, so the file's size does not bound this.
+const MAX_METADATA_BYTES: usize = 4 << 20;
+/// What an entry costs beyond its text, so that empty entries are bounded too.
+const METADATA_ENTRY_COST: usize = 64;
 
 /// What is left of the file-wide limits while tracks and fragments are parsed.
 struct SampleBudget {
@@ -305,6 +310,7 @@ pub fn open<S: ByteSource>(src: S) -> Result<Mp4File> {
     }
     let mut trex: HashMap<u32, Trex> = HashMap::new();
     let mut trak_payloads = Vec::new();
+    let mut meta_left = MAX_METADATA_BYTES;
     for b in boxes(&moov) {
         let b = b?;
         match &b.kind.0 {
@@ -334,8 +340,8 @@ pub fn open<S: ByteSource>(src: S) -> Result<Mp4File> {
                     }
                 }
             }
-            b"udta" => parse_udta(b.payload, &mut file.metadata),
-            b"meta" => parse_meta(b.payload, &mut file.metadata),
+            b"udta" => parse_udta(b.payload, &mut file.metadata, &mut meta_left),
+            b"meta" => parse_meta(b.payload, &mut file.metadata, &mut meta_left),
             _ => {}
         }
     }
@@ -764,7 +770,20 @@ fn parse_stbl(stbl: &[u8], t: &mut Track, is_qt: bool, file_len: u64, budget: &m
         if sample_count > budget.samples {
             return Err(Error::TooLarge("sample count"));
         }
-        samples.reserve(sample_count);
+        // Reserve only what the chunk map can place. A declared count far past it would keep
+        // capacity the budget below never charges, track after track.
+        let mut mapped = 0usize;
+        let mut map_i = 0usize;
+        for ci in 0..chunks.len() {
+            while map_i + 1 < stsc.len() && stsc[map_i + 1].0 as usize <= ci + 1 {
+                map_i += 1;
+            }
+            mapped = mapped.saturating_add(stsc.get(map_i).map_or(0, |e| e.1 as usize));
+            if mapped >= sample_count {
+                break;
+            }
+        }
+        samples.try_reserve(sample_count.min(mapped)).map_err(|_| Error::TooLarge("sample count"))?;
         let mut stss_i = 0usize;
         'chunks: for (ci, &off) in chunks.iter().enumerate() {
             while stsc_i + 1 < stsc.len() && stsc[stsc_i + 1].0 as usize <= ci + 1 {
@@ -903,20 +922,42 @@ fn parse_moof(
     Ok(())
 }
 
-fn parse_udta(p: &[u8], md: &mut Metadata) {
+/// Takes an entry of `key_len` and `value_len` bytes from `left`; false once it would not fit.
+fn charge_metadata(left: &mut usize, key_len: usize, value_len: usize) -> bool {
+    let cost = key_len.saturating_add(value_len).saturating_add(METADATA_ENTRY_COST);
+    match left.checked_sub(cost) {
+        Some(rest) => {
+            *left = rest;
+            true
+        }
+        None => false,
+    }
+}
+
+/// Length of `String::from_utf8_lossy(bytes)` without building it: each invalid sequence
+/// becomes one three-byte U+FFFD.
+fn lossy_len(bytes: &[u8]) -> usize {
+    bytes.utf8_chunks().map(|c| c.valid().len() + if c.invalid().is_empty() { 0 } else { 3 }).sum()
+}
+
+fn parse_udta(p: &[u8], md: &mut Metadata, left: &mut usize) {
     for b in boxes(p).map_while(|b| b.ok()) {
         if b.kind == b"meta" {
-            parse_meta(b.payload, md);
+            parse_meta(b.payload, md, left);
         } else if b.kind.0[0] == 0xA9 && b.payload.len() >= 4 {
             let n = u16::from_be_bytes([b.payload[0], b.payload[1]]) as usize;
-            if let Some(s) = b.payload.get(4..4 + n) {
-                md.entries.push((b.kind.to_string_lossy(), String::from_utf8_lossy(s).into_owned()));
+            // Charged at the stored text's length, before the value is converted.
+            let key = b.kind.to_string_lossy();
+            if let Some(s) = b.payload.get(4..4 + n)
+                && charge_metadata(left, key.len(), lossy_len(s))
+            {
+                md.entries.push((key, String::from_utf8_lossy(s).into_owned()));
             }
         }
     }
 }
 
-fn parse_meta(p: &[u8], md: &mut Metadata) {
+fn parse_meta(p: &[u8], md: &mut Metadata, left: &mut usize) {
     // QuickTime `meta` is a plain box; ISO `meta` is a full box.
     let body = if p.len() >= 8 && &p[4..8] == b"hdlr" { p } else { p.get(4..).unwrap_or(&[]) };
     let mut keys: Vec<String> = Vec::new();
@@ -936,13 +977,20 @@ fn parse_meta(p: &[u8], md: &mut Metadata) {
     let Some(ilst) = find(body, b"ilst") else { return };
     for item in boxes(ilst).map_while(|b| b.ok()) {
         let idx = item.kind.as_u32() as usize;
-        let key = if !keys.is_empty() && idx >= 1 && idx <= keys.len() { keys[idx - 1].clone() } else { item.kind.to_string_lossy() };
+        let named = if idx >= 1 { keys.get(idx - 1) } else { None };
         if let Some(d) = find(item.payload, b"data")
             && d.len() >= 8
         {
             let ty = u32::from_be_bytes([d[0], d[1], d[2], d[3]]) & 0xFF_FFFF;
+            // Charged before the key is copied (every item naming one long key is another copy)
+            // and before the value is converted, at the lengths they will be stored at.
             if ty == 1 || ty == 0 {
-                md.entries.push((key, String::from_utf8_lossy(&d[8..]).into_owned()));
+                let fallback = if named.is_none() { Some(item.kind.to_string_lossy()) } else { None };
+                let key_len = named.or(fallback.as_ref()).map_or(0, String::len);
+                if charge_metadata(left, key_len, lossy_len(&d[8..])) {
+                    let key = named.cloned().or(fallback).unwrap_or_default();
+                    md.entries.push((key, String::from_utf8_lossy(&d[8..]).into_owned()));
+                }
             }
         }
     }
@@ -1002,7 +1050,7 @@ mod tests {
         b.end(il);
         b.end(meta);
         let mut md = Metadata::default();
-        parse_udta(&b.buf, &mut md);
+        parse_udta(&b.buf, &mut md, &mut { MAX_METADATA_BYTES });
         assert_eq!(md.title(), Some("hi"));
         assert_eq!(md.get("©ART"), Some("me"));
     }
@@ -1219,6 +1267,66 @@ mod tests {
         assert_eq!(open(d.as_slice()).unwrap().tracks[0].samples.len(), 4);
     }
 
+    #[test]
+    fn one_long_key_named_by_many_items_stays_within_the_metadata_budget() {
+        let key = vec![b'k'; 64 * 1024];
+        let mut b = BoxBuf::new();
+        b.leaf(b"hdlr", &[0; 24]);
+        let keys = b.start_full(b"keys", 0, 0);
+        b.u32(1);
+        b.u32(8 + key.len() as u32);
+        b.bytes(b"mdta");
+        b.bytes(&key);
+        b.end(keys);
+        let il = b.start(b"ilst");
+        for _ in 0..1000 {
+            let it = b.start(&1u32.to_be_bytes());
+            let d = b.start(b"data");
+            b.u32(1);
+            b.u32(0);
+            b.bytes(b"v");
+            b.end(d);
+            b.end(it);
+        }
+        b.end(il);
+        let mut payload = vec![0; 4];
+        payload.extend_from_slice(&b.buf);
+        let mut md = Metadata::default();
+        parse_meta(&payload, &mut md, &mut { MAX_METADATA_BYTES });
+        assert!(!md.entries.is_empty() && md.entries.len() < 1000, "{}", md.entries.len());
+        assert!(md.entries.iter().map(|(k, v)| k.len() + v.len()).sum::<usize>() <= MAX_METADATA_BYTES);
+        assert_eq!(md.entries[0].0.len(), key.len());
+    }
+
+    #[test]
+    fn metadata_is_charged_at_its_stored_length() {
+        assert_eq!(lossy_len(b"plain"), 5);
+        for bytes in [&b"a\xFFb"[..], b"\xFF\xFE\xFD", b"\xE2\x82", "\u{a9}nam".as_bytes()] {
+            assert_eq!(lossy_len(bytes), String::from_utf8_lossy(bytes).len(), "{bytes:?}");
+        }
+        // Invalid UTF-8 triples in size when converted; the budget sees the converted size.
+        let mut b = BoxBuf::new();
+        let m = b.start(&[0xA9, b'n', b'a', b'm']);
+        b.u16(1000);
+        b.u16(0);
+        b.bytes(&[0xFF; 1000]);
+        b.end(m);
+        let mut md = Metadata::default();
+        let mut left = 3000 + METADATA_ENTRY_COST + 5;
+        parse_udta(&b.buf, &mut md, &mut left);
+        assert_eq!(md.entries.len(), 1);
+        assert_eq!(md.entries[0].0.len() + md.entries[0].1.len(), 3000 + 5 - left);
+    }
+
+    #[test]
+    fn a_declared_sample_count_past_the_chunk_map_reserves_only_what_it_maps() {
+        // stsz claims u32::MAX samples; the two chunks place four.
+        let d = handmade(Some((1, u32::MAX)));
+        let f = open(d.as_slice()).unwrap();
+        assert_eq!(f.tracks[0].samples.len(), 4);
+        assert!(f.tracks[0].samples.capacity() < 64, "{}", f.tracks[0].samples.capacity());
+    }
+
     #[test]
     fn absurd_counts_are_bounded() {
         // Constant-size stsz claiming 4 billion samples is clamped to what fits in the file.

crates/media/src/matroska/meta.rs

@@ -8,6 +8,11 @@ use crate::matroska::track::{AudioInfo, Color, ContentEncoding, MasteringMetadat
 
 // A malformed Matroska file may contain arbitrarily deep ChapterAtom and SimpleTag chains.
 const MAX_META_DEPTH: usize = 64;
+/// Longest header-stripping prefix accepted. It is copied in front of every frame, laced
+/// zero-length frames included, so its length multiplies what a small file can make a reader
+/// produce. Real files strip a few bytes of codec header; a longer one is treated as an
+/// unsupported compression and its frames are not reconstructed.
+const MAX_STRIPPED_HEADER: usize = 256;
 /// Largest accepted `TimestampScale` (about 4.3 s per tick). It keeps tick to nanosecond
 /// arithmetic on 64-bit timestamps within `i128`.
 const MAX_TIMESTAMP_SCALE: u64 = 1 << 32;
@@ -466,7 +471,7 @@ fn parse_encodings(d: &[u8]) -> Vec<ContentEncoding> {
             ContentEncoding::Encryption { scope }
         } else {
             let (algo, settings) = comp.unwrap_or((0, Vec::new()));
-            if algo == 3 {
+            if algo == 3 && settings.len() <= MAX_STRIPPED_HEADER {
                 ContentEncoding::HeaderStripping { bytes: settings, scope }
             } else {
                 ContentEncoding::Compression { algo, settings, scope }
@@ -560,6 +565,26 @@ pub(crate) fn parse_tracks(d: &[u8], timestamp_scale: u64) -> Vec<Track> {
 #[cfg(test)]
 mod metadata_depth_tests {
     use super::*;
+    use crate::matroska::mux::{el, el_uint};
+
+    fn header_stripping(bytes: &[u8]) -> Vec<ContentEncoding> {
+        let mut comp = Vec::new();
+        el_uint(&mut comp, CONTENT_COMP_ALGO, 3);
+        el(&mut comp, CONTENT_COMP_SETTINGS, bytes);
+        let mut enc = Vec::new();
+        el(&mut enc, CONTENT_COMPRESSION, &comp);
+        let mut encs = Vec::new();
+        el(&mut encs, CONTENT_ENCODING, &enc);
+        parse_encodings(&encs)
+    }
+
+    #[test]
+    fn a_long_stripped_header_is_not_put_in_front_of_frames() {
+        let kept = header_stripping(&[7; MAX_STRIPPED_HEADER]);
+        assert!(matches!(&kept[..], [ContentEncoding::HeaderStripping { bytes, .. }] if bytes.len() == MAX_STRIPPED_HEADER));
+        let refused = header_stripping(&[7; MAX_STRIPPED_HEADER + 1]);
+        assert!(matches!(&refused[..], [ContentEncoding::Compression { algo: 3, .. }]), "{refused:?}");
+    }
 
     #[test]
     fn deeply_nested_chapters_stop_at_depth_limit() {

crates/media/src/matroska/mux.rs

@@ -136,7 +136,10 @@ pub struct MkvWriter<W: Write + Seek> {
     cluster_start_ns: i64,
     pending: Option<Pending>,
     cues: Vec<(u64, u64, u64)>, // (time ticks, track number, cluster position)
-    cued_tracks_in_cluster: Vec<bool>,
+    /// The cluster number in which each track last had a cue. Starting a cluster bumps
+    /// `cluster_number` instead of clearing an entry for every declared track.
+    track_cued_in: Vec<u64>,
+    cluster_number: u64,
     end_ns: i64,
     has_video: bool,
 }
@@ -253,7 +256,8 @@ impl<W: Write + Seek> MkvWriter<W> {
             cluster_start_ns: 0,
             pending: None,
             cues: Vec::new(),
-            cued_tracks_in_cluster: vec![false; n],
+            track_cued_in: vec![u64::MAX; n],
+            cluster_number: 0,
             end_ns: 0,
             has_video,
         })
@@ -294,8 +298,8 @@ impl<W: Write + Seek> MkvWriter<W> {
         }
         let cts = self.cluster_ts.unwrap_or(0);
         // cue: video keyframes, or the first frame of each track per cluster when there is no video
-        if (is_video && keyframe) || (!self.has_video && !self.cued_tracks_in_cluster[track]) {
-            self.cued_tracks_in_cluster[track] = true;
+        if (is_video && keyframe) || (!self.has_video && self.track_cued_in[track] != self.cluster_number) {
+            self.track_cued_in[track] = self.cluster_number;
             let pos = self.w.stream_position()? - self.segment_data_start;
             if self.cues.last().is_none_or(|c| c.2 != pos || c.1 != track as u64 + 1) {
                 self.cues.push((ts as u64, track as u64 + 1, pos));
@@ -399,7 +403,7 @@ impl<W: Write + Seek> MkvWriter<W> {
         }
         self.cluster.clear();
         self.cluster_ts = None;
-        self.cued_tracks_in_cluster.iter_mut().for_each(|c| *c = false);
+        self.cluster_number = self.cluster_number.wrapping_add(1);
         Ok(())
     }
 
@@ -568,4 +572,23 @@ mod tests {
         let got: Vec<_> = Demuxer::from_slice(&bytes).unwrap().map(|p| p.unwrap()).map(|p| (p.track, p.pts, p.lace)).collect();
         assert_eq!(got, [(0, 0, 0), (0, 100, 0), (0, 110, 1), (0, 120, 2), (0, 200, 0), (1, 300, 0), (1, 310, 0)]);
     }
+
+    #[test]
+    fn audio_only_files_cue_each_track_once_per_cluster() {
+        let mut a = TrackSpec::new(TrackKind::Audio, "A_PCM/INT/LIT");
+        a.audio = Some((48000.0, 1, Some(16)));
+        let opts = MuxOptions { cluster_duration_ns: 1, lacing: LacingMode::None, ..Default::default() };
+        let mut m = MkvWriter::new(Cursor::new(Vec::new()), vec![a.clone(), a], opts).unwrap();
+        for i in 0..4i64 {
+            m.write_frame(0, i * 10, true, &[1; 4], None).unwrap();
+            m.write_frame(1, i * 10, true, &[2; 4], None).unwrap();
+            m.write_frame(0, i * 10 + 1, true, &[3; 4], None).unwrap();
+        }
+        // Clusters hold [t0, t1] then [t0] at each step: one cue per track per cluster.
+        let mut seen: Vec<(u64, u64)> = m.cues.iter().map(|c| (c.1, c.2)).collect();
+        assert_eq!(seen.len(), 12, "{:?}", m.cues);
+        seen.sort_unstable();
+        seen.dedup();
+        assert_eq!(seen.len(), 12, "{:?}", m.cues);
+    }
 }

crates/media/src/matroska/track.rs

@@ -265,10 +265,16 @@ impl Track {
         p
     }
 
+    /// Length of [`frame_prefix`](Self::frame_prefix), without building it.
+    pub fn frame_prefix_len(&self) -> u64 {
+        let stripped = self.stripped_header().map_or(0, <[u8]>::len) as u64;
+        if matches!(self.codec, Codec::ProRes { .. }) { stripped + 8 } else { stripped }
+    }
+
     /// Frame size of sample `i` as returned by `read_sample` / packets (stored size plus prefix).
     pub fn sample_size(&self, i: usize) -> Option<u64> {
         let s = self.samples.get(i)?;
-        Some(s.size as u64 + self.frame_prefix(s.size as u64).len() as u64)
+        Some(s.size as u64 + self.frame_prefix_len())
     }
 
     /// Indices of keyframe samples, ascending by file order.