Merge pull request #1434 from vangel76/fix-mpegts-ts-resync

#1430: mpegts: resync inside the loaded block, bad adaptation field is not lost sync

crates/mpegts/src/tests.rs

@@ -330,3 +330,99 @@ fn timestamps_unwrap_across_the_33_bit_wrap() {
     assert_eq!(u.unwrap(near + 1000), near + 1000);
     assert_eq!(u.unwrap(4200), (1 << 33) + 4200);
 }
+
+/// A `ByteSource` that counts the bytes requested through `read_at`.
+struct Counting<'a> {
+    data: &'a [u8],
+    read: std::cell::Cell<u64>,
+}
+
+impl ByteSource for Counting<'_> {
+    fn len(&self) -> u64 {
+        self.data.len() as u64
+    }
+    fn read_at(&self, offset: u64, buf: &mut [u8]) -> std::io::Result<()> {
+        self.read.set(self.read.get() + buf.len() as u64);
+        self.data.read_at(offset, buf)
+    }
+}
+
+#[test]
+fn bad_adaptation_fields_do_not_amplify_reads() {
+    // aligned packets whose adaptation_field_length (255) overruns the packet: the packet grid is
+    // intact, so this must cost one pass over the file, not one block read per packet
+    let mut data = Vec::new();
+    for _ in 0..3000 {
+        data.extend([0x47, 0x11, 0x00, 0x30]);
+        data.extend(std::iter::repeat_n(0xFF, 184));
+    }
+    let src = Counting { data: &data, read: Default::default() };
+    let f = open(&src).unwrap();
+    assert!(f.warnings.iter().any(|w| w.contains("adaptation field")), "{:?}", f.warnings);
+    assert!(!f.warnings.iter().any(|w| w.contains("lost packet sync")), "{:?}", f.warnings);
+    let read = src.read.get();
+    assert!(read <= 4 * data.len() as u64, "read {read} bytes of a {}-byte file", data.len());
+}
+
+#[test]
+fn garbage_without_sync_does_not_amplify_reads() {
+    // a sync byte every 189 bytes: five-in-a-row never lines up, every 0x47 is a false candidate
+    let mut data = Vec::new();
+    for _ in 0..3000 {
+        data.extend([0x47]);
+        data.extend(std::iter::repeat_n(0x11, 188));
+    }
+    let head: Vec<u8> = (0..188 * 8).map(|i| if i % 188 == 0 { 0x47 } else { 0 }).collect();
+    let mut all = head;
+    all.extend(data);
+    let src = Counting { data: &all, read: Default::default() };
+    let f = open(&src).unwrap();
+    assert!(f.warnings.iter().any(|w| w.contains("lost packet sync")), "{:?}", f.warnings);
+    let read = src.read.get();
+    assert!(read <= 4 * all.len() as u64, "read {read} bytes of a {}-byte file", all.len());
+}
+
+#[test]
+fn transport_stream_resyncs_after_garbage() {
+    const BLOCK_PACKETS: usize = 4096;
+    for m2ts in [false, true] {
+        let ps = if m2ts { 192 } else { 188 };
+        // garbage in front of the stream (shorter than 8 packets) and, at a few places around the
+        // first block boundary, between two frames; lengths that are not a multiple of the packet
+        for (delta, glen) in [(0usize, 100usize), (1, 37), (2, 188), (3, 301), (5, 5), (8, 1000)] {
+            let edge = (BLOCK_PACKETS - delta) * ps;
+            let mut mux = TsMux::new(m2ts);
+            mux.out.extend([0x11, 0x47, 0x22]); // garbage prefix with a lone 0x47
+            mux.psi(&[(0x02, 0x1011)]);
+            let mut expect = Vec::new();
+            let mut placed = false;
+            let mut k = 0usize;
+            while k < 3000 || !placed {
+                if !placed && mux.out.len() >= edge {
+                    // garbage between two PES packets, with lone sync bytes inside
+                    let mut g: Vec<u8> = (0..glen).map(|i| 0x10 + (i % 7) as u8).collect();
+                    if let Some(b) = g.get_mut(3) {
+                        *b = 0x47;
+                    }
+                    mux.out.extend(g);
+                    placed = true;
+                }
+                let f = mpeg2_frame(if k == 0 { 1 } else { 2 }, (k % 1000) as u16, false, 40 + k % 300);
+                mux.pes(0x1011, &pes(0xE0, Some(90_000 + 3600 * k as i64), None, &f, false));
+                expect.push(f);
+                k += 1;
+            }
+            assert!(mux.out.len() > (BLOCK_PACKETS + 4) * ps, "stream must cross a block boundary");
+            let data = mux.out;
+            let f = open(&data).unwrap();
+            assert!(f.warnings.iter().any(|w| w.contains("lost packet sync")), "{:?}", f.warnings);
+            let v = f.find(Kind::Video, |_| true).unwrap();
+            let s = &f.streams[v];
+            assert_eq!(s.units.len(), expect.len(), "m2ts {m2ts} delta {delta} glen {glen}");
+            for (i, fr) in expect.iter().enumerate() {
+                assert_eq!(&f.read_unit(&data, v, i).unwrap(), fr, "unit {i} (m2ts {m2ts} delta {delta} glen {glen})");
+                assert_eq!(s.units[i].pts, Some(90_000 + 3600 * i as i64));
+            }
+        }
+    }
+}

