diff --git a/crates/mpegts/src/es.rs b/crates/mpegts/src/es.rs index 0b96e148..9735d456 100644 --- a/crates/mpegts/src/es.rs +++ b/crates/mpegts/src/es.rs @@ -91,6 +91,9 @@ pub(crate) struct Splitter { open: Option, pub units: Vec, pub skipped: u64, + /// Units beyond [`MAX_UNITS`] were dropped. + pub truncated: bool, + max_units: usize, // MPEG video state header_start: Option, seq_in_unit: bool, @@ -219,7 +222,8 @@ pub fn frame_info(codec: &Codec, b: &[u8]) -> Option { // E-AC-3: fscod(2) numblkscod(2) acmod(3) lfeon(1) let fscod = b[4] >> 6; let (rate, blocks) = if fscod == 3 { - ([24_000, 22_050, 16_000, 0][((b[4] >> 4) & 3) as usize], 6) + // fscod2 3 is reserved (no sample rate) + (*[24_000, 22_050, 16_000].get(((b[4] >> 4) & 3) as usize)?, 6) } else { ([48_000, 44_100, 32_000][fscod as usize], [1, 2, 3, 6][((b[4] >> 4) & 3) as usize]) }; @@ -229,7 +233,8 @@ pub fn frame_info(codec: &Codec, b: &[u8]) -> Option { let fscod = b[4] >> 6; let rate = *[48_000u32, 44_100, 32_000].get(fscod as usize)?; // bsi: bsid(5) bsmod(3) acmod(3) [cmixlev(2)] [surmixlev(2)] [dsurmod(2)] lfeon(1) - let acmod = b[6] >> 5; + let b6 = *b.get(6)?; + let acmod = b6 >> 5; let mut bit = 3; // bits of byte 6 consumed if acmod & 1 != 0 && acmod != 1 { bit += 2; @@ -240,7 +245,7 @@ pub fn frame_info(codec: &Codec, b: &[u8]) -> Option { if acmod == 2 { bit += 2; } - let word = u16::from_be_bytes([b[6], *b.get(7)?]); + let word = u16::from_be_bytes([b6, *b.get(7)?]); let lfe = (word >> (15 - bit)) & 1; Some(FrameInfo { sample_rate: rate, channels: AC3_CHANNELS[acmod as usize] + lfe as u32, samples: 1536, variant: bsid }) } @@ -257,6 +262,18 @@ fn eac3_joins_unit(b: &[u8]) -> bool { } } +/// Shortest E-AC-3 syncframe accepted. A real syncframe holds the 6 header bytes read here plus +/// bsi and audio blocks (the smallest real ones are tens of bytes), so shorter "frames" +/// (`0B 77 00 01` is 4 bytes) are noise; they would otherwise turn every 4 input bytes into a +/// 64-byte `Unit`. +const EAC3_MIN_FRAME: usize = 8; + +/// Most units one stream may index. A `Unit` is 64 bytes, so this bounds a stream's table to +/// 512 MiB. It is 47 hours of 48 kHz AAC (the densest audio, 46.9 units/s), 9 hours of 240 fps +/// video: far beyond any real file, while a hostile file (a unit per 7 input bytes) cannot grow +/// the table without limit. +pub(crate) const MAX_UNITS: usize = 1 << 23; + /// Full-bandwidth channels per AC-3 audio coding mode (A/52 Table 5.8). const AC3_CHANNELS: [u32; 8] = [2, 1, 2, 3, 3, 4, 4, 5]; @@ -288,7 +305,8 @@ fn frame_len(sync: AudioSync, b: &[u8]) -> Option { ac3_frame_bytes(b[4] >> 6, b[4] & 0x3F) } else if bsid <= 16 { // E-AC-3: frmsiz - Some(((((b[2] & 7) as usize) << 8) | b[3] as usize) * 2 + 2) + let len = ((((b[2] & 7) as usize) << 8) | b[3] as usize) * 2 + 2; + (len >= EAC3_MIN_FRAME).then_some(len) } else { None } @@ -308,6 +326,8 @@ impl Splitter { open: None, units: Vec::new(), skipped: 0, + truncated: false, + max_units: MAX_UNITS, header_start: None, seq_in_unit: false, gop: None, @@ -372,6 +392,10 @@ impl Splitter { if end <= o.start || o.pos == u64::MAX { return; } + if self.units.len() >= self.max_units { + self.truncated = true; + return; + } let size = (end - o.start).min(u32::MAX as u64) as u32; self.units.push(Unit { pts: o.pts, dts: o.dts, key: o.key, disposable: o.disposable, size, pos: o.pos, offset: o.offset, picture: o.picture }); } @@ -381,7 +405,12 @@ impl Splitter { if data.is_empty() { return; } - self.chunks.push_back((pos, self.es_pos)); + // PES-delimited modes resolve a unit's position from `pos` right here and never look a + // chunk up again, so remembering chunks there would only grow without bound (NAL streams + // keep theirs: an access unit can start inside a PES; `keep_from` prunes them) + if matches!(self.mode, Mode::MpegVideo | Mode::Audio(_) | Mode::Nal { .. }) { + self.chunks.push_back((pos, self.es_pos)); + } if let Some(o) = self.open.as_mut() && o.pos == u64::MAX && o.start == self.es_pos @@ -699,4 +728,61 @@ mod tests { // AC-3 frames are never joined assert!(!eac3_joins_unit(&[0x0B, 0x77, 0x40, 0, 0, 8 << 3])); } + + #[test] + fn ac3_header_prefix_of_six_bytes_is_not_indexed_out_of_bounds() { + // exactly the six bytes frame_len() guarantees: byte 6 (acmod) is missing + assert_eq!(frame_info(&Codec::Ac3, &[0x0B, 0x77, 0, 0, 0x1E, 0x40]), None); + assert_eq!(frame_info(&Codec::Ac3, &[0x0B, 0x77, 0, 0, 0x1E, 0x40, 0]), None); + } + + #[test] + fn eac3_reserved_sample_rate_is_refused() { + // fscod 3 and fscod2 3: no sample rate + let mut f = vec![0x0B, 0x77, 0, 7, 0xF0, 16 << 3]; + f.resize(16, 0); + assert_eq!(frame_info(&Codec::Eac3, &f), None); + f[4] = 0xE0; // fscod2 2: 16 kHz + assert_eq!(frame_info(&Codec::Eac3, &f).map(|i| i.sample_rate), Some(16_000)); + } + + #[test] + fn four_byte_eac3_frames_make_no_units() { + let es: Vec = [0x0B, 0x77, 0x00, 0x01].repeat(10_000); + let mut s = Splitter::new(&Codec::Eac3); + s.feed(0, &es); + s.finish(); + assert!(s.units.is_empty(), "{} units from 40000 noise bytes", s.units.len()); + } + + #[test] + fn pes_delimited_modes_do_not_retain_chunks() { + for codec in [Codec::LpcmBluray, Codec::Unknown(0x99)] { + let mut s = Splitter::new(&codec); + for k in 0..5000u64 { + s.pes_start(Some(k as i64), None); + s.feed(k * 188, &[0x11; 184]); + } + s.finish(); + assert!(s.chunks.len() <= 2, "{codec:?}: {} chunks retained", s.chunks.len()); + assert_eq!(s.units.len(), 5000, "{codec:?}"); + assert!(s.units.iter().enumerate().all(|(k, u)| u.pos == k as u64 * 188 && u.offset == 0), "{codec:?}"); + } + } + + #[test] + fn unit_table_is_capped() { + let frame = |k: u8| { + let mut v = vec![0x0B, 0x77, 0, 7, (3 << 4) | (2 << 1), 16 << 3]; + v.resize(16, k); + v + }; + let mut s = Splitter::new(&Codec::Eac3); + s.max_units = 10; + let es: Vec = (0..100).flat_map(frame).collect(); + s.feed(0, &es); + s.finish(); + assert_eq!(s.units.len(), 10); + assert!(s.truncated); + } } diff --git a/crates/mpegts/src/lib.rs b/crates/mpegts/src/lib.rs index cb62b66e..895e3481 100644 --- a/crates/mpegts/src/lib.rs +++ b/crates/mpegts/src/lib.rs @@ -284,6 +284,13 @@ pub struct File { pub warnings: Vec, } +/// Largest unit `read_unit` returns. The biggest real units are intra pictures (a few MB at +/// 4K/8K broadcast rates); a stream without PES timestamps can make one unit of a whole PID. +pub const MAX_UNIT_BYTES: u32 = 128 << 20; + +/// Bytes reserved up front by `read_unit`; the buffer grows beyond this only as data is read. +const UNIT_PREALLOC: usize = 1 << 20; + impl File { /// The first stream of `kind` whose codec satisfies `f`. pub fn find(&self, kind: Kind, f: impl Fn(&Codec) -> bool) -> Option { @@ -294,7 +301,11 @@ impl File { pub fn read_unit(&self, src: &S, s: usize, i: usize) -> Result> { let st = self.streams.get(s).ok_or_else(|| Error::Invalid("no such stream".into()))?; let u = st.units.get(i).ok_or_else(|| Error::Invalid("no such unit".into()))?; - let mut out = Vec::with_capacity(u.size as usize); + if u.size > MAX_UNIT_BYTES { + return Err(Error::Invalid(format!("unit {i} is {} bytes (limit {MAX_UNIT_BYTES})", u.size))); + } + // grow as bytes arrive: a damaged table must not reserve more than the file can supply + let mut out = Vec::with_capacity((u.size as usize).min(UNIT_PREALLOC)); match self.format { Format::Ts { packet_size, sync_offset } => ts::read(src, packet_size, sync_offset, st, u, &mut out)?, Format::Ps { .. } => ps::read(src, st, u, &mut out)?, diff --git a/crates/mpegts/src/tests.rs b/crates/mpegts/src/tests.rs index 5e34bcf6..5cf89db3 100644 --- a/crates/mpegts/src/tests.rs +++ b/crates/mpegts/src/tests.rs @@ -331,6 +331,78 @@ fn timestamps_unwrap_across_the_33_bit_wrap() { assert_eq!(u.unwrap(4200), (1 << 33) + 4200); } +fn adts_frame(k: u8) -> Vec { + let mut v = vec![0xFF, 0xF1, 0x50, 0x80, 0x0C, 0x9F, 0xFC]; + v.extend([k; 93]); + v +} + +#[test] +fn non_pes_payload_is_dropped_until_the_next_unit_start() { + // a unit start whose payload is not a PES header, then continuation packets that happen to + // hold ADTS frames: nothing of it may become a unit + let mut mux = TsMux::new(false); + mux.psi(&[(0x0F, 0x1100)]); + mux.packet(0x1100, true, &[0x55; 184]); + for k in 0..10 { + mux.packet(0x1100, false, &adts_frame(k)); + } + let f = open(&mux.out).unwrap(); + assert_eq!(f.streams.first().map(|s| s.units.len()), Some(0)); + // a real PES packet afterwards is indexed again + let mut mux2 = TsMux::new(false); + mux2.psi(&[(0x0F, 0x1100)]); + mux2.packet(0x1100, true, &[0x55; 184]); + mux2.packet(0x1100, false, &adts_frame(1)); + mux2.pes(0x1100, &pes(0xC0, Some(9000), None, &adts_frame(2), true)); + let f = open(&mux2.out).unwrap(); + let u = &f.streams[0].units; + assert_eq!(u.len(), 1); + assert_eq!(u[0].pts, Some(9000)); +} + +#[test] +fn adaptation_field_only_unit_start_keeps_the_pes_in_progress() { + let mut mux = TsMux::new(false); + mux.psi(&[(0x03, 0x1100)]); + let data = pes(0xC0, Some(9000), None, &[mp2_frame(0), mp2_frame(1)].concat(), true); + mux.packet(0x1100, true, &data[..184]); + // PUSI set, adaptation field only (no payload) + let cc = mux.cc.get(&0x1100).copied().unwrap_or(0); + let mut af = vec![0x47, 0x40 | 0x11, 0x00, 0x20 | cc, 183, 0x00]; + af.resize(188, 0xFF); + mux.out.extend(af); + for chunk in data[184..].chunks(184) { + mux.packet(0x1100, false, chunk); + } + let f = open(&mux.out).unwrap(); + let s = &f.streams[0]; + assert_eq!(s.units.len(), 2); + assert_eq!(s.units[0].pts, Some(9000)); + assert_eq!(f.read_unit(&mux.out, 0, 1).unwrap(), mp2_frame(1)); +} + +#[test] +fn oversized_unit_is_refused_without_reserving_it() { + let f = File { + format: Format::Ts { packet_size: 188, sync_offset: 0 }, + streams: vec![Stream { + id: StreamId::Pid(0x100), + stream_type: 0x1B, + codec: Codec::H264, + language: None, + units: vec![Unit { pts: None, dts: None, key: true, disposable: false, size: u32::MAX, pos: 0, offset: 0, picture: None }], + lpcm: None, + pes_packets: 0, + skipped_bytes: 0, + }], + program: None, + pcr_range: None, + warnings: Vec::new(), + }; + assert!(matches!(f.read_unit(&Vec::::new(), 0, 0), Err(Error::Invalid(_)))); +} + /// A `ByteSource` that counts the bytes requested through `read_at`. struct Counting<'a> { data: &'a [u8], diff --git a/crates/mpegts/src/ts.rs b/crates/mpegts/src/ts.rs index 4a3f567b..b8865aba 100644 --- a/crates/mpegts/src/ts.rs +++ b/crates/mpegts/src/ts.rs @@ -90,14 +90,25 @@ fn packet(p: &[u8]) -> std::result::Result, Bad> { pub(crate) struct PesAsm { hdr: Vec, waiting: bool, + /// The packet that started the current unit was not a PES packet: its continuation packets + /// are dropped until the next unit start. + dropping: bool, } impl PesAsm { /// The PES header completed by this packet (if any) and the packet's ES payload bytes. pub fn packet<'a>(&mut self, pusi: bool, payload: &'a [u8]) -> (Option, &'a [u8]) { if pusi { + if payload.is_empty() { + // an adaptation-field-only packet carries no PES data: it must not cut short + // the PES packet in progress + return (None, &[]); + } self.hdr.clear(); self.waiting = true; + self.dropping = false; + } else if self.dropping { + return (None, &[]); } else if !self.waiting { return (None, payload); } @@ -113,6 +124,7 @@ impl PesAsm { if self.hdr.len() > 1024 || (self.hdr.len() >= 6 && !(self.hdr[0] == 0 && self.hdr[1] == 0 && self.hdr[2] == 1)) { // not a PES packet after all: drop until the next unit start self.waiting = false; + self.dropping = true; self.hdr.clear(); return (None, &[]); } @@ -453,6 +465,9 @@ pub(crate) fn scan(src: &S, packet_size: usize, sync_off e.split.finish(); e.stream.units = std::mem::take(&mut e.split.units); e.stream.skipped_bytes = e.split.skipped; + if e.split.truncated { + warnings.push(format!("PID {pid}: more than {} access units, the rest are not indexed", crate::es::MAX_UNITS)); + } streams.push(e.stream); } Ok(File { format: Format::Ts { packet_size, sync_offset }, streams, program, pcr_range, warnings })