Merge pull request #1433 from vangel76/fix-mpegts-hardening

#1432: mpegts: bounds-check frame_info, cap chunks, units and allocations, drop non-PES data

crates/mpegts/src/es.rs

@@ -91,6 +91,9 @@ pub(crate) struct Splitter {
     open: Option<Open>,
     pub units: Vec<Unit>,
     pub skipped: u64,
+    /// Units beyond [`MAX_UNITS`] were dropped.
+    pub truncated: bool,
+    max_units: usize,
     // MPEG video state
     header_start: Option<u64>,
     seq_in_unit: bool,
@@ -219,7 +222,8 @@ pub fn frame_info(codec: &Codec, b: &[u8]) -> Option<FrameInfo> {
                 // 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<FrameInfo> {
             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<FrameInfo> {
             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<usize> {
                 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<u8> = [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<u8> = (0..100).flat_map(frame).collect();
+        s.feed(0, &es);
+        s.finish();
+        assert_eq!(s.units.len(), 10);
+        assert!(s.truncated);
+    }
 }

crates/mpegts/src/lib.rs

@@ -284,6 +284,13 @@ pub struct File {
     pub warnings: Vec<String>,
 }
 
+/// 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<usize> {
@@ -294,7 +301,11 @@ impl File {
     pub fn read_unit<S: ByteSource + ?Sized>(&self, src: &S, s: usize, i: usize) -> Result<Vec<u8>> {
         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)?,

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<u8> {
+    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::<u8>::new(), 0, 0), Err(Error::Invalid(_))));
+}
+
 /// A `ByteSource` that counts the bytes requested through `read_at`.
 struct Counting<'a> {
     data: &'a [u8],

crates/mpegts/src/ts.rs

@@ -90,14 +90,25 @@ fn packet(p: &[u8]) -> std::result::Result<Packet<'_>, Bad> {
 pub(crate) struct PesAsm {
     hdr: Vec<u8>,
     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<PesHeader>, &'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<S: ByteSource + ?Sized>(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 })