Added reentrancy feedback

src/evm/host.rs

@@ -67,6 +67,9 @@ pub static mut STATE_CHANGE: bool = false;
 pub const RW_SKIPPER_PERCT_IDX: usize = 100;
 pub const RW_SKIPPER_AMT: usize = MAP_SIZE - RW_SKIPPER_PERCT_IDX;
 
+// reentrancy
+pub static mut WRITTEN: bool = false;
+
 // How mant iterations the coverage is the same
 pub static mut COVERAGE_NOT_CHANGED: u32 = 0;
 pub static mut RET_SIZE: usize = 0;
@@ -797,6 +800,9 @@ where
                     JMP_MAP[idx] = if value_changed { 1 } else { 0 };
 
                     STATE_CHANGE |= value_changed;
+
+                    #[cfg(feature = "reentrancy")]
+                    WRITTEN = true;
                 }
 
                 #[cfg(feature = "dataflow")]

src/evm/middlewares/middleware.rs

@@ -1,22 +1,27 @@
-use crate::evm::host::FuzzHost;
+use crate::evm::host::{FuzzHost, JMP_MAP, READ_MAP, WRITE_MAP, CMP_MAP, STATE_CHANGE, 
+    WRITTEN, RET_SIZE, RET_OFFSET};
 use crate::evm::input::{ConciseEVMInput, EVMInput, EVMInputT};
+use crate::evm::types::{as_u64, bytes_to_u64, EVMAddress, EVMU256, generate_random_address, is_zero};
+use crate::generic_vm::vm_executor::MAP_SIZE;
 use crate::generic_vm::vm_state::VMStateT;
 use crate::input::VMInputT;
 use crate::state::{HasCaller, HasItyState};
+use crate::evm::host::FuzzHost;
 
 use bytes::Bytes;
 use libafl::corpus::{Corpus, Testcase};
 use libafl::inputs::Input;
 use libafl::schedulers::Scheduler;
 use libafl::state::{HasCorpus, HasMetadata, State};
-use primitive_types::U512;
+use primitive_types::{H160, U256, U512};
+use revm::{Bytecode, Interpreter, Host};
 use serde::{Deserialize, Serialize};
 
 use std::clone::Clone;
 use std::fmt::Debug;
 
 use std::time::Duration;
-use revm_interpreter::Interpreter;
+use revm_interpreter::{Host, Interpreter};
 use revm_primitives::Bytecode;
 use crate::evm::types::{EVMAddress, EVMU256};
 
