Skip to main content

pdfboss_write/
update.rs

1//! Incremental updates to an existing file (ISO 32000-1 §7.5.6): the base
2//! bytes stay in place and an update section appends new and replaced
3//! objects plus a cross-reference section chained to the base's by `/Prev`,
4//! in the base's own cross-reference style.
5
6use std::io::Write;
7
8use flate2::write::ZlibEncoder;
9use flate2::Compression;
10use pdfboss_core::xref::startxref;
11use pdfboss_core::{Dict, Document, FastMap, Name, ObjRef, Object, Stream};
12
13use crate::error::{Error, Result};
14use crate::ser::{serialize_dict, serialize_object};
15use crate::writer::{WriteOptions, Writer};
16
17/// The resource name every page draws the overlay form under.
18const FORM_NAME: &str = "PdfbossWatermark";
19
20/// Like [`watermark`], but writes a fresh file through the [`Writer`] under
21/// `options` instead of appending an update: every object the base's
22/// catalog reaches is copied over, uncompressed streams are compressed when
23/// `options.compress` is set, and unreachable objects and earlier sections
24/// are left behind, so the result is usually smaller than the base.
25pub fn watermark_with(
26    base: &Document,
27    overlay: &Document,
28    options: WriteOptions,
29) -> Result<Vec<u8>> {
30    let trailer = &base.xref().trailer;
31    if trailer.get("Encrypt").is_some() {
32        return Err(Error::Other(
33            "an encrypted file cannot be rewritten".to_string(),
34        ));
35    }
36    let root = trailer
37        .get_ref("Root")
38        .ok_or_else(|| Error::Other("the base file has no /Root".to_string()))?;
39    let mut writer = Writer::new(options);
40    let prefix = writer.put_stream_raw(Dict::new(), b"q\n".to_vec());
41    let suffix = writer.put_stream_raw(
42        Dict::new(),
43        format!("Q\nq /{FORM_NAME} Do Q\n").into_bytes(),
44    );
45    let mut overlay_copy = Rewrite::new(overlay, options.compress);
46    let form = overlay_copy.form(&mut writer)?;
47    overlay_copy.drain(&mut writer, None)?;
48
49    let pages: FastMap<ObjRef, usize> = (0..base.page_count())
50        .filter_map(|index| {
51            let page = base.page(index).ok()?;
52            page.object_ref().map(|r| (r, index))
53        })
54        .collect();
55    let stamp = Stamp {
56        form,
57        prefix,
58        suffix,
59        pages,
60    };
61    let mut base_copy = Rewrite::new(base, options.compress);
62    let new_root = base_copy.reference(&mut writer, root);
63    if let Some(info) = trailer.get_ref("Info") {
64        let new_info = base_copy.reference(&mut writer, info);
65        writer.set_info(new_info);
66    }
67    base_copy.drain(&mut writer, Some(&stamp))?;
68    writer.finish(new_root)
69}
70
71/// What every stamped page draws: the overlay form and the two content
72/// streams wrapped around the page's own, plus which objects are pages.
73struct Stamp {
74    form: ObjRef,
75    prefix: ObjRef,
76    suffix: ObjRef,
77    pages: FastMap<ObjRef, usize>,
78}
79
80/// Copies one document's object graph into a [`Writer`]: every reference
81/// met is reserved a number once and queued, and the queue is drained
82/// iteratively, so a long chain of references costs no stack.
83struct Rewrite<'a> {
84    source: &'a Document,
85    map: FastMap<ObjRef, ObjRef>,
86    pending: Vec<ObjRef>,
87    compress: bool,
88}
89
90impl<'a> Rewrite<'a> {
91    fn new(source: &'a Document, compress: bool) -> Rewrite<'a> {
92        Rewrite {
93            source,
94            map: FastMap::default(),
95            pending: Vec::new(),
96            compress,
97        }
98    }
99
100    /// The target number for source reference `r`, reserved and queued on
101    /// first sight.
102    fn reference(&mut self, writer: &mut Writer, r: ObjRef) -> ObjRef {
103        if let Some(copied) = self.map.get(&r) {
104            return *copied;
105        }
106        let copied = writer.reserve();
107        self.map.insert(r, copied);
108        self.pending.push(r);
109        copied
110    }
111
112    /// A copy of `obj`'s direct structure with every reference mapped.
113    fn copy(&mut self, writer: &mut Writer, obj: &Object) -> Result<Object> {
114        Ok(match obj {
115            Object::Ref(r) => Object::Ref(self.reference(writer, *r)),
116            Object::Dict(d) => Object::Dict(self.copy_dict(writer, d)?),
117            Object::Array(items) => Object::Array(
118                items
119                    .iter()
120                    .map(|item| self.copy(writer, item))
121                    .collect::<Result<Vec<Object>>>()?,
122            ),
123            Object::Stream(_) => return Err(Error::NestedStream),
124            other => other.clone(),
125        })
126    }
127
128    fn copy_dict(&mut self, writer: &mut Writer, dict: &Dict) -> Result<Dict> {
129        let mut out = Dict::new();
130        for (key, value) in dict.iter() {
131            out.insert(key.clone(), self.copy(writer, value)?);
132        }
133        Ok(out)
134    }
135
136    /// A stream body: its dictionary copied without `/Length` (the writer
137    /// sets it), its data compressed when asked and not already filtered.
138    fn copy_stream(&mut self, writer: &mut Writer, stream: &Stream) -> Result<Object> {
139        let mut dict = stream.dict.clone();
140        dict.remove("Length");
141        let mut dict = self.copy_dict(writer, &dict)?;
142        let data = if self.compress && dict.get("Filter").is_none() {
143            dict.insert(name("Filter"), Object::Name(name("FlateDecode")));
144            deflate(&stream.data)
145        } else {
146            stream.data.clone()
147        };
148        Ok(Object::Stream(Stream { dict, data }))
149    }
150
151    /// Fills every queued object, stamping the pages `stamp` names.
152    fn drain(&mut self, writer: &mut Writer, stamp: Option<&Stamp>) -> Result<()> {
153        while let Some(r) = self.pending.pop() {
154            let target = self.map[&r];
155            let stamped = stamp.and_then(|s| s.pages.get(&r).map(|index| (s, *index)));
156            let body = match stamped {
157                Some((s, index)) => self.stamped_page(writer, index, s)?,
158                None => match self.source.get(r).map_err(core_error)? {
159                    Object::Stream(s) => self.copy_stream(writer, &s)?,
160                    other => self.copy(writer, &other)?,
161                },
162            };
163            writer.fill(target, body)?;
164        }
165        Ok(())
166    }
167
168    /// Page `index`'s dictionary copied with the stamp applied: its
169    /// effective resources gain the form, its content is wrapped in the
170    /// prefix and suffix streams.
171    fn stamped_page(&mut self, writer: &mut Writer, index: usize, stamp: &Stamp) -> Result<Object> {
172        let page = self.source.page(index).map_err(core_error)?;
173        let mut dict = self.copy_dict(writer, page.dict())?;
174        let mut resources = self.copy_dict(writer, &page.resources)?;
175        let mut xobjects = match page.resources.get("XObject") {
176            Some(existing) => {
177                let existing = self.source.resolve(existing).map_err(core_error)?;
178                match existing.as_dict() {
179                    Some(d) => self.copy_dict(writer, d)?,
180                    None => Dict::new(),
181                }
182            }
183            None => Dict::new(),
184        };
185        xobjects.insert(name(FORM_NAME), Object::Ref(stamp.form));
186        resources.insert(name("XObject"), Object::Dict(xobjects));
187        dict.insert(name("Resources"), Object::Dict(resources));
188        let mut contents = vec![Object::Ref(stamp.prefix)];
189        match page.dict().get("Contents") {
190            Some(Object::Array(items)) => {
191                for item in items {
192                    contents.push(self.copy(writer, item)?);
193                }
194            }
195            Some(Object::Ref(r)) => match self.source.get(*r).map_err(core_error)? {
196                Object::Array(items) => {
197                    for item in &items {
198                        contents.push(self.copy(writer, item)?);
199                    }
200                }
201                _ => contents.push(Object::Ref(self.reference(writer, *r))),
202            },
203            _ => {}
204        }
205        contents.push(Object::Ref(stamp.suffix));
206        dict.insert(name("Contents"), Object::Array(contents));
207        Ok(Object::Dict(dict))
208    }
209
210    /// The source's first page as a form XObject, filled into the writer:
211    /// its media box as the bounding box, its decoded content deflated, its
212    /// resources copied.
213    fn form(&mut self, writer: &mut Writer) -> Result<ObjRef> {
214        let page = self.source.page(0).map_err(core_error)?;
215        let content = page.content(self.source).map_err(core_error)?;
216        let resources = self.copy_dict(writer, &page.resources)?;
217        let bbox = page.media_box;
218        let mut dict = Dict::new();
219        dict.insert(name("Type"), Object::Name(name("XObject")));
220        dict.insert(name("Subtype"), Object::Name(name("Form")));
221        dict.insert(name("FormType"), Object::Int(1));
222        dict.insert(
223            name("BBox"),
224            Object::Array(
225                [bbox.x0, bbox.y0, bbox.x1, bbox.y1]
226                    .iter()
227                    .map(|v| Object::Real(f64::from(*v)))
228                    .collect(),
229            ),
230        );
231        dict.insert(name("Resources"), Object::Dict(resources));
232        dict.insert(name("Filter"), Object::Name(name("FlateDecode")));
233        let form = writer.reserve();
234        writer.fill(
235            form,
236            Object::Stream(Stream {
237                dict,
238                data: deflate(&content),
239            }),
240        )?;
241        Ok(form)
242    }
243}
244
245/// Draws the first page of `overlay` over every page of `base`, returning
246/// `base`'s bytes followed by an incremental update: the overlay page as
247/// one form XObject (its resources copied into the base's object space),
248/// and each page's dictionary rewritten with that form in its resources
249/// and its content wrapped in `q … Q` before the form is drawn. Pages
250/// inlined directly into `/Kids`, having no object of their own, are left
251/// as they are. An encrypted base is refused: its new strings and streams
252/// would need encrypting too.
253pub fn watermark(base: &Document, overlay: &Document) -> Result<Vec<u8>> {
254    let mut update = Update::open(base)?;
255    let form = update.import_form(overlay)?;
256    let prefix = update.put(Object::Stream(plain_stream(b"q\n".to_vec())));
257    let suffix = update.put(Object::Stream(plain_stream(
258        format!("Q\nq /{FORM_NAME} Do Q\n").into_bytes(),
259    )));
260    for index in 0..base.page_count() {
261        let page = base.page(index).map_err(core_error)?;
262        let Some(page_ref) = page.object_ref() else {
263            continue;
264        };
265        let mut dict = page.dict().clone();
266        let mut resources = page.resources.clone();
267        let mut xobjects = match resources.get("XObject") {
268            Some(existing) => base
269                .resolve(existing)
270                .map_err(core_error)?
271                .as_dict()
272                .cloned()
273                .unwrap_or_default(),
274            None => Dict::new(),
275        };
276        xobjects.insert(name(FORM_NAME), Object::Ref(form));
277        resources.insert(name("XObject"), Object::Dict(xobjects));
278        dict.insert(name("Resources"), Object::Dict(resources));
279        let mut contents = vec![Object::Ref(prefix)];
280        match dict.get("Contents").cloned() {
281            Some(Object::Array(items)) => contents.extend(items),
282            Some(Object::Ref(r)) => match base.get(r).map_err(core_error)? {
283                Object::Array(items) => contents.extend(items),
284                _ => contents.push(Object::Ref(r)),
285            },
286            _ => {}
287        }
288        contents.push(Object::Ref(suffix));
289        dict.insert(name("Contents"), Object::Array(contents));
290        update.replace(page_ref, Object::Dict(dict));
291    }
292    update.finish()
293}
294
295/// One update section under construction: the objects it will hold, new
296/// ones numbered from the first number the base leaves free.
297struct Update<'a> {
298    base: &'a Document,
299    next: u32,
300    objects: Vec<(ObjRef, Object)>,
301    imported: FastMap<ObjRef, ObjRef>,
302}
303
304impl<'a> Update<'a> {
305    fn open(base: &'a Document) -> Result<Update<'a>> {
306        let trailer = &base.xref().trailer;
307        if trailer.get("Encrypt").is_some() {
308            return Err(Error::Other(
309                "an encrypted file cannot be updated in place".to_string(),
310            ));
311        }
312        let highest = base.xref().iter().map(|(num, _)| num).max().unwrap_or(0);
313        let size = trailer.get_int("Size").unwrap_or(0).max(0) as u32;
314        Ok(Update {
315            base,
316            next: size.max(highest + 1),
317            objects: Vec::new(),
318            imported: FastMap::default(),
319        })
320    }
321
322    /// Adds a new object under the next free number.
323    fn put(&mut self, obj: Object) -> ObjRef {
324        let r = ObjRef {
325            num: self.next,
326            gen: 0,
327        };
328        self.next += 1;
329        self.objects.push((r, obj));
330        r
331    }
332
333    /// Replaces an object of the base under its own number.
334    fn replace(&mut self, r: ObjRef, obj: Object) {
335        self.objects.push((r, obj));
336    }
337
338    /// The overlay's first page as a form XObject in the base's object
339    /// space: its media box as the form's bounding box, its decoded content
340    /// as the form's stream, and its resource graph deep-copied and
341    /// renumbered.
342    fn import_form(&mut self, overlay: &Document) -> Result<ObjRef> {
343        let page = overlay.page(0).map_err(core_error)?;
344        let content = page.content(overlay).map_err(core_error)?;
345        let resources = self.import_object(overlay, &Object::Dict(page.resources.clone()))?;
346        let bbox = page.media_box;
347        let mut dict = Dict::new();
348        dict.insert(name("Type"), Object::Name(name("XObject")));
349        dict.insert(name("Subtype"), Object::Name(name("Form")));
350        dict.insert(name("FormType"), Object::Int(1));
351        dict.insert(
352            name("BBox"),
353            Object::Array(
354                [bbox.x0, bbox.y0, bbox.x1, bbox.y1]
355                    .iter()
356                    .map(|v| Object::Real(f64::from(*v)))
357                    .collect(),
358            ),
359        );
360        dict.insert(name("Resources"), resources);
361        dict.insert(name("Filter"), Object::Name(name("FlateDecode")));
362        Ok(self.put(Object::Stream(Stream {
363            dict,
364            data: deflate(&content),
365        })))
366    }
367
368    /// A deep copy of `obj` from `source` into the update: every reference
369    /// it reaches becomes a new object here, each source object copied once
370    /// however many times it is referenced. Streams keep their encoded
371    /// bytes and filters; their `/Length` is rewritten on emission.
372    fn import_object(&mut self, source: &Document, obj: &Object) -> Result<Object> {
373        Ok(match obj {
374            Object::Ref(r) => {
375                if let Some(copied) = self.imported.get(r) {
376                    return Ok(Object::Ref(*copied));
377                }
378                let copied = ObjRef {
379                    num: self.next,
380                    gen: 0,
381                };
382                self.next += 1;
383                self.imported.insert(*r, copied);
384                let body = source.get(*r).map_err(core_error)?;
385                let body = self.import_object(source, &body)?;
386                self.objects.push((copied, body));
387                Object::Ref(copied)
388            }
389            Object::Dict(d) => Object::Dict(self.import_dict(source, d)?),
390            Object::Array(items) => Object::Array(
391                items
392                    .iter()
393                    .map(|item| self.import_object(source, item))
394                    .collect::<Result<Vec<Object>>>()?,
395            ),
396            Object::Stream(s) => {
397                let mut dict = s.dict.clone();
398                dict.remove("Length");
399                Object::Stream(Stream {
400                    dict: self.import_dict(source, &dict)?,
401                    data: s.data.clone(),
402                })
403            }
404            other => other.clone(),
405        })
406    }
407
408    fn import_dict(&mut self, source: &Document, dict: &Dict) -> Result<Dict> {
409        let mut out = Dict::new();
410        for (key, value) in dict.iter() {
411            out.insert(key.clone(), self.import_object(source, value)?);
412        }
413        Ok(out)
414    }
415
416    /// The base bytes followed by the update section: every object, then a
417    /// cross-reference section in the base's style naming the base's
418    /// section as `/Prev`.
419    fn finish(mut self) -> Result<Vec<u8>> {
420        let base_bytes = self.base.bytes();
421        let prev = startxref(base_bytes)
422            .ok_or_else(|| Error::Other("the base file has no startxref".to_string()))?;
423        let mut out = base_bytes.to_vec();
424        if !out.ends_with(b"\n") {
425            out.push(b'\n');
426        }
427        self.objects.sort_by_key(|(r, _)| r.num);
428        let mut rows: Vec<(ObjRef, usize)> = Vec::with_capacity(self.objects.len() + 1);
429        for (r, obj) in &self.objects {
430            rows.push((*r, out.len()));
431            write_indirect(&mut out, *r, obj)?;
432        }
433        let trailer = &self.base.xref().trailer;
434        let mut section = Dict::new();
435        for key in ["Root", "Info", "ID"] {
436            if let Some(value) = trailer.get(key) {
437                section.insert(name(key), value.clone());
438            }
439        }
440        section.insert(name("Prev"), Object::Int(prev as i64));
441        let stream_style = trailer.get_name("Type").is_some_and(|n| n.0 == "XRef");
442        if stream_style {
443            self.finish_stream(&mut out, rows, section)?;
444        } else {
445            finish_table(&mut out, &rows, section, self.next)?;
446        }
447        Ok(out)
448    }
449
450    /// A cross-reference stream as the section's last object, rows for
451    /// every object of the update and for the stream itself.
452    fn finish_stream(
453        &mut self,
454        out: &mut Vec<u8>,
455        mut rows: Vec<(ObjRef, usize)>,
456        mut section: Dict,
457    ) -> Result<()> {
458        let xref_ref = ObjRef {
459            num: self.next,
460            gen: 0,
461        };
462        self.next += 1;
463        let xref_offset = out.len();
464        rows.push((xref_ref, xref_offset));
465        let mut index = Vec::with_capacity(rows.len() * 2);
466        let mut data = Vec::with_capacity(rows.len() * 7);
467        for (r, offset) in &rows {
468            index.push(Object::Int(i64::from(r.num)));
469            index.push(Object::Int(1));
470            data.push(1);
471            data.extend_from_slice(&field_offset(*offset)?.to_be_bytes());
472            data.extend_from_slice(&r.gen.to_be_bytes());
473        }
474        section.insert(name("Type"), Object::Name(name("XRef")));
475        section.insert(name("Size"), Object::Int(i64::from(self.next)));
476        section.insert(
477            name("W"),
478            Object::Array(vec![Object::Int(1), Object::Int(4), Object::Int(2)]),
479        );
480        section.insert(name("Index"), Object::Array(index));
481        write_indirect(
482            out,
483            xref_ref,
484            &Object::Stream(Stream {
485                dict: section,
486                data,
487            }),
488        )?;
489        out.extend_from_slice(format!("startxref\n{xref_offset}\n%%EOF\n").as_bytes());
490        Ok(())
491    }
492}
493
494/// A classic `xref` table with one subsection per run of consecutive
495/// object numbers, then the `trailer` dictionary.
496fn finish_table(
497    out: &mut Vec<u8>,
498    rows: &[(ObjRef, usize)],
499    mut section: Dict,
500    size: u32,
501) -> Result<()> {
502    let xref_offset = out.len();
503    out.extend_from_slice(b"xref\n");
504    let mut start = 0;
505    while start < rows.len() {
506        let mut end = start + 1;
507        while end < rows.len() && rows[end].0.num == rows[end - 1].0.num + 1 {
508            end += 1;
509        }
510        out.extend_from_slice(format!("{} {}\n", rows[start].0.num, end - start).as_bytes());
511        for (r, offset) in &rows[start..end] {
512            out.extend_from_slice(
513                format!("{:010} {:05} n \n", table_offset(*offset)?, r.gen).as_bytes(),
514            );
515        }
516        start = end;
517    }
518    section.insert(name("Size"), Object::Int(i64::from(size)));
519    out.extend_from_slice(b"trailer\n");
520    serialize_dict(&section, out)?;
521    out.extend_from_slice(format!("\nstartxref\n{xref_offset}\n%%EOF\n").as_bytes());
522    Ok(())
523}
524
525/// Emits `num gen obj` through `endobj`; a stream carries a direct
526/// `/Length` of its stored byte count.
527fn write_indirect(out: &mut Vec<u8>, r: ObjRef, obj: &Object) -> Result<()> {
528    out.extend_from_slice(format!("{} {} obj\n", r.num, r.gen).as_bytes());
529    match obj {
530        Object::Stream(s) => {
531            let mut dict = s.dict.clone();
532            dict.insert(name("Length"), Object::Int(s.data.len() as i64));
533            serialize_dict(&dict, out)?;
534            out.extend_from_slice(b"\nstream\n");
535            out.extend_from_slice(&s.data);
536            out.extend_from_slice(b"\nendstream\nendobj\n");
537        }
538        direct => {
539            serialize_object(direct, out)?;
540            out.extend_from_slice(b"\nendobj\n");
541        }
542    }
543    Ok(())
544}
545
546/// An uncompressed stream with no filter of its own.
547fn plain_stream(data: Vec<u8>) -> Stream {
548    Stream {
549        dict: Dict::new(),
550        data,
551    }
552}
553
554fn deflate(data: &[u8]) -> Vec<u8> {
555    let mut encoder = ZlibEncoder::new(Vec::new(), Compression::default());
556    encoder
557        .write_all(data)
558        .expect("writing into a Vec cannot fail");
559    encoder
560        .finish()
561        .expect("finishing an in-memory zlib stream cannot fail")
562}
563
564/// A byte position as the 4-byte offset field of a cross-reference stream.
565fn field_offset(position: usize) -> Result<u32> {
566    u32::try_from(position)
567        .map_err(|_| Error::Other("file offset exceeds the 4-byte xref field".to_string()))
568}
569
570/// A byte position as the 10-digit offset field of a classic xref table.
571fn table_offset(position: usize) -> Result<usize> {
572    if position as u64 <= 9_999_999_999 {
573        return Ok(position);
574    }
575    Err(Error::Other(
576        "file offset exceeds the 10-digit xref table field".to_string(),
577    ))
578}
579
580fn name(text: &str) -> Name {
581    Name(text.to_string())
582}
583
584fn core_error(error: pdfboss_core::Error) -> Error {
585    Error::Other(error.to_string())
586}