rust_pdf_box/cos/
parser.rs

1use crate::cos::document::COSDocument;
2use crate::cos::lexer::{Lexer, Token};
3use crate::cos::object::{
4    COSArray, COSDictionary, COSName, COSObject, COSStream, COSString, ObjectId,
5};
6use crate::error::{PdfError, PdfResult};
7use std::collections::BTreeMap;
8
9pub struct Parser<'a> {
10    data: &'a [u8],
11}
12
13#[derive(Debug, Clone)]
14#[allow(dead_code)]
15enum XRefEntry {
16    Free,
17    Normal {
18        offset: usize,
19        generation: u16,
20    },
21    Compressed {
22        obj_stm_num: u32,
23        index_in_stream: u32,
24    },
25}
26
27impl<'a> Parser<'a> {
28    pub fn new(data: &'a [u8]) -> Self {
29        Self { data }
30    }
31
32    /// Parses an entire PDF file buffer into a `COSDocument`.
33    pub fn parse(&self) -> PdfResult<COSDocument> {
34        let version = self.parse_header()?;
35        let startxref_offset = self.find_startxref()?;
36
37        let mut doc = COSDocument::new();
38        doc.version = version;
39
40        // Traverse all xref tables / streams starting from the last startxref
41        let mut xref_entries = BTreeMap::new();
42        let mut current_xref_offset = Some(startxref_offset);
43        let mut final_trailer = COSDictionary::new();
44
45        while let Some(offset) = current_xref_offset {
46            let mut lexer = Lexer::with_pos(self.data, offset);
47            lexer.skip_whitespace_and_comments();
48
49            let token = lexer.next_token()?;
50            match token {
51                Token::Keyword(ref k) if k == "xref" => {
52                    let trailer = self.parse_classic_xref_section(&mut lexer, &mut xref_entries)?;
53                    if final_trailer.entries.is_empty() {
54                        final_trailer = trailer.clone();
55                    } else {
56                        // Merge trailer entries without overwriting newer entries
57                        for (k, v) in &trailer.entries {
58                            doc.trailer
59                                .entries
60                                .entry(k.clone())
61                                .or_insert_with(|| v.clone());
62                        }
63                    }
64
65                    if let Some(prev) = trailer.get_integer("Prev") {
66                        current_xref_offset = Some(prev as usize);
67                    } else {
68                        current_xref_offset = None;
69                    }
70                }
71                Token::Integer(_num) => {
72                    // Modern XRef Stream: "num gen obj << /Type /XRef ... >> stream ... endstream"
73                    let gen_token = lexer.next_token()?;
74                    let obj_token = lexer.next_token()?;
75                    if matches!(gen_token, Token::Integer(_))
76                        && matches!(obj_token, Token::Keyword(ref k) if k == "obj")
77                    {
78                        let stream_obj = self.parse_direct_object(&mut lexer)?;
79                        if let COSObject::Stream(s) = stream_obj {
80                            let trailer = s.dict.clone();
81                            self.parse_xref_stream(&s, &mut xref_entries)?;
82                            if final_trailer.entries.is_empty() {
83                                final_trailer = trailer.clone();
84                            }
85                            if let Some(prev) = trailer.get_integer("Prev") {
86                                current_xref_offset = Some(prev as usize);
87                            } else {
88                                current_xref_offset = None;
89                            }
90                        } else {
91                            return Err(PdfError::Parse {
92                                offset: lexer.pos(),
93                                message: "Expected XRef stream object".to_string(),
94                            });
95                        }
96                    } else {
97                        return Err(PdfError::Parse {
98                            offset: lexer.pos(),
99                            message: "Invalid xref table/stream start".to_string(),
100                        });
101                    }
102                }
103                _ => {
104                    return Err(PdfError::Parse {
105                        offset: lexer.pos(),
106                        message: format!("Expected 'xref' or xref stream, found {:?}", token),
107                    });
108                }
109            }
110        }
111
112        doc.trailer = final_trailer;
113
114        // Load all normal objects
115        for (&_obj_id, entry) in &xref_entries {
116            if let XRefEntry::Normal {
117                offset,
118                generation: _,
119            } = entry
120            {
121                if *offset < self.data.len() {
122                    let mut obj_lexer = Lexer::with_pos(self.data, *offset);
123                    if let Ok((id, obj)) = self.parse_indirect_object(&mut obj_lexer) {
124                        doc.set_object(id, obj);
125                    }
126                }
127            }
128        }
129
130        // Load compressed objects from object streams (/ObjStm)
131        let mut obj_streams: BTreeMap<u32, COSStream> = BTreeMap::new();
132        for (id, obj) in &doc.objects {
133            if let COSObject::Stream(s) = obj {
134                if s.dict.get_name("Type") == Some("ObjStm") {
135                    obj_streams.insert(id.num, s.clone());
136                }
137            }
138        }
139
140        for (&_obj_id, entry) in &xref_entries {
141            if let XRefEntry::Compressed {
142                obj_stm_num,
143                index_in_stream,
144            } = entry
145            {
146                if let Some(stm) = obj_streams.get(obj_stm_num) {
147                    let decomp_res: PdfResult<Vec<u8>> = stm.decode();
148                    if let Ok(decompressed) = decomp_res {
149                        let n = stm.dict.get_integer("N").unwrap_or(0) as usize;
150                        let first = stm.dict.get_integer("First").unwrap_or(0) as usize;
151
152                        let mut lex = Lexer::new(&decompressed);
153                        let mut offsets = Vec::new();
154                        for _ in 0..n {
155                            let obj_num_tok = lex.next_token()?;
156                            let offset_tok = lex.next_token()?;
157                            if let (Token::Integer(on), Token::Integer(off)) =
158                                (obj_num_tok, offset_tok)
159                            {
160                                offsets.push((on as u32, first + off as usize));
161                            }
162                        }
163
164                        if let Some(&(target_num, byte_off)) =
165                            offsets.get(*index_in_stream as usize)
166                        {
167                            if byte_off < decompressed.len() {
168                                let mut obj_lex = Lexer::with_pos(&decompressed, byte_off);
169                                if let Ok(obj) = self.parse_direct_object(&mut obj_lex) {
170                                    doc.set_object(ObjectId::new(target_num, 0), obj);
171                                }
172                            }
173                        }
174                    }
175                }
176            }
177        }
178
179        doc.clear_dirty();
180        Ok(doc)
181    }
182
183    fn parse_header(&self) -> PdfResult<String> {
184        let header_prefix = b"%PDF-";
185        if let Some(idx) = self
186            .data
187            .windows(header_prefix.len())
188            .position(|w| w == header_prefix)
189        {
190            let start = idx + header_prefix.len();
191            let mut end = start;
192            while end < self.data.len() && self.data[end] != b'\r' && self.data[end] != b'\n' {
193                end += 1;
194            }
195            let ver = String::from_utf8_lossy(&self.data[start..end]).into_owned();
196            Ok(ver)
197        } else {
198            Err(PdfError::InvalidFormat("Missing %PDF header".to_string()))
199        }
200    }
201
202    pub fn find_startxref(&self) -> PdfResult<usize> {
203        let pattern = b"startxref";
204        let search_window = 2048.min(self.data.len());
205        let start = self.data.len() - search_window;
206        let slice = &self.data[start..];
207
208        if let Some(rel_pos) = slice.windows(pattern.len()).rposition(|w| w == pattern) {
209            let pos = start + rel_pos + pattern.len();
210            let mut lexer = Lexer::with_pos(self.data, pos);
211            let token = lexer.next_token()?;
212            if let Token::Integer(offset) = token {
213                return Ok(offset as usize);
214            }
215        }
216
217        Err(PdfError::InvalidFormat(
218            "Could not locate 'startxref'".to_string(),
219        ))
220    }
221
222    fn parse_classic_xref_section(
223        &self,
224        lexer: &mut Lexer<'a>,
225        entries: &mut BTreeMap<ObjectId, XRefEntry>,
226    ) -> PdfResult<COSDictionary> {
227        loop {
228            let token = lexer.next_token()?;
229            match token {
230                Token::Keyword(ref k) if k == "trailer" => {
231                    let trailer_obj = self.parse_direct_object(lexer)?;
232                    if let COSObject::Dictionary(dict) = trailer_obj {
233                        return Ok(dict);
234                    } else {
235                        return Err(PdfError::Parse {
236                            offset: lexer.pos(),
237                            message: "Expected trailer dictionary".to_string(),
238                        });
239                    }
240                }
241                Token::Integer(start_num) => {
242                    let count_token = lexer.next_token()?;
243                    if let Token::Integer(count) = count_token {
244                        for i in 0..count {
245                            let obj_num = (start_num + i) as u32;
246                            let off_token = lexer.next_token()?;
247                            let gen_token = lexer.next_token()?;
248                            let type_token = lexer.next_token()?;
249
250                            if let (
251                                Token::Integer(off),
252                                Token::Integer(gen_num),
253                                Token::Keyword(t),
254                            ) = (off_token, gen_token, type_token)
255                            {
256                                let obj_id = ObjectId::new(obj_num, gen_num as u16);
257                                if t == "n" {
258                                    entries.entry(obj_id).or_insert(XRefEntry::Normal {
259                                        offset: off as usize,
260                                        generation: gen_num as u16,
261                                    });
262                                } else if t == "f" {
263                                    entries.entry(obj_id).or_insert(XRefEntry::Free);
264                                }
265                            }
266                        }
267                    }
268                }
269                Token::EOF => {
270                    return Err(PdfError::Parse {
271                        offset: lexer.pos(),
272                        message: "Unexpected EOF while parsing xref table".to_string(),
273                    });
274                }
275                _ => {}
276            }
277        }
278    }
279
280    fn parse_xref_stream(
281        &self,
282        stream: &COSStream,
283        entries: &mut BTreeMap<ObjectId, XRefEntry>,
284    ) -> PdfResult<()> {
285        let w_arr = stream
286            .dict
287            .get_array("W")
288            .ok_or_else(|| PdfError::InvalidFormat("Missing /W in xref stream".to_string()))?;
289        if w_arr.len() < 3 {
290            return Err(PdfError::InvalidFormat(
291                "Invalid /W array in xref stream".to_string(),
292            ));
293        }
294
295        let w1 = w_arr.get(0).and_then(|o| o.as_integer()).unwrap_or(0) as usize;
296        let w2 = w_arr.get(1).and_then(|o| o.as_integer()).unwrap_or(0) as usize;
297        let w3 = w_arr.get(2).and_then(|o| o.as_integer()).unwrap_or(0) as usize;
298        let entry_size = w1 + w2 + w3;
299
300        let index_arr = stream.dict.get_array("Index");
301        let ranges: Vec<(u32, u32)> = if let Some(arr) = index_arr {
302            let mut list = Vec::new();
303            let mut iter = arr.iter();
304            while let (Some(a), Some(b)) = (iter.next(), iter.next()) {
305                if let (Some(start), Some(count)) = (a.as_integer(), b.as_integer()) {
306                    list.push((start as u32, count as u32));
307                }
308            }
309            list
310        } else {
311            let size = stream.dict.get_integer("Size").unwrap_or(0) as u32;
312            vec![(0, size)]
313        };
314
315        let decoded = stream.decode()?;
316        let mut byte_offset = 0;
317
318        for (start_num, count) in ranges {
319            for i in 0..count {
320                if byte_offset + entry_size > decoded.len() {
321                    break;
322                }
323                let obj_num = start_num + i;
324                let entry_bytes = &decoded[byte_offset..byte_offset + entry_size];
325                byte_offset += entry_size;
326
327                let read_int = |slice: &[u8]| -> u64 {
328                    let mut val = 0u64;
329                    for &b in slice {
330                        val = (val << 8) | (b as u64);
331                    }
332                    val
333                };
334
335                let t = if w1 > 0 {
336                    read_int(&entry_bytes[0..w1])
337                } else {
338                    1
339                };
340                let f2 = read_int(&entry_bytes[w1..w1 + w2]);
341                let f3 = read_int(&entry_bytes[w1 + w2..w1 + w2 + w3]);
342
343                match t {
344                    0 => {
345                        entries
346                            .entry(ObjectId::new(obj_num, f3 as u16))
347                            .or_insert(XRefEntry::Free);
348                    }
349                    1 => {
350                        entries.entry(ObjectId::new(obj_num, f3 as u16)).or_insert(
351                            XRefEntry::Normal {
352                                offset: f2 as usize,
353                                generation: f3 as u16,
354                            },
355                        );
356                    }
357                    2 => {
358                        entries
359                            .entry(ObjectId::new(obj_num, 0))
360                            .or_insert(XRefEntry::Compressed {
361                                obj_stm_num: f2 as u32,
362                                index_in_stream: f3 as u32,
363                            });
364                    }
365                    _ => {}
366                }
367            }
368        }
369
370        Ok(())
371    }
372
373    pub fn parse_indirect_object(&self, lexer: &mut Lexer<'a>) -> PdfResult<(ObjectId, COSObject)> {
374        let num = match lexer.next_token()? {
375            Token::Integer(i) => i as u32,
376            tok => {
377                return Err(PdfError::Parse {
378                    offset: lexer.pos(),
379                    message: format!("Expected object number, found {:?}", tok),
380                });
381            }
382        };
383
384        let gen_num = match lexer.next_token()? {
385            Token::Integer(i) => i as u16,
386            tok => {
387                return Err(PdfError::Parse {
388                    offset: lexer.pos(),
389                    message: format!("Expected generation number, found {:?}", tok),
390                });
391            }
392        };
393
394        match lexer.next_token()? {
395            Token::Keyword(ref k) if k == "obj" => {}
396            tok => {
397                return Err(PdfError::Parse {
398                    offset: lexer.pos(),
399                    message: format!("Expected 'obj' keyword, found {:?}", tok),
400                });
401            }
402        };
403
404        let obj = self.parse_direct_object(lexer)?;
405        let id = ObjectId::new(num, gen_num);
406        Ok((id, obj))
407    }
408
409    pub fn parse_direct_object(&self, lexer: &mut Lexer<'a>) -> PdfResult<COSObject> {
410        let token = lexer.next_token()?;
411        match token {
412            Token::Keyword(k) => match k.as_str() {
413                "null" => Ok(COSObject::Null),
414                "true" => Ok(COSObject::Boolean(true)),
415                "false" => Ok(COSObject::Boolean(false)),
416                _ => Ok(COSObject::Name(COSName::new(k))),
417            },
418            Token::Integer(i) => {
419                // Lookahead to check if this is an indirect reference "num gen R"
420                let saved_pos = lexer.pos();
421                let next1 = lexer.next_token()?;
422                if let Token::Integer(gen_num) = next1 {
423                    let next2 = lexer.next_token()?;
424                    if let Token::Keyword(ref r) = next2 {
425                        if r == "R" {
426                            return Ok(COSObject::Reference(ObjectId::new(
427                                i as u32,
428                                gen_num as u16,
429                            )));
430                        }
431                    }
432                }
433                // Backtrack if not an indirect reference
434                lexer.set_pos(saved_pos);
435                Ok(COSObject::Integer(i))
436            }
437            Token::Real(f) => Ok(COSObject::Real(f)),
438            Token::String(bytes, fmt) => Ok(COSObject::String(COSString {
439                data: bytes,
440                format: fmt,
441            })),
442            Token::Name(name) => Ok(COSObject::Name(COSName::new(name))),
443            Token::ArrayStart => {
444                let mut arr = COSArray::new();
445                loop {
446                    lexer.skip_whitespace_and_comments();
447                    if lexer.peek_byte() == Some(b']') {
448                        lexer.next_byte();
449                        break;
450                    }
451                    let item = self.parse_direct_object(lexer)?;
452                    arr.push(item);
453                }
454                Ok(COSObject::Array(arr))
455            }
456            Token::DictStart => {
457                let mut dict = COSDictionary::new();
458                loop {
459                    lexer.skip_whitespace_and_comments();
460                    if lexer.peek_byte() == Some(b'>') {
461                        lexer.next_token()?; // consume DictEnd
462                        break;
463                    }
464                    let key_tok = lexer.next_token()?;
465                    let key_name = match key_tok {
466                        Token::Name(n) => n,
467                        Token::DictEnd => break,
468                        _ => {
469                            return Err(PdfError::Parse {
470                                offset: lexer.pos(),
471                                message: format!(
472                                    "Expected dictionary key name, found {:?}",
473                                    key_tok
474                                ),
475                            });
476                        }
477                    };
478
479                    let val = self.parse_direct_object(lexer)?;
480                    dict.set(COSName::new(key_name), val);
481                }
482
483                // Check if this dictionary is followed by 'stream'
484                let saved_pos = lexer.pos();
485                let next_tok = lexer.next_token()?;
486                if let Token::Keyword(ref k) = next_tok {
487                    if k == "stream" {
488                        let stream_len = dict.get_integer("Length").map(|l| l as usize);
489                        let stream_data = lexer.read_stream_data(stream_len)?;
490                        // skip 'endstream'
491                        lexer.skip_whitespace_and_comments();
492                        let end_tok = lexer.next_token()?;
493                        if !matches!(end_tok, Token::Keyword(ref s) if s == "endstream") {
494                            // Some files don't strictly separate endstream, that's fine
495                        }
496                        return Ok(COSObject::Stream(COSStream::new(dict, stream_data)));
497                    }
498                }
499
500                // Not a stream, backtrack
501                lexer.set_pos(saved_pos);
502                Ok(COSObject::Dictionary(dict))
503            }
504            Token::DictEnd | Token::ArrayEnd | Token::EOF => Err(PdfError::Parse {
505                offset: lexer.pos(),
506                message: format!("Unexpected token: {:?}", token),
507            }),
508        }
509    }
510}
rust_pdf_box/cos/parser.rs

rust_pdf_box/cos/
parser.rs
