keyhog_core/

merkle_index.rs

//! Incremental scan support via a persisted file-content index.
//!
//! ## What it does
//!
//! On a fresh scan we compute, for every input chunk, a metadata tuple
//! `(mtime_ns, size, chunk_offset, BLAKE3(content))` and store it under the
//! file's canonical path plus chunk offset. On the next run, files whose
//! `(mtime, size)` match the stored values can be skipped *without re-reading
//! the bytes* -
//! they almost certainly haven't changed (rsync-style trust). When
//! `(mtime, size)` differ but BLAKE3 matches we record the new mtime
//! and still skip - same content, different stat (touched, copied).
//!
//! Tier-B moat innovation #3 from the internal design notes: "10–100×
//! speedup on CI re-runs" by skipping the 99% of files that didn't change.
//!
//! ## Schema versions
//!
//! - **v1 (legacy)** - `path → BLAKE3 hex` only. Loadable but lacks the
//!   metadata short-circuit; treated as cold-start to avoid mixing schemas.
//! - **v2 (legacy)** - `path → (mtime_ns, size, BLAKE3 hex)` plus a
//!   top-level `spec_hash` derived from the loaded detector set. A
//!   spec-hash mismatch invalidates the entire cache; this is the
//!   correctness fix for "added a detector but unchanged files were
//!   silently skipped, missing the new detection forever." Superseded by
//!   v3 and treated as cold-start (it lacks the racy-clean timestamp).
//! - **v3 (legacy)** - v2 plus a top-level `written_at_ns` (wall-clock
//!   nanoseconds when the index was last written). On load, any entry
//!   whose file `mtime_ns` falls in the same clock-second as - or after -
//!   `written_at_ns` is dropped (git's "racy index" guard): a
//!   size-preserving edit made in that window leaves `(mtime, size)`
//!   unchanged on coarse-granularity filesystems (FAT/HFS+/some NFS expose
//!   whole-second mtimes), so trusting the stored hash would skip a
//!   freshly injected secret forever. Dropped entries are simply re-read
//!   and re-hashed on the next scan - slower for those few files, never
//!   unsound.
//! - **v4 (current)** - v3 plus an explicit chunk offset in each persisted
//!   row. Chunked files no longer overwrite every earlier chunk under the
//!   same path, so incremental scans can skip unchanged large files by chunk
//!   instead of re-hashing them on every run. Newer v4 rows also carry a
//!   default-compatible `last_seen_order` so over-cap saves evict oldest cache
//!   rows deterministically instead of depending on hash-map iteration order.
//!
//! ## Serialization
//!
//! JSON, on purpose. The dataset is one row per scanned file (≤ ~1M for
//! any sane repo) and JSON keeps the on-disk format trivial to debug,
//! diff, and version-control if a team wants to.
//!
//! ## Threat model
//!
//! Cached entries do NOT contain credentials. Storing a `(mtime, size,
//! content_hash)` tuple per scanned path leaks that the path *exists*
//! and what its content fingerprint is, which is why `--lockdown`
//! refuses to load or write the cache at all.

use indexmap::{map::Entry, Equivalent, IndexMap};
use std::hash::{Hash, Hasher};
use std::path::{Path, PathBuf};
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::SystemTime;

use parking_lot::RwLock;

// Disk persistence and stale-tmp hygiene are separate filesystem responsibilities;
// the root module owns only the live index and calls those owners through methods.
mod storage;
mod tmp_hygiene;

pub use storage::{default_cache_path, merkle_default_cache_path};

const SCHEMA_VERSION: u32 = 5;

/// Shard count: spreads concurrent `record` / `unchanged` calls across
/// independent locks so tiny-file storms don't serialize all rayon workers.
const MERKLE_SHARDS: usize = 64;

type MerkleShardBuildHasher = ahash::RandomState;
type MerkleShardMap = IndexMap<CacheKey, CacheEntry, MerkleShardBuildHasher>;

#[cfg(target_pointer_width = "64")]
const SHARD_MIX: usize = 0x517c_c1b7_2722_0a95;
#[cfg(target_pointer_width = "32")]
const SHARD_MIX: usize = 0x9e37_79b1;

