pdf-xfa 1.2.0

// Copyright (c) 2026 Innovation Trigger B.V.
//
// PDFluent is available under two licences, at your option: the GNU AGPLv3, or
// the PDFluent Commercial Licence. See the LICENSE file in this repository --
// that file travels with the copy you received, which a URL does not.

//! Content-addressed cache for flattened XFA output (#468).
//!
//! Flattening an XFA form costs 117 to 729 ms on the corpus documents the
//! editor measured, and the render path used to pay it again on every render
//! call. The flatten is a pure function of the input bytes (and of the
//! `XFA_*` environment switches the pipeline reads), so the cache key is a
//! SHA-256 over exactly those, and nothing else:
//!
//! - **Invalidation is by construction.** A mutation that changes any byte of
//!   the document yields a different key and therefore a miss. There is no
//!   invalidation call a caller can forget, and no path by which a stale
//!   flatten can be served for changed input.
//! - **What it does not do.** A mutation that changes bytes is a full flatten
//!   again. Skipping that needs a proof that the changed bytes cannot reach the
//!   flatten, and the flatten reads the whole file (embedded images, fonts, the
//!   static fallback), so no such proof is attempted here.
//!
//! Only successful flattens are stored, so an error is never replayed. The
//! cache is bounded by entry count and by total bytes; the least recently used
//! entry goes first.

use std::collections::HashMap;
use std::sync::{Arc, Mutex, OnceLock};

use sha2::{Digest, Sha256};

use crate::error::XfaError;

/// Environment switches that change what the flatten pipeline produces. They
/// are part of the key: the same bytes under a different switch are a
/// different input.
const KEY_ENV: &[&str] = &[
    "XFA_JS_EXECUTION_MODE",
    "XFA_JS_HARVEST_MODE",
    "XFA_SUPPRESSION_TRUST_LAYOUT",
    "XFA_FORMDOM_ADMIT_DATABOUND",
    "XFA_PRESENCE_PROV",
];

/// Default bound on the number of cached flattens.
pub const DEFAULT_MAX_ENTRIES: usize = 16;
/// Default bound on the total size of cached flattens (128 MiB).
pub const DEFAULT_MAX_BYTES: usize = 128 * 1024 * 1024;

/// Identity of one flatten input.
#[derive(Clone, Copy, PartialEq, Eq, Hash, Debug)]
pub struct FlattenKey([u8; 32]);

impl FlattenKey {
    /// Key for `pdf_bytes` under the current environment switches.
    pub fn for_bytes(pdf_bytes: &[u8]) -> Self {
        Self::for_bytes_with_env(pdf_bytes, |name| std::env::var(name).ok())
    }

    /// As [`for_bytes`](Self::for_bytes) with the switches read through `env`,
    /// so a test can vary one without mutating the process environment.
    pub fn for_bytes_with_env(pdf_bytes: &[u8], env: impl Fn(&str) -> Option<String>) -> Self {
        let mut h = Sha256::new();
        for name in KEY_ENV {
            h.update(name.as_bytes());
            h.update([0]);
            if let Some(v) = env(name) {
                h.update(v.as_bytes());
            }
            h.update([1]);
        }
        h.update((pdf_bytes.len() as u64).to_le_bytes());
        h.update(pdf_bytes);
        Self(h.finalize().into())
    }
}

/// Counters for observing the cache. `hits + misses` is the number of lookups.
#[derive(Clone, Copy, PartialEq, Eq, Debug, Default)]
pub struct FlattenCacheStats {
    /// Lookups answered from the cache.
    pub hits: u64,
    /// Lookups that had to run the flatten.
    pub misses: u64,
    /// Entries currently held.
    pub entries: usize,
    /// Total size in bytes of the entries currently held.
    pub bytes: usize,
}

struct Entry {
    bytes: Arc<Vec<u8>>,
    last_used: u64,
}

#[derive(Default)]
struct Inner {
    map: HashMap<FlattenKey, Entry>,
    tick: u64,
    total_bytes: usize,
    hits: u64,
    misses: u64,
}

/// A bounded, thread-safe flatten cache. Most callers want the process-wide
/// one behind [`flatten_cached`]; the type is public so a host can own a
/// separate instance with its own bounds.
pub struct FlattenCache {
    max_entries: usize,
    max_bytes: usize,
    inner: Mutex<Inner>,
}

impl FlattenCache {
    /// A cache holding at most `max_entries` flattens and `max_bytes` bytes.
    pub fn new(max_entries: usize, max_bytes: usize) -> Self {
        Self {
            max_entries,
            max_bytes,
            inner: Mutex::new(Inner::default()),
        }
    }

