1use 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
17const FORM_NAME: &str = "PdfbossWatermark";
19
20pub 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
71struct Stamp {
74 form: ObjRef,
75 prefix: ObjRef,
76 suffix: ObjRef,
77 pages: FastMap<ObjRef, usize>,
78}
79
80struct 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 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 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 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 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 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 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
245pub 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
295struct 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 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 fn replace(&mut self, r: ObjRef, obj: Object) {
335 self.objects.push((r, obj));
336 }
337
338 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 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 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 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
494fn 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(§ion, out)?;
521 out.extend_from_slice(format!("\nstartxref\n{xref_offset}\n%%EOF\n").as_bytes());
522 Ok(())
523}
524
525fn 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
546fn 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
564fn 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
570fn 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}