@@ -86,12 +91,147 @@ where
     I: VMInputT<VS, EVMAddress, EVMAddress, ConciseEVMInput> + EVMInputT,
     VS: VMStateT,
 {
+    // called on every instruction
     unsafe fn on_step(
         &mut self,
         interp: &mut Interpreter,
         host: &mut FuzzHost<VS, I, S>,
         state: &mut S,
-    );
+    ) {
+        macro_rules! u256_to_u8 {
+            ($key:ident) => {
+                (as_u64($key >> 4) % 254) as u8
+            };
+        }
+        macro_rules! fast_peek {
+            ($idx:expr) => {
+                interp.stack.data()[interp.stack.len() - 1 - $idx]
+            };
+        }
+        unsafe {
+            match *interp.instruction_pointer {
+                0x57 => { // JUMPI
+                    let br = fast_peek!(1);
+                    let jump_dest = if is_zero(br) {
+                        1
+                    } else {
+                        as_u64(fast_peek!(0))
+                    };
+                    let idx = (interp.program_counter() * (jump_dest as usize)) % MAP_SIZE;
+                    if JMP_MAP[idx] == 0 {
+                        *host.coverage_changed = true;
+                    }
+                    if JMP_MAP[idx] < 255 {
+                        JMP_MAP[idx] += 1;
+                    }
+
+                    #[cfg(feature = "cmp")]
+                    {
+                        let idx = (interp.program_counter()) % MAP_SIZE;
+                        CMP_MAP[idx] = br;
+                    }
+                }
+
+                #[cfg(any(feature = "dataflow", feature = "cmp", feature = "reentrancy"))]
+                0x55 => { // SSTORE
+                    #[cfg(feature = "dataflow")]
+                    let value = fast_peek!(1);
+                    {
+                        let mut key = fast_peek!(0);
+                        let v = u256_to_u8!(value) + 1;
+                        WRITE_MAP[process_rw_key!(key)] = v;
+                    }
+                    let res = <FuzzHost<VS, I, S> as Host<S>>::sload(
+                        host,
+                        interp.contract.address, 
+                        fast_peek!(0));
+                    let value_changed = res.expect("sload failed").0 != value;
+
+                    let idx = interp.program_counter() % MAP_SIZE;
+                    JMP_MAP[idx] = if value_changed { 1 } else { 0 };
+
+                    STATE_CHANGE |= value_changed;
+
+                    WRITTEN = true;
+                }
+
+                #[cfg(feature = "dataflow")]
+                0x54 => { // SLOAD
+                    let mut key = fast_peek!(0);
+                    READ_MAP[process_rw_key!(key)] = true;
+                }
+
+                #[cfg(feature = "cmp")]
+                0x10 | 0x12 => { // LT | SLT
+                    let v1 = fast_peek!(0);
+                    let v2 = fast_peek!(1);
+                    let abs_diff = if v1 >= v2 {
+                        if v1 - v2 != EVMU256::ZERO {
+                            v1 - v2
+                        } else {
+                            EVMU256::from(1)
+                        }
+                    } else {
+                        EVMU256::ZERO
+                    };
+                    let idx = interp.program_counter() % MAP_SIZE;
+                    if abs_diff < CMP_MAP[idx] {
+                        CMP_MAP[idx] = abs_diff;
+                    }
+                }
+
+                #[cfg(feature = "cmp")]
+                0x11 | 0x13 => { // GT | SGT
+                    let v1 = fast_peek!(0);
+                    let v2 = fast_peek!(1);
+                    let abs_diff = if v1 <= v2 {
+                        if v2 - v1 != EVMU256::ZERO {
+                            v2 - v1
+                        } else {
+                            U256::from(1)
+                        }
+                    } else {
+                        EVMU256::ZERO
+                    };
+                    let idx = interp.program_counter() % MAP_SIZE;
+                    if abs_diff < CMP_MAP[idx] {
+                        CMP_MAP[idx] = abs_diff;
+                    }
+                }
+
+                #[cfg(feature = "cmp")]
+                0x14 => { // EQ
+                    let v1 = fast_peek!(0);
+                    let v2 = fast_peek!(1);
+                    let abs_diff = 
+                        (if v1 < v2 {v2 - v1} else {v1 - v2}) % (EVMU256::MAX - EVMU256::from(1)) + 1;
+                    let idx = interp.program_counter() % MAP_SIZE;
+                    if abs_diff < CMP_MAP[idx] {
+                        CMP_MAP[idx] = abs_diff;
+                    }
+                }
+
+                0xf1 | 0xf2 | 0xf4 | 0xfa => { // CALL | CALLCODE | DELEGATECALL | STATICCALL
+                    let offset_of_ret_size: usize = match *interp.instruction_pointer {
+                        0xf1 | 0xf2 => 6,
+                        0xf4 | 0xfa => 5,
+                        _ => unreachable!(),
+                    };
+                    unsafe {
+                        RET_OFFSET = as_u64(fast_peek!(offset_of_ret_size - 1)) as usize;
+                        RET_SIZE = as_u64(fast_peek!(offset_of_ret_size)) as usize;
+                    }
+                    host._pc = interp.program_counter();
+                }
+
+                0xf0 | 0xf5 => { // CREATE, CREATE2
+                    host._pc = interp.program_counter();
+                }
+
+                _ => {}
+            }
+        }
+    }
 
     unsafe fn on_return(
         &mut self,

src/evm/vm.rs

@@ -937,6 +937,10 @@ where
         unsafe { &mut CMP_MAP }
     }
 
