Antivirus/crates/hound-defs/src/index.rs

//! The indicator index, and the cuckoo filter in front of it.
//!
//! Checking a project means asking "is this package known bad?" once per
//! dependency, and a real `node_modules` has thousands. Almost every
//! answer is no, so the structure should be optimised for saying no
//! quickly rather than for saying yes well.
//!
//! A cuckoo filter does that: a few hundred kilobytes answers "definitely
//! not in the set" for the overwhelming majority of lookups without ever
//! touching the full index. Only a positive — real or occasional false —
//! costs a map probe, and the map confirms it.
//!
//! Why cuckoo rather than bloom: a cuckoo filter supports **deletion**,
//! which a definitions feed genuinely needs. OSV withdraws records — it
//! once withdrew 157 malware reports after a false-positive incident —
//! and a filter you cannot remove from means a withdrawn record keeps
//! costing a map probe forever, or the whole structure has to be rebuilt.
//!
//! The one property that must never break is **no false negatives**. A
//! false positive costs a hash-map lookup. A false negative is a piece of
//! malware the scanner said was clean.

use crate::osv::{key_for, Indicator};
use std::collections::HashMap;

/// Slots per bucket. Four is the standard choice: enough that the table
/// fills to ~95% before insertion starts failing, small enough that a
/// lookup stays inside one cache line.
const SLOTS: usize = 4;

/// How many times an insert may evict a resident before giving up.
const MAX_KICKS: usize = 500;

/// A cuckoo filter over 16-bit fingerprints.
#[derive(Debug, Clone)]
pub struct CuckooFilter {
    buckets: Vec<[u16; SLOTS]>,
    mask: usize,
    len: usize,
}

impl CuckooFilter {
    /// Size for an expected item count. Rounded up to a power of two,
    /// with headroom so insertion does not start failing near the end.
    pub fn with_capacity(expected: usize) -> Self {
        let needed = (expected.max(1) * 2 / SLOTS).max(16);
        let buckets = needed.next_power_of_two();
        Self {
            buckets: vec![[0u16; SLOTS]; buckets],
            mask: buckets - 1,
            len: 0,
        }
    }

    pub fn len(&self) -> usize {
        self.len
    }

    pub fn is_empty(&self) -> bool {
        self.len == 0
    }

    /// Bytes of memory the filter occupies.
    pub fn size_bytes(&self) -> usize {
        self.buckets.len() * SLOTS * std::mem::size_of::<u16>()
    }

    /// A non-zero 16-bit fingerprint, plus the primary bucket.
    ///
    /// Zero is reserved to mean "empty slot", so a fingerprint that hashes
    /// to zero is nudged to one. Skipping that produces a filter where one
    /// key in 65,536 silently fails to store — a false negative.
    fn fingerprint_and_bucket(&self, key: &str) -> (u16, usize) {
        let h = fnv1a(key.as_bytes());
        let mut fp = (h >> 32) as u16;
        if fp == 0 {
            fp = 1;
        }
        ((fp), (h as usize) & self.mask)
    }

    /// The partner bucket, derived from the fingerprint alone.
    ///
    /// This is what makes the filter work without storing keys: from
    /// either bucket and the fingerprint, the other bucket is computable.
    fn alt_bucket(&self, bucket: usize, fp: u16) -> usize {
        (bucket ^ (fnv1a(&fp.to_le_bytes()) as usize)) & self.mask
    }

    pub fn insert(&mut self, key: &str) -> bool {
        let (fp, b1) = self.fingerprint_and_bucket(key);
        let b2 = self.alt_bucket(b1, fp);

        for b in [b1, b2] {
            if let Some(slot) = self.buckets[b].iter().position(|&s| s == 0) {
                self.buckets[b][slot] = fp;
                self.len += 1;
                return true;
            }
        }

        // Both full: evict a resident and rehome it. The victim slot is
        // chosen deterministically from the fingerprint so the structure
        // is reproducible — a definitions pack must build identically on
        // every machine or its hash is not a version.
        let mut bucket = b2;
        let mut carried = fp;
        for kick in 0..MAX_KICKS {
            let slot = (fnv1a(&[carried.to_le_bytes(), (kick as u16).to_le_bytes()].concat())
                as usize)
                % SLOTS;
            std::mem::swap(&mut carried, &mut self.buckets[bucket][slot]);
            bucket = self.alt_bucket(bucket, carried);
            if let Some(free) = self.buckets[bucket].iter().position(|&s| s == 0) {
                self.buckets[bucket][free] = carried;
                self.len += 1;
                return true;
            }
        }
        false
    }

    /// True when the key *may* be present. False means definitely absent.
    pub fn contains(&self, key: &str) -> bool {
        let (fp, b1) = self.fingerprint_and_bucket(key);
        if self.buckets[b1].contains(&fp) {
            return true;
        }
        let b2 = self.alt_bucket(b1, fp);
        self.buckets[b2].contains(&fp)
    }

