From 68a72f583ce166d5cfcc56d6333407d0b81c5d2f Mon Sep 17 00:00:00 2001 From: Dominik Maier Date: Wed, 25 Nov 2020 21:34:28 +0100 Subject: [PATCH] event manager --- afl/src/engines/mod.rs | 13 +- afl/src/events/mod.rs | 336 ++++++++++++++++++----------------- afl/src/stages/mutational.rs | 7 +- 3 files changed, 187 insertions(+), 169 deletions(-) diff --git a/afl/src/engines/mod.rs b/afl/src/engines/mod.rs index 0f8e723c25..51bedcfd2f 100644 --- a/afl/src/engines/mod.rs +++ b/afl/src/engines/mod.rs @@ -9,7 +9,7 @@ use core::marker::PhantomData; use hashbrown::HashMap; use crate::corpus::{Corpus, Testcase}; -use crate::events::{EventManager, LoadInitialEvent, UpdateStatsEvent}; +use crate::events::{EventManager, Event}; use crate::executors::Executor; use crate::feedbacks::Feedback; use crate::generators::Generator; @@ -17,7 +17,7 @@ use crate::inputs::Input; use crate::observers::Observer; use crate::stages::Stage; use crate::utils::{current_milliseconds, Rand}; -use crate::{fire_event, AflError}; +use crate::AflError; // TODO FeedbackMetadata to store histroy_map @@ -177,9 +177,10 @@ where for _ in 0..num { let input = generator.generate(rand)?; state.add_input(corpus, input)?; - fire_event!(events, LoadInitialEvent)?; + let event = Event::LoadInitial {sender_id: 0, _marker: PhantomData}; + events.fire(event)?; } - events.process(state)?; + events.process(state, corpus)?; Ok(()) } @@ -314,7 +315,7 @@ where stage.perform(rand, state, corpus, events, &input)?; } - events.process(state)?; + events.process(state, corpus)?; Ok(idx) } @@ -331,7 +332,7 @@ where let cur = current_milliseconds(); if cur - last > 60 * 100 { last = cur; - fire_event!(events, UpdateStatsEvent)?; + events.fire(Event::UpdateStats {sender_id: 0, new_execs: 1, _marker: PhantomData})?; // TODO self.new_execs}); } } } diff --git a/afl/src/events/mod.rs b/afl/src/events/mod.rs index 93b02c3bf4..9e6ca4547d 100644 --- a/afl/src/events/mod.rs +++ b/afl/src/events/mod.rs @@ -8,9 +8,8 @@ pub mod shmem_translated; #[cfg(feature = "std")] pub use crate::events::llmp::LLMP; -use core::fmt::Formatter; #[cfg(feature = "std")] -use std::io::Write; +use std::{marker::PhantomData, io::Write}; use crate::corpus::{Corpus, Testcase}; use crate::engines::State; @@ -18,41 +17,127 @@ use crate::executors::Executor; use crate::inputs::Input; use crate::utils::Rand; use crate::AflError; - -pub enum EventDestination { - Main, - Broker, - Clients, +/// Indicate if an event worked or not +enum BrokerEventResult { + /// The broker haneled this. No need to pass it on. + Handled, + /// Pass this message along to the clients. + Forward, } -pub trait Event { - fn name() -> &'static str; +/* - fn destination() -> EventDestination; +/// A custom event, in case a user wants to extend the features (at compile time) +pub trait CustomEvent +where + S: State, + C: Corpus, + E: Executor, + I: Input, + R: Rand, +{ + /// Returns the name of this event + fn name(&self) -> &str; + /// This method will be called in the broker + fn handle_in_broker(&self, broker: &dyn EventManager, state: &mut S, corpus: &mut C) -> Result; + /// This method will be called in the clients after handle_in_broker (unless BrokerEventResult::Handled) was returned in handle_in_broker + fn handle_in_client(&self, client: &dyn EventManager, state: &mut S, corpus: &mut C) -> Result<(), AflError>; +} - fn log(&self, formatter: &mut Formatter, _state: &S) -> Result<(), AflError> - where - S: State, - C: Corpus, - E: Executor, - I: Input, - R: Rand, - { - match write!(formatter, "[{}]", Self::name()) { - Ok(_) => Ok(()), - Err(_) => Err(AflError::Unknown("write error".into())), +struct UnusedCustomEvent {} +impl CustomEvent for UnusedCustomEvent +where + S: State, + C: Corpus, + E: Executor, + I: Input, + R: Rand, +{ + fn name(&self) -> &str {"No custom events"} + fn handle_in_broker(&self, broker: &dyn EventManager, state: &mut S, corpus: &mut C) {Ok(BrokerEventResult::Handled)} + fn handle_in_client(&self, client: &dyn EventManager, state: &mut S, corpus: &mut C) {Ok(())} +} +*/ + +/// Events sent around in the library +pub enum Event +where + S: State, + C: Corpus, + E: Executor, + I: Input, + R: Rand, + // CE: CustomEvent, +{ + LoadInitial {sender_id: u64, _marker: PhantomData<(S, C, E, I, R)>}, + NewTestcase {sender_id: u64, input: I, fitness: u32, _marker: PhantomData<(S, C, E, I, R)>}, + UpdateStats {sender_id: u64, new_execs: usize, _marker: PhantomData<(S, C, E, I, R)>}, + Crash {sender_id: u64, input: I, _marker: PhantomData<(S, C, E, I, R)>}, + Timeout {sender_id: u64, input: I, _marker: PhantomData<(S, C, E, I, R)>}, + Log {sender_id: u64, severity_level: u8, message: String, _marker: PhantomData<(S, C, E, I, R)>}, + //Custom {sender_id: u64, custom_event: CE}, +} + +impl Event +where + S: State, + C: Corpus, + E: Executor, + I: Input, + R: Rand, + //CE: CustomEvent, +{ + fn name(&self) -> &str { + match self { + Event::LoadInitial {sender_id, _marker} => "Initial", + Event::NewTestcase {sender_id, input, fitness, _marker} => "New Testcase", + Event::UpdateStats {sender_id, new_execs, _marker} => "Stats", + Event::Crash {sender_id, input, _marker} => "Crash", + Event::Timeout {sender_id, input, _marker} => "Timeout", + Event::Log {sender_id, severity_level, message, _marker} => "Log", + //Event::Custom {sender_id, custom_event} => custom_event.name(), } } - fn on_recv(&self, _state: &mut S) -> Result<(), AflError> - where - S: State, - C: Corpus, - E: Executor, - I: Input, - R: Rand, - { - Ok(()) + fn handle_in_broker(&self, /*broker: &dyn EventManager,*/ state: &mut S, corpus: &mut C) -> Result { + match self { + Event::LoadInitial {sender_id, _marker} => { + Ok(BrokerEventResult::Handled) + } + Event::NewTestcase {sender_id, input, fitness, _marker} => { + Ok(BrokerEventResult::Forward) + } + Event::UpdateStats {sender_id, new_execs, _marker} => { + // TODO + Ok(BrokerEventResult::Handled) + } + Event::Crash {sender_id, input, _marker} => { + Ok(BrokerEventResult::Handled) + } + Event::Timeout {sender_id, input, _marker} => { + // TODO + Ok(BrokerEventResult::Handled) + }, + Event::Log {sender_id, severity_level, message, _marker} => { + //TODO: broker.log() + println!("{}[{}]: {}", sender_id, severity_level, message); + Ok(BrokerEventResult::Handled) + }, + //Event::Custom {sender_id, custom_event} => custom_event.handle_in_broker(state, corpus), + _ => Ok(BrokerEventResult::Forward) + } + } + + fn handle_in_client(&self, /*client: &dyn EventManager,*/ state: &mut S, corpus: &mut C) -> Result<(), AflError> { + match self { + Event::NewTestcase {sender_id, input, fitness, _marker} => { + let mut testcase = Testcase::new(input.to_owned()); + testcase.set_fitness(*fitness); + corpus.add(testcase); + Ok(()) + } + _ => Err(AflError::Unknown("Received illegal message that message should not have arrived.".into())) + } } // TODO serialize and deserialize, defaults to serde @@ -65,138 +150,56 @@ where E: Executor, I: Input, R: Rand, + //CE: CustomEvent, { /// Check if this EventaManager support a given Event type /// To compare events, use Event::name().as_ptr() - fn enabled(&self) -> bool - where - T: Event; + fn enabled(&self) -> bool; /// Fire an Event - fn fire(&mut self, event: T) -> Result<(), AflError> - where - T: Event; + fn fire(&mut self, event: Event) -> Result<(), AflError>; /// Lookup for incoming events and process them. /// Return the number of processes events or an error - fn process(&mut self, state: &mut S) -> Result; -} + fn process(&mut self, state: &mut S, corpus: &mut C) -> Result; -// e.g. fire_event!(manager, MyEvent, myparam1, ...) -#[macro_export] -macro_rules! fire_event { - ($manager:expr, $event:ty, $( $x:expr ),+ ) => { - { - if $manager.enabled::<$event>() { - $manager.fire(<$event>::new($( $x ),*)) - } else { - Ok(()) - } - } - }; - ($manager:expr, $event:ty) => { - { - if $manager.enabled::<$event>() { - $manager.fire(<$event>::new()) - } else { - Ok(()) - } - } - }; -} -pub struct LoadInitialEvent {} -impl Event for LoadInitialEvent { - fn name() -> &'static str { - "LOAD" - } - - fn destination() -> EventDestination { - EventDestination::Broker - } -} -impl LoadInitialEvent { - pub fn new() -> Self { - LoadInitialEvent {} + fn on_recv(&self, _state: &mut S, corpus: &mut C) -> Result<(), AflError> { + // TODO: Better way to move out of testcase, or get ref + //Ok(corpus.add(self.testcase.take().unwrap())) + Ok(()) } } -pub struct NewTestcaseEvent -where - I: Input, -{ - testcase: Testcase, -} - -impl Event for NewTestcaseEvent -where - I: Input, -{ - fn name() -> &'static str { - "NEW" +/*TODO + fn on_recv(&self, state: &mut S, _corpus: &mut C) -> Result<(), AflError> { + println!( + "#{}\t exec/s: {}", + state.executions(), + //TODO: Count corpus.entries().len(), + state.executions_over_seconds() + ); + Ok(()) } - - fn destination() -> EventDestination { - EventDestination::Clients - } -} - -impl NewTestcaseEvent -where - I: Input, -{ - pub fn new(testcase: Testcase) -> Self { - NewTestcaseEvent { testcase: testcase } - } - - pub fn testcase(&self) -> &Testcase { - &self.testcase - } -} - -pub struct UpdateStatsEvent {} -impl Event for UpdateStatsEvent { - fn name() -> &'static str { - "STATS" - } - - fn destination() -> EventDestination { - EventDestination::Broker - } -} -impl UpdateStatsEvent { - pub fn new() -> Self { - UpdateStatsEvent {} - } -} - -pub struct CrashEvent {} -impl Event for CrashEvent { - fn name() -> &'static str { - "CRASH" - } - - fn destination() -> EventDestination { - EventDestination::Broker - } -} -impl CrashEvent { - pub fn new() -> Self { - CrashEvent {} - } -} +*/ #[cfg(feature = "std")] -pub struct LoggerEventManager +pub struct LoggerEventManager where + S: State, + C: Corpus, + I: Input, + E: Executor, + R: Rand, W: Write, + //CE: CustomEvent, { - events: Vec, + events: Vec>, writer: W, } #[cfg(feature = "std")] -impl EventManager for LoggerEventManager +impl EventManager for LoggerEventManager where S: State, C: Corpus, @@ -204,43 +207,58 @@ where I: Input, R: Rand, W: Write, + //CE: CustomEvent, { - fn enabled(&self) -> bool - where - T: Event, + fn enabled(&self) -> bool { true } - fn fire(&mut self, _event: T) -> Result<(), AflError> - where - T: Event, + fn fire(&mut self, event: Event) -> Result<(), AflError> { - self.events.push(T::name().to_string()); + self.events.push(event); Ok(()) } - fn process(&mut self, state: &mut S) -> Result { - let num = self.events.len(); - for event in &self.events { - writeln!( - &mut self.writer, - "#{}\t[{}] exec/s: {}", - state.executions(), - event, - //TODO: Count corpus.entries().len(), - state.executions_over_seconds() - )?; + fn process(&mut self, state: &mut S, corpus: &mut C) -> Result { + // TODO: iterators + let mut handled = vec!(); + for x in self.events.iter() { + handled.push(x.handle_in_broker(state, corpus)?); } + handled.iter().zip(self.events.iter()).map(|(x, event)| match x { + BrokerEventResult::Forward => { + event.handle_in_client(state, corpus) + }, + // Ignore broker-only events + BrokerEventResult::Handled => Ok(()), + } + ).collect::>(); + let count = self.events.len(); + dbg!("Handled {} events", count); self.events.clear(); - Ok(num) + + /* + let num = self.events.len(); + for event in &self.events {} + + self.events.clear(); + */ + + Ok(count) } } #[cfg(feature = "std")] -impl LoggerEventManager +impl LoggerEventManager where + S: State, + C: Corpus, + I: Input, + E: Executor, + R: Rand, W: Write, + //TODO CE: CustomEvent, { pub fn new(writer: W) -> Self { LoggerEventManager { diff --git a/afl/src/stages/mutational.rs b/afl/src/stages/mutational.rs index 49433ce4cb..4ff0fd5188 100644 --- a/afl/src/stages/mutational.rs +++ b/afl/src/stages/mutational.rs @@ -1,6 +1,6 @@ use core::marker::PhantomData; -use crate::engines::State; +use crate::{events::Event, engines::State}; use crate::events::EventManager; use crate::executors::Executor; use crate::inputs::Input; @@ -9,7 +9,6 @@ use crate::stages::Corpus; use crate::stages::Stage; use crate::utils::Rand; use crate::AflError; -use crate::{events::NewTestcaseEvent, fire_event}; // TODO multi mutators stage @@ -56,8 +55,8 @@ where .post_exec(interesting, &input_mut, i as i32)?; if interesting > 0 { - let new_testcase = state.input_to_testcase(input_mut, interesting)?; - fire_event!(events, NewTestcaseEvent, new_testcase)?; + //let new_testcase = state.input_to_testcase(input_mut, interesting)?; + events.fire(Event::NewTestcase { sender_id: 0, input: input_mut, fitness: interesting , _marker: PhantomData})?; //state.corpus_mut().add(new_testcase); // TODO: Probably no longer needed, once events work } else { state.discard_input(&input_mut)?;