From 56f08338ecab19a56b7fbe8f5c5051779462682b Mon Sep 17 00:00:00 2001 From: Andrea Fioraldi Date: Fri, 27 Nov 2020 15:39:37 +0100 Subject: [PATCH] serialize testcase in newtestcase event (not working) --- afl/src/corpus/mod.rs | 12 ++++--- afl/src/corpus/testcase.rs | 3 +- afl/src/engines/mod.rs | 42 ++++++++++++++++++------- afl/src/events/mod.rs | 61 ++++++++++++++++-------------------- afl/src/stages/mutational.rs | 17 +++++----- 5 files changed, 73 insertions(+), 62 deletions(-) diff --git a/afl/src/corpus/mod.rs b/afl/src/corpus/mod.rs index 8a50907f7a..d92ab7e347 100644 --- a/afl/src/corpus/mod.rs +++ b/afl/src/corpus/mod.rs @@ -35,9 +35,10 @@ where self.entries().len() } - /// Add an entry to the corpus - fn add(&mut self, testcase: Testcase) { + /// Add an entry to the corpus and return its index + fn add(&mut self, testcase: Testcase) -> usize { self.entries_mut().push(testcase); + self.entries().len() - 1 } /// Replaces the testcase at the given idx @@ -199,7 +200,7 @@ where R: Rand, { /// Add an entry and save it to disk - fn add(&mut self, mut entry: Testcase) { + fn add(&mut self, mut entry: Testcase) -> usize { match entry.filename() { None => { // TODO walk entry metadatas to ask for pices of filename (e.g. :havoc in AFL) @@ -210,6 +211,7 @@ where _ => {} } self.entries.push(entry); + self.entries.len() - 1 } fn current_testcase(&self) -> (&Testcase, usize) { @@ -285,8 +287,8 @@ where self.corpus.count() } - fn add(&mut self, entry: Testcase) { - self.corpus.add(entry); + fn add(&mut self, entry: Testcase) -> usize { + self.corpus.add(entry) } /// Removes an entry from the corpus, returning it if it was present. diff --git a/afl/src/corpus/testcase.rs b/afl/src/corpus/testcase.rs index b4c6d6b696..37940f058c 100644 --- a/afl/src/corpus/testcase.rs +++ b/afl/src/corpus/testcase.rs @@ -6,6 +6,7 @@ use core::convert::Into; use core::default::Default; use core::option::Option; use hashbrown::HashMap; +use serde::{Deserialize, Serialize}; use crate::inputs::Input; use crate::AflError; @@ -22,7 +23,7 @@ pub trait TestcaseMetadata { } /// An entry in the Testcase Corpus -#[derive(Default)] +#[derive(Default, Serialize, Deserialize)] pub struct Testcase where I: Input, diff --git a/afl/src/engines/mod.rs b/afl/src/engines/mod.rs index b353e7f2fd..f734aaa7d9 100644 --- a/afl/src/engines/mod.rs +++ b/afl/src/engines/mod.rs @@ -144,16 +144,34 @@ where Ok(testcase) } - /// Adds this input to the corpus, if it's intersting - fn add_input(&mut self, corpus: &mut C, input: I) -> Result<(), AflError> { - let fitness = self.evaluate_input(&input)?; + /// Create a testcase from this input, if it's intersting + fn testcase_if_interesting( + &mut self, + input: I, + fitness: u32, + ) -> Result>, AflError> { if fitness > 0 { - let testcase = self.input_to_testcase(input, fitness)?; - corpus.add(testcase); + Ok(Some(self.input_to_testcase(input, fitness)?)) } else { self.discard_input(&input)?; + Ok(None) + } + } + + /// Adds this input to the corpus, if it's intersting + fn add_if_interesting( + &mut self, + corpus: &mut C, + input: I, + fitness: u32, + ) -> Result, AflError> { + if fitness > 0 { + let testcase = self.input_to_testcase(input, fitness)?; + Ok(Some(corpus.add(testcase))) + } else { + self.discard_input(&input)?; + Ok(None) } - Ok(()) } } @@ -176,12 +194,12 @@ where { for _ in 0..num { let input = generator.generate(rand)?; - state.add_input(corpus, input)?; - let event = Event::LoadInitial { + let fitness = state.evaluate_input(&input)?; + state.add_if_interesting(corpus, input, fitness)?; + events.fire(Event::LoadInitial { sender_id: 0, - _marker: PhantomData, - }; - events.fire(event)?; + phantom: PhantomData, + })?; } events.process(state, corpus)?; Ok(()) @@ -338,7 +356,7 @@ where events.fire(Event::UpdateStats { sender_id: 0, new_execs: 1, - _marker: PhantomData, + phantom: PhantomData, })?; // TODO self.new_execs}); } } diff --git a/afl/src/events/mod.rs b/afl/src/events/mod.rs index 634555e720..cba4968ea3 100644 --- a/afl/src/events/mod.rs +++ b/afl/src/events/mod.rs @@ -1,7 +1,6 @@ #[cfg(feature = "std")] pub mod llmp; -use alloc::borrow::ToOwned; use alloc::string::String; use core::marker::PhantomData; @@ -67,7 +66,7 @@ where */ /// Events sent around in the library -#[derive(Serialize, Deserialize, Debug)] +#[derive(Serialize, Deserialize)] pub enum Event where S: State, @@ -79,34 +78,33 @@ where { LoadInitial { sender_id: u64, - _marker: PhantomData<(S, C, E, I, R)>, + phantom: PhantomData<(S, C, E, I, R)>, }, NewTestcase { sender_id: u64, - input: I, - fitness: u32, - _marker: PhantomData<(S, C, E, I, R)>, + testcase: Testcase, + phantom: PhantomData<(S, C, E, I, R)>, }, UpdateStats { sender_id: u64, new_execs: usize, - _marker: PhantomData<(S, C, E, I, R)>, + phantom: PhantomData<(S, C, E, I, R)>, }, Crash { sender_id: u64, input: I, - _marker: PhantomData<(S, C, E, I, R)>, + phantom: PhantomData<(S, C, E, I, R)>, }, Timeout { sender_id: u64, input: I, - _marker: PhantomData<(S, C, E, I, R)>, + phantom: PhantomData<(S, C, E, I, R)>, }, Log { sender_id: u64, severity_level: u8, message: String, - _marker: PhantomData<(S, C, E, I, R)>, + phantom: PhantomData<(S, C, E, I, R)>, }, //Custom {sender_id: u64, custom_event: CE}, } @@ -124,34 +122,33 @@ where match self { Event::LoadInitial { sender_id: _, - _marker, + phantom, } => "Initial", Event::NewTestcase { sender_id: _, - input: _, - fitness: _, - _marker, + testcase: _, + phantom, } => "New Testcase", Event::UpdateStats { sender_id: _, new_execs: _, - _marker, + phantom, } => "Stats", Event::Crash { sender_id: _, input: _, - _marker, + phantom, } => "Crash", Event::Timeout { sender_id: _, input: _, - _marker, + phantom, } => "Timeout", Event::Log { sender_id: _, severity_level: _, message: _, - _marker, + phantom, } => "Log", //Event::Custom {sender_id, custom_event} => custom_event.name(), } @@ -163,17 +160,16 @@ where _corpus: &mut C, ) -> Result { match self { - Event::LoadInitial { sender_id: _, _marker } => Ok(BrokerEventResult::Handled), + Event::LoadInitial { sender_id: _, phantom } => Ok(BrokerEventResult::Handled), Event::NewTestcase { sender_id: _, - input: _, - fitness: _, - _marker, + testcase: _, + phantom, } => Ok(BrokerEventResult::Forward), Event::UpdateStats { sender_id: _, new_execs: _, - _marker, + phantom, } => { // TODO Ok(BrokerEventResult::Handled) @@ -181,12 +177,12 @@ where Event::Crash { sender_id: _, input: _, - _marker, + phantom, } => Ok(BrokerEventResult::Handled), Event::Timeout { sender_id: _, input: _, - _marker, + phantom, } => { // TODO Ok(BrokerEventResult::Handled) @@ -195,7 +191,7 @@ where sender_id, severity_level, message, - _marker, + phantom, } => { //TODO: broker.log() #[cfg(feature = "std")] @@ -215,13 +211,10 @@ where match self { Event::NewTestcase { sender_id: _, - input, - fitness, - _marker, + testcase, + phantom, } => { - let mut testcase = Testcase::new(input.to_owned()); - testcase.set_fitness(*fitness); - corpus.add(testcase); + corpus.add(*testcase); Ok(()) } _ => Err(AflError::Unknown( @@ -367,7 +360,7 @@ where //CE: CustomEvent, { // TODO... - _marker: PhantomData<(S, C, E, I, R)>, + phantom: PhantomData<(S, C, E, I, R)>, } #[cfg(feature = "std")] @@ -431,7 +424,7 @@ where { pub fn new() -> Self { Self { - _marker: PhantomData, + phantom: PhantomData, } } } diff --git a/afl/src/stages/mutational.rs b/afl/src/stages/mutational.rs index bd96aa15c3..9255e35714 100644 --- a/afl/src/stages/mutational.rs +++ b/afl/src/stages/mutational.rs @@ -49,22 +49,19 @@ where self.mutator_mut() .mutate(rand, corpus, &mut input_mut, i as i32)?; - let interesting = state.evaluate_input(&input_mut)?; + let fitness = state.evaluate_input(&input_mut)?; self.mutator_mut() - .post_exec(interesting, &input_mut, i as i32)?; + .post_exec(fitness, &input_mut, i as i32)?; - if interesting > 0 { - //let new_testcase = state.input_to_testcase(input_mut, interesting)?; + let testcase_maybe = state.testcase_if_interesting(input_mut, fitness)?; + if let Some(testcase) = testcase_maybe { + //corpus.entries()[idx] events.fire(Event::NewTestcase { sender_id: 0, - input: input_mut, - fitness: interesting, - _marker: PhantomData, + testcase: testcase, + phantom: PhantomData, })?; - //state.corpus_mut().add(new_testcase); // TODO: Probably no longer needed, once events work - } else { - state.discard_input(&input_mut)?; } } Ok(())