CachedOnDiskCorpus (#229)

* sized inputs

* push CachedOnDiskCorpus and fix OnDiskCorpus

* comments

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<I>
+where
+    I: Input,
+{
+    inner: OnDiskCorpus<I>,
+    cached_indexes: RefCell<VecDeque<usize>>,
+    cache_max_len: usize,
+}
+
+impl<I> Corpus<I> for CachedOnDiskCorpus<I>
+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<I>) -> Result<usize, Error> {
+        self.inner.add(testcase)
+    }
+
+    /// Replaces the testcase at the given idx
+    #[inline]
+    fn replace(&mut self, idx: usize, testcase: Testcase<I>) -> 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<Option<Testcase<I>>, 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<Testcase<I>>, 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<usize> {
+        self.inner.current()
+    }
+
+    /// Current testcase scheduled (mut)
+    #[inline]
+    fn current_mut(&mut self) -> &mut Option<usize> {
+        self.inner.current_mut()
+    }
+}
+
+impl<I> CachedOnDiskCorpus<I>
+where
+    I: Input,
+{
+    /// Creates the [`CachedOnDiskCorpus`].
+    pub fn new(dir_path: PathBuf, cache_max_len: usize) -> Result<Self, Error> {
+        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<OnDiskMetadataFormat>,
+        cache_max_len: usize,
+    ) -> Result<Self, Error> {
+        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,
+        })
+    }
+}

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;
 

libafl/src/corpus/testcase.rs

@@ -63,21 +63,22 @@ where
 
     /// Store the input to disk if possible
     pub fn store_input(&mut self) -> Result<bool, Error> {
-        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)
-            }
         }
     }