From 6994a912558cb2d74a77501680a17c1a74a2c630 Mon Sep 17 00:00:00 2001
From: Andrea Fioraldi <andreafioraldi@gmail.com>
Date: Tue, 24 Nov 2020 10:19:52 +0100
Subject: [PATCH] try to create better events

---
 afl/src/events/mod.rs         | 56 +++++++++++++++++++++++++----------
 afl/src/mutators/scheduled.rs | 35 ++++++++++++++++++----
 afl/src/stages/mutational.rs  |  2 +-
 3 files changed, 71 insertions(+), 22 deletions(-)

diff --git a/afl/src/events/mod.rs b/afl/src/events/mod.rs
index d5b1b3c95a..1db5674311 100644
--- a/afl/src/events/mod.rs
+++ b/afl/src/events/mod.rs
@@ -3,20 +3,23 @@ pub mod llmp_translated; // TODO: Abstract away.
 pub mod shmem_translated;
 pub use crate::events::llmp::LLMP;
 
+use alloc::rc::Rc;
 use core::any::Any;
+use core::cell::RefCell;
 //use core::any::TypeId;
-use core::fmt::Display;
 // TODO use core version
 use std::io::Write;
 
-use crate::corpus::Corpus;
+use crate::corpus::{Corpus, Testcase};
 use crate::engines::State;
 use crate::executors::Executor;
 use crate::inputs::Input;
 use crate::utils::Rand;
 use crate::AflError;
 
-pub trait Event: Display + Any {}
+pub trait Event: Any {
+    fn name(&self) -> &'static str;
+}
 
 pub trait EventManager<S, C, E, I, R>
 where
@@ -65,10 +68,9 @@ macro_rules! fire_event {
 }
 
 pub struct LoadInitialEvent {}
-impl Event for LoadInitialEvent {}
-impl Display for LoadInitialEvent {
-    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
-        write!(f, "Load")
+impl Event for LoadInitialEvent {
+    fn name(&self) -> &'static str {
+        "LOAD"
     }
 }
 impl LoadInitialEvent {
@@ -77,16 +79,40 @@ impl LoadInitialEvent {
     }
 }
 
-pub struct NewTestcaseEvent {}
-impl Event for NewTestcaseEvent {}
-impl Display for NewTestcaseEvent {
-    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
-        write!(f, "New")
+pub struct NewTestcaseEvent<I>
+where
+    I: Input,
+{
+    testcase: Rc<RefCell<Testcase<I>>>,
+}
+
+impl<I> Event<I> for NewTestcaseEvent<I>
+where
+    I: Input,
+{
+    fn name(&self) -> &'static str {
+        "NEW"
     }
 }
-impl NewTestcaseEvent {
+
+impl<I> NewTestcaseEvent<I>
+where
+    I: Input,
+{
+    pub fn new(testcase: Rc<RefCell<Testcase<I>>>) -> Self {
+        NewTestcaseEvent { testcase: testcase }
+    }
+}
+
+pub struct UpdateStatsEvent {}
+impl Event for UpdateStatsEvent {
+    fn name(&self) -> &'static str {
+        "STATS"
+    }
+}
+impl UpdateStatsEvent {
     pub fn new() -> Self {
-        NewTestcaseEvent {}
+        UpdateStatsEvent {}
     }
 }
 
@@ -136,7 +162,7 @@ where
                 &mut self.writer,
                 "#{}\t[{}] corp: {} exec/s: {}",
                 state.executions(),
-                event,
+                event.name(),
                 state.corpus().entries().len(),
                 state.executions_over_seconds()
             )?;
diff --git a/afl/src/mutators/scheduled.rs b/afl/src/mutators/scheduled.rs
index 5cd51cefc0..b0f2e0e10e 100644
--- a/afl/src/mutators/scheduled.rs
+++ b/afl/src/mutators/scheduled.rs
@@ -248,6 +248,30 @@ where
     }
 }
 
+pub fn mutation_byteneg<M, C, I, R>(
+    _mutator: &mut M,
+    rand: &mut R,
+    _corpus: &mut C,
+    input: &mut I,
+) -> Result<MutationResult, AflError>
+where
+    M: Mutator<C, I, R>,
+    C: Corpus<I, R>,
+    I: Input + HasBytesVec,
+    R: Rand,
+{
+    if input.bytes().len() == 0 {
+        Ok(MutationResult::Skipped)
+    } else {
+        let idx = rand.below(input.bytes().len() as u64) as usize;
+        unsafe {
+            // moar speed, no bound check
+            *input.bytes_mut().get_unchecked_mut(idx) = !(*input.bytes().get_unchecked(idx));
+        }
+        Ok(MutationResult::Mutated)
+    }
+}
+
 /// Returns the first and last diff position between the given vectors, stopping at the min len
 fn locate_diffs(this: &[u8], other: &[u8]) -> (i64, i64) {
     let mut first_diff: i64 = -1;
@@ -374,9 +398,11 @@ where
     pub fn new_default() -> Self {
         let mut scheduled = StdScheduledMutator::<C, I, R>::new();
         scheduled.add_mutation(mutation_bitflip);
-        scheduled.add_mutation(mutation_bitflip);
-        scheduled.add_mutation(mutation_bitflip);
-        scheduled.add_mutation(mutation_bitflip);
+        scheduled.add_mutation(mutation_byteflip);
+        scheduled.add_mutation(mutation_byteinc);
+        scheduled.add_mutation(mutation_bytedec);
+        scheduled.add_mutation(mutation_byteneg);
+
         scheduled.add_mutation(mutation_bitflip);
         scheduled.add_mutation(mutation_bitflip);
         scheduled.add_mutation(mutation_bitflip);
@@ -395,9 +421,6 @@ where
         scheduled.add_mutation(mutation_bitflip);
         scheduled.add_mutation(mutation_bitflip);
 
-        scheduled.add_mutation(mutation_byteflip);
-        scheduled.add_mutation(mutation_byteinc);
-        scheduled.add_mutation(mutation_bytedec);
         scheduled.add_mutation(mutation_splice);
         HavocBytesMutator {
             scheduled: scheduled,
diff --git a/afl/src/stages/mutational.rs b/afl/src/stages/mutational.rs
index 38e7db2eae..6736585290 100644
--- a/afl/src/stages/mutational.rs
+++ b/afl/src/stages/mutational.rs
@@ -54,7 +54,7 @@ where
 
             let (interesting, new_testcase) = state.evaluate_input(input)?;
             if !new_testcase.is_none() {
-                fire_event!(events, NewTestcaseEvent)?;
+                fire_event!(events, NewTestcaseEvent<I>, new_testcase.unwrap())?;
             }
 
             self.mutator_mut()