fix(evm): fuzzing not properly collecting data

evm/src/executor/inspector/fuzzer.rs

@@ -82,12 +82,13 @@ impl Fuzzer {
             state.insert(utils::u256_to_h256_be(*slot).into());
         }
 
-        for index in 0..interpreter.memory.len() / 32 {
-            let mut slot = [0u8; 32];
-            slot.clone_from_slice(interpreter.memory.get_slice(index * 32, 32));
+        // TODO: disabled for now since it's flooding the dictionary
+        // for index in 0..interpreter.memory.len() / 32 {
+        //     let mut slot = [0u8; 32];
+        //     slot.clone_from_slice(interpreter.memory.get_slice(index * 32, 32));
 
-            state.insert(slot);
-        }
+        //     state.insert(slot);
+        // }
     }
 
     /// Overrides an external call and tries to call any method of msg.sender.

evm/src/fuzz/invariant/executor.rs

@@ -23,6 +23,7 @@ use ethers::{
 };
 use eyre::ContextCompat;
 use foundry_common::contracts::{ContractsByAddress, ContractsByArtifact};
+use hashbrown::HashMap;
 use parking_lot::{Mutex, RwLock};
 use proptest::{
     strategy::{BoxedStrategy, Strategy, ValueTree},
@@ -142,12 +143,14 @@ impl<'a> InvariantExecutor<'a> {
                         .expect("could not make raw evm call");
 
                     // Collect data for fuzzing from the state changeset.
-                    let state_changeset =
+                    let mut state_changeset =
                         call_result.state_changeset.to_owned().expect("to have a state changeset.");
 
-                    collect_state_from_call(
-                        &call_result.logs,
-                        &state_changeset,
+                    collect_data(
+                        &mut state_changeset,
+                        sender,
+                        &invariant_contract,
+                        &call_result,
                         fuzz_state.clone(),
                     );
 
@@ -206,6 +209,8 @@ impl<'a> InvariantExecutor<'a> {
             });
         }
 
+        tracing::trace!(target: "forge::test::invariant::dictionary", "{:?}", fuzz_state.read().iter().map(hex::encode));
+
         let (reverts, invariants) = failures.into_inner().into_inner();
 
         Ok(Some(InvariantFuzzTestResult { invariants, cases: fuzz_cases.into_inner(), reverts }))
@@ -230,7 +235,8 @@ impl<'a> InvariantExecutor<'a> {
         }
 
         // Stores fuzz state for use with [fuzz_calldata_from_state].
-        let fuzz_state: EvmFuzzState = build_initial_state(self.executor.backend().mem_db());
+        let fuzz_state: EvmFuzzState =
+            build_initial_state(invariant_contract.address, self.executor.backend().mem_db());
 
         // During execution, any newly created contract is added here and used through the rest of
         // the fuzz run.
@@ -479,6 +485,42 @@ impl<'a> InvariantExecutor<'a> {
     }
 }
 