    /// Remove a key. Only ever call this for a key known to be present —
    /// removing a fingerprint that belongs to a different key would create
    /// a false negative for that other key.
    pub fn remove(&mut self, key: &str) -> bool {
        let (fp, b1) = self.fingerprint_and_bucket(key);
        let b2 = self.alt_bucket(b1, fp);
        for b in [b1, b2] {
            if let Some(slot) = self.buckets[b].iter().position(|&s| s == fp) {
                self.buckets[b][slot] = 0;
                self.len -= 1;
                return true;
            }
        }
        false
    }
}

/// FNV-1a, 64-bit. Not cryptographic and does not need to be: the filter
/// is a performance structure, and every positive is confirmed against the
/// real index before anything is reported.
fn fnv1a(data: &[u8]) -> u64 {
    let mut h: u64 = 0xcbf2_9ce4_8422_2325;
    for b in data {
        h ^= *b as u64;
        h = h.wrapping_mul(0x100_0000_01b3);
    }
    h
}

/// Indicators, with the filter in front.
#[derive(Debug, Default)]
pub struct Index {
    filter: Option<CuckooFilter>,
    entries: HashMap<String, Vec<Indicator>>,
}

impl Index {
    pub fn new() -> Self {
        Self::default()
    }

    /// Build from a set of indicators.
    pub fn build(indicators: Vec<Indicator>) -> Self {
        let mut entries: HashMap<String, Vec<Indicator>> = HashMap::new();
        for ind in indicators {
            entries.entry(ind.key()).or_default().push(ind);
        }
        let mut filter = CuckooFilter::with_capacity(entries.len());
        for key in entries.keys() {
            // A filter that failed to store a key would produce a false
            // negative, so a failed insert abandons the filter rather than
            // shipping one that lies. The map still answers correctly.
            if !filter.insert(key) {
                return Self { filter: None, entries };
            }
        }
        Self { filter: Some(filter), entries }
    }

    pub fn len(&self) -> usize {
        self.entries.len()
    }

    pub fn is_empty(&self) -> bool {
        self.entries.is_empty()
    }

    pub fn filter_bytes(&self) -> usize {
        self.filter.as_ref().map(|f| f.size_bytes()).unwrap_or(0)
    }

    /// Look up a package. `None` for clean.
    ///
    /// The filter short-circuits the common case without touching the map.
    pub fn lookup(&self, ecosystem: &str, name: &str, version: &str) -> Option<&Indicator> {
        let key = key_for(ecosystem, name);
        if let Some(f) = &self.filter {
            if !f.contains(&key) {
                return None;
            }
        }
        self.entries
            .get(&key)?
            .iter()
            .find(|i| i.versions.covers(version))
    }