    fn lock(&self) -> std::sync::MutexGuard<'_, Inner> {
        // A panic inside the flatten closure never holds this lock (it runs
        // outside it), so poison can only mean a panic in the few lines of
        // bookkeeping below; the map is still structurally valid.
        self.inner.lock().unwrap_or_else(|p| p.into_inner())
    }

    /// Return the cached flatten for `key`, or run `flatten` and store its
    /// result. The closure runs without the lock held, so a slow flatten never
    /// blocks other lookups; two threads missing on the same key both flatten
    /// and the later store replaces the earlier one with identical bytes.
    pub fn get_or_flatten<F>(&self, key: FlattenKey, flatten: F) -> Result<Arc<Vec<u8>>, XfaError>
    where
        F: FnOnce() -> Result<Vec<u8>, XfaError>,
    {
        {
            let mut g = self.lock();
            g.tick += 1;
            let tick = g.tick;
            if let Some(e) = g.map.get_mut(&key) {
                e.last_used = tick;
                let hit = Arc::clone(&e.bytes);
                g.hits += 1;
                return Ok(hit);
            }
            g.misses += 1;
        }

        let out = Arc::new(flatten()?);

        if out.len() <= self.max_bytes {
            let mut g = self.lock();
            g.tick += 1;
            let tick = g.tick;
            if let Some(old) = g.map.insert(
                key,
                Entry {
                    bytes: Arc::clone(&out),
                    last_used: tick,
                },
            ) {
                g.total_bytes -= old.bytes.len();
            }
            g.total_bytes += out.len();
            while g.map.len() > self.max_entries || g.total_bytes > self.max_bytes {
                let Some(victim) = g
                    .map
                    .iter()
                    .filter(|(k, _)| **k != key)
                    .min_by_key(|(_, e)| e.last_used)
                    .map(|(k, _)| *k)
                else {
                    break;
                };
                if let Some(e) = g.map.remove(&victim) {
                    g.total_bytes -= e.bytes.len();
                }
            }
        }
        Ok(out)
    }

    /// Drop the entry for `key`, if any. Returns whether one was removed.
    pub fn invalidate(&self, key: &FlattenKey) -> bool {
        let mut g = self.lock();
        match g.map.remove(key) {
            Some(e) => {
                g.total_bytes -= e.bytes.len();
                true
            }
            None => false,
        }
    }

    /// Drop every entry. Counters are kept.
    pub fn clear(&self) {
        let mut g = self.lock();
        g.map.clear();
        g.total_bytes = 0;
    }

    /// Current counters.
    pub fn stats(&self) -> FlattenCacheStats {
        let g = self.lock();
        FlattenCacheStats {
            hits: g.hits,
            misses: g.misses,
            entries: g.map.len(),
            bytes: g.total_bytes,
        }
    }
}

fn global() -> &'static FlattenCache {
    static CACHE: OnceLock<FlattenCache> = OnceLock::new();
    CACHE.get_or_init(|| FlattenCache::new(DEFAULT_MAX_ENTRIES, DEFAULT_MAX_BYTES))
}

/// [`crate::flatten::flatten_xfa_to_pdf`] through the process-wide cache.
///
/// The returned bytes are identical to what the uncached call returns for the
/// same input; only the cost differs. Safe to call from any thread.
pub fn flatten_cached(pdf_bytes: &[u8]) -> Result<Arc<Vec<u8>>, XfaError> {
    global().get_or_flatten(FlattenKey::for_bytes(pdf_bytes), || {
        crate::flatten::flatten_xfa_to_pdf(pdf_bytes)
    })
}

/// Drop the process-wide entry for exactly these input bytes.
///
/// Not needed for correctness (a changed document has a different key); it
/// exists to release memory for a document the caller knows it is done with.
pub fn invalidate(pdf_bytes: &[u8]) -> bool {
    global().invalidate(&FlattenKey::for_bytes(pdf_bytes))
}

/// Drop every process-wide entry.
pub fn clear() {
    global().clear();
}