/// Default upper bound on the number of in-memory cache entries.
///
/// Resident cost per entry is roughly `56 bytes` for the [`CacheEntry`]
/// (`mtime_ns: u64` + `size: u64` + `last_seen_order: u64` + `hash: [u8; 32]`) plus
/// the heap-allocated [`PathBuf`] key (one allocation, length of the
/// canonical path). On a typical repo a path averages ~80-120 bytes, so
/// budget ~150 bytes/entry end-to-end. At the default cap of 8M entries
/// that bounds the index at roughly 1.2 GB resident - large, but bounded,
/// and survivable on the fleet's 32-128 GB boxes. A giant monorepo can
/// raise or lower this via [`MerkleIndex::with_max_entries`] (Tier-A
/// configurability: compiled default, overridable by the caller).
///
/// When the cap is hit we WARN and stop *adding new paths*; updates to
/// paths already in the index are always allowed so an over-cap scan
/// never corrupts an existing entry. An uncached file is simply re-read
/// and re-scanned next run - slower, never unsound. This preserves the
/// module's core guarantee (a file that ever produced a finding is
/// `forget`-ten, never cached) regardless of the cap.
const MERKLE_DEFAULT_MAX_ENTRIES: usize = 8_000_000;

/// Result of loading a persisted [`MerkleIndex`] cache.
#[derive(Debug)]
pub struct MerkleLoadReport {
    index: MerkleIndex,
    status: MerkleLoadStatus,
}

impl MerkleLoadReport {
    /// Status describing whether the cache was loaded or why it cold-started.
    pub fn status(&self) -> &MerkleLoadStatus {
        &self.status
    }

    /// Consume the report and return the live in-memory index.
    pub fn into_index(self) -> MerkleIndex {
        self.index
    }
}

/// Operator-relevant status for a Merkle cache load.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum MerkleLoadStatus {
    /// No cache file exists yet.
    Missing {
        /// Cache path that was probed.
        path: PathBuf,
    },
    /// Cache file loaded and contributed `entries` entries.
    Loaded {
        /// Cache path that loaded successfully.
        path: PathBuf,
        /// Number of entries retained after validation and racy-clean drops.
        entries: usize,
    },
    /// Cache file existed but could not be read.
    ReadFailed {
        /// Cache path that could not be read.
        path: PathBuf,
        /// Filesystem error message.
        error: String,
    },
    /// Cache file existed but was not valid JSON for the current schema.
    ParseFailed {
        /// Cache path that could not be parsed.
        path: PathBuf,
        /// JSON/schema parse error message.
        error: String,
    },
    /// Cache file used an incompatible schema version.
    SchemaMismatch {
        /// Cache path that used the incompatible schema.
        path: PathBuf,
        /// Schema version found in the cache file.
        version: u32,
        /// Schema version required by this binary.
        expected: u32,
    },
    /// Cache file was built for a different detector corpus or scan config.
    SpecChanged {
        /// Cache path whose stored detector-spec hash did not match.
        path: PathBuf,
    },
    /// One persisted entry carried an invalid BLAKE3 hex digest.
    InvalidEntryHash {
        /// Cache path containing the invalid entry.
        path: PathBuf,
        /// Persisted source path key for the invalid entry.
        entry_path: String,
        /// Invalid hash string from the cache file.
        hash: String,
    },
}

fn shard_index(path: &Path) -> usize {
    #[cfg(unix)]
    {
        use std::os::unix::ffi::OsStrExt;
        shard_index_bytes(path.as_os_str().as_bytes())
    }

    #[cfg(not(unix))]
    {
        shard_index_bytes(path.as_os_str().to_string_lossy().as_bytes())
    }
}

#[inline]
fn shard_index_bytes(bytes: &[u8]) -> usize {
    debug_assert!(MERKLE_SHARDS.is_power_of_two());
    let mut hash = SHARD_MIX ^ bytes.len();
    for &byte in bytes {
        hash ^= usize::from(byte);
        hash = hash.rotate_left(5).wrapping_mul(SHARD_MIX);
    }
    hash & (MERKLE_SHARDS - 1)
}

#[inline]
fn shard_capacity(max_entries: usize) -> usize {
    if max_entries == 0 {
        0
    } else {
        (max_entries / MERKLE_SHARDS) + usize::from(max_entries % MERKLE_SHARDS != 0)
    }
}

/// Live-cache key. Sharding stays path-based so all chunks from one file live
/// in the same shard; `forget(path)` can then evict the whole file cheaply.
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
struct CacheKey {
    path: PathBuf,
    chunk_offset: u64,
}

impl CacheKey {
    fn file(path: PathBuf) -> Self {
        Self::chunk(path, 0)
    }

    fn chunk(path: PathBuf, chunk_offset: u64) -> Self {
        Self { path, chunk_offset }
    }
}

struct CacheKeyRef<'a> {
    path: &'a Path,
    chunk_offset: u64,
}

impl Hash for CacheKeyRef<'_> {
    fn hash<H: Hasher>(&self, state: &mut H) {
        self.path.hash(state);
        self.chunk_offset.hash(state);
    }
}

