From d6ee2dbe122df61978ecb1a4704811012f8f3ae6 Mon Sep 17 00:00:00 2001
From: "Dongjia \"toka\" Zhang" <tokazerkje@outlook.com>
Date: Fri, 17 Mar 2023 23:25:01 +0900
Subject: [PATCH] IfStage (#1157)

* macro

* bracket

* IfStage

* remove macro

* revert
---
 libafl/src/stages/logics.rs      | 87 +++++++++++++++++++++++++++++---
 libafl_frida/src/asan/asan_rt.rs |  4 +-
 2 files changed, 82 insertions(+), 9 deletions(-)

diff --git a/libafl/src/stages/logics.rs b/libafl/src/stages/logics.rs
index 49347910f7..6f50712ce1 100644
--- a/libafl/src/stages/logics.rs
+++ b/libafl/src/stages/logics.rs
@@ -77,9 +77,8 @@ where
     }
 }
 
-/// Perform the stage if closure evaluates to true
 #[derive(Debug)]
-pub struct IfElseStage<CB, E, EM, ST, Z>
+pub struct IfStage<CB, E, EM, ST, Z>
 where
     CB: FnMut(&mut Z, &mut E, &mut E::State, &mut EM, CorpusId) -> Result<bool, Error>,
     E: UsesState,
@@ -89,11 +88,10 @@ where
 {
     closure: CB,
     if_stages: ST,
-    else_stages: ST,
     phantom: PhantomData<(E, EM, Z)>,
 }
 
-impl<CB, E, EM, ST, Z> UsesState for IfElseStage<CB, E, EM, ST, Z>
+impl<CB, E, EM, ST, Z> UsesState for IfStage<CB, E, EM, ST, Z>
 where
     CB: FnMut(&mut Z, &mut E, &mut E::State, &mut EM, CorpusId) -> Result<bool, Error>,
     E: UsesState,
@@ -104,13 +102,85 @@ where
     type State = E::State;
 }
 
-impl<CB, E, EM, ST, Z> Stage<E, EM, Z> for IfElseStage<CB, E, EM, ST, Z>
+impl<CB, E, EM, ST, Z> Stage<E, EM, Z> for IfStage<CB, E, EM, ST, Z>
 where
     CB: FnMut(&mut Z, &mut E, &mut E::State, &mut EM, CorpusId) -> Result<bool, Error>,
     E: UsesState,
     EM: UsesState<State = E::State>,
     ST: StagesTuple<E, EM, E::State, Z>,
     Z: UsesState<State = E::State>,
+{
+    fn perform(
+        &mut self,
+        fuzzer: &mut Z,
+        executor: &mut E,
+        state: &mut E::State,
+        manager: &mut EM,
+        corpus_idx: CorpusId,
+    ) -> Result<(), Error> {
+        if (self.closure)(fuzzer, executor, state, manager, corpus_idx)? {
+            self.if_stages
+                .perform_all(fuzzer, executor, state, manager, corpus_idx)?;
+        }
+        Ok(())
+    }
+}
+
+impl<CB, E, EM, ST, Z> IfStage<CB, E, EM, ST, Z>
+where
+    CB: FnMut(&mut Z, &mut E, &mut E::State, &mut EM, CorpusId) -> Result<bool, Error>,
+    E: UsesState,
+    EM: UsesState<State = E::State>,
+    ST: StagesTuple<E, EM, E::State, Z>,
+    Z: UsesState<State = E::State>,
+{
+    /// Constructor
+    pub fn new(closure: CB, if_stages: ST) -> Self {
+        Self {
+            closure,
+            if_stages,
+            phantom: PhantomData,
+        }
+    }
+}
+
+/// Perform the stage if closure evaluates to true
+#[derive(Debug)]
+pub struct IfElseStage<CB, E, EM, ST1, ST2, Z>
+where
+    CB: FnMut(&mut Z, &mut E, &mut E::State, &mut EM, CorpusId) -> Result<bool, Error>,
+    E: UsesState,
+    EM: UsesState<State = E::State>,
+    ST1: StagesTuple<E, EM, E::State, Z>,
+    ST2: StagesTuple<E, EM, E::State, Z>,
+    Z: UsesState<State = E::State>,
+{
+    closure: CB,
+    if_stages: ST1,
+    else_stages: ST2,
+    phantom: PhantomData<(E, EM, Z)>,
+}
+
+impl<CB, E, EM, ST1, ST2, Z> UsesState for IfElseStage<CB, E, EM, ST1, ST2, Z>
+where
+    CB: FnMut(&mut Z, &mut E, &mut E::State, &mut EM, CorpusId) -> Result<bool, Error>,
+    E: UsesState,
+    EM: UsesState<State = E::State>,
+    ST1: StagesTuple<E, EM, E::State, Z>,
+    ST2: StagesTuple<E, EM, E::State, Z>,
+    Z: UsesState<State = E::State>,
+{
+    type State = E::State;
+}
+
+impl<CB, E, EM, ST1, ST2, Z> Stage<E, EM, Z> for IfElseStage<CB, E, EM, ST1, ST2, Z>
+where
+    CB: FnMut(&mut Z, &mut E, &mut E::State, &mut EM, CorpusId) -> Result<bool, Error>,
+    E: UsesState,
+    EM: UsesState<State = E::State>,
+    ST1: StagesTuple<E, EM, E::State, Z>,
+    ST2: StagesTuple<E, EM, E::State, Z>,
+    Z: UsesState<State = E::State>,
 {
     fn perform(
         &mut self,
@@ -131,16 +201,17 @@ where
     }
 }
 
-impl<CB, E, EM, ST, Z> IfElseStage<CB, E, EM, ST, Z>
+impl<CB, E, EM, ST1, ST2, Z> IfElseStage<CB, E, EM, ST1, ST2, Z>
 where
     CB: FnMut(&mut Z, &mut E, &mut E::State, &mut EM, CorpusId) -> Result<bool, Error>,
     E: UsesState,
     EM: UsesState<State = E::State>,
-    ST: StagesTuple<E, EM, E::State, Z>,
+    ST1: StagesTuple<E, EM, E::State, Z>,
+    ST2: StagesTuple<E, EM, E::State, Z>,
     Z: UsesState<State = E::State>,
 {
     /// Constructor
-    pub fn new(closure: CB, if_stages: ST, else_stages: ST) -> Self {
+    pub fn new(closure: CB, if_stages: ST1, else_stages: ST2) -> Self {
         Self {
             closure,
             if_stages,
diff --git a/libafl_frida/src/asan/asan_rt.rs b/libafl_frida/src/asan/asan_rt.rs
index 82504c643b..4b1136c193 100644
--- a/libafl_frida/src/asan/asan_rt.rs
+++ b/libafl_frida/src/asan/asan_rt.rs
@@ -1672,7 +1672,9 @@ impl AsanRuntime {
 
         // apple aarch64 requires MAP_JIT to allocates WX pages
         #[cfg(all(target_vendor = "apple", target_arch = "aarch64"))]
-        map_flags |= MapFlags::MAP_JIT;
+        {
+            map_flags |= MapFlags::MAP_JIT;
+        }
 
         unsafe {
             let mapping = mmap(