+/// Collects data from call for fuzzing. However, it first verifies that the sender is not an EOA
+/// before inserting it into the dictionary. Otherwise, we flood the dictionary with
+/// randomly generated addresses.
+fn collect_data(
+    state_changeset: &mut HashMap<Address, revm::Account>,
+    sender: &Address,
+    invariant_contract: &InvariantContract,
+    call_result: &RawCallResult,
+    fuzz_state: EvmFuzzState,
+) {
+    // Verify it has no code.
+    let mut has_code = false;
+    if let Some(Some(code)) = state_changeset.get(sender).map(|account| account.info.code.as_ref())
+    {
+        has_code = !code.is_empty();
+    }
+
+    // We keep the nonce changes to apply later.
+    let mut sender_changeset = None;
+    if !has_code {
+        sender_changeset = state_changeset.remove(sender);
+    }
+
+    collect_state_from_call(
+        invariant_contract.address,
+        &call_result.logs,
+        &*state_changeset,
+        fuzz_state,
+    );
+
+    // Re-add changes
+    if let Some(changed) = sender_changeset {
+        state_changeset.insert(*sender, changed);
+    }
+}
+
 /// Verifies that the invariant run execution can continue.
 fn can_continue(
     invariant_contract: &InvariantContract,

evm/src/fuzz/mod.rs

@@ -62,9 +62,9 @@ impl<'a> FuzzedExecutor<'a> {
 
         // Stores fuzz state for use with [fuzz_calldata_from_state]
         let state: EvmFuzzState = if let Some(fork_db) = self.executor.backend().active_fork_db() {
-            build_initial_state(fork_db)
+            build_initial_state(address, fork_db)
         } else {
-            build_initial_state(self.executor.backend().mem_db())
+            build_initial_state(address, self.executor.backend().mem_db())
         };
 
         // TODO: We should have a `FuzzerOpts` struct where we can configure the fuzzer. When we
@@ -83,7 +83,7 @@ impl<'a> FuzzedExecutor<'a> {
                 call.state_changeset.as_ref().expect("We should have a state changeset.");
 
             // Build fuzzer state
-            collect_state_from_call(&call.logs, state_changeset, state.clone());
+            collect_state_from_call(address, &call.logs, state_changeset, state.clone());
 
             // When assume cheat code is triggered return a special string "FOUNDRY::ASSUME"
             if call.result.as_ref() == ASSUME_MAGIC_RETURN_CODE {
@@ -123,6 +123,7 @@ impl<'a> FuzzedExecutor<'a> {
                 ))
             }
         });
