pdf_xfa/
flatten_cache.rs

1// Copyright (c) 2026 Innovation Trigger B.V.
2//
3// PDFluent is available under two licences, at your option: the GNU AGPLv3, or
4// the PDFluent Commercial Licence. See the LICENSE file in this repository --
5// that file travels with the copy you received, which a URL does not.
6
7//! Content-addressed cache for flattened XFA output (#468).
8//!
9//! Flattening an XFA form costs 117 to 729 ms on the corpus documents the
10//! editor measured, and the render path used to pay it again on every render
11//! call. The flatten is a pure function of the input bytes (and of the
12//! `XFA_*` environment switches the pipeline reads), so the cache key is a
13//! SHA-256 over exactly those, and nothing else:
14//!
15//! - **Invalidation is by construction.** A mutation that changes any byte of
16//!   the document yields a different key and therefore a miss. There is no
17//!   invalidation call a caller can forget, and no path by which a stale
18//!   flatten can be served for changed input.
19//! - **What it does not do.** A mutation that changes bytes is a full flatten
20//!   again. Skipping that needs a proof that the changed bytes cannot reach the
21//!   flatten, and the flatten reads the whole file (embedded images, fonts, the
22//!   static fallback), so no such proof is attempted here.
23//!
24//! Only successful flattens are stored, so an error is never replayed. The
25//! cache is bounded by entry count and by total bytes; the least recently used
26//! entry goes first.
27
28use std::collections::HashMap;
29use std::sync::{Arc, Mutex, OnceLock};
30
31use sha2::{Digest, Sha256};
32
33use crate::error::XfaError;
34
35/// Environment switches that change what the flatten pipeline produces. They
36/// are part of the key: the same bytes under a different switch are a
37/// different input.
38const KEY_ENV: &[&str] = &[
39    "XFA_JS_EXECUTION_MODE",
40    "XFA_JS_HARVEST_MODE",
41    "XFA_SUPPRESSION_TRUST_LAYOUT",
42    "XFA_FORMDOM_ADMIT_DATABOUND",
43    "XFA_PRESENCE_PROV",
44];
45
46/// Default bound on the number of cached flattens.
47pub const DEFAULT_MAX_ENTRIES: usize = 16;
48/// Default bound on the total size of cached flattens (128 MiB).
49pub const DEFAULT_MAX_BYTES: usize = 128 * 1024 * 1024;
50
51/// Identity of one flatten input.
52#[derive(Clone, Copy, PartialEq, Eq, Hash, Debug)]
53pub struct FlattenKey([u8; 32]);
54
55impl FlattenKey {
56    /// Key for `pdf_bytes` under the current environment switches.
57    pub fn for_bytes(pdf_bytes: &[u8]) -> Self {
58        Self::for_bytes_with_env(pdf_bytes, |name| std::env::var(name).ok())
59    }
60
61    /// As [`for_bytes`](Self::for_bytes) with the switches read through `env`,
62    /// so a test can vary one without mutating the process environment.
63    pub fn for_bytes_with_env(pdf_bytes: &[u8], env: impl Fn(&str) -> Option<String>) -> Self {
64        let mut h = Sha256::new();
65        for name in KEY_ENV {
66            h.update(name.as_bytes());
67            h.update([0]);
68            if let Some(v) = env(name) {
69                h.update(v.as_bytes());
70            }
71            h.update([1]);
72        }
73        h.update((pdf_bytes.len() as u64).to_le_bytes());
74        h.update(pdf_bytes);
75        Self(h.finalize().into())
76    }
77}
78
79/// Counters for observing the cache. `hits + misses` is the number of lookups.
80#[derive(Clone, Copy, PartialEq, Eq, Debug, Default)]
81pub struct FlattenCacheStats {
82    /// Lookups answered from the cache.
83    pub hits: u64,
84    /// Lookups that had to run the flatten.
85    pub misses: u64,
86    /// Entries currently held.
87    pub entries: usize,
88    /// Total size in bytes of the entries currently held.
89    pub bytes: usize,
90}
91
92struct Entry {
93    bytes: Arc<Vec<u8>>,
94    last_used: u64,
95}
96
97#[derive(Default)]
98struct Inner {
99    map: HashMap<FlattenKey, Entry>,
100    tick: u64,
101    total_bytes: usize,
102    hits: u64,
103    misses: u64,
104}
105
106/// A bounded, thread-safe flatten cache. Most callers want the process-wide
107/// one behind [`flatten_cached`]; the type is public so a host can own a
108/// separate instance with its own bounds.
109pub struct FlattenCache {
110    max_entries: usize,
111    max_bytes: usize,
112    inner: Mutex<Inner>,
113}
114
115impl FlattenCache {
116    /// A cache holding at most `max_entries` flattens and `max_bytes` bytes.
117    pub fn new(max_entries: usize, max_bytes: usize) -> Self {
118        Self {
119            max_entries,
120            max_bytes,
121            inner: Mutex::new(Inner::default()),
122        }
123    }
124
125    fn lock(&self) -> std::sync::MutexGuard<'_, Inner> {
126        // A panic inside the flatten closure never holds this lock (it runs
127        // outside it), so poison can only mean a panic in the few lines of
128        // bookkeeping below; the map is still structurally valid.
129        self.inner.lock().unwrap_or_else(|p| p.into_inner())
130    }
131
132    /// Return the cached flatten for `key`, or run `flatten` and store its
133    /// result. The closure runs without the lock held, so a slow flatten never
134    /// blocks other lookups; two threads missing on the same key both flatten
135    /// and the later store replaces the earlier one with identical bytes.
136    pub fn get_or_flatten<F>(&self, key: FlattenKey, flatten: F) -> Result<Arc<Vec<u8>>, XfaError>
137    where
138        F: FnOnce() -> Result<Vec<u8>, XfaError>,
139    {
140        {
141            let mut g = self.lock();
142            g.tick += 1;
143            let tick = g.tick;
144            if let Some(e) = g.map.get_mut(&key) {
145                e.last_used = tick;
146                let hit = Arc::clone(&e.bytes);
147                g.hits += 1;
148                return Ok(hit);
149            }
150            g.misses += 1;
151        }
152
153        let out = Arc::new(flatten()?);
154
155        if out.len() <= self.max_bytes {
156            let mut g = self.lock();
157            g.tick += 1;
158            let tick = g.tick;
159            if let Some(old) = g.map.insert(
160                key,
161                Entry {
162                    bytes: Arc::clone(&out),
163                    last_used: tick,
164                },
165            ) {
166                g.total_bytes -= old.bytes.len();
167            }
168            g.total_bytes += out.len();
169            while g.map.len() > self.max_entries || g.total_bytes > self.max_bytes {
170                let Some(victim) = g
171                    .map
172                    .iter()
173                    .filter(|(k, _)| **k != key)
174                    .min_by_key(|(_, e)| e.last_used)
175                    .map(|(k, _)| *k)
176                else {
177                    break;
178                };
179                if let Some(e) = g.map.remove(&victim) {
180                    g.total_bytes -= e.bytes.len();
181                }
182            }
183        }
184        Ok(out)
185    }
186
187    /// Drop the entry for `key`, if any. Returns whether one was removed.
188    pub fn invalidate(&self, key: &FlattenKey) -> bool {
189        let mut g = self.lock();
190        match g.map.remove(key) {
191            Some(e) => {
192                g.total_bytes -= e.bytes.len();
193                true
194            }
195            None => false,
196        }
197    }
198
199    /// Drop every entry. Counters are kept.
200    pub fn clear(&self) {
201        let mut g = self.lock();
202        g.map.clear();
203        g.total_bytes = 0;
204    }
205
206    /// Current counters.
207    pub fn stats(&self) -> FlattenCacheStats {
208        let g = self.lock();
209        FlattenCacheStats {
210            hits: g.hits,
211            misses: g.misses,
212            entries: g.map.len(),
213            bytes: g.total_bytes,
214        }
215    }
216}
217
218fn global() -> &'static FlattenCache {
219    static CACHE: OnceLock<FlattenCache> = OnceLock::new();
220    CACHE.get_or_init(|| FlattenCache::new(DEFAULT_MAX_ENTRIES, DEFAULT_MAX_BYTES))
221}
222
223/// [`crate::flatten::flatten_xfa_to_pdf`] through the process-wide cache.
224///
225/// The returned bytes are identical to what the uncached call returns for the
226/// same input; only the cost differs. Safe to call from any thread.
227pub fn flatten_cached(pdf_bytes: &[u8]) -> Result<Arc<Vec<u8>>, XfaError> {
228    global().get_or_flatten(FlattenKey::for_bytes(pdf_bytes), || {
229        crate::flatten::flatten_xfa_to_pdf(pdf_bytes)
230    })
231}
232
233/// Drop the process-wide entry for exactly these input bytes.
234///
235/// Not needed for correctness (a changed document has a different key); it
236/// exists to release memory for a document the caller knows it is done with.
237pub fn invalidate(pdf_bytes: &[u8]) -> bool {
238    global().invalidate(&FlattenKey::for_bytes(pdf_bytes))
239}
240
241/// Drop every process-wide entry.
242pub fn clear() {
243    global().clear();
244}
245
246/// Counters of the process-wide cache.
247pub fn stats() -> FlattenCacheStats {
248    global().stats()
249}
250
251#[cfg(test)]
252mod tests {
253    use super::*;
254    use std::cell::Cell;
255
256    fn key(n: u8) -> FlattenKey {
257        FlattenKey::for_bytes(&[n; 8])
258    }
259
260    #[test]
261    fn second_lookup_is_a_hit_and_does_not_run_the_flatten() {
262        let c = FlattenCache::new(4, 1 << 20);
263        let runs = Cell::new(0);
264        let f = || {
265            runs.set(runs.get() + 1);
266            Ok(vec![1, 2, 3])
267        };
268        let a = c.get_or_flatten(key(1), f).unwrap();
269        let b = c.get_or_flatten(key(1), f).unwrap();
270        assert_eq!(runs.get(), 1);
271        assert!(Arc::ptr_eq(&a, &b));
272        let s = c.stats();
273        assert_eq!((s.hits, s.misses, s.entries, s.bytes), (1, 1, 1, 3));
274    }
275
276    #[test]
277    fn changed_input_is_a_different_key() {
278        // A mutation flips one byte; it must never be served the old flatten.
279        let mut bytes = vec![7u8; 4096];
280        let before = FlattenKey::for_bytes(&bytes);
281        bytes[2048] ^= 1;
282        assert_ne!(before, FlattenKey::for_bytes(&bytes));
283        // Same length, same prefix, different tail: the length alone is not
284        // the key.
285        assert_ne!(
286            FlattenKey::for_bytes(&[1, 2, 3]),
287            FlattenKey::for_bytes(&[1, 2, 4])
288        );
289    }
290
291    #[test]
292    fn explicit_invalidation_forces_a_new_flatten() {
293        let c = FlattenCache::new(4, 1 << 20);
294        let runs = Cell::new(0);
295        let f = || {
296            runs.set(runs.get() + 1);
297            Ok(vec![9])
298        };
299        c.get_or_flatten(key(2), f).unwrap();
300        assert!(c.invalidate(&key(2)));
301        assert!(!c.invalidate(&key(2)));
302        assert_eq!(c.stats().bytes, 0);
303        c.get_or_flatten(key(2), f).unwrap();
304        assert_eq!(runs.get(), 2);
305    }
306
307    #[test]
308    fn errors_are_not_cached() {
309        let c = FlattenCache::new(4, 1 << 20);
310        let r = c.get_or_flatten(key(3), || Err(XfaError::LayoutFailed("boom".into())));
311        assert!(r.is_err());
312        assert_eq!(c.stats().entries, 0);
313        let ok = c.get_or_flatten(key(3), || Ok(vec![1])).unwrap();
314        assert_eq!(ok.as_slice(), &[1]);
315    }
316
317    #[test]
318    fn evicts_least_recently_used_by_count() {
319        let c = FlattenCache::new(2, 1 << 20);
320        c.get_or_flatten(key(1), || Ok(vec![1])).unwrap();
321        c.get_or_flatten(key(2), || Ok(vec![2])).unwrap();
322        // Touch 1 so 2 is the oldest.
323        c.get_or_flatten(key(1), || Ok(vec![0])).unwrap();
324        c.get_or_flatten(key(3), || Ok(vec![3])).unwrap();
325        let runs = Cell::new(0);
326        let f = |v: u8| {
327            let runs = &runs;
328            move || {
329                runs.set(runs.get() + 1);
330                Ok(vec![v])
331            }
332        };
333        c.get_or_flatten(key(1), f(1)).unwrap();
334        c.get_or_flatten(key(3), f(3)).unwrap();
335        assert_eq!(runs.get(), 0, "1 and 3 must still be cached");
336        c.get_or_flatten(key(2), f(2)).unwrap();
337        assert_eq!(runs.get(), 1, "2 must have been evicted");
338    }
339
340    #[test]
341    fn byte_budget_bounds_total_and_skips_oversized_entries() {
342        let c = FlattenCache::new(10, 10);
343        c.get_or_flatten(key(1), || Ok(vec![0; 6])).unwrap();
344        c.get_or_flatten(key(2), || Ok(vec![0; 6])).unwrap();
345        let s = c.stats();
346        assert!(s.bytes <= 10, "budget exceeded: {}", s.bytes);
347        assert_eq!(s.entries, 1);
348        // Larger than the whole budget: returned, never stored.
349        let big = c.get_or_flatten(key(3), || Ok(vec![0; 64])).unwrap();
350        assert_eq!(big.len(), 64);
351        assert!(c.stats().bytes <= 10);
352    }
353
354    #[test]
355    fn every_environment_switch_changes_the_key() {
356        let bytes = [5u8; 32];
357        let base = FlattenKey::for_bytes_with_env(&bytes, |_| None);
358        for name in KEY_ENV {
359            let with = FlattenKey::for_bytes_with_env(&bytes, |n| {
360                (n == *name).then(|| "sandboxed".to_string())
361            });
362            assert_ne!(base, with, "{name} does not reach the key");
363        }
364        // A value moved from one switch to another is not the same input.
365        let a =
366            FlattenKey::for_bytes_with_env(&bytes, |n| (n == KEY_ENV[0]).then(|| "1".to_string()));
367        let b =
368            FlattenKey::for_bytes_with_env(&bytes, |n| (n == KEY_ENV[1]).then(|| "1".to_string()));
369        assert_ne!(a, b);
370    }
371}
pdf_xfa/flatten_cache.rs

pdf_xfa/
flatten_cache.rs
