1use std::collections::HashMap;
39use std::sync::{Arc, Mutex, OnceLock};
40
41use pdfrum_common::{DiagKind, Diagnostics, Limits, Severity};
42use pdfrum_crypt::{CryptClass, SecurityHandler};
43use pdfrum_object::{
44 Array, ByteSpan, Dict, Name, ObjRef, Object, PdfString, Resolve, Stream, names,
45};
46
47use crate::error::Error;
48use crate::lexer::Lexer;
49use crate::objstm::ObjStm;
50use crate::syntax::{Context, Strictness, indirect};
51use crate::xref::{Entry, Xref};
52
53#[derive(Debug)]
59pub struct ObjectStore {
60 bytes: ByteSpan,
62 xref: Arc<Xref>,
70 limits: Limits,
72 cells: Mutex<Cells>,
82 in_progress: Mutex<Vec<u32>>,
84 containers: Mutex<Containers>,
86 security: SecurityHandler,
88 metadata_exempt: Option<u32>,
91 diags: Mutex<Diagnostics>,
93}
94
95type Cells = HashMap<u32, Arc<OnceLock<Arc<Object>>>, pdfrum_common::FxBuildHasher>;
101
102type Containers = HashMap<u32, Option<Arc<ObjStm>>, pdfrum_common::FxBuildHasher>;
107
108impl ObjectStore {
109 pub(crate) fn new(
111 bytes: ByteSpan,
112 xref: Arc<Xref>,
113 limits: Limits,
114 security: SecurityHandler,
115 ) -> Self {
116 Self {
117 bytes,
118 xref,
119 limits,
120 cells: Mutex::new(HashMap::default()),
121 in_progress: Mutex::new(Vec::new()),
122 containers: Mutex::new(HashMap::default()),
123 security,
124 metadata_exempt: None,
125 diags: Mutex::new(Diagnostics::default()),
126 }
127 }
128
129 pub(crate) fn exempt_from_decryption(&mut self, num: u32) {
132 self.metadata_exempt = Some(num);
133 }
134
135 pub(crate) fn xref(&self) -> &Xref {
137 &self.xref
138 }
139
140 pub(crate) fn limits(&self) -> &Limits {
142 &self.limits
143 }
144
145 pub(crate) fn security(&self) -> &SecurityHandler {
147 &self.security
148 }
149
150 pub(crate) fn drain_diags(&self) -> Diagnostics {
152 match self.diags.lock() {
153 Ok(mut guard) => std::mem::take(&mut *guard),
154 Err(_) => Diagnostics::default(),
155 }
156 }
157
158 pub(crate) fn peek_diags(&self) -> Diagnostics {
164 match self.diags.lock() {
165 Ok(guard) => guard.clone(),
166 Err(_) => Diagnostics::default(),
167 }
168 }
169
170 pub(crate) fn note(&self, severity: Severity, what: DiagKind, at: Option<u64>) {
177 if let Ok(mut guard) = self.diags.lock() {
178 guard.record(severity, what, at);
179 }
180 }
181
182 fn with_diags<T>(&self, f: impl FnOnce(&mut Diagnostics) -> T) -> T {
184 let mut local = Diagnostics::default();
185 let out = f(&mut local);
186 if let Ok(mut guard) = self.diags.lock() {
187 guard.extend(&local);
188 }
189 out
190 }
191
192 pub fn get(&self, num: u32) -> Result<Arc<Object>, Error> {
200 let reference = ObjRef::new(num, self.xref.generation(num));
201 if num == 0 || reference.is_invalid() {
202 return Err(Error::Unresolved(reference));
203 }
204 if !self.xref.is_valid_object_number(num) {
206 return Err(Error::Unresolved(reference));
207 }
208
209 let cell = self.cell(num);
210 if let Some(object) = cell.get() {
211 return Ok(Arc::clone(object));
212 }
213
214 let guard = Guard::enter(self, num).ok_or(Error::Cycle(reference))?;
215 if let Some(object) = cell.get() {
217 return Ok(Arc::clone(object));
218 }
219
220 let object = self
225 .parse(num, guard.nesting())
226 .ok_or(Error::Unresolved(reference))?;
227 let object = Arc::new(object);
228 let _ = cell.set(Arc::clone(&object));
230 Ok(cell.get().map_or(object, Arc::clone))
231 }
232
233 fn cell(&self, num: u32) -> Arc<OnceLock<Arc<Object>>> {
235 match self.cells.lock() {
236 Ok(mut cells) => Arc::clone(cells.entry(num).or_default()),
237 Err(_) => Arc::new(OnceLock::new()),
238 }
239 }
240
241 fn parse(&self, num: u32, depth: u32) -> Option<Object> {
247 match self.xref.entry(num)? {
248 Entry::Offset(pos) if pos > 0 => self.parse_at(num, pos, depth),
249 Entry::Free | Entry::Offset(_) => None,
252 Entry::InObjStream { stream, index } => {
253 self.parse_member(num, stream.num, index, depth)
254 }
255 }
256 }
257
258 fn parse_at(&self, num: u32, pos: u64, depth: u32) -> Option<Object> {
263 let pos = usize::try_from(pos).ok()?;
264 if pos >= self.bytes.len() {
265 return None;
266 }
267 let parsed = self.with_diags(|diags| {
268 let mut ctx = Context {
269 limits: &self.limits,
270 diags,
271 file: Some(&self.bytes),
272 store: Some(self),
273 };
274 let mut lx = Lexer::at(&self.bytes, pos);
275 indirect(&mut lx, &mut ctx, Strictness::Loose, depth).ok()
276 })?;
277
278 if parsed.num != num {
279 self.note(
280 Severity::Suspicious,
281 DiagKind::ObjNumMismatch,
282 Some(pos as u64),
283 );
284 return None;
285 }
286 Some(self.decrypted(ObjRef::new(parsed.num, parsed.generation), parsed.object))
287 }
288
289 fn parse_member(&self, num: u32, archive: u32, index: u32, depth: u32) -> Option<Object> {
291 let container = self.container(archive)?;
292 self.with_diags(|diags| container.member(num, index, &self.limits, diags, self, depth))
294 }
295
296 fn container(&self, archive: u32) -> Option<Arc<ObjStm>> {
298 if !self.xref.is_object_stream(archive) {
300 return None;
301 }
302 if let Ok(cache) = self.containers.lock()
303 && let Some(hit) = cache.get(&archive)
304 {
305 return hit.clone();
306 }
307
308 let built = self.build_container(archive);
309 if let Ok(mut cache) = self.containers.lock() {
310 cache.insert(archive, built.clone());
311 }
312 built
313 }
314
315 fn build_container(&self, archive: u32) -> Option<Arc<ObjStm>> {
317 let Object::Stream(stream) = &*self.get(archive).ok()? else {
318 return None;
319 };
320 self.with_diags(|diags| ObjStm::build(stream, &self.limits, diags, self).map(Arc::new))
321 }
322
323 fn decrypted(&self, obj: ObjRef, object: Object) -> Object {
328 if matches!(self.security, SecurityHandler::Identity)
329 || self.metadata_exempt == Some(obj.num)
330 {
331 return object;
332 }
333 let mut deferred = Vec::new();
334 let out = decrypt_node(&self.security, obj, object, &mut deferred, false);
335 resolve_deferred(&self.security, obj, out, &deferred)
340 }
341}
342
343impl Resolve for ObjectStore {
344 fn fetch(&self, r: ObjRef) -> Result<Arc<Object>, pdfrum_object::Error> {
345 self.get(r.num).map_err(Into::into)
346 }
347}
348
349struct Guard<'s> {
351 store: &'s ObjectStore,
352 num: u32,
353 nesting: u32,
355}
356
357impl<'s> Guard<'s> {
358 fn enter(store: &'s ObjectStore, num: u32) -> Option<Self> {
360 let mut running = store.in_progress.lock().ok()?;
361 if running.contains(&num) {
362 return None;
363 }
364 let nesting = u32::try_from(running.len()).unwrap_or(u32::MAX);
365 running.push(num);
366 drop(running);
367 Some(Self {
368 store,
369 num,
370 nesting,
371 })
372 }
373
374 fn nesting(&self) -> u32 {
376 self.nesting
377 }
378}
379
380impl Drop for Guard<'_> {
381 fn drop(&mut self) {
382 if let Ok(mut running) = self.store.in_progress.lock() {
383 running.retain(|&n| n != self.num);
384 }
385 }
386}
387
388#[derive(Debug, Clone)]
391struct Deferred {
392 parent: Dict,
394 path: Vec<Step>,
396}
397
398#[derive(Debug, Clone)]
400enum Step {
401 Key(Name),
403 Index(usize),
405}
406
407fn decrypt_node(
412 handler: &SecurityHandler,
413 obj: ObjRef,
414 object: Object,
415 deferred: &mut Vec<Deferred>,
416 in_sig: bool,
417) -> Object {
418 match object {
419 Object::Str(s) => {
420 if in_sig {
421 Object::Str(s)
422 } else {
423 let plain = handler.decrypt(obj, CryptClass::String, &s.bytes);
424 Object::Str(PdfString::new(plain, syntax_of(&s)))
425 }
426 }
427 Object::Array(a) => Object::Array(
428 a.iter()
429 .map(|v| decrypt_node(handler, obj, v.clone(), deferred, in_sig))
430 .collect(),
431 ),
432 Object::Dict(d) => Object::Dict(decrypt_dict(handler, obj, &d, deferred, in_sig, &[])),
433 Object::Stream(s) => {
434 let dict = decrypt_dict(handler, obj, &s.dict, deferred, in_sig, &[]);
435 let class = if s.dict.name(names::TYPE) == Some(names::EMBEDDED_FILE) {
442 CryptClass::Embedded
443 } else {
444 CryptClass::Stream
445 };
446 let plain = handler.decrypt(obj, class, &s.data);
447 Object::Stream(Box::new(Stream::new(
448 dict,
449 pdfrum_object::ByteSpan::from(plain),
450 )))
451 }
452 other => other,
453 }
454}
455
456fn decrypt_dict(
459 handler: &SecurityHandler,
460 obj: ObjRef,
461 dict: &Dict,
462 deferred: &mut Vec<Deferred>,
463 in_sig: bool,
464 path: &[Step],
465) -> Dict {
466 let suspicious = dict.contains_key(names::TYPE) || dict.contains_key(names::FT);
469 let mut out = Dict::new();
470 let mut skipped: Vec<(Name, Object)> = Vec::new();
471
472 for (key, value) in dict.iter() {
473 if suspicious && key == names::CONTENTS && !in_sig {
474 skipped.push((key.clone(), value.clone()));
475 out.push(key.clone(), value.clone());
476 continue;
477 }
478 let mut child = path.to_vec();
479 child.push(Step::Key(key.clone()));
480 out.push(
481 key.clone(),
482 decrypt_value(handler, obj, value.clone(), deferred, in_sig, &child),
483 );
484 }
485
486 for (key, _) in skipped {
487 let mut child = path.to_vec();
488 child.push(Step::Key(key));
489 deferred.push(Deferred {
490 parent: out.clone(),
491 path: child,
492 });
493 }
494 out
495}
496
497fn decrypt_value(
500 handler: &SecurityHandler,
501 obj: ObjRef,
502 value: Object,
503 deferred: &mut Vec<Deferred>,
504 in_sig: bool,
505 path: &[Step],
506) -> Object {
507 match value {
508 Object::Dict(d) => Object::Dict(decrypt_dict(handler, obj, &d, deferred, in_sig, path)),
509 Object::Array(a) => {
510 let mut out = Array::new();
511 for (i, v) in a.iter().enumerate() {
512 let mut child = path.to_vec();
513 child.push(Step::Index(i));
514 out.push(decrypt_value(
515 handler,
516 obj,
517 v.clone(),
518 deferred,
519 in_sig,
520 &child,
521 ));
522 }
523 Object::Array(out)
524 }
525 other => decrypt_node(handler, obj, other, deferred, in_sig),
526 }
527}
528
529fn resolve_deferred(
532 handler: &SecurityHandler,
533 obj: ObjRef,
534 object: Object,
535 deferred: &[Deferred],
536) -> Object {
537 let mut out = object;
538 for entry in deferred {
539 if pdfrum_crypt::is_signature_dict(&entry.parent) {
540 continue;
543 }
544 out = rewrite_at(handler, obj, out, &entry.path);
545 }
546 out
547}
548
549fn rewrite_at(handler: &SecurityHandler, obj: ObjRef, object: Object, path: &[Step]) -> Object {
551 let Some((step, rest)) = path.split_first() else {
552 let mut ignored = Vec::new();
553 return decrypt_node(handler, obj, object, &mut ignored, false);
554 };
555 match (object, step) {
556 (Object::Dict(d), Step::Key(key)) => {
557 Object::Dict(Dict::from_pairs(d.iter().map(|(k, v)| {
558 if k == key {
559 (k.clone(), rewrite_at(handler, obj, v.clone(), rest))
560 } else {
561 (k.clone(), v.clone())
562 }
563 })))
564 }
565 (Object::Stream(s), Step::Key(key)) => {
566 let dict = Dict::from_pairs(s.dict.iter().map(|(k, v)| {
567 if k == key {
568 (k.clone(), rewrite_at(handler, obj, v.clone(), rest))
569 } else {
570 (k.clone(), v.clone())
571 }
572 }));
573 Object::Stream(Box::new(Stream::new(dict, s.data)))
574 }
575 (Object::Array(a), Step::Index(index)) => Object::Array(
576 a.iter()
577 .enumerate()
578 .map(|(i, v)| {
579 if i == *index {
580 rewrite_at(handler, obj, v.clone(), rest)
581 } else {
582 v.clone()
583 }
584 })
585 .collect(),
586 ),
587 (other, _) => other,
588 }
589}
590
591fn syntax_of(s: &PdfString) -> pdfrum_object::StringSyntax {
594 if s.hex {
595 pdfrum_object::StringSyntax::Hex
596 } else {
597 pdfrum_object::StringSyntax::Literal
598 }
599}
600
601#[cfg(test)]
602mod tests {
603 use super::ObjectStore;
604 use crate::error::Error;
605 use crate::xref::Xref;
606 use pdfrum_common::Limits;
607 use pdfrum_crypt::SecurityHandler;
608 use pdfrum_object::ByteSpan;
609 use pdfrum_object::{ObjRef, Resolve, names};
610 use std::sync::Arc;
611
612 fn store(file: &[u8], build: impl FnOnce(&mut Xref)) -> ObjectStore {
613 let mut xref = Xref::new();
614 build(&mut xref);
615 ObjectStore::new(
616 ByteSpan::from(file.to_vec()),
617 Arc::new(xref),
618 Limits::default(),
619 SecurityHandler::Identity,
620 )
621 }
622
623 #[test]
624 fn fetches_and_caches() {
625 let file = b"%PDF-1.7\n1 0 obj << /Type /Page >> endobj\n";
626 let s = store(file, |x| {
627 x.add_normal(1, 0, false, 9, &Limits::default());
628 });
629 let first = s.get(1).expect("object");
630 assert!(first.as_dict().is_some());
631 let second = s.get(1).expect("object");
633 assert!(Arc::ptr_eq(&first, &second));
634 }
635
636 #[test]
637 fn nested_fetches_share_one_nesting_budget() {
638 let count: u32 = 80;
644 let mut file = b"%PDF-1.7\n".to_vec();
645 let mut offsets = Vec::new();
646 for i in 1..=count {
647 offsets.push(file.len());
648 if i == count {
649 file.extend_from_slice(format!("{i} 0 obj 4 endobj\n").as_bytes());
650 } else {
651 file.extend_from_slice(
652 format!(
653 "{i} 0 obj << /Length {} 0 R >> stream\nDATA\nendstream endobj\n",
654 i + 1
655 )
656 .as_bytes(),
657 );
658 }
659 }
660 let s = store(&file, |x| {
661 for (i, offset) in offsets.iter().enumerate() {
662 let num = u32::try_from(i).unwrap_or(0) + 1;
663 x.add_normal(num, 0, false, *offset as u64, &Limits::default());
664 }
665 });
666
667 let first = s.get(1).expect("object 1");
671 assert_eq!(&*first.as_stream().expect("stream").data, b"DATA");
672 }
673
674 #[test]
675 fn object_zero_never_resolves() {
676 let s = store(b"", |x| {
677 x.add_normal(1, 0, false, 0, &Limits::default());
678 });
679 assert!(matches!(s.get(0), Err(Error::Unresolved(_))));
680 }
681
682 #[test]
683 fn numbers_past_the_table_are_unfetchable() {
684 let file = b"%PDF-1.7\n1 0 obj 5 endobj\n9 0 obj 7 endobj\n";
685 let s = store(file, |x| {
686 x.add_normal(1, 0, false, 9, &Limits::default());
687 });
688 assert!(matches!(s.get(9), Err(Error::Unresolved(_))));
690 }
691
692 #[test]
693 fn a_free_entry_resolves_to_nothing() {
694 let s = store(b"1 0 obj 5 endobj", |x| {
695 x.add_normal(2, 0, false, 0, &Limits::default());
696 x.set_free(1, 1);
697 });
698 assert!(matches!(s.get(1), Err(Error::Unresolved(_))));
699 }
700
701 #[test]
702 fn a_header_naming_another_object_fails_the_fetch() {
703 let file = b"%PDF-1.7\n7 0 obj << >> endobj\n";
704 let s = store(file, |x| {
705 x.add_normal(1, 0, false, 9, &Limits::default());
707 });
708 assert!(matches!(s.get(1), Err(Error::Unresolved(_))));
709 assert!(
710 s.drain_diags()
711 .contains(&pdfrum_common::DiagKind::ObjNumMismatch)
712 );
713 }
714
715 #[test]
716 fn a_self_referential_length_ends_as_a_keyword_scan() {
717 let file = b"%PDF-1.7\n1 0 obj << /Length 1 0 R >> stream\nDATA\nendstream endobj\n";
719 let s = store(file, |x| {
720 x.add_normal(1, 0, false, 9, &Limits::default());
721 });
722 let obj = s.get(1).expect("object");
723 assert_eq!(&*obj.as_stream().expect("stream").data, b"DATA");
724 }
725
726 #[test]
727 fn a_failed_fetch_is_retried_rather_than_remembered() {
728 let file = b"%PDF-1.7\nnot an object\n";
729 let s = store(file, |x| {
730 x.add_normal(1, 0, false, 9, &Limits::default());
731 });
732 assert!(s.get(1).is_err());
733 assert!(s.get(1).is_err());
734 }
735
736 #[test]
737 fn resolving_reads_through_the_trait() {
738 let file = b"%PDF-1.7\n1 0 obj << /Count 4 >> endobj\n";
739 let s = store(file, |x| {
740 x.add_normal(1, 0, false, 9, &Limits::default());
741 });
742 let fetched = s.fetch(ObjRef::new(1, 0)).expect("object");
743 assert_eq!(
744 fetched.as_dict().and_then(|d| d.direct_int(names::COUNT)),
745 Some(4)
746 );
747 }
748
749 #[test]
750 fn the_store_is_send_and_sync() {
751 fn assert_both<T: Send + Sync>() {}
752 assert_both::<ObjectStore>();
753 }
754}