/// Counters of the process-wide cache.
pub fn stats() -> FlattenCacheStats {
    global().stats()
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::cell::Cell;

    fn key(n: u8) -> FlattenKey {
        FlattenKey::for_bytes(&[n; 8])
    }

    #[test]
    fn second_lookup_is_a_hit_and_does_not_run_the_flatten() {
        let c = FlattenCache::new(4, 1 << 20);
        let runs = Cell::new(0);
        let f = || {
            runs.set(runs.get() + 1);
            Ok(vec![1, 2, 3])
        };
        let a = c.get_or_flatten(key(1), f).unwrap();
        let b = c.get_or_flatten(key(1), f).unwrap();
        assert_eq!(runs.get(), 1);
        assert!(Arc::ptr_eq(&a, &b));
        let s = c.stats();
        assert_eq!((s.hits, s.misses, s.entries, s.bytes), (1, 1, 1, 3));
    }

    #[test]
    fn changed_input_is_a_different_key() {
        // A mutation flips one byte; it must never be served the old flatten.
        let mut bytes = vec![7u8; 4096];
        let before = FlattenKey::for_bytes(&bytes);
        bytes[2048] ^= 1;
        assert_ne!(before, FlattenKey::for_bytes(&bytes));
        // Same length, same prefix, different tail: the length alone is not
        // the key.
        assert_ne!(
            FlattenKey::for_bytes(&[1, 2, 3]),
            FlattenKey::for_bytes(&[1, 2, 4])
        );
    }

    #[test]
    fn explicit_invalidation_forces_a_new_flatten() {
        let c = FlattenCache::new(4, 1 << 20);
        let runs = Cell::new(0);
        let f = || {
            runs.set(runs.get() + 1);
            Ok(vec![9])
        };
        c.get_or_flatten(key(2), f).unwrap();
        assert!(c.invalidate(&key(2)));
        assert!(!c.invalidate(&key(2)));
        assert_eq!(c.stats().bytes, 0);
        c.get_or_flatten(key(2), f).unwrap();
        assert_eq!(runs.get(), 2);
    }

    #[test]
    fn errors_are_not_cached() {
        let c = FlattenCache::new(4, 1 << 20);
        let r = c.get_or_flatten(key(3), || Err(XfaError::LayoutFailed("boom".into())));
        assert!(r.is_err());
        assert_eq!(c.stats().entries, 0);
        let ok = c.get_or_flatten(key(3), || Ok(vec![1])).unwrap();
        assert_eq!(ok.as_slice(), &[1]);
    }

    #[test]
    fn evicts_least_recently_used_by_count() {
        let c = FlattenCache::new(2, 1 << 20);
        c.get_or_flatten(key(1), || Ok(vec![1])).unwrap();
        c.get_or_flatten(key(2), || Ok(vec![2])).unwrap();
        // Touch 1 so 2 is the oldest.
        c.get_or_flatten(key(1), || Ok(vec![0])).unwrap();
        c.get_or_flatten(key(3), || Ok(vec![3])).unwrap();
        let runs = Cell::new(0);
        let f = |v: u8| {
            let runs = &runs;
            move || {
                runs.set(runs.get() + 1);
                Ok(vec![v])
            }
        };
        c.get_or_flatten(key(1), f(1)).unwrap();
        c.get_or_flatten(key(3), f(3)).unwrap();
        assert_eq!(runs.get(), 0, "1 and 3 must still be cached");
        c.get_or_flatten(key(2), f(2)).unwrap();
        assert_eq!(runs.get(), 1, "2 must have been evicted");
    }

    #[test]
    fn byte_budget_bounds_total_and_skips_oversized_entries() {
        let c = FlattenCache::new(10, 10);
        c.get_or_flatten(key(1), || Ok(vec![0; 6])).unwrap();
        c.get_or_flatten(key(2), || Ok(vec![0; 6])).unwrap();
        let s = c.stats();
        assert!(s.bytes <= 10, "budget exceeded: {}", s.bytes);
        assert_eq!(s.entries, 1);
        // Larger than the whole budget: returned, never stored.
        let big = c.get_or_flatten(key(3), || Ok(vec![0; 64])).unwrap();
        assert_eq!(big.len(), 64);
        assert!(c.stats().bytes <= 10);
    }

    #[test]
    fn every_environment_switch_changes_the_key() {
        let bytes = [5u8; 32];
        let base = FlattenKey::for_bytes_with_env(&bytes, |_| None);
        for name in KEY_ENV {
            let with = FlattenKey::for_bytes_with_env(&bytes, |n| {
                (n == *name).then(|| "sandboxed".to_string())
            });
            assert_ne!(base, with, "{name} does not reach the key");
        }
        // A value moved from one switch to another is not the same input.
        let a =
            FlattenKey::for_bytes_with_env(&bytes, |n| (n == KEY_ENV[0]).then(|| "1".to_string()));
        let b =
            FlattenKey::for_bytes_with_env(&bytes, |n| (n == KEY_ENV[1]).then(|| "1".to_string()));
        assert_ne!(a, b);
    }
}