+    fn get_written(&self) -> &'static mut bool {
+        unsafe { &mut WRITTEN }
+    }
+
     fn state_changed(&self) -> bool {
         unsafe { STATE_CHANGE }
     }

src/feedback.rs

@@ -468,6 +468,15 @@ where
     {
         let mut cmp_interesting = false;
         let mut cov_interesting = false;
+        let mut reentrancy_interesting = false;
+
+        // checking for reentrancy
+        let new_state = _state.get_execution_result().new_state.state;
+        while new_state.has_post_execution() {
+            // now execute again, checking for writes to EVM memory.
+            self.vm.deref().borrow_mut().execute(_input, _state);
+        }
+        reentrancy_interesting = *self.vm.deref().borrow_mut().get_written();
 
         // check if the current distance is smaller than the min_map
         for i in 0..MAP_SIZE {
@@ -492,7 +501,16 @@ where
                 .vote(state.get_infant_state_state(), input.get_state_idx(), 3);
         }
 
+        // if reentrancy has occurred, vote for the state (tentative)
+        if reentrancy_interesting {
+            println!("Voted for {} because of REENTRANCY", input.get_state_idx());
+            self.scheduler
+                .vote(state.get_infant_state_state(), input.get_state_idx(), 3);
+        }
+
         unsafe {
+            let cur_read_map = self.vm.deref().borrow_mut().get_read();
+            let cur_write_map = self.vm.deref().borrow_mut().get_write();
             // hack to account for saving reentrancy without dataflow
             let pc_interesting = state
                 .get_execution_result()
@@ -506,7 +524,22 @@ where
                 if self.known_states.contains(&hash) {
                     return Ok(false);
                 }
-                return Ok(true);
+
+                // ensure dataflow / coverage-wise is interesting for new encountered VM state
+                let mut df_interesting = false;
+                for i in 0..MAP_SIZE {
+                    if cur_read_map[i] && cur_write_map[i] != 0 {
+                        df_interesting = true;
+                        break;
+                    }
+                }
+                for i in 0..MAP_SIZE {
+                    cur_write_map[i] = 0;
+                }
+                if df_interesting || pc_interesting { // || reentrancy_interesting {
+                    self.known_states.insert(hash);
+                    return Ok(true);
+                }
             }
         }
         Ok(false)
@@ -524,3 +557,128 @@ where
         Ok(())
     }
 }
