Skip to main content

rightkit_docs/
lib.rs

1//! RightKit documents: read and write DOCX, XLSX and PDF behind one typed API, for agents and
2//! apps that must not trust the files they open.
3//!
4//! ```no_run
5//! use std::path::Path;
6//! use rightkit_docs::{read, write, Block, Op, ReadOptions, Source, WriteOptions};
7//!
8//! let path = Path::new("report.docx");
9//! write(path, &[Op::CreateDocx { blocks: vec![Block::Heading { level: 1, text: "Q3".into() }] }],
10//!       &WriteOptions::default())?;
11//! let doc = read(Source::Path(path), &ReadOptions::default())?;
12//! // Edit against exactly what was read: refuses if the file changed in between.
13//! write(path, &[Op::ReplaceText { find: "Q3".into(), replace: "Q4".into() }],
14//!       &WriteOptions { expect_sha256: Some(doc.sha256), ..Default::default() })?;
15//! # Ok::<(), rightkit_docs::DocsError>(())
16//! ```
17//!
18//! | module   | responsibility                                                              |
19//! |----------|-----------------------------------------------------------------------------|
20//! | `detect` | format from content: PDF, DOCX, XLSX; refuses macros, encryption, the rest   |
21//! | `package`| OOXML zip preflight (entries, sizes, ratio, names, DOCTYPE) and raw rewrite  |
22//! | `docx`   | read, create, append blocks, replace text by byte-range patch                |
23//! | `xlsx`   | read (calamine), create (rust_xlsxwriter), set cells and append rows in place|
24//! | `pdf`    | facade over `rightkit-pdf`: extract, replace text, add text, select pages    |
25//!
26//! Every write is built in memory, validated by re-detecting it, and published once through
27//! `rightkit_fs::AtomicWriter` after a last source-hash recheck.
28
29mod detect;
30mod docx;
31mod error;
32mod limits;
33mod model;
34mod package;
35mod pdf;
36mod xlsx;
37mod xmlutil;
38
39use std::collections::BTreeSet;
40use std::fs;
41use std::io::Read as _;
42use std::path::{Path, PathBuf};
43use std::sync::atomic::{AtomicBool, Ordering};
44
45use sha2::{Digest, Sha256};
46
47pub use error::{DocsError, Result};
48pub use limits::Limits;
49pub use model::*;
50
51use detect::{detect, Opened};
52
53fn sha256_hex(bytes: &[u8]) -> String {
54    Sha256::digest(bytes)
55        .iter()
56        .map(|b| format!("{b:02x}"))
57        .collect()
58}
59
60fn read_file(path: &Path, limit: u64) -> Result<Vec<u8>> {
61    let meta = fs::metadata(path).map_err(|e| DocsError::Io(format!("{}: {e}", path.display())))?;
62    if !meta.is_file() {
63        return Err(DocsError::Io(format!(
64            "{} is not a regular file",
65            path.display()
66        )));
67    }
68    if meta.len() > limit {
69        return Err(DocsError::TooLarge {
70            limit,
71            actual: meta.len(),
72        });
73    }
74    let file =
75        fs::File::open(path).map_err(|e| DocsError::Io(format!("{}: {e}", path.display())))?;
76    let mut buf = Vec::with_capacity(meta.len() as usize);
77    file.take(limit + 1)
78        .read_to_end(&mut buf)
79        .map_err(|e| DocsError::Io(format!("{}: {e}", path.display())))?;
80    if buf.len() as u64 > limit {
81        return Err(DocsError::TooLarge {
82            limit,
83            actual: buf.len() as u64,
84        });
85    }
86    Ok(buf)
87}
88
89/// Cut `s` to at most `max` characters on a character boundary.
90fn cap(mut s: String, max: usize) -> (String, bool) {
91    match s.char_indices().nth(max) {
92        Some((i, _)) => {
93            s.truncate(i);
94            (s, true)
95        }
96        None => (s, false),
97    }
98}
99
100/// Read a DOCX, XLSX or PDF. The format comes from the bytes; the file name is never consulted.
101pub fn read(source: Source<'_>, options: &ReadOptions) -> Result<DocumentRead> {
102    let owned;
103    let bytes: &[u8] = match source {
104        Source::Bytes(b) => b,
105        Source::Path(p) => {
106            owned = read_file(p, options.limits.max_input_bytes)?;
107            &owned
108        }
109    };
110    let sha = sha256_hex(bytes);
111    let size = bytes.len() as u64;
112    match detect(bytes, &options.limits)? {
113        Opened::Pdf => {
114            let r = pdf::read(bytes, &options.limits, &options.pages, options.max_chars)?;
115            Ok(DocumentRead {
116                format: Format::Pdf,
117                sha256: sha,
118                size_bytes: size,
119                text: r.text,
120                markdown: options.markdown.then_some(r.markdown),
121                truncated: r.truncated,
122                content: Content::Pdf(r.content),
123                omitted: r.omitted,
124            })
125        }
126        Opened::Docx(o) => {
127            let r = docx::read(&o, options.max_chars)?;
128            let (text, t1) = cap(docx::plain_blocks(&r.content.blocks), options.max_chars);
129            let (md, t2) = if options.markdown {
130                let (m, t) = cap(docx::markdown_blocks(&r.content.blocks), options.max_chars);
131                (Some(m), t)
132            } else {
133                (None, false)
134            };
135            Ok(DocumentRead {
136                format: Format::Docx,
137                sha256: sha,
138                size_bytes: size,
139                text,
140                markdown: md,
141                truncated: r.truncated || t1 || t2,
142                content: Content::Docx(r.content),
143                omitted: r.omitted,
144            })
145        }
146        Opened::Xlsx(o) => {
147            let r = xlsx::read(&o, bytes, options)?;
148            let (text, md) = xlsx::render(&r.content);
149            let (text, t1) = cap(text, options.max_chars);
150            let (md, t2) = if options.markdown {
151                let (m, t) = cap(md, options.max_chars);
152                (Some(m), t)
153            } else {
154                (None, false)
155            };
156            Ok(DocumentRead {
157                format: Format::Xlsx,
158                sha256: sha,
159                size_bytes: size,
160                text,
161                markdown: md,
162                truncated: r.truncated || t1 || t2,
163                content: Content::Xlsx(r.content),
164                omitted: r.omitted,
165            })
166        }
167    }
168}
169
170fn op_name(op: &Op) -> &'static str {
171    match op {
172        Op::CreateDocx { .. } => "createDocx",
173        Op::CreateXlsx { .. } => "createXlsx",
174        Op::ReplaceText { .. } => "replaceText",
175        Op::AppendBlocks { .. } => "appendBlocks",
176        Op::SetCells { .. } => "setCells",
177        Op::AppendRows { .. } => "appendRows",
178        Op::AddPdfText { .. } => "addPdfText",
179        Op::SelectPdfPages { .. } => "selectPdfPages",
180    }
181}
182
183#[derive(Default)]
184struct Totals {
185    replacements: usize,
186    cells_set: usize,
187    rows_appended: usize,
188    touched: BTreeSet<String>,
189    losses: Vec<Loss>,
190    warnings: Vec<String>,
191}
192
193impl Totals {
194    fn absorb(&mut self, o: docx::EditOutcome) -> Vec<u8> {
195        self.touched.extend(o.touched);
196        for l in o.losses {
197            if !self.losses.contains(&l) {
198                self.losses.push(l);
199            }
200        }
201        for w in o.warnings {
202            if !self.warnings.contains(&w) {
203                self.warnings.push(w);
204            }
205        }
206        o.bytes
207    }
208}
209
210fn wrong(op: &Op, fmt: Format) -> DocsError {
211    DocsError::WrongFormat {
212        op: op_name(op),
213        format: fmt.as_str(),
214    }
215}
216
217/// Apply one edit op to `current`, whose format is detected fresh.
218fn apply(
219    op: &Op,
220    current: &[u8],
221    limits: &Limits,
222    totals: &mut Totals,
223) -> Result<(Vec<u8>, Format)> {
224    match (op, detect(current, limits)?) {
225        (Op::ReplaceText { find, replace }, Opened::Docx(o)) => {
226            let out = docx::replace_text(&o, find, replace)?;
227            totals.replacements += out.count;
228            Ok((totals.absorb(out), Format::Docx))
229        }
230        (Op::AppendBlocks { blocks }, Opened::Docx(o)) => {
231            let out = docx::append_blocks(&o, blocks)?;
232            Ok((totals.absorb(out), Format::Docx))
233        }
234        (Op::SetCells { sheet, cells }, Opened::Xlsx(o)) => {
235            let (out, n) = xlsx::set_cells(&o, sheet, cells)?;
236            totals.cells_set += n;
237            Ok((totals.absorb(out), Format::Xlsx))
238        }
239        (Op::AppendRows { sheet, rows }, Opened::Xlsx(o)) => {
240            let (out, n) = xlsx::append_rows(&o, sheet, rows)?;
241            totals.rows_appended += n;
242            Ok((totals.absorb(out), Format::Xlsx))
243        }
244        (Op::ReplaceText { find, replace }, Opened::Pdf) => {
245            let out = pdf::replace_text(current, limits, find, replace)?;
246            totals.replacements += out.count;
247            Ok((totals.absorb(out), Format::Pdf))
248        }
249        (
250            Op::AddPdfText {
251                page,
252                x,
253                y,
254                text,
255                size_pt,
256            },
257            Opened::Pdf,
258        ) => {
259            let out = pdf::add_text(current, limits, *page, *x, *y, text, *size_pt)?;
260            Ok((totals.absorb(out), Format::Pdf))
261        }
262        (Op::SelectPdfPages { pages }, Opened::Pdf) => {
263            let out = pdf::select_pages(current, limits, pages)?;
264            Ok((totals.absorb(out), Format::Pdf))
265        }
266        (Op::CreateDocx { .. } | Op::CreateXlsx { .. }, _) => Err(DocsError::InvalidOp(
267            "a create op is only valid as the first op".into(),
268        )),
269        (op, Opened::Pdf) => Err(wrong(op, Format::Pdf)),
270        (op, Opened::Docx(_)) => Err(wrong(op, Format::Docx)),
271        (op, Opened::Xlsx(_)) => Err(wrong(op, Format::Xlsx)),
272    }
273}
274
275fn format_of(bytes: &[u8], limits: &Limits) -> Result<Format> {
276    Ok(match detect(bytes, limits)? {
277        Opened::Pdf => Format::Pdf,
278        Opened::Docx(_) => Format::Docx,
279        Opened::Xlsx(_) => Format::Xlsx,
280    })
281}
282
283fn fs_err(e: rightkit_fs::FsError) -> DocsError {
284    if matches!(e, rightkit_fs::FsError::AlreadyExists(_)) {
285        return DocsError::TargetExists;
286    }
287    if e.target_replaced() {
288        return DocsError::SavedWithWarning(e.to_string());
289    }
290    DocsError::Io(e.to_string())
291}
292
293/// Apply typed ops to `target` and publish the result atomically.
294///
295/// A leading `CreateDocx` or `CreateXlsx` makes a new file (the target must not exist unless
296/// `options.overwrite`); any other first op edits the existing file. Ops apply in order to one
297/// in-memory document; nothing touches the disk until every op has succeeded and the result has
298/// been re-detected as the same format. An edit then rechecks the target's SHA-256 (against
299/// `options.expect_sha256` when given, always against the bytes this call started from) as the
300/// last step before the rename, so a file changed meanwhile is never overwritten
301/// (`source_changed`). The receipt lists everything the edit could not keep.
302pub fn write(target: &Path, ops: &[Op], options: &WriteOptions) -> Result<WriteReceipt> {
303    let limits = &options.limits;
304    let first = ops
305        .first()
306        .ok_or_else(|| DocsError::InvalidOp("no ops".into()))?;
307    let creating = matches!(first, Op::CreateDocx { .. } | Op::CreateXlsx { .. });
308
309    // One resolution of the path: a symlink target is replaced where it points, the link stays.
310    let existing = fs::symlink_metadata(target).is_ok();
311    let real: PathBuf = if existing {
312        fs::canonicalize(target).map_err(|e| DocsError::Io(format!("{}: {e}", target.display())))?
313    } else {
314        target.to_path_buf()
315    };
316
317    let mut totals = Totals::default();
318    let (mut current, mut format, source_sha): (Vec<u8>, Format, Option<String>) = if creating {
319        let source_sha = if existing {
320            if !options.overwrite {
321                return Err(DocsError::TargetExists);
322            }
323            let old = read_file(&real, limits.max_input_bytes)?;
324            let h = sha256_hex(&old);
325            if let Some(expect) = &options.expect_sha256 {
326                if *expect != h {
327                    return Err(DocsError::SourceChanged);
328                }
329            }
330            Some(h)
331        } else {
332            None
333        };
334        match first {
335            Op::CreateDocx { blocks } => (docx::create(blocks)?, Format::Docx, source_sha),
336            Op::CreateXlsx { sheets } => (xlsx::create(sheets)?, Format::Xlsx, source_sha),
337            _ => unreachable!("creating implies a create op"),
338        }
339    } else {
340        if !existing {
341            return Err(DocsError::TargetMissing);
342        }
343        let bytes = read_file(&real, limits.max_input_bytes)?;
344        let h = sha256_hex(&bytes);
345        if let Some(expect) = &options.expect_sha256 {
346            if *expect != h {
347                return Err(DocsError::SourceChanged);
348            }
349        }
350        let fmt = format_of(&bytes, limits)?;
351        (bytes, fmt, Some(h))
352    };
353
354    let start = usize::from(creating);
355    for op in &ops[start..] {
356        let (next, fmt) = apply(op, &current, limits, &mut totals)?;
357        current = next;
358        format = fmt;
359    }
360    if creating {
361        totals.touched.insert("(new document)".into());
362    }
363
364    // The output must be a document this crate itself accepts, in the format it started as.
365    if format_of(&current, limits)? != format {
366        return Err(DocsError::Malformed(
367            "the edit produced a different kind of document; nothing was written".into(),
368        ));
369    }
370    if current.len() as u64 > limits.max_input_bytes {
371        return Err(DocsError::TooLarge {
372            limit: limits.max_input_bytes,
373            actual: current.len() as u64,
374        });
375    }
376
377    // Publish. Last source recheck runs inside the writer, after the temp file is complete and
378    // synced and immediately before the rename.
379    let write_opts = rightkit_fs::atomic::AtomicWriteOptions {
380        replace: if existing {
381            rightkit_fs::atomic::ReplaceMode::Rename
382        } else {
383            rightkit_fs::atomic::ReplaceMode::NoReplace
384        },
385        ..Default::default()
386    };
387    let mut writer = rightkit_fs::AtomicWriter::create_with(&real, write_opts).map_err(fs_err)?;
388    if existing {
389        if let Ok(meta) = fs::metadata(&real) {
390            // Keep the permission bits (and the read-only attribute) of the file being replaced.
391            let _ = writer.file_mut().set_permissions(meta.permissions());
392        }
393    }
394    writer.write_all(&current).map_err(fs_err)?;
395    #[cfg(test)]
396    tests::run_before_publish_hook();
397    let changed = AtomicBool::new(false);
398    let expected = source_sha.clone();
399    let recheck_path = real.clone();
400    let limit = limits.max_input_bytes;
401    let publish = writer.commit_checked(|_temp| {
402        if let Some(h0) = &expected {
403            let now = read_file(&recheck_path, limit).map(|b| sha256_hex(&b)).ok();
404            if now.as_deref() != Some(h0.as_str()) {
405                changed.store(true, Ordering::SeqCst);
406                return Err(std::io::Error::other("source changed"));
407            }
408        } else if recheck_path.exists() {
409            changed.store(true, Ordering::SeqCst);
410            return Err(std::io::Error::other("target appeared"));
411        }
412        Ok(())
413    });
414    if let Err(e) = publish {
415        if changed.load(Ordering::SeqCst) {
416            return Err(if source_sha.is_some() {
417                DocsError::SourceChanged
418            } else {
419                DocsError::TargetExists
420            });
421        }
422        return Err(fs_err(e));
423    }
424
425    Ok(WriteReceipt {
426        path: target.to_path_buf(),
427        format,
428        created: !existing,
429        source_sha256: source_sha,
430        sha256: sha256_hex(&current),
431        bytes_written: current.len() as u64,
432        ops_applied: ops.len(),
433        replacements: totals.replacements,
434        cells_set: totals.cells_set,
435        rows_appended: totals.rows_appended,
436        touched_parts: totals.touched.into_iter().collect(),
437        losses: totals.losses,
438        warnings: totals.warnings,
439    })
440}
441
442#[cfg(test)]
443mod tests {
444    use super::*;
445    use std::cell::RefCell;
446
447    type Hook = Box<dyn Fn()>;
448    thread_local! {
449        static BEFORE_PUBLISH: RefCell<Option<Hook>> = const { RefCell::new(None) };
450    }
451
452    pub(super) fn run_before_publish_hook() {
453        BEFORE_PUBLISH.with(|h| {
454            if let Some(f) = h.borrow().as_ref() {
455                f();
456            }
457        });
458    }
459
460    fn scratch(tag: &str) -> PathBuf {
461        let dir =
462            std::env::temp_dir().join(format!("rightkit-docs-unit-{tag}-{}", std::process::id()));
463        let _ = fs::remove_dir_all(&dir);
464        fs::create_dir_all(&dir).unwrap();
465        dir
466    }
467
468    /// Another process rewrites the file after this call has read it and built its result, and
469    /// before the rename: the last recheck sees it, nothing is published, and the other writer's
470    /// bytes survive.
471    #[test]
472    fn a_file_changed_between_read_and_publish_is_never_overwritten() {
473        let dir = scratch("race");
474        let path = dir.join("a.docx");
475        let make = |t: &str| Op::CreateDocx {
476            blocks: vec![Block::Paragraph { text: t.into() }],
477        };
478        write(&path, &[make("original")], &WriteOptions::default()).unwrap();
479
480        let theirs = docx::create(&[Block::Paragraph {
481            text: "someone else".into(),
482        }])
483        .unwrap();
484        let (p, t) = (path.clone(), theirs.clone());
485        BEFORE_PUBLISH
486            .with(|h| *h.borrow_mut() = Some(Box::new(move || fs::write(&p, &t).unwrap())));
487        let e = write(
488            &path,
489            &[Op::ReplaceText {
490                find: "original".into(),
491                replace: "mine".into(),
492            }],
493            &WriteOptions::default(),
494        )
495        .unwrap_err();
496        BEFORE_PUBLISH.with(|h| *h.borrow_mut() = None);
497
498        assert_eq!(e.code(), "source_changed");
499        assert_eq!(
500            fs::read(&path).unwrap(),
501            theirs,
502            "the concurrent writer's file is intact"
503        );
504        let names: Vec<_> = fs::read_dir(&dir)
505            .unwrap()
506            .map(|e| e.unwrap().file_name())
507            .collect();
508        assert_eq!(names.len(), 1, "no temp file left behind: {names:?}");
509    }
510
511    /// The target file is created by someone else between the existence check and the rename of a
512    /// create: the create does not clobber it.
513    #[test]
514    fn a_create_never_clobbers_a_file_that_appears_first() {
515        let dir = scratch("create-race");
516        let path = dir.join("new.docx");
517        let (p, t) = (path.clone(), b"theirs".to_vec());
518        BEFORE_PUBLISH
519            .with(|h| *h.borrow_mut() = Some(Box::new(move || fs::write(&p, &t).unwrap())));
520        let e = write(
521            &path,
522            &[Op::CreateDocx {
523                blocks: vec![Block::Paragraph { text: "x".into() }],
524            }],
525            &WriteOptions::default(),
526        )
527        .unwrap_err();
528        BEFORE_PUBLISH.with(|h| *h.borrow_mut() = None);
529        assert_eq!(e.code(), "target_exists");
530        assert_eq!(fs::read(&path).unwrap(), b"theirs");
531    }
532}