diff --git a/crates/media/src/isobmff/demux.rs b/crates/media/src/isobmff/demux.rs index 323eac35..8926ddbd 100644 --- a/crates/media/src/isobmff/demux.rs +++ b/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(src: S) -> Result { } let mut trex: HashMap = 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(src: S) -> Result { } } } - 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 = 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::() <= 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. diff --git a/crates/media/src/matroska/meta.rs b/crates/media/src/matroska/meta.rs index 3843c91f..1b559bca 100644 --- a/crates/media/src/matroska/meta.rs +++ b/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::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 { #[cfg(test)] mod metadata_depth_tests { use super::*; + use crate::matroska::mux::{el, el_uint}; + + fn header_stripping(bytes: &[u8]) -> Vec { + 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() { diff --git a/crates/media/src/matroska/mux.rs b/crates/media/src/matroska/mux.rs index 2ba663a4..8f90ed4e 100644 --- a/crates/media/src/matroska/mux.rs +++ b/crates/media/src/matroska/mux.rs @@ -136,7 +136,10 @@ pub struct MkvWriter { cluster_start_ns: i64, pending: Option, cues: Vec<(u64, u64, u64)>, // (time ticks, track number, cluster position) - cued_tracks_in_cluster: Vec, + /// 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, + cluster_number: u64, end_ns: i64, has_video: bool, } @@ -253,7 +256,8 @@ impl MkvWriter { 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 MkvWriter { } 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 MkvWriter { } 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); + } } diff --git a/crates/media/src/matroska/track.rs b/crates/media/src/matroska/track.rs index 2b0b4494..9a5dfd7a 100644 --- a/crates/media/src/matroska/track.rs +++ b/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 { 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.