use std::sync::{Arc, Mutex};
use pdfrum_common::{
DiagKind, Diagnostics, LimitExceeded, Limits, Operation, PageIndex, PdfVersion, Severity,
};
use pdfrum_crypt::{Permissions, SecurityHandler};
use pdfrum_object::{ByteSpan, Dict, NoResolve, ObjRef, Object, Resolve, names};
use crate::error::Error;
use crate::store::ObjectStore;
use crate::xref::{Trailer, Xref, XrefShape};
const HEADER_SIZE: usize = 9;
#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
#[non_exhaustive]
pub enum LoadError {
#[error("not a PDF file")]
NotPdf,
#[error("wrong password")]
WrongPassword,
#[error("unsupported encryption: {0}")]
UnsupportedEncryption(String),
#[error("damaged beyond recovery: {0}")]
Broken(String),
#[error(transparent)]
Limit(LimitExceeded),
}
#[derive(Debug, Clone, Default)]
pub struct LoadOptions {
pub password: Option<Vec<u8>>,
pub limits: Limits,
}
#[derive(Debug, Clone, PartialEq)]
pub struct PageDict {
pub dict: Dict,
pub reference: Option<ObjRef>,
}
impl PageDict {
#[must_use]
pub fn inherited(&self, key: &pdfrum_object::Name, r: &impl Resolve) -> Option<Object> {
let mut node = self.dict.clone();
let mut seen: Vec<Dict> = Vec::new();
for _ in 0..64 {
if let Some(value) = node.get(key, r) {
return Some(value.get().clone());
}
if seen.contains(&node) {
return None;
}
seen.push(node.clone());
node = node.dict(names::PARENT, r)?;
}
None
}
}
#[derive(Debug)]
pub struct Document {
bytes: ByteSpan,
trailer: Trailer,
store: Arc<ObjectStore>,
version: Option<PdfVersion>,
header_offset: u64,
xref_shape: XrefShape,
encrypt: Option<(Dict, bool)>,
page_count: u32,
pages: Mutex<PageCache>,
pub diags: Diagnostics,
}
#[derive(Debug, Default)]
struct PageCache {
slots: Vec<Option<PageDict>>,
poisoned: bool,
}
pub fn load(bytes: impl Into<ByteSpan>, opts: &LoadOptions) -> Result<Document, LoadError> {
opts.limits
.check_deadline(Operation::Open)
.map_err(LoadError::Limit)?;
let mut diags = Diagnostics::default();
let bytes: ByteSpan = bytes.into();
let header_offset = find_header(&bytes, &opts.limits).ok_or(LoadError::NotPdf)?;
if bytes.len() < header_offset.saturating_add(HEADER_SIZE) {
return Err(LoadError::NotPdf);
}
if header_offset > 0 {
diags.record(
Severity::Recovered,
DiagKind::HeaderOffset,
Some(header_offset as u64),
);
}
let body = bytes
.subspan(header_offset..bytes.len())
.unwrap_or_else(|_| ByteSpan::empty());
let version = read_version(&body);
let (xref, mut trailer, mut xref_shape) =
crate::xref::read_xref_full(&body, &opts.limits, &mut diags).map_err(|e| match e {
crate::Error::Limit(limit) => LoadError::Limit(limit),
other => LoadError::Broken(other.to_string()),
})?;
let mut shared_xref = Arc::new(xref);
let mut security = build_security(&body, &shared_xref, &trailer.dict, opts, &mut diags)?;
let mut store = build_store(&body, &shared_xref, opts, &trailer.dict, security);
let mut page_count = catalog_page_count(&store, &trailer.dict, &opts.limits);
if page_count.is_none() {
if xref_shape.rebuilt {
return Err(LoadError::Broken("no document catalog".into()));
}
diags.record(Severity::Recovered, DiagKind::RootRecovered, None);
let mut fresh = Xref::new();
let mut fresh_trailer = Trailer::default();
let rebuilt = crate::xref::rebuild(
&body,
&mut fresh,
&mut fresh_trailer,
&opts.limits,
&mut diags,
&NoResolve,
)
.map_err(LoadError::Limit)?;
if !rebuilt {
return Err(LoadError::Broken("no document catalog".into()));
}
let mut merged = (*shared_xref).clone();
merged.merge_up(&fresh);
crate::xref::merge_trailers(&mut trailer, &fresh_trailer);
xref_shape = XrefShape::rebuilt();
shared_xref = Arc::new(merged);
security = build_security(&body, &shared_xref, &trailer.dict, opts, &mut diags)?;
store = build_store(&body, &shared_xref, opts, &trailer.dict, security);
if catalog(&store, &trailer.dict).is_none() {
return Err(LoadError::Broken("no document catalog".into()));
}
page_count = catalog_page_count(&store, &trailer.dict, &opts.limits);
}
let page_count = page_count.unwrap_or(0);
let encrypt = encrypt_dict_located(&body, &shared_xref, &trailer.dict, &opts.limits);
diags.extend(&store.drain_diags());
Ok(Document {
bytes: body,
trailer,
store,
version,
header_offset: header_offset as u64,
xref_shape,
encrypt,
page_count,
pages: Mutex::new(PageCache {
slots: vec![None; usize::try_from(page_count).unwrap_or(0)],
poisoned: false,
}),
diags,
})
}
fn find_header(bytes: &[u8], limits: &Limits) -> Option<usize> {
let window = usize::try_from(limits.header_scan).unwrap_or(usize::MAX);
let last = bytes.len().checked_sub(4)?.min(window);
(0..=last).find(|&i| bytes.get(i..i + 4) == Some(b"%PDF"))
}
fn read_version(body: &[u8]) -> Option<PdfVersion> {
let digit = |i: usize| -> u8 {
body.get(i)
.filter(|b| b.is_ascii_digit())
.map_or(0, |b| b - b'0')
};
match digit(5).saturating_mul(10).saturating_add(digit(7)) {
0 => None,
packed => Some(PdfVersion::new(packed / 10, packed % 10)),
}
}
fn build_security(
body: &ByteSpan,
xref: &Arc<Xref>,
trailer: &Dict,
opts: &LoadOptions,
diags: &mut Diagnostics,
) -> Result<SecurityHandler, LoadError> {
let Some(encrypt) = encrypt_dict(body, xref, trailer, &opts.limits) else {
return Ok(SecurityHandler::Identity);
};
if encrypt.name(names::FILTER) != Some(names::STANDARD) {
return Err(LoadError::UnsupportedEncryption(
encrypt
.name(names::FILTER)
.map_or_else(|| "unnamed".to_owned(), |n| n.as_text().into_owned()),
));
}
let file_id = trailer
.array(names::ID, &NoResolve)
.and_then(|a| a.string_at(0).map(|s| s.bytes.to_vec()))
.unwrap_or_default();
let password = opts.password.clone().unwrap_or_default();
match SecurityHandler::from_encrypt_dict(&encrypt, &file_id, &password, &NoResolve) {
Ok(handler) => {
if handler.password_encoding() != pdfrum_crypt::PasswordEncoding::AsGiven {
diags.record(Severity::Recovered, DiagKind::PasswordReencoded, None);
}
Ok(handler)
}
Err(pdfrum_crypt::Error::WrongPassword) => Err(LoadError::WrongPassword),
Err(pdfrum_crypt::Error::UnsupportedHandler(name)) => Err(
LoadError::UnsupportedEncryption(String::from_utf8_lossy(&name).into_owned()),
),
Err(e) => Err(LoadError::UnsupportedEncryption(e.to_string())),
}
}
fn encrypt_dict(
body: &ByteSpan,
xref: &Arc<Xref>,
trailer: &Dict,
limits: &Limits,
) -> Option<Dict> {
encrypt_dict_located(body, xref, trailer, limits).map(|(d, _)| d)
}
fn encrypt_dict_located(
body: &ByteSpan,
xref: &Arc<Xref>,
trailer: &Dict,
limits: &Limits,
) -> Option<(Dict, bool)> {
match trailer.raw(names::ENCRYPT)? {
Object::Dict(d) => Some((d.clone(), true)),
Object::Ref(r) => {
let plain = ObjectStore::new(
body.clone(),
Arc::clone(xref),
limits.clone(),
SecurityHandler::Identity,
);
Some((plain.get(r.num).ok()?.as_dict().cloned()?, false))
}
_ => None,
}
}
fn exempt_metadata(store: &mut ObjectStore, trailer: &Dict) {
if store.security().encrypt_metadata() {
return;
}
let Some(root) = trailer.reference(names::ROOT) else {
return;
};
let Ok(catalog) = store.get(root.num) else {
return;
};
if let Some(metadata) = catalog.as_dict().and_then(|d| d.reference(names::METADATA)) {
store.exempt_from_decryption(metadata.num);
}
}
fn build_store(
body: &ByteSpan,
xref: &Arc<Xref>,
opts: &LoadOptions,
trailer: &Dict,
security: SecurityHandler,
) -> Arc<ObjectStore> {
let mut store = ObjectStore::new(
body.clone(),
Arc::clone(xref),
opts.limits.clone(),
security,
);
exempt_metadata(&mut store, trailer);
Arc::new(store)
}
fn catalog(store: &ObjectStore, trailer: &Dict) -> Option<Dict> {
let root = trailer.reference(names::ROOT)?;
store.get(root.num).ok()?.as_dict().cloned()
}
fn catalog_page_count(store: &ObjectStore, trailer: &Dict, limits: &Limits) -> Option<u32> {
let catalog = catalog(store, trailer)?;
let count = page_count_of(store, &catalog, limits);
(count > 0).then_some(count)
}
fn page_count_of(store: &ObjectStore, catalog: &Dict, limits: &Limits) -> u32 {
let Some(pages) = catalog.dict(names::PAGES, store) else {
return 0;
};
if pages.raw(names::KIDS).is_none() {
return 1;
}
let mut ancestors = vec![pages.clone()];
count_subtree(store, &pages, limits, &mut ancestors).unwrap_or(0)
}
fn count_subtree(
store: &ObjectStore,
node: &Dict,
limits: &Limits,
ancestors: &mut Vec<Dict>,
) -> Option<u32> {
if let Some(count) = node.int(names::COUNT, store)
&& count > 0
&& count < i64::from(limits.max_page_count)
&& let Ok(count) = u32::try_from(count)
{
return Some(count);
}
let Some(kids) = node.array(names::KIDS, store) else {
return Some(0);
};
let mut total: u32 = 0;
for kid in kids.iter() {
let Some(kid) = kid.resolve(store).ok().and_then(|k| k.as_dict().cloned()) else {
continue;
};
if ancestors.contains(&kid) {
continue;
}
total = total.saturating_add(match node_kind(&kid) {
NodeKind::Branch => {
ancestors.push(kid.clone());
let under = count_subtree(store, &kid, limits, ancestors);
ancestors.pop();
under?
}
NodeKind::Leaf => 1,
});
if total >= limits.max_page_count {
return None;
}
}
Some(total)
}
enum NodeKind {
Branch,
Leaf,
}
fn node_kind(node: &Dict) -> NodeKind {
match node.name(names::TYPE) {
Some(t) if t == names::PAGES => NodeKind::Branch,
Some(t) if t == names::PAGE => NodeKind::Leaf,
_ => {
if node.contains_key(names::KIDS) {
NodeKind::Branch
} else {
NodeKind::Leaf
}
}
}
}
impl Document {
#[must_use]
pub fn page_count(&self) -> u32 {
self.page_count
}
pub fn page(&self, index: impl Into<PageIndex>) -> Result<PageDict, Error> {
let index = index.into();
let slot = usize::try_from(index.get()).unwrap_or(usize::MAX);
if index.get() >= self.page_count {
return Err(Error::NoPage(index));
}
let Ok(mut pages) = self.pages.lock() else {
return Err(Error::NoPage(index));
};
if let Some(Some(found)) = pages.slots.get(slot) {
return Ok(found.clone());
}
if pages.poisoned {
return Err(Error::NoPage(index));
}
self.walk_pages(&mut pages);
pages
.slots
.get(slot)
.and_then(Clone::clone)
.ok_or(Error::NoPage(index))
}
fn walk_pages(&self, pages: &mut PageCache) {
let Some(root) = self.trailer.dict.reference(names::ROOT) else {
return;
};
let Ok(catalog) = self.store.get(root.num) else {
return;
};
let Some(node) = catalog
.as_dict()
.and_then(|d| d.dict(names::PAGES, &*self.store))
else {
return;
};
let mut next: usize = 0;
let mut ancestors = Vec::new();
let objref = catalog.as_dict().and_then(|d| d.reference(names::PAGES));
self.visit(&node, objref, pages, &mut next, 0, &mut ancestors);
}
fn visit(
&self,
node: &Dict,
reference: Option<ObjRef>,
pages: &mut PageCache,
next: &mut usize,
depth: u32,
ancestors: &mut Vec<Dict>,
) {
if *next >= pages.slots.len() {
return;
}
if node.raw(names::KIDS).is_none() {
if matches!(node_kind(node), NodeKind::Branch) {
return;
}
if let Some(slot) = pages.slots.get_mut(*next) {
*slot = Some(PageDict {
dict: node.clone(),
reference,
});
}
*next += 1;
return;
}
if depth >= self.store.limits().max_page_tree_depth {
self.store
.note(Severity::Suspicious, DiagKind::PageTreeDepthExceeded, None);
pages.poisoned = true;
return;
}
let Some(kids) = node.array(names::KIDS, &*self.store) else {
return;
};
ancestors.push(node.clone());
for kid in kids.iter() {
let kid_ref = kid.as_ref_id();
let loaded = kid
.resolve(&*self.store)
.ok()
.and_then(|k| k.as_dict().cloned());
let Some(loaded) = loaded else {
*next += 1;
continue;
};
if ancestors.contains(&loaded) {
self.store
.note(Severity::Recovered, DiagKind::PageTreeRepaired, None);
continue;
}
self.visit(&loaded, kid_ref, pages, next, depth + 1, ancestors);
if *next >= pages.slots.len() {
break;
}
}
ancestors.pop();
}
#[must_use]
pub fn trailer(&self) -> &Dict {
&self.trailer.dict
}
#[must_use]
pub fn trailer_object_number(&self) -> u32 {
self.trailer.object_number
}
pub fn catalog(&self) -> Result<Dict, Error> {
let root = self
.trailer
.dict
.reference(names::ROOT)
.ok_or(Error::NoCatalog)?;
self.store
.get(root.num)
.ok()
.and_then(|c| c.as_dict().cloned())
.ok_or(Error::NoCatalog)
}
#[must_use]
pub fn version(&self) -> Option<PdfVersion> {
self.version
}
#[must_use]
pub fn header_offset(&self) -> u64 {
self.header_offset
}
#[must_use]
pub fn lazy_diagnostics(&self) -> Diagnostics {
self.store.peek_diags()
}
#[must_use]
pub fn xref_was_rebuilt(&self) -> bool {
self.xref_shape.rebuilt
}
#[must_use]
pub fn last_xref_offset(&self) -> u64 {
self.xref_shape.last_offset
}
#[must_use]
pub fn main_xref_is_stream(&self) -> bool {
self.xref_shape.main_is_stream
}
#[must_use]
pub fn encrypt_dict(&self) -> Option<(&Dict, bool)> {
self.encrypt.as_ref().map(|(d, inline)| (d, *inline))
}
#[must_use]
pub fn permissions(&self) -> Permissions {
self.store.security().permissions()
}
#[must_use]
pub fn owner_permissions(&self) -> Permissions {
self.store.security().owner_permissions()
}
#[must_use]
pub fn is_encrypted(&self) -> bool {
!matches!(self.store.security(), SecurityHandler::Identity)
}
#[must_use]
pub fn security_handler(&self) -> &SecurityHandler {
self.store.security()
}
#[must_use]
pub fn bytes(&self) -> &[u8] {
&self.bytes
}
#[must_use]
pub fn store(&self) -> &Arc<ObjectStore> {
&self.store
}
#[must_use]
pub fn xref(&self) -> &Xref {
self.store.xref()
}
}
impl Resolve for Document {
fn fetch(&self, r: ObjRef) -> Result<Arc<Object>, pdfrum_object::Error> {
self.store.fetch(r)
}
}
#[cfg(test)]
mod tests {
use super::{LoadError, LoadOptions, find_header, load, read_version};
use pdfrum_common::{DiagKind, Limits, PdfVersion};
use pdfrum_crypt::Permissions;
use pdfrum_object::{Name, names};
use std::sync::Arc;
fn open(bytes: &[u8]) -> Result<super::Document, LoadError> {
load(Arc::from(bytes), &LoadOptions::default())
}
fn build(count: usize) -> Vec<u8> {
let mut out = Vec::new();
out.extend_from_slice(b"%PDF-1.7\n");
let mut offsets = vec![0usize];
offsets.push(out.len());
out.extend_from_slice(b"1 0 obj\n<< /Type /Catalog /Pages 2 0 R >>\nendobj\n");
offsets.push(out.len());
let kids: Vec<String> = (0..count).map(|i| format!("{} 0 R", i + 3)).collect();
out.extend_from_slice(
format!(
"2 0 obj\n<< /Type /Pages /Count {count} /Kids [{}] /MediaBox [0 0 612 792] >>\nendobj\n",
kids.join(" ")
)
.as_bytes(),
);
for i in 0..count {
offsets.push(out.len());
out.extend_from_slice(
format!(
"{} 0 obj\n<< /Type /Page /Parent 2 0 R /PageNumber {i} >>\nendobj\n",
i + 3
)
.as_bytes(),
);
}
let xref_at = out.len();
out.extend_from_slice(format!("xref\n0 {}\n", offsets.len()).as_bytes());
out.extend_from_slice(b"0000000000 65535 f \n");
for offset in offsets.iter().skip(1) {
out.extend_from_slice(format!("{offset:010} 00000 n \n").as_bytes());
}
out.extend_from_slice(
format!(
"trailer\n<< /Size {} /Root 1 0 R >>\nstartxref\n{xref_at}\n%%EOF\n",
offsets.len()
)
.as_bytes(),
);
out
}
#[test]
fn finds_a_header_at_the_start_or_after_junk() {
assert_eq!(find_header(b"%PDF-1.7\n", &Limits::default()), Some(0));
assert_eq!(find_header(b"junk%PDF-1.7\n", &Limits::default()), Some(4));
assert_eq!(find_header(b"no header", &Limits::default()), None);
}
#[test]
fn version_digits_are_read_not_validated() {
assert_eq!(read_version(b"%PDF-1.7\n"), Some(PdfVersion::PDF_1_7));
assert_eq!(read_version(b"%PDF-2.0\n"), Some(PdfVersion::PDF_2_0));
assert_eq!(read_version(b"%PDF-x.y\n"), None);
for major in 0..=9u8 {
for minor in 0..=9u8 {
let header = format!("%PDF-{major}.{minor}\n");
let expected = (major, minor) != (0, 0);
assert_eq!(
read_version(header.as_bytes()),
expected.then(|| PdfVersion::new(major, minor)),
"header {header:?}"
);
}
}
}
#[test]
fn a_file_without_a_header_is_not_a_pdf() {
assert_eq!(open(b"just some bytes").err(), Some(LoadError::NotPdf));
assert_eq!(open(b"%PDF").err(), Some(LoadError::NotPdf));
}
#[test]
fn opens_a_document_and_counts_its_pages() {
let doc = open(&build(3)).expect("document");
assert_eq!(doc.page_count(), 3);
assert_eq!(doc.version(), Some(PdfVersion::PDF_1_7));
assert!(!doc.xref_was_rebuilt());
assert!(!doc.is_encrypted());
assert_eq!(doc.permissions(), Permissions::ALL);
}
#[test]
fn reads_pages_in_order() {
let doc = open(&build(5)).expect("document");
let number = Name::from("PageNumber");
for i in 0..5u32 {
let page = doc.page(i).expect("page");
assert_eq!(page.dict.direct_int(&number), Some(i64::from(i)));
}
assert!(doc.page(5).is_err());
}
#[test]
fn reads_pages_in_reverse_and_out_of_order() {
let doc = open(&build(5)).expect("document");
let number = Name::from("PageNumber");
for i in (0..5u32).rev() {
assert_eq!(
doc.page(i).expect("page").dict.direct_int(&number),
Some(i64::from(i))
);
}
assert!(doc.page(99).is_err());
assert_eq!(doc.page(3).expect("page").dict.direct_int(&number), Some(3));
}
#[test]
fn a_count_larger_than_the_tree_reports_the_lie() {
let text = String::from_utf8_lossy(&build(3)).replace("/Count 3", "/Count 9");
let doc = open(text.as_bytes()).expect("document");
assert_eq!(doc.page_count(), 9);
assert!(doc.page(0).is_ok());
assert!(doc.page(2).is_ok());
assert!(doc.page(3).is_err());
assert!(doc.page(8).is_err());
assert!(doc.page(2).is_ok());
}
#[test]
fn a_kids_less_pages_node_counts_as_a_page_it_cannot_produce() {
let file = b"%PDF-1.7\n\
1 0 obj\n<< /Type /Catalog /Pages 2 0 R >>\nendobj\n\
2 0 obj\n<< /Type /Pages /Count 3 >>\nendobj\n\
trailer\n<< /Root 1 0 R >>\nstartxref\n0\n%%EOF\n";
let doc = open(file).expect("document");
assert_eq!(doc.page_count(), 1);
assert!(doc.page(0).is_err());
assert!(doc.page(1).is_err());
}
#[test]
fn a_kids_less_node_that_does_not_claim_to_be_a_branch_is_a_page() {
let file = b"%PDF-1.7\n\
1 0 obj\n<< /Type /Catalog /Pages 2 0 R >>\nendobj\n\
2 0 obj\n<< /MediaBox [0 0 10 10] >>\nendobj\n\
trailer\n<< /Root 1 0 R >>\nstartxref\n0\n%%EOF\n";
let doc = open(file).expect("document");
assert_eq!(doc.page_count(), 1);
assert!(doc.page(0).is_ok());
}
#[test]
fn a_catalog_without_pages_will_not_open() {
let file = b"%PDF-1.7\n\
1 0 obj\n<< /Type /Catalog >>\nendobj\n\
trailer\n<< /Root 1 0 R >>\nstartxref\n0\n%%EOF\n";
assert!(matches!(open(file), Err(LoadError::Broken(_))));
}
#[test]
fn a_root_written_inline_does_not_name_a_catalog() {
let file = b"%PDF-1.7\n\
1 0 obj\n<< /Type /Page >>\nendobj\n\
trailer\n<< /Root << /Type /Catalog /Pages 2 0 R >> >>\n\
startxref\n0\n%%EOF\n";
assert!(matches!(open(file), Err(LoadError::Broken(_))));
}
#[test]
fn a_broken_start_xref_still_opens_the_document() {
let text = String::from_utf8_lossy(&build(2)).into_owned();
let broken = text
.replace("%%EOF", "")
.replace("startxref\n", "startxref\n1\n");
let doc = open(broken.as_bytes()).expect("document");
assert!(doc.xref_was_rebuilt());
assert_eq!(doc.page_count(), 2);
assert!(doc.diags.contains(&DiagKind::XrefRebuilt));
}
#[test]
fn a_header_after_junk_shifts_every_offset() {
let mut file = vec![b'x'; 100];
file.extend_from_slice(&build(2));
let doc = open(&file).expect("document");
assert_eq!(doc.header_offset(), 100);
assert_eq!(doc.page_count(), 2);
assert!(doc.diags.contains(&DiagKind::HeaderOffset));
}
#[test]
fn inheritable_attributes_come_from_the_parent() {
let doc = open(&build(2)).expect("document");
let page = doc.page(0).expect("page");
let inherited = page
.inherited(names::MEDIA_BOX, &doc)
.expect("inherited media box");
let array = inherited.as_array().expect("array");
assert_eq!(array.number_at(2), Some(612.0));
assert_eq!(array.number_at(3), Some(792.0));
assert!(page.inherited(names::ROTATE, &doc).is_none());
}
#[test]
fn a_pages_reference_is_reported() {
let doc = open(&build(1)).expect("document");
assert_eq!(doc.page(0).expect("page").reference.map(|r| r.num), Some(3));
}
#[test]
fn documents_are_send_and_sync() {
fn assert_both<T: Send + Sync>() {}
assert_both::<super::Document>();
}
#[test]
fn never_panics_on_arbitrary_bytes() {
let seeds: &[&[u8]] = &[
b"",
b"%PDF",
b"%PDF-1.7",
b"%PDF-1.7\nstartxref\n0\n%%EOF",
b"%PDF-1.7\ntrailer<</Root 1 0 R>>",
b"%PDF-1.7\n1 0 obj<</Length 1 0 R>>stream\n",
b"%PDF-1.7\n\x00\xff\x80\x0b",
];
for seed in seeds {
let _ = open(seed);
}
}
}