impl Equivalent<CacheKey> for CacheKeyRef<'_> {
    fn equivalent(&self, key: &CacheKey) -> bool {
        self.path == key.path.as_path() && self.chunk_offset == key.chunk_offset
    }
}

/// In-memory per-entry record. Mirrors the durable storage entry but holds the hash as
/// a fixed-size array - saves the per-lookup hex-decode cost on the
/// `unchanged` hot path.
#[derive(Debug, Clone, Copy)]
struct CacheEntry {
    mtime_ns: u64,
    /// Inode change time. `0` means "unknown" (legacy v4 entry or a platform
    /// with no change time); such entries are never trusted by the read-free
    /// fast-path skip.
    ctime_ns: u64,
    size: u64,
    last_seen_order: u64,
    hash: [u8; 32],
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
struct CacheFileFingerprint {
    modified: SystemTime,
    len: u64,
}

/// In-memory file-hash index loaded from / saved to a JSON cache file.
///
/// Concurrency model: the orchestrator holds an `Arc<MerkleIndex>` and
/// records new entries as chunks arrive from rayon-parallel sources.
/// Paths are sharded across [`MERKLE_SHARDS`] mutex-protected maps so
/// concurrent updates rarely contend.
#[derive(Debug)]
pub struct MerkleIndex {
    shards: Vec<RwLock<MerkleShardMap>>,
    /// Upper bound on the number of retained entries across all shards.
    /// Defaults to [`MERKLE_DEFAULT_MAX_ENTRIES`]. Once reached, only
    /// updates to existing paths are accepted; new paths are dropped
    /// (with a one-shot WARN) so a giant monorepo can't silently grow
    /// the index without bound.
    max_entries: usize,
    /// Set once the cap is first hit so we WARN at most once per index
    /// rather than once per dropped entry (which would be a log storm
    /// on a multi-million-file overflow).
    cap_warned: std::sync::atomic::AtomicBool,
    /// Approximate live entry count, maintained on the insert hot path so
    /// the cap check is O(1) instead of summing all 64 shard lengths per
    /// insert (that scan would dominate a multi-million-file scan). It is
    /// incremented only on a NEW-path insert and never decremented (the
    /// `forget` path is for found-secret invalidation, not bulk eviction),
    /// so it is a monotonic upper bound on live entries - exactly the
    /// conservative side for a "stop growing" budget. Exact counts use
    /// [`Self::len`].
    approx_count: std::sync::atomic::AtomicUsize,
    /// Fingerprint of the cache file that populated this index, or the cache
    /// file written by the most recent successful save. Save uses this to skip
    /// the expensive read/parse merge when disk has not changed under us.
    cache_file_fingerprint: RwLock<Option<CacheFileFingerprint>>,
    /// Monotonic access order used for deterministic persisted-cache eviction.
    /// New or updated entries receive a larger value than entries loaded from
    /// disk, so over-cap saves evict stale rows before fresh scan work.
    access_order: AtomicU64,
}

impl MerkleIndex {
    /// Construct a fresh, empty [`MerkleIndex`] with no cached entries and
    /// the default entry cap ([`MERKLE_DEFAULT_MAX_ENTRIES`]).
    pub fn empty() -> Self {
        Self::with_max_entries(MERKLE_DEFAULT_MAX_ENTRIES)
    }

    /// Construct a fresh, empty [`MerkleIndex`].
    pub fn new() -> Self {
        Self::empty()
    }

    /// Construct a fresh, empty [`MerkleIndex`] with an explicit entry cap.
    /// A cap of `0` is treated as "unbounded" for callers that genuinely
    /// want the old behavior, but the documented resident cost still
    /// applies (~150 bytes/entry).
    pub(crate) fn with_max_entries(max_entries: usize) -> Self {
        let shard_capacity = shard_capacity(max_entries);
        Self {
            shards: (0..MERKLE_SHARDS)
                .map(|_| {
                    RwLock::new(IndexMap::with_capacity_and_hasher(
                        shard_capacity,
                        MerkleShardBuildHasher::default(),
                    ))
                })
                .collect(),
            max_entries,
            cap_warned: std::sync::atomic::AtomicBool::new(false),
            approx_count: std::sync::atomic::AtomicUsize::new(0),
            cache_file_fingerprint: RwLock::new(None),
            access_order: AtomicU64::new(0),
        }
    }

    /// The configured maximum number of retained entries (`0` = unbounded).
    pub(crate) fn max_entries(&self) -> usize {
        self.max_entries
    }

