From 08a60aa950a042a817fa3d05ae0530df6bee89e5 Mon Sep 17 00:00:00 2001 From: Andrea Fioraldi Date: Tue, 15 Dec 2020 12:19:44 +0100 Subject: [PATCH] stats as trait --- afl/src/engines/mod.rs | 4 +- afl/src/events/mod.rs | 241 ++++++++++++++++++++++++----------------- 2 files changed, 145 insertions(+), 100 deletions(-) diff --git a/afl/src/engines/mod.rs b/afl/src/engines/mod.rs index dc9b44e7f3..3f0e6cdeef 100644 --- a/afl/src/engines/mod.rs +++ b/afl/src/engines/mod.rs @@ -410,7 +410,7 @@ mod tests { use crate::corpus::{Corpus, InMemoryCorpus, Testcase}; use crate::engines::{Engine, Fuzzer, State, StdFuzzer}; #[cfg(feature = "std")] - use crate::events::LoggerEventManager; + use crate::events::{LoggerEventManager, SimpleStats}; use crate::executors::inmemory::InMemoryExecutor; use crate::executors::{Executor, ExitKind}; use crate::inputs::bytes::BytesInput; @@ -434,7 +434,7 @@ mod tests { let executor = InMemoryExecutor::::new("main", harness, tuple_list!()); let mut state = State::new(tuple_list!()); - let mut events_manager = LoggerEventManager::new(stderr()); + let mut events_manager = LoggerEventManager::new(SimpleStats::new(|s| { println!("{}", s); })); let mut engine = Engine::new(executor); let mut mutator = StdScheduledMutator::new(); mutator.add_mutation(mutation_bitflip); diff --git a/afl/src/events/mod.rs b/afl/src/events/mod.rs index cfed5e27a5..a9d1c279e0 100644 --- a/afl/src/events/mod.rs +++ b/afl/src/events/mod.rs @@ -42,60 +42,21 @@ pub struct ClientStats { executions: u64, } -/* -/// A custom event, for own messages, with own handler. -pub trait CustomEvent: SerdeAny -where - I: Input, -{ - /// Returns the name of this event - fn name(&self) -> &str; - /// This method will be called in the broker - fn handle_in_broker(&self) -> 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) -> Result<(), AflError>; -} -*/ +// TODO Stats as a trait in order to implement logging in an abstract way +// Something like stats.show() that will render the AFL status screen -/// Events sent around in the library -pub trait Event -where - I: Input, -{ - /// Returns the name of this event - fn name(&self) -> &str; - /// This method will be called in the broker - fn handle_in_broker(&self, stats: &mut Stats) -> 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, - state: &mut State, - corpus: &mut C, - ) -> Result<(), AflError> - where - C: Corpus, - OT: ObserversTuple, - FT: FeedbacksTuple, - R: Rand; -} - -#[derive(Debug, Clone, Default)] -pub struct Stats { - start_time: Duration, - corpus_size: usize, - client_stats: Vec, -} - -impl Stats { - /// the client stats, mutable - fn client_stats_mut(&mut self) -> &mut Vec { - &mut self.client_stats - } +pub trait Stats { + /// the client stats (mut) + fn client_stats_mut(&mut self) -> &mut Vec; /// the client stats - fn client_stats(&self) -> &[ClientStats] { - &self.client_stats - } + fn client_stats(&self) -> &[ClientStats]; + + /// creation time + fn start_time(&mut self) -> time::Duration; + + /// show the stats to the user + fn show(&mut self, event_msg: String); /// Amount of elements in the corpus (combined for all children) fn corpus_size(&self) -> u64 { @@ -104,11 +65,6 @@ impl Stats { .fold(0u64, |acc, x| acc + x.corpus_size) } - /// Time this fuzzing run stated - fn start_time(&mut self) -> time::Duration { - self.start_time - } - /// Total executions #[inline] fn total_execs(&mut self) -> u64 { @@ -140,6 +96,102 @@ impl Stats { } } +#[derive(Debug)] +pub struct SimpleStats +where + F: FnMut(String) +{ + print_fn: F, + start_time: Duration, + corpus_size: usize, + client_stats: Vec, +} + +impl Stats for SimpleStats +where + F: FnMut(String) +{ + /// the client stats, mutable + fn client_stats_mut(&mut self) -> &mut Vec { + &mut self.client_stats + } + + /// the client stats + fn client_stats(&self) -> &[ClientStats] { + &self.client_stats + } + + /// Time this fuzzing run stated + fn start_time(&mut self) -> time::Duration { + self.start_time + } + + fn show(&mut self, event_msg: String) { + let fmt = format!("[{}] corpus: {}, executions: {}, exec/sec: {}", event_msg, self.corpus_size(), self.total_execs(), self.execs_per_sec()); + (self.print_fn)(fmt); + } +} + +impl SimpleStats +where + F: FnMut(String) +{ + pub fn new(print_fn: F) -> Self { + Self { + print_fn: print_fn, + start_time: utils::current_time(), + corpus_size: 0, + client_stats: vec![] + } + } + + pub fn with_time(print_fn: F, start_time: time::Duration) -> Self { + Self { + print_fn: print_fn, + start_time: start_time, + corpus_size: 0, + client_stats: vec![] + } + } +} + +/* +/// A custom event, for own messages, with own handler. +pub trait CustomEvent: SerdeAny +where + I: Input, +{ + /// Returns the name of this event + fn name(&self) -> &str; + /// This method will be called in the broker + fn handle_in_broker(&self) -> 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) -> Result<(), AflError>; +} +*/ + +/// Events sent around in the library +pub trait Event +where + I: Input, +{ + /// Returns the name of this event + fn name(&self) -> &str; + /// This method will be called in the broker + fn handle_in_broker(&self, stats: &mut ST) -> Result where ST: Stats; + /// 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, + state: &mut State, + corpus: &mut C, + ) -> Result<(), AflError> + where + C: Corpus, + OT: ObserversTuple, + FT: FeedbacksTuple, + R: Rand; +} + pub trait EventManager where C: Corpus, @@ -258,7 +310,7 @@ where /// Broker fun #[inline] - fn handle_in_broker(&self, stats: &mut Stats) -> Result { + fn handle_in_broker(&self, stats: &mut ST) -> Result where ST: Stats { match self { LoggerEvent::NewTestcase { corpus_size, @@ -321,7 +373,7 @@ where } } -pub struct LoggerEventManager +pub struct LoggerEventManager where C: Corpus, E: Executor, @@ -329,19 +381,17 @@ where FT: FeedbacksTuple, I: Input, R: Rand, - W: Write, + ST: Stats, //CE: CustomEvent, { - writer: W, - - stats: Stats, + stats: ST, events: Vec>, // stats (maybe we need a separated struct?) phantom: PhantomData<(C, E, I, R, OT, FT)>, } -impl EventManager - for LoggerEventManager +impl EventManager + for LoggerEventManager where C: Corpus, E: Executor, @@ -349,7 +399,7 @@ where FT: FeedbacksTuple, I: Input, R: Rand, - W: Write, + ST: Stats, //CE: CustomEvent, { fn process( @@ -431,7 +481,7 @@ where } } -impl LoggerEventManager +impl LoggerEventManager where C: Corpus, I: Input, @@ -439,16 +489,12 @@ where OT: ObserversTuple, FT: FeedbacksTuple, R: Rand, - W: Write, + ST: Stats, //TODO CE: CustomEvent, { - pub fn new(writer: W) -> Self { + pub fn new(stats: ST) -> Self { Self { - stats: Stats { - start_time: utils::current_time(), - ..Default::default() - }, - writer: writer, + stats: stats, phantom: PhantomData, events: vec![], } @@ -534,7 +580,7 @@ where /// Broker fun #[inline] - fn handle_in_broker(&self, stats: &mut Stats) -> Result { + fn handle_in_broker(&self, stats: &mut ST) -> Result where ST: Stats { match &self.kind { LLMPEventKind::NewTestcase { input: _, @@ -633,20 +679,7 @@ const _LLMP_TAG_EVENT_TO_BROKER: llmp::Tag = 0x2B80438; const LLMP_TAG_EVENT_TO_BOTH: llmp::Tag = 0x2B0741; #[cfg(feature = "std")] -pub struct LlmpEventManager -where - W: Write, - //CE: CustomEvent, -{ - writer: W, - - // stats (maybe we need a separated struct?) - llmp: llmp::LlmpConnection, - stats: Stats, - phantom: PhantomData<(C, E, OT, FT, I, R)>, -} - -impl LlmpEventManager +pub struct LlmpEventManager where C: Corpus, E: Executor, @@ -654,20 +687,32 @@ where FT: FeedbacksTuple, I: Input, R: Rand, - W: Write, + ST: Stats, + //CE: CustomEvent, +{ + llmp: llmp::LlmpConnection, + stats: ST, + phantom: PhantomData<(C, E, OT, FT, I, R)>, +} + +impl LlmpEventManager +where + C: Corpus, + E: Executor, + OT: ObserversTuple, + FT: FeedbacksTuple, + I: Input, + R: Rand, + ST: Stats, { /// Create llmp on a port /// If the port is not yet bound, it will act as broker /// Else, it will act as client. - pub fn new_on_port(port: u16, writer: W) -> Result { + pub fn new_on_port(port: u16, stats: ST) -> Result { let mgr = Self { llmp: llmp::LlmpConnection::on_port(port)?, - stats: Stats { - start_time: utils::current_time(), - ..Default::default() - }, + stats: stats, phantom: PhantomData, - writer, }; Ok(mgr) } @@ -723,8 +768,8 @@ where } #[cfg(feature = "std")] -impl EventManager - for LlmpEventManager +impl EventManager + for LlmpEventManager where C: Corpus, E: Executor, @@ -732,7 +777,7 @@ where OT: ObserversTuple, I: Input, R: Rand, - W: Write, + ST: Stats, //CE: CustomEvent, { fn process(