Add peek function to Corpus (#2238)

* add peek function to Corpus

* send 0 as next peek in NopCorpus

* rename peek to peek_next_free_id

* fix clippy

* add peek_next_free_id to libafl_libfuzzer

* impl peek_next_free_id for ArtifactCorpus

libafl/src/corpus/cached.rs

@@ -147,6 +147,12 @@ where
         self.inner.next(idx)
     }
 
+    /// Peek the next free corpus id
+    #[inline]
+    fn peek_next_free_id(&self) -> CorpusId {
+        self.inner.peek_next_free_id()
+    }
+
     #[inline]
     fn prev(&self, idx: CorpusId) -> Option<CorpusId> {
         self.inner.prev(idx)

libafl/src/corpus/inmemory.rs

@@ -264,6 +264,12 @@ where
         self._insert(testcase, false)
     }
 
+    #[must_use]
+    /// Peek the next free corpus id
+    pub fn peek_next_free_id(&self) -> CorpusId {
+        CorpusId::from(self.progressive_idx)
+    }
+
     /// Insert a testcase assigning a `CorpusId` to it
     pub fn insert_disabled(&mut self, testcase: RefCell<Testcase<I>>) -> CorpusId {
         self._insert(testcase, true)
@@ -430,6 +436,12 @@ where
         &mut self.current
     }
 
+    /// Peek the next free corpus id
+    #[inline]
+    fn peek_next_free_id(&self) -> CorpusId {
+        self.storage.peek_next_free_id()
+    }
+
     #[inline]
     fn next(&self, idx: CorpusId) -> Option<CorpusId> {
         self.storage.enabled.next(idx)

libafl/src/corpus/inmemory_ondisk.rs

@@ -149,6 +149,12 @@ where
         self.inner.next(idx)
     }
 
+    /// Peek the next free corpus id
+    #[inline]
+    fn peek_next_free_id(&self) -> CorpusId {
+        self.inner.peek_next_free_id()
+    }
+
     #[inline]
     fn prev(&self, idx: CorpusId) -> Option<CorpusId> {
         self.inner.prev(idx)

libafl/src/corpus/mod.rs

@@ -131,6 +131,9 @@ pub trait Corpus: UsesInput + Serialize + for<'de> Deserialize<'de> {
     /// Get the next corpus id
     fn next(&self, id: CorpusId) -> Option<CorpusId>;
 
+    /// Peek the next free corpus id
+    fn peek_next_free_id(&self) -> CorpusId;
+
     /// Get the prev corpus id
     fn prev(&self, id: CorpusId) -> Option<CorpusId>;
 

libafl/src/corpus/nop.rs

@@ -87,6 +87,12 @@ where
         &self.empty
     }
 
+    /// Peek the next free corpus id
+    #[inline]
+    fn peek_next_free_id(&self) -> CorpusId {
+        CorpusId::from(0_usize)
+    }
+
     /// Current testcase scheduled (mutable)
     #[inline]
     fn current_mut(&mut self) -> &mut Option<CorpusId> {

libafl/src/corpus/ondisk.rs

@@ -106,6 +106,12 @@ where
         self.inner.replace(idx, testcase)
     }
 
+    /// Peek the next free corpus id
+    #[inline]
+    fn peek_next_free_id(&self) -> CorpusId {
+        self.inner.peek_next_free_id()
+    }
+
     /// Removes an entry from the corpus, returning it if it was present.
     #[inline]
     fn remove(&mut self, idx: CorpusId) -> Result<Testcase<I>, Error> {

libafl_libfuzzer/libafl_libfuzzer_runtime/src/corpus.rs

@@ -203,6 +203,9 @@ where
     fn next(&self, id: CorpusId) -> Option<CorpusId> {
         self.mapping.enabled.next(id)
     }
+    fn peek_next_free_id(&self) -> CorpusId {
+        self.mapping.peek_next_free_id()
+    }
 
     fn prev(&self, id: CorpusId) -> Option<CorpusId> {
         self.mapping.enabled.prev(id)
@@ -356,6 +359,10 @@ where
         maybe_last.ok_or_else(|| Error::illegal_argument("Can only get the last corpus ID."))
     }
 
+    fn peek_next_free_id(&self) -> CorpusId {
+        CorpusId::from(self.count)
+    }
+
     // This just calls Self::get as ArtifactCorpus disregards disabled entries
     fn get_from_all(&self, id: CorpusId) -> Result<&RefCell<Testcase<Self::Input>>, Error> {
         self.get(id)