    /// Whether any version of a package is known bad, regardless of the
    /// version in hand. Used when a lockfile does not pin one.
    pub fn any_version(&self, ecosystem: &str, name: &str) -> Option<&Indicator> {
        let key = key_for(ecosystem, name);
        if let Some(f) = &self.filter {
            if !f.contains(&key) {
                return None;
            }
        }
        self.entries.get(&key)?.first()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::osv::Versions;

    fn ind(eco: &str, name: &str, versions: Versions) -> Indicator {
        Indicator {
            ecosystem: eco.into(),
            name: name.into(),
            versions,
            id: "MAL-2024-1339".into(),
            summary: "Malicious code".into(),
        }
    }

    // ── the property that must never break ──

    #[test]
    fn the_filter_never_produces_a_false_negative() {
        // A false positive costs a hash lookup. A false negative is
        // malware reported as clean. Every inserted key must be found.
        let mut f = CuckooFilter::with_capacity(20_000);
        let keys: Vec<String> = (0..20_000).map(|i| format!("npm:package-{i}")).collect();
        for k in &keys {
            assert!(f.insert(k), "insert failed at {k}");
        }
        for k in &keys {
            assert!(f.contains(k), "false negative for {k}");
        }
        assert_eq!(f.len(), 20_000);
    }

    #[test]
    fn a_fingerprint_of_zero_is_still_stored() {
        // Zero marks an empty slot, so a key whose fingerprint hashes to
        // zero would silently vanish. Exercised across a wide key space
        // because it is rare by construction.
        let mut f = CuckooFilter::with_capacity(70_000);
        let keys: Vec<String> = (0..65_600).map(|i| format!("k{i}")).collect();
        for k in &keys {
            f.insert(k);
        }
        let missing: Vec<&String> = keys.iter().filter(|k| !f.contains(k)).collect();
        assert!(missing.is_empty(), "{} keys vanished", missing.len());
    }

    #[test]
    fn the_false_positive_rate_is_low_enough_to_be_worth_it() {
        let mut f = CuckooFilter::with_capacity(10_000);
        for i in 0..10_000 {
            f.insert(&format!("npm:real-{i}"));
        }
        let probes = 100_000;
        let fp = (0..probes)
            .filter(|i| f.contains(&format!("npm:absent-{i}")))
            .count();
        let rate = fp as f64 / probes as f64;
        assert!(rate < 0.01, "false-positive rate {rate:.4} is too high to save work");
    }

    #[test]
    fn removal_works_so_withdrawn_records_can_be_dropped() {
        // OSV withdraws records. Without deletion, a withdrawn indicator
        // costs a map probe forever or forces a full rebuild.
        let mut f = CuckooFilter::with_capacity(100);
        f.insert("npm:withdrawn");
        assert!(f.contains("npm:withdrawn"));
        assert!(f.remove("npm:withdrawn"));
        assert!(!f.contains("npm:withdrawn"));
        assert_eq!(f.len(), 0);
    }

    #[test]
    fn building_the_filter_is_deterministic() {
        // A pack's hash is its version. If the same input produced a
        // different filter on two machines, the hash would be meaningless.
        let keys: Vec<String> = (0..5_000).map(|i| format!("pypi:pkg{i}")).collect();
        let build = || {
            let mut f = CuckooFilter::with_capacity(5_000);
            for k in &keys {
                f.insert(k);
            }
            f.buckets.clone()
        };
        assert_eq!(build(), build());
    }

    #[test]
    fn the_filter_is_small() {
        let f = CuckooFilter::with_capacity(226_000);
        let mb = f.size_bytes() as f64 / (1024.0 * 1024.0);
        assert!(mb < 4.0, "226k indicators should fit in a few MB, got {mb:.1} MB");
    }

    // ── the index ──

    #[test]
    fn finds_a_known_malicious_package() {
        let idx = Index::build(vec![ind("npm", "test-poc2", Versions::All)]);
        let hit = idx.lookup("npm", "test-poc2", "1.0.0").expect("should match");
        assert_eq!(hit.id, "MAL-2024-1339");
    }

    #[test]
    fn a_clean_package_is_not_reported() {
        let idx = Index::build(vec![ind("npm", "test-poc2", Versions::All)]);
        assert!(idx.lookup("npm", "react", "18.2.0").is_none());
    }

    #[test]
    fn the_ecosystem_has_to_match() {
        // A malicious npm package named "requests" says nothing about the
        // PyPI package of the same name.
        let idx = Index::build(vec![ind("npm", "requests", Versions::All)]);
        assert!(idx.lookup("npm", "requests", "1.0.0").is_some());
        assert!(idx.lookup("pypi", "requests", "2.31.0").is_none());
    }

    #[test]
    fn lookups_are_case_insensitive() {
        let idx = Index::build(vec![ind("npm", "EvilPkg", Versions::All)]);
        assert!(idx.lookup("NPM", "evilpkg", "1.0.0").is_some());
    }

    #[test]
    fn an_exact_version_record_only_matches_that_version() {
        let idx = Index::build(vec![ind(
            "pypi",
            "thing",
            Versions::Exact(vec!["1.0.0".into()]),
        )]);
        assert!(idx.lookup("pypi", "thing", "1.0.0").is_some());
        assert!(
            idx.lookup("pypi", "thing", "2.0.0").is_none(),
            "condemning a version the feed did not is how a scanner loses trust"
        );
    }

    #[test]
    fn any_version_ignores_the_version_when_a_lockfile_does_not_pin_one() {
        let idx = Index::build(vec![ind(
            "pypi",
            "thing",
            Versions::Exact(vec!["1.0.0".into()]),
        )]);
        assert!(idx.any_version("pypi", "thing").is_some());
        assert!(idx.any_version("pypi", "other").is_none());
    }

    #[test]
    fn several_records_for_one_package_all_stay_reachable() {
        let idx = Index::build(vec![
            ind("npm", "x", Versions::Exact(vec!["1.0.0".into()])),
            ind("npm", "x", Versions::Exact(vec!["2.0.0".into()])),
        ]);
        assert!(idx.lookup("npm", "x", "1.0.0").is_some());
        assert!(idx.lookup("npm", "x", "2.0.0").is_some());
        assert!(idx.lookup("npm", "x", "3.0.0").is_none());
    }

    #[test]
    fn an_empty_index_answers_cleanly() {
        let idx = Index::new();
        assert!(idx.is_empty());
        assert!(idx.lookup("npm", "anything", "1.0.0").is_none());
    }

    #[test]
    fn a_large_index_still_answers_correctly_through_the_filter() {
        let mut inds: Vec<Indicator> = (0..50_000)
            .map(|i| ind("npm", &format!("bad-{i}"), Versions::All))
            .collect();
        inds.push(ind("pypi", "needle", Versions::All));
        let idx = Index::build(inds);

        assert!(idx.lookup("pypi", "needle", "1.0.0").is_some());
        assert!(idx.lookup("npm", "bad-49999", "1.0.0").is_some());
        assert!(idx.lookup("npm", "definitely-clean", "1.0.0").is_none());
        assert!(idx.filter_bytes() > 0, "the filter should have been built");
    }
}