Skip to content

Commit

Permalink
Add peek function to Corpus (AFLplusplus#2238)
Browse files Browse the repository at this point in the history
* 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
  • Loading branch information
R9295 authored May 22, 2024
1 parent 4b67b55 commit 1ed1c4e
Show file tree
Hide file tree
Showing 7 changed files with 46 additions and 0 deletions.
6 changes: 6 additions & 0 deletions libafl/src/corpus/cached.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
12 changes: 12 additions & 0 deletions libafl/src/corpus/inmemory.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down Expand Up @@ -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)
Expand Down
6 changes: 6 additions & 0 deletions libafl/src/corpus/inmemory_ondisk.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
3 changes: 3 additions & 0 deletions libafl/src/corpus/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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>;

Expand Down
6 changes: 6 additions & 0 deletions libafl/src/corpus/nop.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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> {
Expand Down
6 changes: 6 additions & 0 deletions libafl/src/corpus/ondisk.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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> {
Expand Down
7 changes: 7 additions & 0 deletions libafl_libfuzzer/libafl_libfuzzer_runtime/src/corpus.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down Expand Up @@ -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)
Expand Down

0 comments on commit 1ed1c4e

Please sign in to comment.