diff --git a/libafl/src/corpus/cached.rs b/libafl/src/corpus/cached.rs new file mode 100644 index 0000000000..ecb3970242 --- /dev/null +++ b/libafl/src/corpus/cached.rs @@ -0,0 +1,136 @@ +//! The cached ondisk corpus stores testcases to disk keeping a part of them in memory. + +use alloc::collections::vec_deque::VecDeque; +use core::cell::RefCell; +use serde::{Deserialize, Serialize}; +use std::path::PathBuf; + +use crate::{ + corpus::{ + ondisk::{OnDiskCorpus, OnDiskMetadataFormat}, + Corpus, Testcase, + }, + inputs::Input, + Error, +}; + +/// A corpus that keep in memory a maximun number of testcases. The eviction policy is FIFO. +#[cfg(feature = "std")] +#[derive(Default, Serialize, Deserialize, Clone, Debug)] +#[serde(bound = "I: serde::de::DeserializeOwned")] +pub struct CachedOnDiskCorpus +where + I: Input, +{ + inner: OnDiskCorpus, + cached_indexes: RefCell>, + cache_max_len: usize, +} + +impl Corpus for CachedOnDiskCorpus +where + I: Input, +{ + /// Returns the number of elements + #[inline] + fn count(&self) -> usize { + self.inner.count() + } + + /// Add an entry to the corpus and return its index + #[inline] + fn add(&mut self, testcase: Testcase) -> Result { + self.inner.add(testcase) + } + + /// Replaces the testcase at the given idx + #[inline] + fn replace(&mut self, idx: usize, testcase: Testcase) -> Result<(), Error> { + // TODO finish + self.inner.replace(idx, testcase) + } + + /// Removes an entry from the corpus, returning it if it was present. + #[inline] + fn remove(&mut self, idx: usize) -> Result>, Error> { + let testcase = self.inner.remove(idx)?; + if testcase.is_some() { + self.cached_indexes.borrow_mut().retain(|e| *e != idx); + } + Ok(testcase) + } + + /// Get by id + #[inline] + fn get(&self, idx: usize) -> Result<&RefCell>, Error> { + let testcase = { self.inner.get(idx)? }; + if testcase.borrow().input().is_none() { + let _ = testcase.borrow_mut().load_input()?; + let current = *self.current(); + while self.cached_indexes.borrow().len() >= self.cache_max_len { + let removed = self.cached_indexes.borrow_mut().pop_front().unwrap(); + if let Some(cur) = current { + if cur == removed { + self.cached_indexes.borrow_mut().push_back(cur); + if self.cache_max_len == 1 { + break; + } + continue; + } + } + *self.inner.get(removed)?.borrow_mut().input_mut() = None; + } + self.cached_indexes.borrow_mut().push_back(idx); + } + Ok(testcase) + } + + /// Current testcase scheduled + #[inline] + fn current(&self) -> &Option { + self.inner.current() + } + + /// Current testcase scheduled (mut) + #[inline] + fn current_mut(&mut self) -> &mut Option { + self.inner.current_mut() + } +} + +impl CachedOnDiskCorpus +where + I: Input, +{ + /// Creates the [`CachedOnDiskCorpus`]. + pub fn new(dir_path: PathBuf, cache_max_len: usize) -> Result { + if cache_max_len == 0 { + return Err(Error::IllegalArgument( + "The max cache len in CachedOnDiskCorpus cannot be 0".into(), + )); + } + Ok(Self { + inner: OnDiskCorpus::new(dir_path)?, + cached_indexes: RefCell::new(VecDeque::new()), + cache_max_len, + }) + } + + /// Creates the [`CachedOnDiskCorpus`] specifying the type of `Metadata` to be saved to disk. + pub fn new_save_meta( + dir_path: PathBuf, + meta_format: Option, + cache_max_len: usize, + ) -> Result { + if cache_max_len == 0 { + return Err(Error::IllegalArgument( + "The max cache len in CachedOnDiskCorpus cannot be 0".into(), + )); + } + Ok(Self { + inner: OnDiskCorpus::new_save_meta(dir_path, meta_format)?, + cached_indexes: RefCell::new(VecDeque::new()), + cache_max_len, + }) + } +} diff --git a/libafl/src/corpus/mod.rs b/libafl/src/corpus/mod.rs index 15e6a8315d..ac638c1de6 100644 --- a/libafl/src/corpus/mod.rs +++ b/libafl/src/corpus/mod.rs @@ -11,6 +11,11 @@ pub mod ondisk; #[cfg(feature = "std")] pub use ondisk::OnDiskCorpus; +#[cfg(feature = "std")] +pub mod cached; +#[cfg(feature = "std")] +pub use cached::CachedOnDiskCorpus; + pub mod queue; pub use queue::QueueCorpusScheduler; diff --git a/libafl/src/corpus/testcase.rs b/libafl/src/corpus/testcase.rs index c75cbc76f9..6906101a6b 100644 --- a/libafl/src/corpus/testcase.rs +++ b/libafl/src/corpus/testcase.rs @@ -63,21 +63,22 @@ where /// Store the input to disk if possible pub fn store_input(&mut self) -> Result { - let fname; match self.filename() { - Some(f) => { - fname = f.clone(); + Some(fname) => { + let saved = match self.input() { + None => false, + Some(i) => { + i.to_file(fname)?; + true + } + }; + if saved { + // remove the input from memory + *self.input_mut() = None; + } + Ok(saved) } - None => { - return Ok(false); - } - }; - match self.input_mut() { None => Ok(false), - Some(i) => { - i.to_file(fname)?; - Ok(true) - } } }