print: charge skipped whitespace and comments to the lexer's token budget (#547)

* print: charge skipped whitespace and comments to the lexer's token budget

Lexer::skip_ws scanned whitespace and comments without charging them,
so an annotation object behind a long comment was rescanned in full for
every reference to it, outside the per-object and job-wide budgets. The
skipped bytes now count like tokens, and the scan stops at the budget.

* address review: fail objstm load when header lexer exhausts, restore token charge on failed reference lookahead

* address review: charge objstm header lexer tokens to the job budget

crates/print/src/sheets.rs

@@ -138,13 +138,22 @@ impl<'a> Lexer<'a> {
         self.s.get(self.pos..).unwrap_or(&[])
     }
 
+    /// Skips whitespace and comments, charged to the token budget like the tokens they sit
+    /// between: an object behind a long comment is otherwise scanned again for every reference
+    /// to it, outside every budget.
     fn skip_ws(&mut self) {
+        let start = self.pos;
+        let end = self.pos.saturating_add(self.max_tokens.saturating_sub(self.tokens));
         while let Some(b) = self.peek() {
+            if self.pos >= end {
+                self.exhausted = true;
+                break;
+            }
             if is_white(b) {
                 self.pos += 1;
             } else if b == b'%' {
                 while let Some(b) = self.peek() {
-                    if b == b'\n' || b == b'\r' {
+                    if b == b'\n' || b == b'\r' || self.pos >= end {
                         break;
                     }
                     self.pos += 1;
@@ -153,6 +162,7 @@ impl<'a> Lexer<'a> {
                 break;
             }
         }
+        self.tokens = self.tokens.saturating_add(self.pos - start);
     }
 
     /// A run of regular characters (a number, keyword or name body).
@@ -328,14 +338,15 @@ impl<'a> Lexer<'a> {
             return None;
         }
         if word.iter().all(u8::is_ascii_digit) {
-            let after = self.pos;
+            let (after, tokens, exhausted) = (self.pos, self.tokens, self.exhausted);
             if let (Some(num), Some(generation)) = (std::str::from_utf8(word).ok().and_then(|w| w.parse::<u32>().ok()), self.unsigned())
                 && let Ok(generation) = u16::try_from(generation)
                 && self.keyword(b"R")
             {
                 return Some(Obj::Ref(num, generation));
             }
-            self.pos = after;
+            // The caller rescans what the lookahead skipped: it must not be charged twice.
+            (self.pos, self.tokens, self.exhausted) = (after, tokens, exhausted);
         }
         self.tokens = self.tokens.checked_add(word.len())?;
         if self.tokens > self.max_tokens {
@@ -818,6 +829,7 @@ impl Reader<'_> {
             let (Ok(num), Some(at)) = (u32::try_from(num), usize::try_from(at).ok().and_then(|a| a.checked_add(first))) else { break };
             objects.push((num, at));
         }
+        self.budget.parsed(&lx)?;
         Some(ObjStm { data, objects })
     }
 

crates/print/src/tests/pdf_bounds.rs

@@ -284,6 +284,18 @@ fn an_oversized_inherited_media_box_is_not_copied_for_every_page() {
     assert_eq!(sheets.sheets.len(), pages);
 }
 
+#[test]
+fn a_comment_before_a_shared_annotation_is_not_rescanned_outside_the_budget() {
+    // One annotation behind a 600 KB comment, referenced 20,000 times: 12 GB of scanning if
+    // every reference re-reads the comment for nothing.
+    let refs = vec!["6 0 R"; 20_000].join(" ");
+    let mut objects = three_page_objects(&format!("/Annots [{refs}]"));
+    objects[5] = format!("%{}\n<< /Type /Annot /F 4 >>", "x".repeat(600_000));
+    let started = std::time::Instant::now();
+    let _ = read(&classic(&objects));
+    assert!(started.elapsed() < std::time::Duration::from_secs(10), "{:?}", started.elapsed());
+}
+
 /// A classic-xref PDF of these objects, numbered from 1.
 fn classic(objects: &[String]) -> Vec<u8> {
     let mut pdf = b"%PDF-1.7\n".to_vec();
@@ -476,3 +488,25 @@ fn an_over_limit_flate_stream_sets_exhausted() {
     assert!(budget.exhausted);
     assert_eq!(budget.decoded, MAX_DECODED - 100);
 }
+
+#[test]
+fn a_failed_reference_lookahead_does_not_charge_skipped_bytes_twice() {
+    // `1`, 600 KB of blanks, `2`: the lookahead for `1 2 R` skips the blanks and rewinds.
+    let array = format!("[1{}2]", " ".repeat(600_000));
+    let obj = Lexer::new(array.as_bytes(), 0).object(0);
+    assert!(matches!(obj, Some(Obj::Array(a)) if a.len() == 2));
+}
+
+#[test]
+fn an_object_stream_header_past_the_token_limit_exhausts_the_budget() {
+    let mut raw = b"1 ".to_vec();
+    raw.resize(MAX_OBJECT_TOKENS + 16, b' ');
+    raw.extend_from_slice(b"0 null");
+    let packed = miniz_oxide::deflate::compress_to_vec_zlib(&raw, 6);
+    let mut pdf = format!("6 0 obj\n<< /Type /ObjStm /N 1 /First 4 /Filter /FlateDecode /Length {} >>\nstream\n", packed.len()).into_bytes();
+    pdf.extend_from_slice(&packed);
+    pdf.extend_from_slice(b"\nendstream\nendobj\n");
+    let mut reader = Reader { data: &pdf, entries: HashMap::from([(6, Entry::At(0))]), objstms: HashMap::new(), budget: Budget::default() };
+    assert!(reader.load_objstm(6).is_none());
+    assert!(reader.budget.exhausted);
+}