From 398ff1fd09e6ec67129ea9b3f52f0fc3c50dd550 Mon Sep 17 00:00:00 2001 From: Arsenii Kulikov Date: Thu, 4 Apr 2024 22:43:20 +0400 Subject: [PATCH] fix: do not flood dictionary with data dependent on fuzz inputs (#7552) * fix dictionary * clippy + fmt * fix --- crates/evm/evm/src/executors/fuzz/mod.rs | 15 +- crates/evm/evm/src/executors/invariant/mod.rs | 24 +- crates/evm/fuzz/src/inspector.rs | 7 +- crates/evm/fuzz/src/strategies/calldata.rs | 2 +- crates/evm/fuzz/src/strategies/mod.rs | 3 +- crates/evm/fuzz/src/strategies/param.rs | 4 +- crates/evm/fuzz/src/strategies/state.rs | 230 +++++++++++------- crates/forge/src/runner.rs | 14 +- 8 files changed, 170 insertions(+), 129 deletions(-) diff --git a/crates/evm/evm/src/executors/fuzz/mod.rs b/crates/evm/evm/src/executors/fuzz/mod.rs index 41beb7e43317..ae63d2386bae 100644 --- a/crates/evm/evm/src/executors/fuzz/mod.rs +++ b/crates/evm/evm/src/executors/fuzz/mod.rs @@ -10,10 +10,7 @@ use foundry_evm_core::{ }; use foundry_evm_coverage::HitMaps; use foundry_evm_fuzz::{ - strategies::{ - build_initial_state, collect_state_from_call, fuzz_calldata, fuzz_calldata_from_state, - EvmFuzzState, - }, + strategies::{build_initial_state, fuzz_calldata, fuzz_calldata_from_state, EvmFuzzState}, BaseCounterExample, CounterExample, FuzzCase, FuzzError, FuzzTestResult, }; use foundry_evm_traces::CallTraceArena; @@ -89,7 +86,7 @@ impl FuzzedExecutor { ]; debug!(func=?func.name, should_fail, "fuzzing"); let run_result = self.runner.clone().run(&strat, |calldata| { - let fuzz_res = self.single_fuzz(&state, address, should_fail, calldata)?; + let fuzz_res = self.single_fuzz(address, should_fail, calldata)?; match fuzz_res { FuzzOutcome::Case(case) => { @@ -195,7 +192,6 @@ impl FuzzedExecutor { /// or a `CounterExampleOutcome` pub fn single_fuzz( &self, - state: &EvmFuzzState, address: Address, should_fail: bool, calldata: alloy_primitives::Bytes, @@ -206,9 +202,6 @@ impl FuzzedExecutor { .map_err(|_| TestCaseError::fail(FuzzError::FailedContractCall))?; let state_changeset = call.state_changeset.take().unwrap(); - // Build fuzzer state - collect_state_from_call(&call.logs, &state_changeset, state, &self.config.dictionary); - // When the `assume` cheatcode is called it returns a special string if call.result.as_ref() == MAGIC_ASSUME { return Err(TestCaseError::reject(FuzzError::AssumeReject)) @@ -247,9 +240,9 @@ impl FuzzedExecutor { /// Stores fuzz state for use with [fuzz_calldata_from_state] pub fn build_fuzz_state(&self) -> EvmFuzzState { if let Some(fork_db) = self.executor.backend.active_fork_db() { - build_initial_state(fork_db, &self.config.dictionary) + build_initial_state(fork_db, self.config.dictionary) } else { - build_initial_state(self.executor.backend.mem_db(), &self.config.dictionary) + build_initial_state(self.executor.backend.mem_db(), self.config.dictionary) } } } diff --git a/crates/evm/evm/src/executors/invariant/mod.rs b/crates/evm/evm/src/executors/invariant/mod.rs index 5ed7998ac3e1..835a2bc877bf 100644 --- a/crates/evm/evm/src/executors/invariant/mod.rs +++ b/crates/evm/evm/src/executors/invariant/mod.rs @@ -6,7 +6,7 @@ use alloy_primitives::{Address, FixedBytes, U256}; use alloy_sol_types::{sol, SolCall}; use eyre::{eyre, ContextCompat, Result}; use foundry_common::contracts::{ContractsByAddress, ContractsByArtifact}; -use foundry_config::{FuzzDictionaryConfig, InvariantConfig}; +use foundry_config::InvariantConfig; use foundry_evm_core::{ constants::{CALLER, CHEATCODE_ADDRESS, HARDHAT_CONSOLE_ADDRESS, MAGIC_ASSUME}, utils::{get_function, StateChangeset}, @@ -17,8 +17,8 @@ use foundry_evm_fuzz::{ RandomCallGenerator, SenderFilters, TargetedContracts, }, strategies::{ - build_initial_state, collect_created_contracts, collect_state_from_call, invariant_strat, - override_call_strat, CalldataFuzzDictionary, EvmFuzzState, + build_initial_state, collect_created_contracts, invariant_strat, override_call_strat, + CalldataFuzzDictionary, EvmFuzzState, }, FuzzCase, FuzzedCases, }; @@ -246,13 +246,7 @@ impl<'a> InvariantExecutor<'a> { let mut state_changeset = call_result.state_changeset.to_owned().expect("no changesets"); - collect_data( - &mut state_changeset, - sender, - &call_result, - &fuzz_state, - &self.config.dictionary, - ); + collect_data(&mut state_changeset, sender, &call_result, &fuzz_state); if let Err(error) = collect_created_contracts( &state_changeset, @@ -325,11 +319,14 @@ impl<'a> InvariantExecutor<'a> { } fuzz_cases.borrow_mut().push(FuzzedCases::new(fuzz_runs)); + // Revert state to not persist values between runs. + fuzz_state.revert(); + Ok(()) }); trace!(target: "forge::test::invariant::calldata_address_fuzz_dictionary", "{:?}", calldata_fuzz_dictionary.inner.addresses); - trace!(target: "forge::test::invariant::dictionary", "{:?}", fuzz_state.read().values().iter().map(hex::encode).collect::>()); + trace!(target: "forge::test::invariant::dictionary", "{:?}", fuzz_state.dictionary_read().values().iter().map(hex::encode).collect::>()); let (reverts, error) = failures.into_inner().into_inner(); @@ -361,7 +358,7 @@ 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(), &self.config.dictionary); + build_initial_state(self.executor.backend.mem_db(), self.config.dictionary); // During execution, any newly created contract is added here and used through the rest of // the fuzz run. @@ -668,7 +665,6 @@ fn collect_data( sender: &Address, call_result: &RawCallResult, fuzz_state: &EvmFuzzState, - config: &FuzzDictionaryConfig, ) { // Verify it has no code. let mut has_code = false; @@ -683,7 +679,7 @@ fn collect_data( sender_changeset = state_changeset.remove(sender); } - collect_state_from_call(&call_result.logs, &*state_changeset, fuzz_state, config); + fuzz_state.collect_state_from_call(&call_result.logs, &*state_changeset); // Re-add changes if let Some(changed) = sender_changeset { diff --git a/crates/evm/fuzz/src/inspector.rs b/crates/evm/fuzz/src/inspector.rs index e58afb214fee..5e7e644b11d6 100644 --- a/crates/evm/fuzz/src/inspector.rs +++ b/crates/evm/fuzz/src/inspector.rs @@ -1,4 +1,5 @@ use crate::{invariant::RandomCallGenerator, strategies::EvmFuzzState}; +use alloy_primitives::U256; use revm::{ interpreter::{CallInputs, CallOutcome, CallScheme, Interpreter}, Database, EvmContext, Inspector, @@ -61,11 +62,7 @@ impl Inspector for Fuzzer { impl Fuzzer { /// Collects `stack` and `memory` values into the fuzz dictionary. fn collect_data(&mut self, interpreter: &Interpreter) { - let mut state = self.fuzz_state.write(); - - for slot in interpreter.stack().data() { - state.values_mut().insert(slot.to_be_bytes()); - } + self.fuzz_state.collect_values(interpreter.stack().data().iter().map(U256::to_be_bytes)); // TODO: disabled for now since it's flooding the dictionary // for index in 0..interpreter.shared_memory.len() / 32 { diff --git a/crates/evm/fuzz/src/strategies/calldata.rs b/crates/evm/fuzz/src/strategies/calldata.rs index b3d41bfadbbd..ff3bb5713446 100644 --- a/crates/evm/fuzz/src/strategies/calldata.rs +++ b/crates/evm/fuzz/src/strategies/calldata.rs @@ -41,7 +41,7 @@ impl CalldataFuzzDictionaryConfig { if dict_size > 0 { addresses.extend(std::iter::repeat_with(Address::random).take(dict_size)); // Add all addresses that already had their PUSH bytes collected. - addresses.extend(state.read().addresses()); + addresses.extend(state.dictionary_read().addresses()); } Self { addresses: addresses.into_iter().collect() } diff --git a/crates/evm/fuzz/src/strategies/mod.rs b/crates/evm/fuzz/src/strategies/mod.rs index 63e008ec0c96..0e82a4d4b8d6 100644 --- a/crates/evm/fuzz/src/strategies/mod.rs +++ b/crates/evm/fuzz/src/strategies/mod.rs @@ -14,8 +14,7 @@ pub use calldata::{ mod state; pub use state::{ - build_initial_state, collect_created_contracts, collect_state_from_call, - fuzz_calldata_from_state, EvmFuzzState, + build_initial_state, collect_created_contracts, fuzz_calldata_from_state, EvmFuzzState, }; mod invariants; diff --git a/crates/evm/fuzz/src/strategies/param.rs b/crates/evm/fuzz/src/strategies/param.rs index f287c75dfde1..20e69a27e725 100644 --- a/crates/evm/fuzz/src/strategies/param.rs +++ b/crates/evm/fuzz/src/strategies/param.rs @@ -89,7 +89,7 @@ pub fn fuzz_param_from_state( let state = state.clone(); // Use `Index` instead of `Selector` to not iterate over the entire dictionary. any::().prop_map(move |index| { - let state = state.read(); + let state = state.dictionary_read(); let values = state.values(); let index = index.index(values.len()); *values.iter().nth(index).unwrap() @@ -184,7 +184,7 @@ mod tests { let f = "testArray(uint64[2] calldata values)"; let func = get_func(f).unwrap(); let db = CacheDB::new(EmptyDB::default()); - let state = build_initial_state(&db, &FuzzDictionaryConfig::default()); + let state = build_initial_state(&db, FuzzDictionaryConfig::default()); let strat = proptest::prop_oneof![ 60 => fuzz_calldata(func.clone()), 40 => fuzz_calldata_from_state(func, &state), diff --git a/crates/evm/fuzz/src/strategies/state.rs b/crates/evm/fuzz/src/strategies/state.rs index a8965a884557..8c98c472e1ab 100644 --- a/crates/evm/fuzz/src/strategies/state.rs +++ b/crates/evm/fuzz/src/strategies/state.rs @@ -7,7 +7,7 @@ use foundry_common::contracts::{ContractsByAddress, ContractsByArtifact}; use foundry_config::FuzzDictionaryConfig; use foundry_evm_core::utils::StateChangeset; use indexmap::IndexSet; -use parking_lot::RwLock; +use parking_lot::{lock_api::RwLockReadGuard, RawRwLock, RwLock}; use proptest::prelude::{BoxedStrategy, Strategy}; use revm::{ db::{CacheDB, DatabaseRef}, @@ -19,16 +19,103 @@ use std::{fmt, sync::Arc}; /// A set of arbitrary 32 byte data from the VM used to generate values for the strategy. /// /// Wrapped in a shareable container. -pub type EvmFuzzState = Arc>; +#[derive(Clone, Debug)] +pub struct EvmFuzzState { + inner: Arc>, +} + +impl EvmFuzzState { + pub fn new(dictionary: FuzzDictionary) -> Self { + Self { inner: Arc::new(RwLock::new(dictionary)) } + } + + pub fn collect_values(&self, values: impl IntoIterator) { + let mut dict = self.inner.write(); + for value in values { + dict.insert_value(value); + } + } + + /// Collects state changes from a [StateChangeset] and logs into an [EvmFuzzState] according to + /// the given [FuzzDictionaryConfig]. + pub fn collect_state_from_call(&self, logs: &[Log], state_changeset: &StateChangeset) { + let mut dict = self.inner.write(); + + // Insert log topics and data. + for log in logs { + for topic in log.topics() { + dict.insert_value(topic.0); + } + let chunks = log.data.data.chunks_exact(32); + let rem = chunks.remainder(); + for chunk in chunks { + dict.insert_value(chunk.try_into().unwrap()); + } + if !rem.is_empty() { + dict.insert_value(B256::right_padding_from(rem).0); + } + } + + for (address, account) in state_changeset { + // Insert basic account information + dict.insert_value(address.into_word().into()); + + if dict.config.include_push_bytes { + // Insert push bytes + if let Some(code) = &account.info.code { + dict.insert_address(*address); + for push_byte in collect_push_bytes(code.bytes()) { + dict.insert_value(push_byte); + } + } + } + + if dict.config.include_storage { + // Insert storage + for (slot, value) in &account.storage { + let value = value.present_value; + dict.insert_value(B256::from(*slot).0); + dict.insert_value(B256::from(value).0); + // also add the value below and above the storage value to the dictionary. + if value != U256::ZERO { + let below_value = value - U256::from(1); + dict.insert_value(B256::from(below_value).0); + } + if value != U256::MAX { + let above_value = value + U256::from(1); + dict.insert_value(B256::from(above_value).0); + } + } + } + } + } + + /// Removes all newly added entries from the dictionary. + /// + /// Should be called between fuzz/invariant runs to avoid accumumlating data derived from fuzz + /// inputs. + pub fn revert(&self) { + self.inner.write().revert(); + } + + pub fn dictionary_read(&self) -> RwLockReadGuard<'_, RawRwLock, FuzzDictionary> { + self.inner.read() + } +} // We're using `IndexSet` to have a stable element order when restoring persisted state, as well as // for performance when iterating over the sets. -#[derive(Default)] pub struct FuzzDictionary { /// Collected state values. state_values: IndexSet<[u8; 32]>, /// Addresses that already had their PUSH bytes collected. addresses: IndexSet
, + /// Configuration for the dictionary. + config: FuzzDictionaryConfig, + /// New keys added to the dictionary since container initialization. + new_values: IndexSet<[u8; 32]>, + /// New addresses added to the dictionary since container initialization. + new_addreses: IndexSet
, } impl fmt::Debug for FuzzDictionary { @@ -41,14 +128,39 @@ impl fmt::Debug for FuzzDictionary { } impl FuzzDictionary { - #[inline] - pub fn values(&self) -> &IndexSet<[u8; 32]> { - &self.state_values + pub fn new( + initial_values: IndexSet<[u8; 32]>, + initial_addresses: IndexSet
, + config: FuzzDictionaryConfig, + ) -> Self { + Self { + state_values: initial_values, + addresses: initial_addresses, + config, + new_values: IndexSet::new(), + new_addreses: IndexSet::new(), + } + } + + pub fn insert_value(&mut self, value: [u8; 32]) { + if self.state_values.len() < self.config.max_fuzz_dictionary_values && + self.state_values.insert(value) + { + self.new_values.insert(value); + } + } + + pub fn insert_address(&mut self, address: Address) { + if self.addresses.len() < self.config.max_fuzz_dictionary_addresses && + self.addresses.insert(address) + { + self.new_addreses.insert(address); + } } #[inline] - pub fn values_mut(&mut self) -> &mut IndexSet<[u8; 32]> { - &mut self.state_values + pub fn values(&self) -> &IndexSet<[u8; 32]> { + &self.state_values } #[inline] @@ -56,9 +168,16 @@ impl FuzzDictionary { &self.addresses } - #[inline] - pub fn addresses_mut(&mut self) -> &mut IndexSet
{ - &mut self.addresses + pub fn revert(&mut self) { + for key in self.new_values.iter() { + self.state_values.swap_remove(key); + } + for address in self.new_addreses.iter() { + self.addresses.swap_remove(address); + } + + self.new_values.clear(); + self.new_addreses.clear(); } } @@ -88,9 +207,10 @@ pub fn fuzz_calldata_from_state(func: Function, state: &EvmFuzzState) -> BoxedSt /// Builds the initial [EvmFuzzState] from a database. pub fn build_initial_state( db: &CacheDB, - config: &FuzzDictionaryConfig, + config: FuzzDictionaryConfig, ) -> EvmFuzzState { - let mut state = FuzzDictionary::default(); + let mut values = IndexSet::new(); + let mut addresses = IndexSet::new(); // Sort accounts to ensure deterministic dictionary generation from the same setUp state. let mut accs = db.accounts.iter().collect::>(); @@ -99,15 +219,14 @@ pub fn build_initial_state( for (address, account) in accs { let address: Address = *address; // Insert basic account information - state.values_mut().insert(address.into_word().into()); + values.insert(address.into_word().into()); // Insert push bytes if config.include_push_bytes { if let Some(code) = &account.info.code { - if state.addresses_mut().insert(address) { - for push_byte in collect_push_bytes(code.bytes()) { - state.values_mut().insert(push_byte); - } + addresses.insert(address); + for push_byte in collect_push_bytes(code.bytes()) { + values.insert(push_byte); } } } @@ -115,16 +234,16 @@ pub fn build_initial_state( if config.include_storage { // Insert storage for (slot, value) in &account.storage { - state.values_mut().insert(B256::from(*slot).0); - state.values_mut().insert(B256::from(*value).0); + values.insert(B256::from(*slot).0); + values.insert(B256::from(*value).0); // also add the value below and above the storage value to the dictionary. if *value != U256::ZERO { let below_value = value - U256::from(1); - state.values_mut().insert(B256::from(below_value).0); + values.insert(B256::from(below_value).0); } if *value != U256::MAX { let above_value = value + U256::from(1); - state.values_mut().insert(B256::from(above_value).0); + values.insert(B256::from(above_value).0); } } } @@ -132,73 +251,12 @@ pub fn build_initial_state( // need at least some state data if db is empty otherwise we can't select random data for state // fuzzing - if state.values().is_empty() { + if values.is_empty() { // prefill with a random addresses - state.values_mut().insert(Address::random().into_word().into()); + values.insert(Address::random().into_word().into()); } - Arc::new(RwLock::new(state)) -} - -/// Collects state changes from a [StateChangeset] and logs into an [EvmFuzzState] according to the -/// given [FuzzDictionaryConfig]. -pub fn collect_state_from_call( - logs: &[Log], - state_changeset: &StateChangeset, - state: &EvmFuzzState, - config: &FuzzDictionaryConfig, -) { - let mut state = state.write(); - - // Insert log topics and data. - for log in logs { - for topic in log.topics() { - state.values_mut().insert(topic.0); - } - let chunks = log.data.data.chunks_exact(32); - let rem = chunks.remainder(); - for chunk in chunks { - state.values_mut().insert(chunk.try_into().unwrap()); - } - if !rem.is_empty() { - state.values_mut().insert(B256::right_padding_from(rem).0); - } - } - - for (address, account) in state_changeset { - // Insert basic account information - state.values_mut().insert(address.into_word().into()); - - if config.include_push_bytes && state.addresses.len() < config.max_fuzz_dictionary_addresses - { - // Insert push bytes - if let Some(code) = &account.info.code { - if state.addresses_mut().insert(*address) { - for push_byte in collect_push_bytes(code.bytes()) { - state.values_mut().insert(push_byte); - } - } - } - } - - if config.include_storage && state.state_values.len() < config.max_fuzz_dictionary_values { - // Insert storage - for (slot, value) in &account.storage { - let value = value.present_value; - state.values_mut().insert(B256::from(*slot).0); - state.values_mut().insert(B256::from(value).0); - // also add the value below and above the storage value to the dictionary. - if value != U256::ZERO { - let below_value = value - U256::from(1); - state.values_mut().insert(B256::from(below_value).0); - } - if value != U256::MAX { - let above_value = value + U256::from(1); - state.values_mut().insert(B256::from(above_value).0); - } - } - } - } + EvmFuzzState::new(FuzzDictionary::new(values, addresses, config)) } /// The maximum number of bytes we will look at in bytecodes to find push bytes (24 KiB). diff --git a/crates/forge/src/runner.rs b/crates/forge/src/runner.rs index 7f3258e359fa..508ee9234423 100644 --- a/crates/forge/src/runner.rs +++ b/crates/forge/src/runner.rs @@ -611,7 +611,6 @@ impl<'a> ContractRunner<'a> { self.sender, fuzz_config.clone(), ); - let state = fuzzed_executor.build_fuzz_state(); let result = fuzzed_executor.fuzz(func, address, should_fail, self.revert_decoder); let mut debug = Default::default(); @@ -650,13 +649,12 @@ impl<'a> ContractRunner<'a> { result.first_case.calldata.clone() }; // rerun the last relevant test with traces - let debug_result = FuzzedExecutor::new( - debug_executor, - runner, - self.sender, - fuzz_config, - ) - .single_fuzz(&state, address, should_fail, calldata); + let debug_result = + FuzzedExecutor::new(debug_executor, runner, self.sender, fuzz_config).single_fuzz( + address, + should_fail, + calldata, + ); (debug, breakpoints) = match debug_result { Ok(fuzz_outcome) => match fuzz_outcome {