crates/mpegts/src/ts.rs

@@ -29,9 +29,19 @@ struct Packet<'a> {
     payload: &'a [u8],
 }
 
-fn packet(p: &[u8]) -> Option<Packet<'_>> {
+/// Why a packet could not be parsed.
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+enum Bad {
+    /// No sync byte (or too short): the packet grid is lost.
+    Sync,
+    /// Sync byte present but the adaptation field overruns the packet: the grid is intact, the
+    /// packet carries nothing usable.
+    Adaptation,
+}
+
+fn packet(p: &[u8]) -> std::result::Result<Packet<'_>, Bad> {
     if p.len() < 188 || p[0] != 0x47 {
-        return None;
+        return Err(Bad::Sync);
     }
     let pusi = p[1] & 0x40 != 0;
     let pid = (((p[1] & 0x1F) as u16) << 8) | p[2] as u16;
@@ -44,7 +54,7 @@ fn packet(p: &[u8]) -> Option<Packet<'_>> {
         let al = p[4] as usize;
         start = 5 + al;
         if start > 188 {
-            return None;
+            return Err(Bad::Adaptation);
         }
         if al > 0 {
             let flags = p[5];
@@ -57,7 +67,7 @@ fn packet(p: &[u8]) -> Option<Packet<'_>> {
         }
     }
     let has_payload = afc & 1 != 0;
-    Some(Packet { pusi, pid, cc, has_payload, discontinuity, pcr, payload: if has_payload { &p[start..188] } else { &[] } })
+    Ok(Packet { pusi, pid, cc, has_payload, discontinuity, pcr, payload: if has_payload { &p[start..188] } else { &[] } })
 }
 
 /// Reassembles a PES header that may span packets; yields each packet's elementary-stream bytes.
@@ -216,6 +226,39 @@ fn first_sync(head: &[u8], packet_size: usize, sync_offset: usize) -> Option<usi
     (0..packet_size.min(head.len())).find(|&s| (0..5).all(|k| head.get(s + sync_offset + k * packet_size) == Some(&0x47)))
 }
 
