1use std::collections::HashMap;
2
3use zpdf_core::{Error, ObjectId, ParseLimits, PdfDict, PdfObject, Result};
4
5use crate::lexer::Lexer;
6
7#[derive(Debug, Clone)]
8pub enum XrefEntry {
9 InUse {
10 offset: u64,
11 gen: u16,
12 },
13 Free {
14 next: u32,
15 gen: u16,
16 },
17 Compressed {
18 stream_obj: u32,
19 index_in_stream: u32,
20 },
21}
22
23#[derive(Debug, Clone, Default)]
24pub struct XrefTable {
25 entries: HashMap<ObjectId, XrefEntry>,
26}
27
28impl XrefTable {
29 pub fn new() -> Self {
30 Self::default()
31 }
32
33 pub fn get(&self, id: ObjectId) -> Option<&XrefEntry> {
34 self.entries.get(&id)
35 }
36
37 pub fn insert(&mut self, id: ObjectId, entry: XrefEntry) {
38 self.entries.entry(id).or_insert(entry);
39 }
40
41 pub fn insert_overwrite(&mut self, id: ObjectId, entry: XrefEntry) {
46 self.entries.insert(id, entry);
47 }
48
49 pub fn len(&self) -> usize {
50 self.entries.len()
51 }
52
53 pub fn is_empty(&self) -> bool {
54 self.entries.is_empty()
55 }
56
57 pub fn object_ids(&self) -> impl Iterator<Item = ObjectId> + '_ {
58 self.entries.keys().copied()
59 }
60}
61
62pub fn parse_xref_and_trailer(data: &[u8], limits: &ParseLimits) -> Result<(XrefTable, PdfDict)> {
63 let startxref_offset = find_startxref(data)?;
64 let mut table = XrefTable::new();
65
66 let xref_offset = startxref_offset;
67 let (trailer, next_prev) = parse_xref_section(data, xref_offset, &mut table, limits)?;
68
69 let mut visited = std::collections::HashSet::new();
73 visited.insert(xref_offset);
74 parse_hybrid_xrefstm(data, &trailer, &mut table, limits, &mut visited);
77 let mut prev = next_prev;
78 while let Some(prev_offset) = prev {
79 let prev_offset_usize =
80 usize::try_from(prev_offset).map_err(|_| Error::InvalidXref(prev_offset))?;
81 if !visited.insert(prev_offset_usize) {
82 break;
83 }
84 let (section_trailer, next) =
85 parse_xref_section(data, prev_offset_usize, &mut table, limits)?;
86 parse_hybrid_xrefstm(data, §ion_trailer, &mut table, limits, &mut visited);
87 prev = next;
88 }
89
90 Ok((table, trailer))
91}
92
93fn parse_hybrid_xrefstm(
101 data: &[u8],
102 trailer: &PdfDict,
103 table: &mut XrefTable,
104 limits: &ParseLimits,
105 visited: &mut std::collections::HashSet<usize>,
106) {
107 let Some(PdfObject::Integer(off)) = trailer.get("XRefStm") else {
108 return;
109 };
110 let Ok(off) = usize::try_from(*off) else {
111 tracing::warn!("/XRefStm offset {off} is negative; ignoring");
112 return;
113 };
114 if !visited.insert(off) {
116 return;
117 }
118 if let Err(e) = parse_xref_stream(data, off, table, limits) {
119 tracing::warn!("failed to parse /XRefStm at offset {off}: {e}");
120 }
121}
122
123fn parse_xref_section(
124 data: &[u8],
125 offset: usize,
126 table: &mut XrefTable,
127 limits: &ParseLimits,
128) -> Result<(PdfDict, Option<u64>)> {
129 if offset >= data.len() {
133 return Err(Error::InvalidXref(offset as u64));
134 }
135 if data[offset..].starts_with(b"xref") {
136 parse_traditional_xref(data, offset, table, limits)
137 } else {
138 parse_xref_stream(data, offset, table, limits)
139 }
140}
141
142fn parse_xref_stream(
143 data: &[u8],
144 offset: usize,
145 table: &mut XrefTable,
146 limits: &ParseLimits,
147) -> Result<(PdfDict, Option<u64>)> {
148 use crate::filters;
149 use crate::object_parser::ObjectParser;
150
151 let parser = ObjectParser::new(data, limits);
152 let obj = parser.parse_indirect_at(offset)?;
153 let stream = match obj {
154 PdfObject::Stream(s) => s,
155 _ => return Err(Error::InvalidXref(offset as u64)),
156 };
157
158 let dict = &stream.dict;
159 if dict.get_name("Type").unwrap_or("") != "XRef" {
160 return Err(Error::InvalidXref(offset as u64));
161 }
162
163 let size =
164 u32::try_from(dict.get_i64("Size")?).map_err(|_| Error::InvalidXref(offset as u64))?;
165
166 let w_arr = dict.get_array("W")?;
170 if w_arr.len() != 3 {
171 return Err(Error::InvalidXref(offset as u64));
172 }
173 let field_width = |obj: &PdfObject| -> Result<usize> {
174 let w = obj.as_i64()?;
175 if !(0..=8).contains(&w) {
176 return Err(Error::InvalidXref(offset as u64));
177 }
178 Ok(w as usize)
179 };
180 let w1 = field_width(&w_arr[0])?;
181 let w2 = field_width(&w_arr[1])?;
182 let w3 = field_width(&w_arr[2])?;
183 let entry_size = w1 + w2 + w3;
184 if entry_size == 0 {
185 return Err(Error::InvalidXref(offset as u64));
186 }
187
188 let decoded = filters::decode_stream_with_limits(&stream.data, dict, limits)?;
190
191 let index_ranges: Vec<(u32, u32)> = if let Ok(idx_arr) = dict.get_array("Index") {
193 if !idx_arr.len().is_multiple_of(2) {
194 return Err(Error::InvalidXref(offset as u64));
195 }
196 let mut ranges = Vec::with_capacity(idx_arr.len() / 2);
197 for pair in idx_arr.chunks_exact(2) {
198 let start =
199 u32::try_from(pair[0].as_i64()?).map_err(|_| Error::InvalidXref(offset as u64))?;
200 let count =
201 u32::try_from(pair[1].as_i64()?).map_err(|_| Error::InvalidXref(offset as u64))?;
202 ranges.push((start, count));
203 }
204 ranges
205 } else {
206 vec![(0, size)]
207 };
208
209 let mut pos = 0usize;
210 for &(start, count) in &index_ranges {
211 let _range_end = start
213 .checked_add(count)
214 .ok_or(Error::InvalidXref(offset as u64))?;
215
216 let new_total = table
218 .len()
219 .checked_add(count as usize)
220 .ok_or(Error::InvalidXref(offset as u64))?;
221 if new_total > limits.max_objects as usize {
222 return Err(Error::InvalidXref(offset as u64));
223 }
224
225 for i in 0..count {
226 if pos + entry_size > decoded.len() {
227 break;
228 }
229 let obj_num = start + i;
230
231 let field1 = read_field(&decoded[pos..], w1);
232 let field2 = read_field(&decoded[pos + w1..], w2);
233 let field3 = read_field(&decoded[pos + w1 + w2..], w3);
234 pos += entry_size;
235
236 let entry_type = if w1 == 0 {
237 1
238 } else {
239 u8::try_from(field1).map_err(|_| Error::InvalidXref(offset as u64))?
240 };
241
242 match entry_type {
243 0 => {
244 let next =
245 u32::try_from(field2).map_err(|_| Error::InvalidXref(offset as u64))?;
246 let gen =
247 u16::try_from(field3).map_err(|_| Error::InvalidXref(offset as u64))?;
248 table.insert(ObjectId(obj_num, gen), XrefEntry::Free { next, gen });
249 }
250 1 => {
251 let gen =
252 u16::try_from(field3).map_err(|_| Error::InvalidXref(offset as u64))?;
253 table.insert(
254 ObjectId(obj_num, gen),
255 XrefEntry::InUse {
256 offset: field2,
257 gen,
258 },
259 );
260 }
261 2 => {
262 let stream_obj =
263 u32::try_from(field2).map_err(|_| Error::InvalidXref(offset as u64))?;
264 let index_in_stream =
265 u32::try_from(field3).map_err(|_| Error::InvalidXref(offset as u64))?;
266 table.insert(
267 ObjectId(obj_num, 0),
268 XrefEntry::Compressed {
269 stream_obj,
270 index_in_stream,
271 },
272 );
273 }
274 _ => {}
275 }
276 }
277 }
278
279 let trailer = dict.clone();
281 let prev = trailer.get("Prev").and_then(|obj| match obj {
282 PdfObject::Integer(n) => u64::try_from(*n).ok(),
283 _ => None,
284 });
285
286 Ok((trailer, prev))
287}
288
289fn read_field(data: &[u8], width: usize) -> u64 {
290 let mut val = 0u64;
291 for &byte in &data[..width] {
292 val = (val << 8) | byte as u64;
293 }
294 val
295}
296
297fn parse_traditional_xref(
298 data: &[u8],
299 offset: usize,
300 table: &mut XrefTable,
301 limits: &ParseLimits,
302) -> Result<(PdfDict, Option<u64>)> {
303 let mut pos = offset + 4; skip_eol(data, &mut pos);
305
306 loop {
308 skip_whitespace(data, &mut pos);
309
310 if pos >= data.len() {
314 return Err(Error::InvalidXref(pos as u64));
315 }
316
317 if data[pos..].starts_with(b"trailer") {
318 pos += 7;
319 break;
320 }
321
322 let (first_obj, count) = parse_subsection_header(data, &mut pos)?;
324
325 let _range_end = first_obj
327 .checked_add(count)
328 .ok_or(Error::InvalidXref(pos as u64))?;
329
330 let new_total = table
335 .len()
336 .checked_add(count as usize)
337 .ok_or(Error::InvalidXref(pos as u64))?;
338 if new_total > limits.max_objects as usize {
339 return Err(Error::InvalidXref(pos as u64));
340 }
341
342 for i in 0..count {
343 skip_whitespace(data, &mut pos);
344 let (entry_offset, gen, in_use) = parse_xref_entry_at(data, &mut pos)?;
350 let id = ObjectId(first_obj + i, gen);
351
352 if in_use {
353 table.insert(
354 id,
355 XrefEntry::InUse {
356 offset: entry_offset,
357 gen,
358 },
359 );
360 } else {
361 table.insert(
362 id,
363 XrefEntry::Free {
364 next: u32::try_from(entry_offset)
365 .map_err(|_| Error::InvalidXref(pos as u64))?,
366 gen,
367 },
368 );
369 }
370 }
371 }
372
373 let mut lex = Lexer::new(data, pos, limits);
375 let trailer_obj = lex.next_token()?;
376 let trailer = match trailer_obj {
377 PdfObject::Dict(d) => d,
378 _ => return Err(Error::InvalidXref(pos as u64)),
379 };
380
381 let prev = trailer.get("Prev").and_then(|obj| match obj {
382 PdfObject::Integer(n) => u64::try_from(*n).ok(),
383 _ => None,
384 });
385
386 Ok((trailer, prev))
387}
388
389fn find_startxref(data: &[u8]) -> Result<usize> {
390 let marker = b"startxref";
397 let marker_pos = data
398 .windows(marker.len())
399 .rposition(|w| w == marker)
400 .ok_or(Error::InvalidXref(0))?;
401
402 let after_marker = marker_pos + marker.len();
403 let num_start = data[after_marker..]
404 .iter()
405 .position(|b| b.is_ascii_digit())
406 .ok_or(Error::InvalidXref(0))?;
407
408 let num_bytes = &data[after_marker + num_start..];
409 let num_end = num_bytes
410 .iter()
411 .position(|b| !b.is_ascii_digit())
412 .unwrap_or(num_bytes.len());
413
414 let offset_str =
415 std::str::from_utf8(&num_bytes[..num_end]).map_err(|_| Error::InvalidXref(0))?;
416 let offset: usize = offset_str.parse().map_err(|_| Error::InvalidXref(0))?;
417
418 Ok(offset)
419}
420
421fn parse_subsection_header(data: &[u8], pos: &mut usize) -> Result<(u32, u32)> {
422 let start = *pos;
423 while *pos < data.len() && data[*pos].is_ascii_digit() {
424 *pos += 1;
425 }
426 let first: u32 = std::str::from_utf8(&data[start..*pos])
427 .map_err(|_| Error::InvalidXref(start as u64))?
428 .parse()
429 .map_err(|_| Error::InvalidXref(start as u64))?;
430
431 skip_whitespace(data, pos);
432
433 let count_start = *pos;
434 while *pos < data.len() && data[*pos].is_ascii_digit() {
435 *pos += 1;
436 }
437 let count: u32 = std::str::from_utf8(&data[count_start..*pos])
438 .map_err(|_| Error::InvalidXref(count_start as u64))?
439 .parse()
440 .map_err(|_| Error::InvalidXref(count_start as u64))?;
441
442 skip_eol(data, pos);
443 Ok((first, count))
444}
445
446fn parse_xref_entry_at(data: &[u8], pos: &mut usize) -> Result<(u64, u16, bool)> {
451 let start = *pos as u64;
452 let offset = read_decimal(data, pos).ok_or(Error::InvalidXref(start))?;
453 skip_whitespace(data, pos);
454 let gen = read_decimal(data, pos)
455 .and_then(|g| u16::try_from(g).ok())
456 .ok_or(Error::InvalidXref(start))?;
457 skip_whitespace(data, pos);
458 let in_use = match data.get(*pos) {
459 Some(b'n') => true,
460 Some(b'f') => false,
461 _ => return Err(Error::InvalidXref(start)),
462 };
463 *pos += 1;
464 Ok((offset, gen, in_use))
465}
466
467fn read_decimal(data: &[u8], pos: &mut usize) -> Option<u64> {
470 let start = *pos;
471 while *pos < data.len() && data[*pos].is_ascii_digit() {
472 *pos += 1;
473 }
474 if *pos == start {
475 return None;
476 }
477 std::str::from_utf8(&data[start..*pos]).ok()?.parse().ok()
478}
479
480fn skip_whitespace(data: &[u8], pos: &mut usize) {
481 while *pos < data.len() && matches!(data[*pos], b' ' | b'\t' | b'\r' | b'\n') {
482 *pos += 1;
483 }
484}
485
486fn skip_eol(data: &[u8], pos: &mut usize) {
487 while *pos < data.len() && matches!(data[*pos], b' ' | b'\t' | b'\r' | b'\n') {
488 *pos += 1;
489 }
490}
491
492#[cfg(test)]
493mod tests {
494 use super::*;
495
496 #[test]
497 fn find_startxref_offset() {
498 let data = b"%PDF-1.4\n...lots of content...\nstartxref\n1234\n%%EOF";
499 let offset = find_startxref(data).unwrap();
500 assert_eq!(offset, 1234);
501 }
502
503 #[test]
504 fn parse_xref_entry_in_use() {
505 let mut pos = 0usize;
506 let (offset, gen, in_use) =
507 parse_xref_entry_at(b"0000000010 00000 n \r\n", &mut pos).unwrap();
508 assert_eq!(offset, 10);
509 assert_eq!(gen, 0);
510 assert!(in_use);
511 assert_eq!(pos, 18, "cursor stops just past the type letter");
512 }
513
514 #[test]
515 fn parse_xref_entry_free() {
516 let mut pos = 0usize;
517 let (offset, gen, in_use) =
518 parse_xref_entry_at(b"0000000000 65535 f \r\n", &mut pos).unwrap();
519 assert_eq!(offset, 0);
520 assert_eq!(gen, 65535);
521 assert!(!in_use);
522 }
523
524 #[test]
525 fn parse_xref_entry_truncated_errors() {
526 let mut pos = 0usize;
527 assert!(parse_xref_entry_at(b"0000000010 000", &mut pos).is_err());
528 }
529
530 #[test]
531 fn traditional_xref_with_19_byte_entries() {
532 let mut d = Vec::new();
534 d.extend_from_slice(b"%PDF-1.4\n");
535 let off1 = d.len();
536 d.extend_from_slice(b"1 0 obj\n<< /Type /Catalog /Pages 2 0 R >>\nendobj\n");
537 let off2 = d.len();
538 d.extend_from_slice(b"2 0 obj\n<< /Type /Pages /Kids [] /Count 0 >>\nendobj\n");
539 let xref_off = d.len();
540 d.extend_from_slice(b"xref\n0 3\n");
541 d.extend_from_slice(b"0000000000 65535 f\n"); d.extend_from_slice(format!("{off1:010} 00000 n\n").as_bytes()); d.extend_from_slice(format!("{off2:010} 00000 n\n").as_bytes()); d.extend_from_slice(
545 format!("trailer\n<< /Size 3 /Root 1 0 R >>\nstartxref\n{xref_off}\n%%EOF\n")
546 .as_bytes(),
547 );
548
549 let (table, trailer) = parse_xref_and_trailer(&d, &ParseLimits::default()).unwrap();
550 assert_eq!(trailer.get_ref("Root").unwrap(), ObjectId(1, 0));
551 match table.get(ObjectId(1, 0)).unwrap() {
552 XrefEntry::InUse { offset, .. } => assert_eq!(*offset as usize, off1),
553 other => panic!("expected InUse, got {other:?}"),
554 }
555 match table.get(ObjectId(2, 0)).unwrap() {
556 XrefEntry::InUse { offset, .. } => assert_eq!(*offset as usize, off2),
557 other => panic!("expected InUse, got {other:?}"),
558 }
559 }
560
561 fn xref_stream_bytes(w: &str, size: u32, index: &str, body: &[u8]) -> Vec<u8> {
564 let mut d = format!(
565 "9 0 obj\n<< /Type /XRef /Size {size} /W {w} {index} /Length {} >>\nstream\n",
566 body.len()
567 )
568 .into_bytes();
569 d.extend_from_slice(body);
570 d.extend_from_slice(b"\nendstream\nendobj\n");
571 d
572 }
573
574 #[test]
575 fn xref_stream_rejects_negative_w_width() {
576 let d = xref_stream_bytes("[1 -2 2]", 1, "", &[]);
577 let mut table = XrefTable::new();
578 assert!(parse_xref_stream(&d, 0, &mut table, &ParseLimits::default()).is_err());
579 }
580
581 #[test]
582 fn xref_stream_rejects_oversized_w_width() {
583 let d = xref_stream_bytes("[9 4 2]", 1, "", &[]);
584 let mut table = XrefTable::new();
585 assert!(parse_xref_stream(&d, 0, &mut table, &ParseLimits::default()).is_err());
586 }
587
588 #[test]
589 fn xref_stream_rejects_zero_entry_size() {
590 let d = xref_stream_bytes("[0 0 0]", 1, "", &[]);
591 let mut table = XrefTable::new();
592 assert!(parse_xref_stream(&d, 0, &mut table, &ParseLimits::default()).is_err());
593 }
594
595 #[test]
596 fn xref_stream_rejects_negative_index_values() {
597 let d = xref_stream_bytes("[1 1 1]", 1, "/Index [-1 1]", &[1, 0, 0]);
598 let mut table = XrefTable::new();
599 assert!(parse_xref_stream(&d, 0, &mut table, &ParseLimits::default()).is_err());
600 }
601
602 #[test]
603 fn xref_stream_rejects_truncating_entry_fields() {
604 let d = xref_stream_bytes("[2 0 0]", 1, "/Index [0 1]", &[1, 1]);
606 let mut table = XrefTable::new();
607 assert!(parse_xref_stream(&d, 0, &mut table, &ParseLimits::default()).is_err());
608
609 let d = xref_stream_bytes("[1 1 3]", 1, "/Index [0 1]", &[1, 0, 1, 0, 0]);
611 let mut table = XrefTable::new();
612 assert!(parse_xref_stream(&d, 0, &mut table, &ParseLimits::default()).is_err());
613 }
614
615 #[test]
616 fn sparse_xref_stream_limits_entry_count_not_object_number() {
617 let d = xref_stream_bytes("[1 1 1]", 1_000_001, "/Index [1000000 1]", &[1, 0, 0]);
618 let limits = ParseLimits {
619 max_objects: 1,
620 ..ParseLimits::default()
621 };
622 let mut table = XrefTable::new();
623 parse_xref_stream(&d, 0, &mut table, &limits).unwrap();
624 assert!(matches!(
625 table.get(ObjectId(1_000_000, 0)),
626 Some(XrefEntry::InUse { .. })
627 ));
628 }
629
630 #[test]
631 fn sparse_traditional_xref_limits_entry_count_not_object_number() {
632 let d = b"xref\n1000000 1\n0000000000 00000 n \ntrailer\n<< /Size 1000001 >>\n";
633 let limits = ParseLimits {
634 max_objects: 1,
635 ..ParseLimits::default()
636 };
637 let mut table = XrefTable::new();
638 parse_traditional_xref(d, 0, &mut table, &limits).unwrap();
639 assert!(matches!(
640 table.get(ObjectId(1_000_000, 0)),
641 Some(XrefEntry::InUse { .. })
642 ));
643 }
644
645 #[test]
646 fn hybrid_xrefstm_is_parsed_with_correct_precedence() {
647 let mut d = Vec::new();
652 d.extend_from_slice(b"%PDF-1.4\n");
653 let off1 = d.len();
654 d.extend_from_slice(b"1 0 obj\n<< /Type /Catalog /Pages 2 0 R >>\nendobj\n");
655 let off4_table = d.len();
656 d.extend_from_slice(b"4 0 obj\n<< /Marker /FromTable >>\nendobj\n");
657 let off4_stm = d.len();
658 d.extend_from_slice(b"4 0 obj\n<< /Marker /FromStm >>\nendobj\n");
659 let off5 = d.len();
660 d.extend_from_slice(b"5 0 obj\n<< /Marker /StmOnly >>\nendobj\n");
661
662 let mut body = Vec::new();
664 for (off, gen) in [(off4_stm as u32, 0u16), (off5 as u32, 0)] {
665 body.push(1u8); body.extend_from_slice(&off.to_be_bytes());
667 body.extend_from_slice(&gen.to_be_bytes());
668 }
669 let off6 = d.len();
670 d.extend_from_slice(
671 format!(
672 "6 0 obj\n<< /Type /XRef /Size 7 /W [1 4 2] /Index [4 2] /Length {} >>\nstream\n",
673 body.len()
674 )
675 .as_bytes(),
676 );
677 d.extend_from_slice(&body);
678 d.extend_from_slice(b"\nendstream\nendobj\n");
679
680 let xref_off = d.len();
681 d.extend_from_slice(b"xref\n0 2\n0000000000 65535 f \n");
682 d.extend_from_slice(format!("{off1:010} 00000 n \n").as_bytes());
683 d.extend_from_slice(b"4 1\n");
684 d.extend_from_slice(format!("{off4_table:010} 00000 n \n").as_bytes());
685 d.extend_from_slice(
686 format!(
687 "trailer\n<< /Size 7 /Root 1 0 R /XRefStm {off6} >>\nstartxref\n{xref_off}\n%%EOF\n"
688 )
689 .as_bytes(),
690 );
691
692 let (table, trailer) = parse_xref_and_trailer(&d, &ParseLimits::default()).unwrap();
693 assert_eq!(trailer.get_ref("Root").unwrap(), ObjectId(1, 0));
694 match table.get(ObjectId(5, 0)).unwrap() {
696 XrefEntry::InUse { offset, .. } => assert_eq!(*offset as usize, off5),
697 other => panic!("expected InUse from XRefStm, got {other:?}"),
698 }
699 match table.get(ObjectId(4, 0)).unwrap() {
701 XrefEntry::InUse { offset, .. } => assert_eq!(*offset as usize, off4_table),
702 other => panic!("expected InUse, got {other:?}"),
703 }
704
705 let file = crate::PdfFile::parse(d).unwrap();
707 let o5 = file.resolve(ObjectId(5, 0)).unwrap();
708 assert_eq!(o5.as_dict().unwrap().get_name("Marker").unwrap(), "StmOnly");
709 let o4 = file.resolve(ObjectId(4, 0)).unwrap();
710 assert_eq!(
711 o4.as_dict().unwrap().get_name("Marker").unwrap(),
712 "FromTable"
713 );
714 }
715}