+        tracing::trace!(target: "forge::test::fuzz::dictionary", "{:?}", state.read().iter().map(hex::encode));
 
         let (calldata, call) = counterexample.into_inner();
         let mut result = FuzzTestResult {

evm/src/fuzz/strategies/invariants.rs

@@ -13,6 +13,8 @@ use proptest::prelude::*;
 pub use proptest::test_runner::Config as FuzzConfig;
 use std::sync::Arc;
 
+use super::fuzz_param_from_state;
+
 /// Given a target address, we generate random calldata.
 pub fn override_call_strat(
     fuzz_state: EvmFuzzState,
@@ -75,7 +77,7 @@ fn generate_call(
             let senders = senders.clone();
             let fuzz_state = fuzz_state.clone();
             func.prop_flat_map(move |func| {
-                let sender = select_random_sender(senders.clone());
+                let sender = select_random_sender(fuzz_state.clone(), senders.clone());
                 (sender, fuzz_contract_with_calldata(fuzz_state.clone(), contract, func))
             })
         })
@@ -86,9 +88,25 @@ fn generate_call(
 /// * If `senders` is empty, then it's a completely random address.
 /// * If `senders` is not empty, then there's an 80% chance that one from the list is selected. The
 ///   remaining 20% will be random.
-fn select_random_sender(senders: Vec<Address>) -> impl Strategy<Value = Address> {
-    let fuzz_strategy =
-        fuzz_param(&ParamType::Address).prop_map(move |addr| addr.into_address().unwrap()).boxed();
+fn select_random_sender(
+    fuzz_state: EvmFuzzState,
+    senders: Vec<Address>,
+) -> impl Strategy<Value = Address> {
+    let fuzz_strategy = proptest::strategy::Union::new_weighted(vec![
+        (
+            10,
+            fuzz_param(&ParamType::Address)
+                .prop_map(move |addr| addr.into_address().unwrap())
+                .boxed(),
+        ),
+        (
+            90,
+            fuzz_param_from_state(&ParamType::Address, fuzz_state)
+                .prop_map(move |addr| addr.into_address().unwrap())
+                .boxed(),
+        ),
+    ])
+    .boxed();
 
     if !senders.is_empty() {
         let selector =

evm/src/fuzz/strategies/param.rs

@@ -99,7 +99,11 @@ pub fn fuzz_param_from_state(param: &ParamType, arc_state: EvmFuzzState) -> Boxe
         },
         ParamType::Bool => value.prop_map(move |value| Token::Bool(value[31] == 1)).boxed(),
         ParamType::String => value
-            .prop_map(move |value| Token::String(String::from_utf8_lossy(&value[..]).to_string()))
+            .prop_map(move |value| {
+                Token::String(
+                    String::from_utf8_lossy(&value[..]).trim().trim_end_matches('\0').to_string(),
+                )
+            })
             .boxed(),
         ParamType::Array(param) => proptest::collection::vec(
             fuzz_param_from_state(param, arc_state.clone()),
@@ -128,7 +132,10 @@ pub fn fuzz_param_from_state(param: &ParamType, arc_state: EvmFuzzState) -> Boxe
 
 #[cfg(test)]
 mod tests {
-    use crate::fuzz::strategies::{build_initial_state, fuzz_calldata, fuzz_calldata_from_state};
+    use crate::{
+        fuzz::strategies::{build_initial_state, fuzz_calldata, fuzz_calldata_from_state},
+        CALLER,
+    };
     use ethers::abi::HumanReadableParser;
     use revm::db::{CacheDB, EmptyDB};
 
@@ -138,7 +145,7 @@ mod tests {
         let func = HumanReadableParser::parse_function(f).unwrap();
 
         let db = CacheDB::new(EmptyDB());
-        let state = build_initial_state(&db);
+        let state = build_initial_state(CALLER, &db);
 
         let strat = proptest::strategy::Union::new_weighted(vec![
             (60, fuzz_calldata(func.clone())),

evm/src/fuzz/strategies/state.rs

@@ -53,20 +53,35 @@ This is a bug, please open an issue: https://github.com/foundry-rs/foundry/issue
 }
 
 /// Builds the initial [EvmFuzzState] from a database.
-pub fn build_initial_state<DB: DatabaseRef>(db: &CacheDB<DB>) -> EvmFuzzState {
+pub fn build_initial_state<DB: DatabaseRef>(
+    test_address: Address,
+    db: &CacheDB<DB>,
+) -> EvmFuzzState {
     let mut state: BTreeSet<[u8; 32]> = BTreeSet::new();
     for (address, account) in db.accounts.iter() {
+        // We don't want to collect data from the test contract.
+        if *address == test_address {
+            continue
+        }
+
         let info = db.basic(*address);
 
         // Insert basic account information
         state.insert(H256::from(*address).into());
-        state.insert(utils::u256_to_h256_le(info.balance).into());
-        state.insert(utils::u256_to_h256_le(U256::from(info.nonce)).into());
+        state.insert(utils::u256_to_h256_be(info.balance).into());
+        state.insert(utils::u256_to_h256_be(U256::from(info.nonce)).into());
 
         // Insert storage
         for (slot, value) in &account.storage {
-            state.insert(utils::u256_to_h256_le(*slot).into());
-            state.insert(utils::u256_to_h256_le(*value).into());
+            state.insert(utils::u256_to_h256_be(*slot).into());
+            state.insert(utils::u256_to_h256_be(*value).into());
+        }
+
+        // Insert push bytes
+        if let Some(code) = &account.info.code {
+            for push_byte in collect_push_bytes(code.bytes().clone()) {
+                state.insert(push_byte);
+            }
         }
     }
 
@@ -82,22 +97,28 @@ pub fn build_initial_state<DB: DatabaseRef>(db: &CacheDB<DB>) -> EvmFuzzState {
 
 /// Collects state changes from a [StateChangeset] and logs into an [EvmFuzzState].
 pub fn collect_state_from_call(
+    test_address: Address,
     logs: &[Log],
     state_changeset: &StateChangeset,
     state: EvmFuzzState,
 ) {
     let mut state = state.write();
 
     for (address, account) in state_changeset {
+        // We don't want to collect data from the test contract.
+        if *address == test_address {
+            continue
+        }
+
         // Insert basic account information
         state.insert(H256::from(*address).into());
-        state.insert(utils::u256_to_h256_le(account.info.balance).into());
-        state.insert(utils::u256_to_h256_le(U256::from(account.info.nonce)).into());
+        state.insert(utils::u256_to_h256_be(account.info.balance).into());
+        state.insert(utils::u256_to_h256_be(U256::from(account.info.nonce)).into());
 
         // Insert storage
         for (slot, value) in &account.storage {
-            state.insert(utils::u256_to_h256_le(*slot).into());
-            state.insert(utils::u256_to_h256_le(*value).into());
+            state.insert(utils::u256_to_h256_be(*slot).into());
+            state.insert(utils::u256_to_h256_be(*value).into());
         }
 
         // Insert push bytes
@@ -151,11 +172,8 @@ fn collect_push_bytes(code: Bytes) -> Vec<[u8; 32]> {
                 return bytes
             }
 
-            let mut buffer: [u8; 32] = [0; 32];
-            let _ = (&mut buffer[..])
-                .write(&code[push_start..push_end])
-                .expect("push was larger than 32 bytes");
-            bytes.push(buffer);
+            bytes.push(U256::from_big_endian(&code[push_start..push_end]).into());
+
             i += push_size;
         }
         i += 1;

forge/src/lib.rs

@@ -51,11 +51,19 @@ pub struct TestOptions {
 }
 
 impl TestOptions {
+    pub fn invariant_fuzzer(&self) -> TestRunner {
+        self.fuzzer_with_cases(self.invariant_runs)
+    }
+
     pub fn fuzzer(&self) -> TestRunner {
+        self.fuzzer_with_cases(self.fuzz_runs)
+    }
+
+    pub fn fuzzer_with_cases(&self, cases: u32) -> TestRunner {
         // TODO: Add Options to modify the persistence
         let cfg = proptest::test_runner::Config {
             failure_persistence: None,
-            cases: self.fuzz_runs,
+            cases,
             max_local_rejects: self.fuzz_max_local_rejects,
             max_global_rejects: self.fuzz_max_global_rejects,
             ..Default::default()

forge/src/runner.rs

@@ -302,7 +302,7 @@ impl<'a> ContractRunner<'a> {
                 .collect();
 
             let results = self.run_invariant_test(
-                test_options.fuzzer(),
+                test_options.invariant_fuzzer(),
                 setup,
                 test_options,
                 functions.clone(),

forge/tests/it/fuzz.rs

@@ -2,8 +2,8 @@
 
 use crate::{config::*, test_helpers::filter::Filter};
 use forge::result::SuiteResult;
-
 use foundry_evm::decode::decode_console_logs;
+use std::collections::BTreeMap;
 
 #[test]
 fn test_fuzz() {
@@ -40,3 +40,34 @@ fn test_fuzz() {
         }
     }
 }
+
+#[test]
+fn test_fuzz_collection() {
+    let mut runner = runner();
+
+    let mut opts = TEST_OPTS;
+    opts.invariant_depth = 200;
+    opts.fuzz_runs = 1000;
+    runner.test_options = opts;
+
+    let results =
+        runner.test(&Filter::new(".*", ".*", ".*fuzz/FuzzCollection.t.sol"), None, opts).unwrap();
+
+    assert_multiple(
+        &results,
+        BTreeMap::from([(
+            "fuzz/FuzzCollection.t.sol:SampleContractTest",
+            vec![
+                ("invariantCounter", false, Some("broken counter.".into()), None, None),
+                (
+                    "testIncrement(address)",
+                    false,
+                    Some("Call did not revert as expected".into()),
+                    None,
+                    None,
+                ),
+                ("testNeedle(uint256)", false, Some("needle found.".into()), None, None),
+            ],
+        )]),
+    );
+}