dragonfly_client_rs/
scan_cache.rs

1use std::{
2    collections::HashMap,
3    fs::{self, File},
4    io::Read,
5    path::{Path, PathBuf},
6};
7
8#[cfg(test)]
9use color_eyre::eyre::ensure;
10use color_eyre::Result;
11use sha2::{Digest, Sha256};
12use tempfile::TempDir;
13use xxhash_rust::xxh3::Xxh3;
14use yara::Rules;
15
16use crate::{exts::RuleExt, scanner::RuleScore};
17
18#[derive(Clone, Debug, PartialEq, serde::Serialize, serde::Deserialize)]
19struct ContentMatch {
20    score: RuleScore,
21    filetypes: Vec<String>,
22}
23
24pub(crate) struct FileHashes {
25    pub fast: u128,
26    pub sha256: String,
27}
28
29struct CachedFile {
30    path: PathBuf,
31    matches: Vec<ContentMatch>,
32}
33
34/// Reuse successful content scans within one package and one compiled ruleset.
35/// Representative bytes live on bounded temporary storage, not in memory.
36pub struct ScanCache<'a> {
37    rules: &'a Rules,
38    reuse: Option<&'a crate::reuse_cache::ReuseCache>,
39    pub stats: crate::reuse_cache::CacheStats,
40    directory: TempDir,
41    files: HashMap<(u128, u64), CachedFile>,
42    max_entries: usize,
43    remaining_bytes: u64,
44    pub scanned_files: usize,
45    pub reused_files: usize,
46}
47
48impl<'a> ScanCache<'a> {
49    pub fn new(rules: &'a Rules, max_entries: usize, max_bytes: u64) -> Result<Self> {
50        Ok(Self {
51            rules,
52            reuse: None,
53            stats: crate::reuse_cache::CacheStats::new("yara", crate::reuse_cache::CacheMode::Off),
54            directory: tempfile::tempdir()?,
55            files: HashMap::new(),
56            max_entries,
57            remaining_bytes: max_bytes,
58            scanned_files: 0,
59            reused_files: 0,
60        })
61    }
62
63    pub fn set_reuse(&mut self, reuse: Option<&'a crate::reuse_cache::ReuseCache>) {
64        self.reuse = reuse;
65        self.stats.mode = reuse.map_or(crate::reuse_cache::CacheMode::Off, |cache| cache.mode);
66    }
67
68    #[cfg(test)]
69    pub fn scan(&mut self, path: &Path, max_scan_size: u64) -> Result<Vec<RuleScore>> {
70        let size = path.metadata()?.len();
71        ensure!(
72            size <= max_scan_size,
73            "file {} is {size} bytes, exceeding the {max_scan_size}-byte scan limit",
74            path.display()
75        );
76        self.scan_hashed(path, size, &hash_file(path)?)
77    }
78
79    pub(crate) fn prefetch(&mut self, files: &[(PathBuf, u64, FileHashes)]) {
80        if let Some(reuse) = self.reuse.filter(|cache| cache.uses_database()) {
81            let keys = files
82                .iter()
83                .filter(|(_, size, hash)| !self.files.contains_key(&(hash.fast, *size)))
84                .map(|(_, _, hash)| {
85                    (
86                        format!("sha256:{}", hash.sha256),
87                        crate::durable_cache::Key {
88                            file_digest: hash.sha256.clone(),
89                            language: String::new(),
90                        },
91                    )
92                })
93                .collect::<Vec<_>>();
94            reuse.prefetch(&keys, &mut self.stats);
95        }
96    }
97
98    pub(crate) fn scan_hashed(
99        &mut self,
100        path: &Path,
101        size: u64,
102        hashes: &FileHashes,
103    ) -> Result<Vec<RuleScore>> {
104        let identity = (hashes.fast, size);
105        if let Some(cached) = self.files.get(&identity) {
106            // A hash collision must never suppress a scan of different bytes.
107            if files_equal(path, &cached.path, size)? {
108                self.reused_files += 1;
109                return Ok(matches_for_path(&cached.matches, path));
110            }
111        }
112
113        let key = if self
114            .reuse
115            .is_some_and(crate::reuse_cache::ReuseCache::uses_database)
116        {
117            format!("sha256:{}", hashes.sha256)
118        } else {
119            format!("{:032x}:{}", identity.0, size)
120        };
121        let mut cached = self
122            .reuse
123            .and_then(|cache| cache.lookup::<Vec<ContentMatch>>(&key, path, &mut self.stats));
124        if let Some(matches) = &mut cached {
125            canonicalize(matches);
126        }
127        let reuse_enabled = cached.is_some()
128            && self
129                .reuse
130                .is_some_and(crate::reuse_cache::ReuseCache::should_reuse);
131        let matches = if let Some(matches) = cached.as_ref().filter(|_| reuse_enabled) {
132            self.stats.reused_files += 1;
133            self.stats.reused_bytes += size;
134            matches.clone()
135        } else {
136            let started = std::time::Instant::now();
137            self.stats.engine_files += 1;
138            self.stats.engine_bytes += size;
139            let scanned = self.rules.scan_file(path, 10);
140            self.stats.engine_us += started.elapsed().as_micros();
141            let mut matches: Vec<_> = scanned?
142                .into_iter()
143                .map(|rule| ContentMatch {
144                    filetypes: rule
145                        .get_filetypes()
146                        .into_iter()
147                        .map(str::to_owned)
148                        .collect(),
149                    score: RuleScore::from(rule),
150                })
151                .collect();
152            canonicalize(&mut matches);
153            self.scanned_files += 1;
154            if let Some(previous) = cached {
155                self.stats.validated_files += 1;
156                if previous != matches {
157                    self.stats.mismatched_files += 1;
158                    if let Some(cache) = self.reuse {
159                        cache.quarantine(&key, &mut self.stats);
160                    }
161                    report_mismatch(path, &hashes.sha256, &previous, &matches);
162                    if self.stats.reused_files > 0 {
163                        color_eyre::eyre::bail!("Cache validation failed; cached results for this job must be discarded");
164                    }
165                }
166            }
167            if let Some(cache) = self.reuse {
168                cache.insert(key, path, &matches, &mut self.stats);
169            }
170            matches
171        };
172        let result = matches_for_path(&matches, path);
173        if self.files.len() < self.max_entries
174            && size <= self.remaining_bytes
175            && !self.files.contains_key(&identity)
176        {
177            let destination = self.directory.path().join(self.files.len().to_string());
178            match fs::copy(path, &destination) {
179                Ok(_) => {
180                    self.remaining_bytes -= size;
181                    self.files.insert(
182                        identity,
183                        CachedFile {
184                            path: destination,
185                            matches,
186                        },
187                    );
188                }
189                Err(error) => {
190                    // The package result is already valid. Stop new writes; any
191                    // partial representative is removed with the temporary directory.
192                    self.max_entries = 0;
193                    tracing::warn!(
194                        event = "content_scan_cache_write_failed",
195                        %error,
196                        "Disabling new cache entries; continuing with successful YARA results"
197                    );
198                }
199            }
200        }
201        Ok(result)
202    }
203}
204
205fn report_mismatch(path: &Path, digest: &str, previous: &[ContentMatch], matches: &[ContentMatch]) {
206    tracing::error!(
207        event = "scan_reuse_mismatch",
208        file_sha256 = %digest,
209        file_path = %path.display(),
210        cached_count = previous.len(),
211        fresh_count = matches.len(),
212        cached_matches = ?previous.iter().take(16).collect::<Vec<_>>(),
213        fresh_matches = ?matches.iter().take(16).collect::<Vec<_>>(),
214        "Cached YARA findings differ from fresh scan; quarantining this file"
215    );
216}
217
218fn canonicalize(matches: &mut [ContentMatch]) {
219    for matched in matches.iter_mut() {
220        matched.filetypes.sort_unstable();
221        matched.filetypes.dedup();
222    }
223    matches.sort_unstable_by(|a, b| {
224        (&a.score.name, a.score.score, &a.filetypes).cmp(&(
225            &b.score.name,
226            b.score.score,
227            &b.filetypes,
228        ))
229    });
230}
231
232fn matches_for_path(matches: &[ContentMatch], path: &Path) -> Vec<RuleScore> {
233    matches
234        .iter()
235        .filter(|matched| {
236            matched.filetypes.is_empty()
237                || matched
238                    .filetypes
239                    .iter()
240                    .any(|suffix| path.to_string_lossy().ends_with(suffix))
241        })
242        .map(|matched| matched.score.clone())
243        .collect()
244}
245
246pub(crate) fn hash_file(path: &Path) -> Result<FileHashes> {
247    let mut file = File::open(path)?;
248    let mut hasher = Xxh3::new();
249    let mut sha256 = Sha256::new();
250    let mut buffer = [0_u8; 8192];
251    loop {
252        let read = file.read(&mut buffer)?;
253        if read == 0 {
254            return Ok(FileHashes {
255                fast: hasher.digest128(),
256                sha256: format!("{:x}", sha256.finalize()),
257            });
258        }
259        hasher.update(&buffer[..read]);
260        sha256.update(&buffer[..read]);
261    }
262}
263
264fn files_equal(left: &Path, right: &Path, mut remaining: u64) -> Result<bool> {
265    let mut left = File::open(left)?;
266    let mut right = File::open(right)?;
267    let mut left_buffer = [0_u8; 8192];
268    let mut right_buffer = [0_u8; 8192];
269    while remaining > 0 {
270        let length = usize::try_from(remaining.min(8192))?;
271        left.read_exact(&mut left_buffer[..length])?;
272        right.read_exact(&mut right_buffer[..length])?;
273        if left_buffer[..length] != right_buffer[..length] {
274            return Ok(false);
275        }
276        remaining -= u64::try_from(length)?;
277    }
278    Ok(true)
279}
280
281#[cfg(test)]
282mod tests {
283    use super::{hash_file, ScanCache};
284    use std::fs;
285    use tempfile::tempdir;
286    use yara::{Compiler, Rules};
287
288    fn rules() -> Rules {
289        Compiler::new()
290            .unwrap()
291            .add_rules_str(
292                r#"rule python {
293                    meta: filetype = ".py .pyi" weight = 5
294                    strings: $a = "danger"
295                    condition: $a
296                }
297                rule compound_suffix {
298                    meta: filetype = "special.txt" weight = 3
299                    strings: $a = "danger"
300                    condition: $a
301                }"#,
302            )
303            .unwrap()
304            .compile_rules()
305            .unwrap()
306    }
307
308    #[test]
309    fn cache_validation_ignores_order_but_preserves_findings_and_filters() {
310        use super::{canonicalize, ContentMatch};
311        use crate::{
312            reuse_cache::{CacheMode, CacheStats, ReuseCache},
313            scanner::RuleScore,
314        };
315        let rules = rules();
316        let dir = tempdir().unwrap();
317        let path = dir.path().join("danger.py");
318        fs::write(&path, b"danger").unwrap();
319        let old = vec![
320            ContentMatch {
321                score: RuleScore {
322                    name: "python".into(),
323                    score: 5,
324                },
325                filetypes: vec![".pyi".into(), ".py".into()],
326            },
327            ContentMatch {
328                score: RuleScore {
329                    name: "compound_suffix".into(),
330                    score: 3,
331                },
332                filetypes: vec!["special.txt".into()],
333            },
334        ];
335        let shared = ReuseCache::new(CacheMode::Observe, 10, 4096);
336        let key = format!("{:032x}:6", hash_file(&path).unwrap().fast);
337        shared.insert(
338            key,
339            &path,
340            &old,
341            &mut CacheStats::new("yara", CacheMode::Observe),
342        );
343        let mut scan = ScanCache::new(&rules, 10, 4096).unwrap();
344        scan.set_reuse(Some(&shared));
345        assert_eq!(scan.scan(&path, 1024).unwrap().len(), 1);
346        assert_eq!(scan.stats.validated_files, 1);
347        assert_eq!(scan.stats.mismatched_files, 0);
348        let mut canonical = old;
349        canonicalize(&mut canonical);
350        let mut changed = canonical.clone();
351        changed[0].score.score += 1;
352        canonicalize(&mut changed);
353        assert_ne!(canonical, changed);
354        changed = canonical.clone();
355        changed[0].filetypes.push(".sh".into());
356        canonicalize(&mut changed);
357        assert_ne!(canonical, changed);
358    }
359
360    #[test]
361    fn a_sample_mismatch_preserves_fresh_results_and_quarantines_only_that_key() {
362        use crate::reuse_cache::{CacheMode, CacheStats, ReuseCache};
363        let rules = rules();
364        let dir = tempdir().unwrap();
365        let path = dir.path().join("danger.py");
366        fs::write(&path, b"danger").unwrap();
367        let shared = ReuseCache::new(CacheMode::Reuse, 10, 1024);
368        let mut setup = CacheStats::new("yara", CacheMode::Reuse);
369        shared.insert(
370            format!("{:032x}:6", hash_file(&path).unwrap().fast),
371            &path,
372            &Vec::<super::ContentMatch>::new(),
373            &mut setup,
374        );
375        for _ in 0..99 {
376            assert!(shared.should_reuse());
377        }
378        let mut cache = ScanCache::new(&rules, 10, 1024).unwrap();
379        cache.set_reuse(Some(&shared));
380        assert_eq!(cache.scan(&path, 1024).unwrap().len(), 1);
381        assert_eq!(cache.stats.mismatched_files, 1);
382        assert!(!shared.is_disabled());
383        let key = format!("{:032x}:6", hash_file(&path).unwrap().fast);
384        assert!(shared
385            .lookup::<Vec<super::ContentMatch>>(&key, &path, &mut setup)
386            .is_none());
387        shared.insert(
388            key.clone(),
389            &path,
390            &Vec::<super::ContentMatch>::new(),
391            &mut setup,
392        );
393        assert!(shared
394            .lookup::<Vec<super::ContentMatch>>(&key, &path, &mut setup)
395            .is_none());
396        shared.insert("unrelated".into(), &path, &vec![7_u8], &mut setup);
397        assert_eq!(
398            shared.lookup::<Vec<u8>>("unrelated", &path, &mut setup),
399            Some(vec![7])
400        );
401    }
402
403    #[test]
404    fn cross_job_reuse_observation_and_rules_reset_preserve_results() {
405        use crate::reuse_cache::{CacheMode, ReuseCache};
406        let rules = rules();
407        let dir = tempdir().unwrap();
408        let path = dir.path().join("first.py");
409        fs::write(&path, b"danger").unwrap();
410        for mode in [CacheMode::Off, CacheMode::Observe, CacheMode::Reuse] {
411            let mut shared = ReuseCache::new(mode, 10, 1024);
412            for iteration in 0..2 {
413                let mut job = ScanCache::new(&rules, 10, 1024).unwrap();
414                job.set_reuse(Some(&shared));
415                assert_eq!(job.scan(&path, 1024).unwrap()[0].name, "python");
416                assert_eq!(
417                    job.stats.reused_files,
418                    u64::from(iteration == 1 && mode == CacheMode::Reuse)
419                );
420                assert_eq!(
421                    job.stats.validated_files,
422                    u64::from(iteration == 1 && mode == CacheMode::Observe)
423                );
424                assert_eq!(job.stats.mismatched_files, 0);
425            }
426            shared.clear();
427            let changed = Compiler::new()
428                .unwrap()
429                .add_rules_str("rule changed { condition: true }")
430                .unwrap()
431                .compile_rules()
432                .unwrap();
433            let mut job = ScanCache::new(&changed, 10, 1024).unwrap();
434            job.set_reuse(Some(&shared));
435            assert_eq!(job.scan(&path, 1024).unwrap()[0].name, "changed");
436            assert_eq!(job.stats.reused_files, 0);
437        }
438    }
439
440    #[test]
441    fn reuse_filters_each_original_path_including_compound_suffixes() {
442        let rules = rules();
443        let directory = tempdir().unwrap();
444        let mut cache = ScanCache::new(&rules, 10, 1024).unwrap();
445        for (name, expected) in [
446            ("first.txt", None),
447            ("second.py", Some("python")),
448            ("third.pyi", Some("python")),
449            ("special.txt", Some("compound_suffix")),
450        ] {
451            let path = directory.path().join(name);
452            fs::write(&path, b"danger").unwrap();
453            let matches = cache.scan(&path, 1024).unwrap();
454            assert_eq!(
455                matches.first().map(|matched| matched.name.as_str()),
456                expected
457            );
458        }
459        assert_eq!((cache.scanned_files, cache.reused_files), (1, 3));
460    }
461
462    #[test]
463    fn cache_limits_fall_back_to_scanning_and_keep_existing_hits() {
464        let rules = rules();
465        let directory = tempdir().unwrap();
466        for (entries, bytes) in [(1, 1024), (10, 6)] {
467            let mut cache = ScanCache::new(&rules, entries, bytes).unwrap();
468            let original = directory.path().join("original.py");
469            let other = directory.path().join("other.py");
470            fs::write(&original, b"danger").unwrap();
471            fs::write(&other, b"danger too").unwrap();
472            assert_eq!(cache.scan(&original, 1024).unwrap().len(), 1);
473            assert_eq!(cache.scan(&other, 1024).unwrap().len(), 1);
474            assert_eq!(cache.scan(&other, 1024).unwrap().len(), 1);
475            assert_eq!(cache.scan(&original, 1024).unwrap().len(), 1);
476            assert_eq!((cache.scanned_files, cache.reused_files), (3, 1));
477            assert_eq!(cache.files.len(), 1);
478        }
479    }
480
481    #[test]
482    fn cache_write_failure_preserves_successful_results_and_stops_new_writes() {
483        let rules = rules();
484        let directory = tempdir().unwrap();
485        let path = directory.path().join("module.py");
486        fs::write(&path, b"danger").unwrap();
487        let mut cache = ScanCache::new(&rules, 10, 1024).unwrap();
488        fs::remove_dir(cache.directory.path()).unwrap();
489        assert_eq!(cache.scan(&path, 1024).unwrap().len(), 1);
490        assert_eq!(cache.scan(&path, 1024).unwrap().len(), 1);
491        assert_eq!((cache.scanned_files, cache.reused_files), (2, 0));
492        assert_eq!(cache.max_entries, 0);
493        assert!(cache.files.is_empty());
494    }
495
496    #[test]
497    fn hash_collision_does_not_reuse_a_clean_result() {
498        let rules = rules();
499        let directory = tempdir().unwrap();
500        let clean = directory.path().join("clean.py");
501        let malicious = directory.path().join("malicious.py");
502        fs::write(&clean, b"benign").unwrap();
503        fs::write(&malicious, b"danger").unwrap();
504        let mut cache = ScanCache::new(&rules, 10, 1024).unwrap();
505        assert!(cache.scan(&clean, 1024).unwrap().is_empty());
506        let cached = cache
507            .files
508            .remove(&(hash_file(&clean).unwrap().fast, 6))
509            .unwrap();
510        cache
511            .files
512            .insert((hash_file(&malicious).unwrap().fast, 6), cached);
513        assert_eq!(cache.scan(&malicious, 1024).unwrap().len(), 1);
514        assert_eq!((cache.scanned_files, cache.reused_files), (2, 0));
515    }
516
517    #[test]
518    fn clean_empty_files_reuse_but_size_checks_still_apply() {
519        let rules = rules();
520        let directory = tempdir().unwrap();
521        let path = directory.path().join("empty.py");
522        fs::write(&path, b"").unwrap();
523        let mut cache = ScanCache::new(&rules, 10, 1024).unwrap();
524        assert!(cache.scan(&path, 0).unwrap().is_empty());
525        assert!(cache.scan(&path, 0).unwrap().is_empty());
526        assert_eq!((cache.scanned_files, cache.reused_files), (1, 1));
527        fs::write(&path, b"danger").unwrap();
528        cache.scan(&path, 6).unwrap();
529        assert!(cache.scan(&path, 5).is_err());
530        assert_eq!((cache.scanned_files, cache.reused_files), (2, 1));
531    }
532}
dragonfly_client_rs/scan_cache.rs

dragonfly_client_rs/
scan_cache.rs
