Merge pull request #1582

Ogg: bound the work and the warnings of damage recovery

crates/ogg/src/lib.rs

@@ -196,6 +196,21 @@ pub struct OggFile {
     pub warnings: Vec<String>,
 }
 
+/// Source bytes damage recovery may read per file: the windows the scan for the next capture
+/// pattern fills and the pages read only to find their CRC wrong.
+const RECOVERY_BUDGET: u64 = 32 << 20;
+/// Warnings kept per file (not counting the one about the budget); the rest are counted in one last entry.
+const MAX_WARNINGS: usize = 256;
+
+/// Record a warning, keeping at most `MAX_WARNINGS` and one entry saying the rest were dropped.
+fn warn(warnings: &mut Vec<String>, msg: String) {
+    match warnings.len().cmp(&MAX_WARNINGS) {
+        std::cmp::Ordering::Less => warnings.push(msg),
+        std::cmp::Ordering::Equal => warnings.push("further warnings omitted".into()),
+        std::cmp::Ordering::Greater => {}
+    }
+}
+
 struct Partial {
     parts: Vec<(u64, u32)>,
     size: u32,
@@ -219,7 +234,7 @@ pub fn open(src: &(impl ByteSource + ?Sized)) -> Result<OggFile> {
             match scanner.find(src, pos + 1, len)? {
                 Some(p) => {
                     if found_any {
-                        warnings.push(format!("lost sync at byte {pos}; resumed at {p}"));
+                        warn(&mut warnings, format!("lost sync at byte {pos}; resumed at {p}"));
                     }
                     pos = p;
                     continue;
@@ -229,14 +244,14 @@ pub fn open(src: &(impl ByteSource + ?Sized)) -> Result<OggFile> {
         }
         let nseg = hdr[26] as usize;
         if pos + 27 + nseg as u64 > len {
-            warnings.push(format!("truncated page header at byte {pos}"));
+            warn(&mut warnings, format!("truncated page header at byte {pos}"));
             break;
         }
         src.read_at(pos + 27, &mut hdr[27..27 + nseg])?;
         let body_len: u64 = hdr[27..27 + nseg].iter().map(|&b| b as u64).sum();
         let page_len = 27 + nseg as u64 + body_len;
         if pos + page_len > len {
-            warnings.push(format!("truncated page at byte {pos}"));
+            warn(&mut warnings, format!("truncated page at byte {pos}"));
             break;
         }
         let mut page = vec![0u8; page_len as usize];
@@ -244,7 +259,11 @@ pub fn open(src: &(impl ByteSource + ?Sized)) -> Result<OggFile> {
         let stored_crc = u32::from_le_bytes([page[22], page[23], page[24], page[25]]);
         page[22..26].fill(0);
         if crc32(&page) != stored_crc {
-            warnings.push(format!("bad CRC on page at byte {pos}"));
+            warn(&mut warnings, format!("bad CRC on page at byte {pos}"));
+            // a page read only to find it damaged is recovery work, charged like the scan's windows
+            if scanner.charge(page_len) {
+                break;
+            }
             pos += 1;
             continue;
         }
@@ -275,7 +294,7 @@ pub fn open(src: &(impl ByteSource + ?Sized)) -> Result<OggFile> {
         // a continued page without a pending packet (lost a page): drop the fragment
         let continued = flags & 0x01 != 0;
         if !continued && partial.remove(&serial).is_some() {
-            warnings.push(format!("incomplete packet before page at byte {pos}"));
+            warn(&mut warnings, format!("incomplete packet before page at byte {pos}"));
         }
         let mut completed: Vec<Packet> = Vec::new();
         let mut off = 0u64;
@@ -332,7 +351,11 @@ pub fn open(src: &(impl ByteSource + ?Sized)) -> Result<OggFile> {
         return Err(Error::NotOgg);
     }
     if !partial.is_empty() {
-        warnings.push("file ends inside a packet".into());
+        warn(&mut warnings, "file ends inside a packet".into());
+    }
+    if scanner.spent() {
+        // once per file and last, so it is not subject to the cap on the others
+        warnings.push("recovery budget spent; the rest of the damaged data was not searched".into());
     }
     Ok(OggFile { streams, warnings })
 }
@@ -340,12 +363,23 @@ pub fn open(src: &(impl ByteSource + ?Sized)) -> Result<OggFile> {
 struct CaptureScanner {
     buf: Vec<u8>,
     buf_start: u64,
-    budget: usize,
+    budget: u64,
 }
 
 impl CaptureScanner {
     fn new() -> Self {
-        Self { buf: Vec::new(), buf_start: 0, budget: 32 * 1024 * 1024 }
+        Self { buf: Vec::new(), buf_start: 0, budget: RECOVERY_BUDGET }
+    }
+
+    /// Charge `bytes` of recovery work; true once the budget is spent.
+    fn charge(&mut self, bytes: u64) -> bool {
+        self.budget = self.budget.saturating_sub(bytes);
+        self.spent()
+    }
+
+    /// Nothing left to search with: fewer bytes than a capture pattern can still be charged for.
+    fn spent(&self) -> bool {
+        self.budget < 4
     }
 
     fn find(&mut self, src: &(impl ByteSource + ?Sized), from: u64, len: u64) -> Result<Option<u64>> {
@@ -360,17 +394,17 @@ impl CaptureScanner {
                 at = buf_end.saturating_sub(3);
                 continue;
             }
-            if self.budget == 0 {
+            if self.spent() {
                 return Ok(None);
             }
-            let n = ((len - at) as usize).min(64 * 1024).min(self.budget);
+            let n = ((len - at) as usize).min(64 * 1024).min(usize::try_from(self.budget).unwrap_or(usize::MAX));
             if n < 4 {
                 break;
             }
             self.buf.resize(n, 0);
             src.read_at(at, &mut self.buf)?;
             self.buf_start = at;
-            self.budget = self.budget.saturating_sub(n);
+            self.budget = self.budget.saturating_sub(n as u64);
             if let Some(i) = self.buf.windows(4).position(|w| w == b"OggS") {
                 return Ok(Some(at + i as u64));
             }

crates/ogg/src/tests.rs

@@ -1,6 +1,9 @@
 //! Synthetic Ogg streams: a small page writer (RFC 3533) and Opus-like packets.
 
-use super::{Codec, Error, OpusTiming, Packet, Stream, crc32, open, opus_packet_samples, sniff};
+use std::cell::Cell;
+use std::io;
+
+use super::{ByteSource, CaptureScanner, Codec, Error, OpusTiming, Packet, Stream, crc32, open, opus_packet_samples, sniff};
 
 /// Write one page holding `segments` (lacing already computed by `lace`).
 fn page(out: &mut Vec<u8>, flags: u8, granule: i64, serial: u32, seq: u32, lacing: &[u8], body: &[u8]) {
@@ -196,3 +199,78 @@ fn opus_timing_rejects_a_granule_that_overflows() {
     assert_eq!(ok.total, 480);
     assert!(OpusTiming::of(&packet(i64::MIN), 0).is_err());
 }
+
+/// A source that counts the bytes read from it.
+struct Counting<'a> {
+    data: &'a [u8],
+    read: Cell<u64>,
+}
+
+impl ByteSource for Counting<'_> {
+    fn len(&self) -> u64 {
+        self.data.len() as u64
+    }
+    fn read_at(&self, offset: u64, buf: &mut [u8]) -> io::Result<()> {
+        self.read.set(self.read.get() + buf.len() as u64);
+        self.data.read_at(offset, buf)
+    }
+}
+
+/// A valid file followed by `len` bytes of 0xFF with a capture pattern every 64 bytes. Each one
+/// claims 255 full lacing values, so it reads as a page of about 64 KiB that fails its CRC.
+fn dense_fake_pages(len: usize) -> Vec<u8> {
+    let mut f = opus_file(6, 312, 0);
+    let start = f.len();
+    f.resize(start + len, 0xFF);
+    for at in (start..start + len - 64).step_by(64) {
+        f[at..at + 5].copy_from_slice(b"OggS\0");
+        f[at + 26] = 255;
+    }
+    f
+}
+
+#[test]
+fn recovery_reads_a_bounded_amount_of_a_damaged_file() {
+    let f = dense_fake_pages(1 << 20);
+    let src = Counting { data: &f, read: Cell::new(0) };
+    let o = open(&src).unwrap();
+    assert_eq!(o.streams[0].packets.len(), 6, "the good pages are kept");
+    let read = src.read.get();
+    assert!(read <= 40 << 20, "read {read} bytes of a {} byte file", f.len());
+    assert!(o.warnings.iter().any(|w| w.contains("recovery budget")), "{:?}", o.warnings);
+}
+
+#[test]
+fn warnings_are_capped() {
+    // 100,000 pages with no segments and a wrong CRC, 27 bytes each
+    let mut f = opus_file(6, 312, 0);
+    for _ in 0..100_000 {
+        let at = f.len();
+        f.extend_from_slice(b"OggS\0\0");
+        f.resize(at + 27, 0);
+        f[at + 22] = 1;
+    }
+    let o = open(&f).unwrap();
+    assert!(o.warnings.len() <= 257, "{} warnings", o.warnings.len());
+    assert!(o.warnings.iter().any(|w| w.contains("omitted")), "{:?}", o.warnings.last());
+}
+
+#[test]
+fn a_file_of_good_pages_is_not_charged_for_recovery() {
+    let f = opus_file(2000, 312, 0);
+    let src = Counting { data: &f, read: Cell::new(0) };
+    let o = open(&src).unwrap();
+    assert_eq!(o.streams[0].packets.len(), 2000);
+    assert!(o.warnings.is_empty(), "{:?}", o.warnings);
+}
+
+#[test]
+fn a_budget_too_small_to_search_is_spent() {
+    // 1 to 3 bytes left cannot hold a capture pattern: the scan would stop without saying why
+    let mut scanner = CaptureScanner::new();
+    assert!(!scanner.spent());
+    assert!(scanner.charge(super::RECOVERY_BUDGET - 3));
+    assert!(scanner.spent());
+    let f = opus_file(6, 312, 0);
+    assert_eq!(scanner.find(&f, 0, f.len() as u64).unwrap(), None);
+}