+
+//////////////////////////////////////////////////////////////////////////////
+
+// #[cfg(feature = "reentrancy")]
+pub struct ReentrancyFeedback<I, S, VS, Loc, Addr, Out, Code, By, SlotTy, CI>
+where
+    I: VMInputT<VS, Loc, Addr, CI>,
+    VS: Default + VMStateT,
+    Addr: Serialize + DeserializeOwned + Debug + Clone,
+    Loc: Serialize + DeserializeOwned + Debug + Clone,
+    CI: Serialize + DeserializeOwned + Debug + Clone + ConciseSerde,
+{
+    /// write status of the current execution
+    written: bool,
+    vm: Rc<RefCell<dyn GenericVM<VS, Code, By, Loc, Addr, SlotTy, Out, I, S, CI>>>,
+    phantom: PhantomData<(VS, Loc, Addr, Out)>,
+}
+
+// #[cfg(feature = "reentrancy")]
+impl<I, S, VS, Loc, Addr, Out, Code, By, SlotTy, CI> Debug for ReentrancyFeedback<I, S, VS, Loc, Addr, Out, Code, By, SlotTy, CI> 
+where
+    I: VMInputT<VS, Loc, Addr, CI>,
+    VS: Default + VMStateT,
+    Addr: Serialize + DeserializeOwned + Debug + Clone,
+    Loc: Serialize + DeserializeOwned + Debug + Clone,
+    CI: Serialize + DeserializeOwned + Debug + Clone + ConciseSerde,
+{
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("ReentrancyFeedback").finish()
+    }
+}
+
+// #[cfg(feature = "reentrancy")]
+impl<I, S, VS, Loc, Addr, Out, Code, By, SlotTy, CI> Named for ReentrancyFeedback<I, S, VS, Loc, Addr, Out, Code, By, SlotTy, CI> 
+where
+    I: VMInputT<VS, Loc, Addr, CI>,
+    VS: Default + VMStateT,
+    Addr: Serialize + DeserializeOwned + Debug + Clone,
+    Loc: Serialize + DeserializeOwned + Debug + Clone,
+    CI: Serialize + DeserializeOwned + Debug + Clone + ConciseSerde,
+{
+    fn name(&self) -> &str {
+        "ReentrancyFeedback"
+    }
+}
+
+// #[cfg(feature = "reentrancy")]
+impl<I, S, VS, Loc, Addr, Out, Code, By, SlotTy, CI> ReentrancyFeedback<I, S, VS, Loc, Addr, Out, Code, By, SlotTy, CI> 
+where
+    I: VMInputT<VS, Loc, Addr, CI>,
+    VS: Default + VMStateT,
+    Addr: Serialize + DeserializeOwned + Debug + Clone,
+    Loc: Serialize + DeserializeOwned + Debug + Clone,
+    CI: Serialize + DeserializeOwned + Debug + Clone + ConciseSerde,
+{
+    pub fn new(
+        written: bool, 
+        vm: Rc<RefCell<dyn GenericVM<VS, Code, By, Loc, Addr, SlotTy, Out, I, S, CI>>>) -> Self {
+        Self {
+            written,
+            vm,
+            phantom: Default::default(),
+        }
+    }
+}
+
+// #[cfg(feature = "reentrancy")]
+impl<I, S, VS, Loc, Addr, Out, Code, By, SlotTy, CI> Feedback<I, S> for ReentrancyFeedback<I, S, VS, Loc, Addr, Out, Code, By, SlotTy, CI>
+where
+    S: State + HasClientPerfMonitor + HasExecutionResult<Loc, Addr, VS, Out, CI>,
+    I: VMInputT<VS, Loc, Addr, CI>,
+    VS: Default + VMStateT,
+    Addr: Serialize + DeserializeOwned + Debug + Clone,
+    Loc: Serialize + DeserializeOwned + Debug + Clone,
+    Out: Default,
+    CI: Serialize + DeserializeOwned + Debug + Clone + ConciseSerde,
+{
+    fn init_state(&mut self, _state: &mut S) -> Result<(), Error> {
+        Ok(())
+    }
+
+    // simply returns true if written is true, i.e. there has been a write after a control leak.
+    fn is_interesting<EMI, OT>(
+        &mut self,
+        _state: &mut S,
+        _manager: &mut EMI,
+        _input: &I,
+        _observers: &OT,
+        _exit_kind: &ExitKind,
+    ) -> Result<bool, Error>
+    where
+        EMI: EventFirer<I>,
+        OT: ObserversTuple<I, S>,
+    {
+        self.written = false;
+
+        // state after executing. should have been executed before is_interesting was called.
+        let new_state = _state.get_execution_result().new_state.state;
+
+        // if there is a post execution context, there has been a new control leak.
+        while new_state.has_post_execution() {
+            // now execute again, checking for writes to EVM memory.
+            self.vm.deref().borrow().execute(_input, _state);
+            if self.written {
+                return Ok(true);
+            }
+        }  
+        return Ok(false);
+    }
+
+    fn append_metadata(
+        &mut self,
+        _state: &mut S,
+        _testcase: &mut Testcase<I>,
+    ) -> Result<(), Error> {
+        Ok(())
+    }
+
+    fn discard_metadata(&mut self, _state: &mut S, _input: &I) -> Result<(), Error> {
+        Ok(())
+    }
+}
+
+/////////////////////////////////////////////////////////////////////////////////////////////
+