    /// Hash the given content with BLAKE3 (32-byte output).
    pub(crate) fn hash_content(content: &[u8]) -> [u8; 32] {
        *blake3::hash(content).as_bytes()
    }

    /// Record one observed chunk at its absolute byte offset and return `true`
    /// when the same `(path, chunk_offset)` previously held the same content.
    pub fn record_chunk_at_offset_and_check_unchanged(
        &self,
        path: PathBuf,
        chunk_offset: u64,
        mtime_ns: u64,
        ctime_ns: u64,
        size: u64,
        content: &[u8],
    ) -> bool {
        self.record_chunk_path_at_offset_and_check_unchanged(
            path.as_path(),
            chunk_offset,
            mtime_ns,
            ctime_ns,
            size,
            content,
        )
    }

    /// Borrowing variant for hot dispatch loops that already hold a path
    /// string/reference. The index still owns the persisted key, but callers do
    /// not need to allocate a temporary `PathBuf` before handing it off.
    pub fn record_chunk_path_at_offset_and_check_unchanged(
        &self,
        path: &Path,
        chunk_offset: u64,
        mtime_ns: u64,
        ctime_ns: u64,
        size: u64,
        content: &[u8],
    ) -> bool {
        let content_hash = Self::hash_content(content);
        self.record_borrowed_key_at_offset_with_metadata(
            path,
            chunk_offset,
            CacheEntry {
                mtime_ns,
                ctime_ns,
                size,
                last_seen_order: self.next_access_order(),
                hash: content_hash,
            },
        )
    }

    /// Returns `true` when `path` was previously indexed with the SAME
    /// content hash. Kept for callers that already have the hash in hand
    /// (e.g. the orchestrator's chunk-level skip path).
    pub(crate) fn unchanged(&self, path: &Path, content_hash: &[u8; 32]) -> bool {
        let _profile = keyhog_profile::span(keyhog_profile::Stage::IncrementalLookup);
        // perf: borrow the path via CacheKeyRef (zero-alloc Equivalent lookup,
        // same as the write path at `record_*`) instead of allocating a
        // CacheKey::file(path.to_path_buf()) on every chunk-level skip check.
        let i = shard_index(path);
        self.shards[i]
            .read()
            .get(&CacheKeyRef {
                path,
                chunk_offset: 0,
            })
            .is_some_and(|prev| &prev.hash == content_hash)
    }

    /// Returns `true` when the stored entry for `path` carries exactly this
    /// `(mtime_ns, ctime_ns, size)` identity. This is the **fast-path skip** -
    /// it avoids reading the file at all, which is the dominant cost on
    /// cold-cache disk. A `false` return means "either we've never seen this
    /// path, or metadata differs - caller must read + hash to decide."
    ///
    /// The change time is part of the identity: a userspace tamper can restore
    /// mtime and size over modified content, but the kernel-owned inode change
    /// time moves on any write and cannot be set back. An entry recorded
    /// without one (`0`: legacy v4 cache row, or a platform whose stat has no
    /// change time) is never trusted here.
    pub fn metadata_unchanged(&self, path: &Path, mtime_ns: u64, ctime_ns: u64, size: u64) -> bool {
        let _profile = keyhog_profile::span(keyhog_profile::Stage::IncrementalLookup);
        // perf: this is the per-file fast-path skip (dominant cold-cache cost);
        // borrow the path via CacheKeyRef so it never allocates a PathBuf.
        let i = shard_index(path);
        self.shards[i]
            .read()
            .get(&CacheKeyRef {
                path,
                chunk_offset: 0,
            })
            .is_some_and(|prev| {
                prev.mtime_ns == mtime_ns
                    && prev.size == size
                    && prev.ctime_ns != 0
                    && prev.ctime_ns == ctime_ns
            })
    }

    /// Returns the stored `(mtime_ns, size, content_hash)` for `path`,
    /// or `None` if the index hasn't seen it. Used by paranoid-mode
    /// verifiers that want to confirm content didn't change even when
    /// metadata happens to match.
    pub(crate) fn lookup(&self, path: &Path) -> Option<(u64, u64, [u8; 32])> {
        let _profile = keyhog_profile::span(keyhog_profile::Stage::IncrementalLookup);
        // perf: zero-alloc borrowed lookup (see `metadata_unchanged`).
        let i = shard_index(path);
        self.shards[i]
            .read()
            .get(&CacheKeyRef {
                path,
                chunk_offset: 0,
            })
            .map(|e| (e.mtime_ns, e.size, e.hash))
    }
    pub(crate) fn seed_for_testing(
        &self,
        path: PathBuf,
        mtime_ns: u64,
        ctime_ns: u64,
        size: u64,
        content_hash: [u8; 32],
    ) {
        self.record_key_with_metadata(
            CacheKey::file(path),
            CacheEntry {
                mtime_ns,
                ctime_ns,
                size,
                last_seen_order: self.next_access_order(),
                hash: content_hash,
            },
        );
    }

