cos: charge the object-stream pre-scan to one budget per rebuild

Overlapping unterminated dictionaries could each be read up to 1 MiB. A shared
budget (the file length) is now taken off on every path; a spent budget keeps the
object a candidate, as before.

Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>

crates/cos/src/document/reader.rs

@@ -623,6 +623,9 @@ impl Document {
         // `/ObjStm` type; the others are not parsed for the object-stream check below.
         let mut candidates: HashSet<u32> = HashSet::new();
         let headers = object_headers(&data);
+        // What the dictionary pre-check may read in all, so overlapping unterminated dictionaries
+        // cost no more than one more pass over the file.
+        let mut scan_budget = data.len();
         // Where `/XRef` occurs in the file, found in one pass.
         let mut xref_marks = Vec::new();
         let mut mark_from = 0;
@@ -633,7 +636,7 @@ impl Document {
         for &(num, offset, generation, body) in &headers {
             entries.insert(num, XrefEntry::InFile { offset, generation });
             objstms.push(num);
-            if may_be_object_stream(&data, body) {
+            if may_be_object_stream(&data, body, &mut scan_budget) {
                 candidates.insert(num);
             } else {
                 candidates.remove(&num);
@@ -758,83 +761,86 @@ const DICTIONARY_SCAN_LIMIT: usize = 1 << 20;
 /// parser does: strings, hex strings, comments, names and nested dictionaries. `Some(None)`: the
 /// object is not a dictionary, so it is never a stream. `None`: the end is not certain (no end
 /// within the limit, a hex string with a bad digit, a `>>` inside an array); the caller must
-/// then treat the object as a possible object stream.
-fn dictionary_end(data: &[u8], from: usize) -> Option<Option<usize>> {
-    let limit = data.len().min(from.saturating_add(DICTIONARY_SCAN_LIMIT));
+/// then treat the object as a possible object stream. `limit` ends the read; the second value is
+/// how far the read got, for the caller's budget.
+fn dictionary_end(data: &[u8], from: usize, limit: usize) -> (Option<Option<usize>>, usize) {
     let mut i = from;
-    while i < limit && (is_whitespace(data[i]) || data[i] == b'%') {
-        if data[i] == b'%' {
-            while i < limit && !matches!(data[i], b'\n' | b'\r') {
-                i += 1;
-            }
-        } else {
-            i += 1;
-        }
-    }
-    if data.get(i..i + 2) != Some(b"<<") {
-        return if i < limit { Some(None) } else { None };
-    }
-    i += 2;
-    let (mut dicts, mut arrays) = (1usize, 0usize);
-    while i < limit {
-        match data[i] {
-            b'(' => {
-                let mut depth = 1usize;
-                i += 1;
-                while i < limit && depth > 0 {
-                    match data[i] {
-                        b'\\' => i += 1,
-                        b'(' => depth += 1,
-                        b')' => depth -= 1,
-                        _ => {}
-                    }
-                    i += 1;
-                }
-                if depth > 0 {
-                    return None;
-                }
-            }
-            b'<' if data.get(i + 1) == Some(&b'<') => {
-                dicts += 1;
-                i += 2;
-            }
-            b'<' => {
-                i += 1;
-                while i < limit && data[i] != b'>' {
-                    if !data[i].is_ascii_hexdigit() && !is_whitespace(data[i]) {
-                        return None;
-                    }
-                    i += 1;
-                }
-                i += 1;
-            }
-            b'>' if data.get(i + 1) == Some(&b'>') => {
-                if arrays > 0 {
-                    return None;
-                }
-                dicts -= 1;
-                i += 2;
-                if dicts == 0 {
-                    return Some(Some(i));
-                }
-            }
-            b'%' => {
+    let result = 'scan: {
+        while i < limit && (is_whitespace(data[i]) || data[i] == b'%') {
+            if data[i] == b'%' {
                 while i < limit && !matches!(data[i], b'\n' | b'\r') {
                     i += 1;
                 }
-            }
-            b'[' => {
-                arrays += 1;
+            } else {
                 i += 1;
             }
-            b']' => {
-                arrays = arrays.saturating_sub(1);
-                i += 1;
-            }
-            _ => i += 1,
         }
-    }
-    None
+        if data.get(i..i + 2) != Some(b"<<") {
+            break 'scan if i < limit { Some(None) } else { None };
+        }
+        i += 2;
+        let (mut dicts, mut arrays) = (1usize, 0usize);
+        while i < limit {
+            match data[i] {
+                b'(' => {
+                    let mut depth = 1usize;
+                    i += 1;
+                    while i < limit && depth > 0 {
+                        match data[i] {
+                            b'\\' => i += 1,
+                            b'(' => depth += 1,
+                            b')' => depth -= 1,
+                            _ => {}
+                        }
+                        i += 1;
+                    }
+                    if depth > 0 {
+                        break 'scan None;
+                    }
+                }
+                b'<' if data.get(i + 1) == Some(&b'<') => {
+                    dicts += 1;
+                    i += 2;
+                }
+                b'<' => {
+                    i += 1;
+                    while i < limit && data[i] != b'>' {
+                        if !data[i].is_ascii_hexdigit() && !is_whitespace(data[i]) {
+                            break 'scan None;
+                        }
+                        i += 1;
+                    }
+                    i += 1;
+                }
+                b'>' if data.get(i + 1) == Some(&b'>') => {
+                    if arrays > 0 {
+                        break 'scan None;
+                    }
+                    dicts -= 1;
+                    i += 2;
+                    if dicts == 0 {
+                        break 'scan Some(Some(i));
+                    }
+                }
+                b'%' => {
+                    while i < limit && !matches!(data[i], b'\n' | b'\r') {
+                        i += 1;
+                    }
+                }
+                b'[' => {
+                    arrays += 1;
+                    i += 1;
+                }
+                b']' => {
+                    arrays = arrays.saturating_sub(1);
+                    i += 1;
+                }
+                _ => i += 1,
+            }
+        }
+        None
+    };
+    (result, i.min(limit))
 }
 
 /// Whether the object whose body starts at `data[from..]` could be an object stream: its
@@ -843,8 +849,19 @@ fn dictionary_end(data: &[u8], from: usize) -> Option<Option<usize>> {
 /// dictionary ends is read by [`dictionary_end`], so text inside a string (an object header, a
 /// `stream` keyword) is not taken for syntax; when that is not certain the object counts as a
 /// candidate, so a doubtful object is parsed as it always was.
-pub(super) fn may_be_object_stream(data: &[u8], from: usize) -> bool {
-    match dictionary_end(data, from) {
+///
+/// `budget` is the bytes all the calls of one rebuild may still read. A call reads at most
+/// [`DICTIONARY_SCAN_LIMIT`] bytes or what is left, takes what it read off the budget on every
+/// path, and says "maybe" once the budget cannot cover it. Dictionaries that do not overlap
+/// add up to no more than the file, so only unterminated ones can use the budget up.
+pub(super) fn may_be_object_stream(data: &[u8], from: usize, budget: &mut usize) -> bool {
+    let limit = data.len().min(from.saturating_add(DICTIONARY_SCAN_LIMIT.min(*budget)));
+    if *budget == 0 {
+        return true;
+    }
+    let (end, read) = dictionary_end(data, from, limit);
+    *budget = budget.saturating_sub(read.saturating_sub(from).max(1));
+    match end {
         Some(None) => false,
         None => true,
         Some(Some(end)) => {

crates/cos/src/document/tests.rs

@@ -1193,6 +1193,11 @@ fn parses_as_object_stream(file: &[u8]) -> bool {
     matches!(parse_indirect(file, 0, &|_| None), Ok((_, Object::Stream(s))) if s.dict.name(b"Type") == Some(b"ObjStm"))
 }
 
+/// A scan budget that never runs out.
+fn unlimited() -> usize {
+    usize::MAX
+}
+
 #[test]
 fn the_object_stream_pre_check_never_rejects_what_parses_as_one() {
     // Two inputs from review: header-shaped text, and `>>` + `stream` + newline, inside a
@@ -1205,14 +1210,26 @@ fn the_object_stream_pre_check_never_rejects_what_parses_as_one() {
     ] {
         let file = format!("1 0 obj\n{dict}\nstream\nxx\nendstream\nendobj\n").into_bytes();
         assert!(parses_as_object_stream(&file), "test input must be an object stream: {dict}");
-        assert!(may_be_object_stream(&file, 7), "{dict}");
+        assert!(may_be_object_stream(&file, 7, &mut unlimited()), "{dict}");
     }
     // Ordinary streams stay out; their data is never read.
     let plain = b"1 0 obj\n<< /Length 9 >>\nstream\nObjStm # \nendstream\nendobj\n";
-    assert!(!may_be_object_stream(plain, 7));
-    assert!(!may_be_object_stream(b"1 0 obj\n[ /ObjStm ]\nendobj\n", 7));
+    assert!(!may_be_object_stream(plain, 7, &mut unlimited()));
+    assert!(!may_be_object_stream(b"1 0 obj\n[ /ObjStm ]\nendobj\n", 7, &mut unlimited()));
     // Unterminated: not certain, so a candidate.
-    assert!(may_be_object_stream(b"1 0 obj\n<< /Note (open ", 7));
+    assert!(may_be_object_stream(b"1 0 obj\n<< /Note (open ", 7, &mut unlimited()));
+    // The scan is charged to the shared budget on every path, and a spent budget says "maybe"
+    // without reading.
+    let open = b"1 0 obj\n<< /Note (open ";
+    let mut budget = 1000;
+    assert!(may_be_object_stream(open, 7, &mut budget));
+    assert_eq!(budget, 1000 - (open.len() - 7));
+    let mut budget = 3;
+    assert!(may_be_object_stream(plain, 7, &mut budget), "a budget too small for the dictionary is a candidate");
+    assert_eq!(budget, 0);
+    let mut budget = 0;
+    assert!(may_be_object_stream(plain, 7, &mut budget));
+    assert_eq!(budget, 0);
     // Random dictionaries built from tricky pieces: whatever parses as an object stream is
     // a candidate.
     let pieces: [&str; 22] = [
@@ -1257,7 +1274,7 @@ fn the_object_stream_pre_check_never_rejects_what_parses_as_one() {
         let file = format!("1 0 obj\n{dict}\nstream\nxx\nendstream\nendobj\n").into_bytes();
         if parses_as_object_stream(&file) {
             hits += 1;
-            assert!(may_be_object_stream(&file, 7), "rejected an object stream: {dict:?}");
+            assert!(may_be_object_stream(&file, 7, &mut unlimited()), "rejected an object stream: {dict:?}");
         }
     }
     assert!(hits > 1000, "the generator must produce object streams ({hits})");