address review: skip arrays and dictionaries in a calculator program without building them

crates/pdf/src/color.rs

@@ -220,7 +220,18 @@ pub enum PsOp {
 /// (the caller rejects it when `ops` ends past the cap).
 fn ps_parse(lx: &mut crate::object::Lexer, depth: usize, ops: &mut usize) -> Vec<PsOp> {
     let mut out = vec![];
-    while let Some(o) = lx.next() {
+    loop {
+        // A calculator holds no arrays or dictionaries: one is skipped as a single token without
+        // being built (the lexer would allocate every element before this counted it).
+        if lx.at_compound() {
+            if *ops > MAX_CALC_OPS {
+                break;
+            }
+            lx.skip_compound();
+            *ops += 1;
+            continue;
+        }
+        let Some(o) = lx.next() else { break };
         if *ops > MAX_CALC_OPS {
             break;
         }
@@ -805,6 +816,26 @@ mod tests {
         assert_eq!(ps_len(&code), MAX_CALC_OPS + 1);
     }
 
+    #[test]
+    fn arrays_and_dictionaries_in_a_calculator_are_skipped_whole() {
+        // Strings with brackets, nested arrays and a dictionary are one skipped token each; the
+        // operators around them still parse.
+        let src = b"{ 1 [ (a]b) [1 2] << /K [3 4] /S (>>) >> 5 ] 2 add << /A 1 >> mul }";
+        let mut lx = crate::object::Lexer::new(src, 0);
+        let Some(PsOp::Proc(code)) = ps_parse(&mut lx, 0, &mut 0).into_iter().next() else { panic!() };
+        assert_eq!(code, vec![PsOp::Num(1.0), PsOp::Num(2.0), PsOp::Op("add".into()), PsOp::Op("mul".into())]);
+        // A huge array is skipped once the cap is passed, not read to its end.
+        let src = format!("{{ [ {}] 1 }}", "0 ".repeat(MAX_CALC_OPS * 100));
+        let mut lx = crate::object::Lexer::new(src.as_bytes(), 0);
+        let mut scanned = 0;
+        let Some(PsOp::Proc(code)) = ps_parse(&mut lx, 0, &mut scanned).into_iter().next() else { panic!() };
+        assert_eq!(code, vec![PsOp::Num(1.0)]);
+        // An unclosed one ends at the end of the data.
+        let mut lx = crate::object::Lexer::new(b"[ 1 2 ", 0);
+        lx.skip_compound();
+        assert_eq!(lx.pos, 6);
+    }
+
     #[test]
     fn skipped_calculator_tokens_count_toward_the_cap() {
         // `null` is not a calculator token: skipped, but still scanned.

crates/pdf/src/object.rs

@@ -82,6 +82,48 @@ impl<'a> Lexer<'a> {
         }
     }
 
+    /// Whether the next token opens an array or a dictionary (skipping white space first).
+    pub fn at_compound(&mut self) -> bool {
+        self.skip_ws();
+        match self.data.get(self.pos) {
+            Some(b'[') => true,
+            Some(b'<') => self.data.get(self.pos + 1) == Some(&b'<'),
+            _ => false,
+        }
+    }
+
+    /// Skip the array or dictionary at the current position (see [`Lexer::at_compound`]) as
+    /// [`Lexer::next`] would read it, without building it: for readers that have no use for one
+    /// and must not allocate a vector per element of a hostile stream.
+    pub fn skip_compound(&mut self) {
+        let mut open = 0usize;
+        loop {
+            self.skip_ws();
+            match self.data.get(self.pos) {
+                None => return,
+                Some(b'[') => {
+                    open += 1;
+                    self.pos += 1;
+                }
+                Some(b'<') if self.data.get(self.pos + 1) == Some(&b'<') => {
+                    open += 1;
+                    self.pos += 2;
+                }
+                Some(b']') | Some(b'>') if self.data[self.pos] == b']' || self.data.get(self.pos + 1) == Some(&b'>') => {
+                    open = open.saturating_sub(1);
+                    self.pos += if self.data[self.pos] == b']' { 1 } else { 2 };
+                    if open == 0 {
+                        return;
+                    }
+                }
+                // Anything else is one token (a string, name, number or word).
+                Some(_) => {
+                    self.next_depth(65);
+                }
+            }
+        }
+    }
+
     fn word(&mut self) -> &'a [u8] {
         let s = self.pos;
         while self.pos < self.data.len() && !is_white(self.data[self.pos]) && !is_delim(self.data[self.pos]) {