    fn next_access_order(&self) -> u64 {
        self.access_order
            .fetch_add(1, Ordering::Relaxed)
            .saturating_add(1)
    }

    fn observe_loaded_access_order(&self, loaded_order: u64) {
        self.access_order.fetch_max(loaded_order, Ordering::Relaxed);
    }

    fn record_key_with_metadata(&self, key: CacheKey, entry: CacheEntry) {
        self.try_insert(key, entry);
    }

    fn record_borrowed_key_at_offset_with_metadata(
        &self,
        path: &Path,
        chunk_offset: u64,
        entry: CacheEntry,
    ) -> bool {
        let _profile = keyhog_profile::span(keyhog_profile::Stage::ResultMerge);
        let i = shard_index(path);
        let mut shard = self.shards[i].write();
        let lookup = CacheKeyRef { path, chunk_offset };
        if let Some(previous) = shard.get_mut(&lookup) {
            let unchanged = previous.hash == entry.hash;
            *previous = entry;
            return unchanged;
        }

        if self.entry_cap_reached() {
            return false;
        }
        shard.insert(CacheKey::chunk(path.to_path_buf(), chunk_offset), entry);
        self.approx_count.fetch_add(1, Ordering::Relaxed);
        false
    }

    /// Insert or update one entry, honoring [`Self::max_entries`].
    ///
    /// Returns `true` if the entry is now present (inserted or updated),
    /// `false` if it was a NEW path dropped because the cap is reached.
    /// Updates to an already-present path always succeed (they don't grow
    /// the working set) so an over-cap scan never corrupts existing state.
    /// The first drop emits a single WARN; subsequent drops are silent to
    /// avoid a log storm on a multi-million-file overflow.
    fn try_insert(&self, key: CacheKey, entry: CacheEntry) -> bool {
        let _profile = keyhog_profile::span(keyhog_profile::Stage::ResultMerge);
        let i = shard_index(&key.path);
        let mut shard = self.shards[i].write();
        let slot = match shard.entry(key) {
            Entry::Occupied(mut slot) => {
                slot.insert(entry);
                return true;
            }
            Entry::Vacant(slot) => slot,
        };

        // `max_entries == 0` means unbounded (opt-in legacy behavior).
        // The cap is a soft budget checked against `approx_count` (O(1),
        // no shard scan). Concurrent new-path inserts across shards can
        // overshoot by at most the number of in-flight `record` calls -
        // bounded and harmless (a few entries over budget, never
        // unbounded growth).
        if self.entry_cap_reached() {
            return false;
        }
        slot.insert(entry);
        self.approx_count.fetch_add(1, Ordering::Relaxed);
        true
    }

    fn entry_cap_reached(&self) -> bool {
        if self.max_entries == 0 || self.approx_count.load(Ordering::Relaxed) < self.max_entries {
            return false;
        }
        if !self.cap_warned.swap(true, Ordering::Relaxed) {
            tracing::warn!(
                cap = self.max_entries,
                "merkle index entry cap reached; new paths will not be \
                 cached this run (they are re-scanned next run). Raise \
                 the cap for very large trees if the rescan cost matters."
            );
        }
        true
    }

    /// Remove `path` from the index so the next scan treats it as new and
    /// re-reads + re-scans it.
    ///
    /// This is how incremental mode keeps its core safety guarantee: a file
    /// that produced ANY finding is never cached, so a secret in an otherwise
    /// unchanged file still surfaces on every later run instead of being
    /// silently skipped (the failure this module's own header warns about).
    /// Clean files - the 99% - stay cached, so the 10-100x speedup is
    /// unaffected, and because we store the ABSENCE of an entry rather than the
    /// finding, no secret value ever touches the on-disk index.
    pub fn forget(&self, path: &Path) {
        let _profile = keyhog_profile::span(keyhog_profile::Stage::ResultMerge);
        let i = shard_index(path);
        self.shards[i].write().retain(|key, _| key.path != path);
    }

    /// Number of indexed entries.
    pub(crate) fn len(&self) -> usize {
        self.shards.iter().map(|s| s.read().len()).sum()
    }

    /// Returns true if no cached entries are present across any shard.
    pub(crate) fn is_empty(&self) -> bool {
        self.shards.iter().all(|s| s.read().is_empty())
    }
}

impl Default for MerkleIndex {
    fn default() -> Self {
        Self::empty()
    }
}