+/// Find where packets line up again, searching `block` from byte `from`: the first packet start
+/// `i` whose sync byte and the next four packets' sync bytes are all 0x47. Returns `(i, false)`
+/// when found. Returns `(i, true)` when the block ran out first: `i` is where to refill from (a
+/// candidate whose five-packet window does not fit yet, or the last place a whole packet fits).
+/// At the end of the file (`at_end`) a candidate is accepted with as many sync bytes as exist.
+fn resync(block: &[u8], from: usize, packet_size: usize, sync_offset: usize, at_end: bool) -> (usize, bool) {
+    let mut i = from;
+    while i + packet_size <= block.len() {
+        if block.get(i + sync_offset) == Some(&0x47) {
+            let mut aligned = true;
+            let mut complete = true;
+            for k in 1..5 {
+                match block.get(i + sync_offset + k * packet_size) {
+                    Some(&0x47) => {}
+                    Some(_) => {
+                        aligned = false;
+                        break;
+                    }
+                    None => {
+                        complete = false;
+                        break;
+                    }
+                }
+            }
+            if aligned {
+                return (i, !complete && !at_end);
+            }
+        }
+        i += 1;
+    }
+    (i, true)
+}
+
 pub(crate) fn scan<S: ByteSource + ?Sized>(src: &S, packet_size: usize, sync_offset: usize) -> Result<File> {
     let len = src.len();
     let mut head = vec![0u8; packet_size * 8];
@@ -236,30 +279,39 @@ pub(crate) fn scan<S: ByteSource + ?Sized>(src: &S, packet_size: usize, sync_off
     let (mut lost_sync, mut cc_errors, mut crc_errors) = (0u64, 0u64, 0u64);
     const BLOCK: usize = 4096;
     let mut block = vec![0u8; BLOCK * packet_size];
+    let mut bad_adaptation = 0u64;
     while off + packet_size as u64 <= len {
         let got = read_some(src, off, &mut block)?;
-        let whole = got / packet_size;
-        if whole == 0 {
+        if got < packet_size {
             break;
         }
-        let mut k = 0;
-        while k < whole {
-            let pos = off + (k * packet_size) as u64;
-            let raw = &block[k * packet_size..(k + 1) * packet_size];
-            let Some(pk) = packet(&raw[sync_offset..]) else {
-                // lost sync: find the next offset where packets line up again
-                lost_sync += 1;
-                let rest = &block[k * packet_size + 1..got];
-                match first_sync(rest, packet_size, sync_offset).or_else(|| rest.iter().position(|&b| b == 0x47).map(|p| p.saturating_sub(sync_offset))) {
-                    Some(s) => {
-                        off = pos + 1 + s as u64;
+        let at_end = off + got as u64 >= len;
+        // `p`: byte index of the next packet in `block`. Packets are parsed in place; a lost
+        // sync is repaired inside the loaded block, and the block is refilled from the file
+        // (at `off + p`) only when it is used up, so the bytes read stay close to the file size.
+        let mut p = 0usize;
+        while p + packet_size <= got {
+            let pos = off + p as u64;
+            let raw = &block[p..p + packet_size];
+            let pk = match packet(&raw[sync_offset..]) {
+                Ok(pk) => pk,
+                Err(Bad::Adaptation) => {
+                    // the packet grid is intact: this packet just has no usable payload
+                    bad_adaptation += 1;
+                    p += packet_size;
+                    continue;
+                }
+                Err(Bad::Sync) => {
+                    lost_sync += 1;
+                    let (next, exhausted) = resync(&block[..got], p + 1, packet_size, sync_offset, at_end);
+                    p = next;
+                    if exhausted {
+                        break;
                     }
-                    None => off = pos + got as u64,
+                    continue;
                 }
-                k = usize::MAX;
-                break;
             };
-            k += 1;
+            p += packet_size;
             if let Some(pcr) = pk.pcr {
                 // 33-bit base × 300 wraps at 2^33·300
                 let wrap = (1i64 << 33) * 300;
@@ -362,13 +414,15 @@ pub(crate) fn scan<S: ByteSource + ?Sized>(src: &S, packet_size: usize, sync_off
             }
             e.split.feed(pos, bytes);
         }
-        if k != usize::MAX {
-            off += (whole * packet_size) as u64;
-        }
+        // `p >= 1` here (the loop ran at least once or `resync` moved past the failed packet)
+        off += p.max(1) as u64;
     }
     if lost_sync > 0 {
         warnings.push(format!("lost packet sync {lost_sync} time(s)"));
     }
+    if bad_adaptation > 0 {
+        warnings.push(format!("{bad_adaptation} packet(s) with an invalid adaptation field"));
+    }
     if cc_errors > 0 {
         warnings.push(format!("{cc_errors} continuity counter error(s)"));
     }
@@ -406,7 +460,11 @@ pub(crate) fn read<S: ByteSource + ?Sized>(src: &S, packet_size: usize, sync_off
         }
         for k in 0..whole {
             let raw = &block[k * packet_size..(k + 1) * packet_size];
-            let Some(pk) = packet(&raw[sync_offset..]) else { return Ok(()) };
+            let pk = match packet(&raw[sync_offset..]) {
+                Ok(pk) => pk,
+                Err(Bad::Sync) => return Ok(()),
+                Err(Bad::Adaptation) => continue,
+            };
             if pk.pid != pid {
                 continue;
             }