From 1ed1c4eb31b86c45c66cdd3fc23e2ff676941e23 Mon Sep 17 00:00:00 2001
From: Aarnav <aarnavbos@gmail.com>
Date: Wed, 22 May 2024 04:30:47 -0700
Subject: [PATCH] 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                          |  6 ++++++
 libafl/src/corpus/inmemory.rs                        | 12 ++++++++++++
 libafl/src/corpus/inmemory_ondisk.rs                 |  6 ++++++
 libafl/src/corpus/mod.rs                             |  3 +++
 libafl/src/corpus/nop.rs                             |  6 ++++++
 libafl/src/corpus/ondisk.rs                          |  6 ++++++
 .../libafl_libfuzzer_runtime/src/corpus.rs           |  7 +++++++
 7 files changed, 46 insertions(+)

diff --git a/libafl/src/corpus/cached.rs b/libafl/src/corpus/cached.rs
index d585871c0e..c38edaf007 100644
--- a/libafl/src/corpus/cached.rs
+++ b/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)
diff --git a/libafl/src/corpus/inmemory.rs b/libafl/src/corpus/inmemory.rs
index f0ce03cfc2..0e6169614c 100644
--- a/libafl/src/corpus/inmemory.rs
+++ b/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)
diff --git a/libafl/src/corpus/inmemory_ondisk.rs b/libafl/src/corpus/inmemory_ondisk.rs
index 90a7d170a5..f77f376d5c 100644
--- a/libafl/src/corpus/inmemory_ondisk.rs
+++ b/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)
diff --git a/libafl/src/corpus/mod.rs b/libafl/src/corpus/mod.rs
index 21bb9f3259..80ef4e427c 100644
--- a/libafl/src/corpus/mod.rs
+++ b/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>;
 
diff --git a/libafl/src/corpus/nop.rs b/libafl/src/corpus/nop.rs
index ea6211a23d..afec997bdc 100644
--- a/libafl/src/corpus/nop.rs
+++ b/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> {
diff --git a/libafl/src/corpus/ondisk.rs b/libafl/src/corpus/ondisk.rs
index 926b7269c7..0b8e878eb1 100644
--- a/libafl/src/corpus/ondisk.rs
+++ b/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> {
diff --git a/libafl_libfuzzer/libafl_libfuzzer_runtime/src/corpus.rs b/libafl_libfuzzer/libafl_libfuzzer_runtime/src/corpus.rs
index 20f0ad3092..0a3665046f 100644
--- a/libafl_libfuzzer/libafl_libfuzzer_runtime/src/corpus.rs
+++ b/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)