Skip to content

Commit

Permalink
Browse files Browse the repository at this point in the history
- added `FuzzDictionaryConfig.max_calldata_fuzz_dictionary_addresses` option to specify how many random addresses to generate and to randomly select from when fuzzing calldata. If option is not specified then current behavior applies
- to narrow down number of runs / addresses involved in invariant test the `CalldataFuzzDictionaryConfig` is populated with random addresses plus all accounts from db (from `EvmFuzzState`)
- added `fuzz_calldata_with_config` fn that accepts `Option<CalldataFuzzDictionary>` as param. Non invariants tests use existing `fuzz_calldata` fn and pass None as config arg
  • Loading branch information
grandizzy committed Feb 26, 2024
1 parent 474591a commit 4991196
Show file tree
Hide file tree
Showing 7 changed files with 140 additions and 27 deletions.
4 changes: 4 additions & 0 deletions crates/config/src/fuzz.rs
Original file line number Diff line number Diff line change
Expand Up @@ -86,6 +86,9 @@ pub struct FuzzDictionaryConfig {
/// Once the fuzzer exceeds this limit, it will start evicting random entries
#[serde(deserialize_with = "crate::deserialize_usize_or_max")]
pub max_fuzz_dictionary_values: usize,
/// How many random addresses to use and to recycle when fuzzing calldata.
/// If not specified then `max_fuzz_dictionary_addresses` value applies.
pub max_calldata_fuzz_dictionary_addresses: Option<usize>,
}

impl Default for FuzzDictionaryConfig {
Expand All @@ -98,6 +101,7 @@ impl Default for FuzzDictionaryConfig {
max_fuzz_dictionary_addresses: (300 * 1024 * 1024) / 20,
// limit this to 200MB
max_fuzz_dictionary_values: (200 * 1024 * 1024) / 32,
max_calldata_fuzz_dictionary_addresses: None,
}
}
}
Expand Down
21 changes: 17 additions & 4 deletions crates/evm/evm/src/executors/invariant/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -33,13 +33,18 @@ use std::{cell::RefCell, collections::BTreeMap, sync::Arc};

mod error;
pub use error::{InvariantFailures, InvariantFuzzError, InvariantFuzzTestResult};
use foundry_evm_fuzz::strategies::{CalldataFuzzDictionary, CalldataFuzzDictionaryConfig};

mod funcs;
pub use funcs::{assert_invariants, replay_run};

/// Alias for (Dictionary for fuzzing, initial contracts to fuzz and an InvariantStrategy).
type InvariantPreparation =
(EvmFuzzState, FuzzRunIdentifiedContracts, BoxedStrategy<Vec<BasicTxDetails>>);
type InvariantPreparation = (
EvmFuzzState,
FuzzRunIdentifiedContracts,
BoxedStrategy<Vec<BasicTxDetails>>,
CalldataFuzzDictionary,
);

/// Enriched results of an invariant run check.
///
Expand Down Expand Up @@ -104,7 +109,8 @@ impl<'a> InvariantExecutor<'a> {
return Err(eyre!("Invariant test function should have no inputs"))
}

let (fuzz_state, targeted_contracts, strat) = self.prepare_fuzzing(&invariant_contract)?;
let (fuzz_state, targeted_contracts, strat, calldata_fuzz_dictionary) =
self.prepare_fuzzing(&invariant_contract)?;

// Stores the consumed gas and calldata of every successful fuzz call.
let fuzz_cases: RefCell<Vec<FuzzedCases>> = RefCell::new(Default::default());
Expand Down Expand Up @@ -239,6 +245,7 @@ impl<'a> InvariantExecutor<'a> {
Ok(())
});

trace!(target: "forge::test::invariant::calldata_address_fuzz_dictionary", "{:?}", calldata_fuzz_dictionary.clone().addresses);
trace!(target: "forge::test::invariant::dictionary", "{:?}", fuzz_state.read().values().iter().map(hex::encode).collect::<Vec<_>>());

let (reverts, error) = failures.into_inner().into_inner();
Expand Down Expand Up @@ -277,12 +284,17 @@ impl<'a> InvariantExecutor<'a> {
let targeted_contracts: FuzzRunIdentifiedContracts =
Arc::new(Mutex::new(targeted_contracts));

let calldata_fuzz_config: CalldataFuzzDictionary = Arc::new(
CalldataFuzzDictionaryConfig::new(&self.config.dictionary, fuzz_state.clone()),
);

// Creates the invariant strategy.
let strat = invariant_strat(
fuzz_state.clone(),
targeted_senders,
targeted_contracts.clone(),
self.config.dictionary.dictionary_weight,
calldata_fuzz_config.clone(),
)
.no_shrink()
.boxed();
Expand All @@ -300,6 +312,7 @@ impl<'a> InvariantExecutor<'a> {
fuzz_state.clone(),
targeted_contracts.clone(),
target_contract_ref.clone(),
calldata_fuzz_config.clone(),
),
target_contract_ref,
));
Expand All @@ -308,7 +321,7 @@ impl<'a> InvariantExecutor<'a> {
self.executor.inspector.fuzzer =
Some(Fuzzer { call_generator, fuzz_state: fuzz_state.clone(), collect: true });

Ok((fuzz_state, targeted_contracts, strat))
Ok((fuzz_state, targeted_contracts, strat, calldata_fuzz_config))
}

/// Fills the `InvariantExecutor` with the artifact identifier filters (in `path:name` string
Expand Down
52 changes: 49 additions & 3 deletions crates/evm/fuzz/src/strategies/calldata.rs
Original file line number Diff line number Diff line change
@@ -1,18 +1,64 @@
use super::fuzz_param;
use crate::strategies::{fuzz_param, EvmFuzzState};
use alloy_dyn_abi::JsonAbiExt;
use alloy_json_abi::Function;
use alloy_primitives::Bytes;
use alloy_primitives::{Address, Bytes};
use foundry_config::FuzzDictionaryConfig;
use hashbrown::HashSet;
use proptest::prelude::{BoxedStrategy, Strategy};
use std::{fmt, sync::Arc};

pub type CalldataFuzzDictionary = Arc<CalldataFuzzDictionaryConfig>;

#[derive(Clone)]
pub struct CalldataFuzzDictionaryConfig {
/// Addresses that can be used for fuzzing calldata.
pub addresses: Vec<Address>,
}

impl fmt::Debug for CalldataFuzzDictionaryConfig {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("CalldataFuzzDictionaryConfig").field("addresses", &self.addresses).finish()
}
}

impl CalldataFuzzDictionaryConfig {
pub fn new(config: &FuzzDictionaryConfig, state: EvmFuzzState) -> Self {
let mut addresses: HashSet<Address> = HashSet::new();

if let Some(len) = config.max_calldata_fuzz_dictionary_addresses {
loop {
if addresses.len() == len {
break
}
addresses.insert(Address::random());
}

// add any state address calldata fuzz dictionary, in addition to random generated
// addresses
let mut state = state.write();
addresses.extend(state.addresses());
}

Self { addresses: Vec::from_iter(addresses) }
}
}

/// Given a function, it returns a strategy which generates valid calldata
/// for that function's input types.
pub fn fuzz_calldata(func: Function) -> BoxedStrategy<Bytes> {
fuzz_calldata_with_config(func, None)
}

pub fn fuzz_calldata_with_config(
func: Function,
config: Option<CalldataFuzzDictionary>,
) -> BoxedStrategy<Bytes> {
// We need to compose all the strategies generated for each parameter in all
// possible combinations
let strats = func
.inputs
.iter()
.map(|input| fuzz_param(&input.selector_type().parse().unwrap()))
.map(|input| fuzz_param(&input.selector_type().parse().unwrap(), config.clone()))
.collect::<Vec<_>>();

strats
Expand Down
34 changes: 27 additions & 7 deletions crates/evm/fuzz/src/strategies/invariants.rs
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
use super::fuzz_param_from_state;
use super::{fuzz_calldata_with_config, fuzz_param_from_state, CalldataFuzzDictionary};
use crate::{
invariant::{BasicTxDetails, FuzzRunIdentifiedContracts, SenderFilters},
strategies::{fuzz_calldata, fuzz_calldata_from_state, fuzz_param, EvmFuzzState},
strategies::{fuzz_calldata_from_state, fuzz_param, EvmFuzzState},
};
use alloy_json_abi::{Function, JsonAbi};
use alloy_primitives::{Address, Bytes};
Expand All @@ -14,6 +14,7 @@ pub fn override_call_strat(
fuzz_state: EvmFuzzState,
contracts: FuzzRunIdentifiedContracts,
target: Arc<RwLock<Address>>,
calldata_fuzz_config: CalldataFuzzDictionary,
) -> SBoxedStrategy<(Address, Bytes)> {
let contracts_ref = contracts.clone();

Expand All @@ -27,10 +28,16 @@ pub fn override_call_strat(
])
.prop_flat_map(move |target_address| {
let fuzz_state = fuzz_state.clone();
let calldata_fuzz_config = calldata_fuzz_config.clone();
let (_, abi, functions) = contracts.lock().get(&target_address).unwrap().clone();
let func = select_random_function(abi, functions);
func.prop_flat_map(move |func| {
fuzz_contract_with_calldata(fuzz_state.clone(), target_address, func)
fuzz_contract_with_calldata(
fuzz_state.clone(),
calldata_fuzz_config.clone(),
target_address,
func,
)
})
})
.sboxed()
Expand All @@ -51,10 +58,12 @@ pub fn invariant_strat(
senders: SenderFilters,
contracts: FuzzRunIdentifiedContracts,
dictionary_weight: u32,
calldata_fuzz_config: CalldataFuzzDictionary,
) -> impl Strategy<Value = Vec<BasicTxDetails>> {
// We only want to seed the first value, since we want to generate the rest as we mutate the
// state
generate_call(fuzz_state, senders, contracts, dictionary_weight).prop_map(|x| vec![x])
generate_call(fuzz_state, senders, contracts, dictionary_weight, calldata_fuzz_config)
.prop_map(|x| vec![x])
}

/// Strategy to generate a transaction where the `sender`, `target` and `calldata` are all generated
Expand All @@ -64,6 +73,7 @@ fn generate_call(
senders: SenderFilters,
contracts: FuzzRunIdentifiedContracts,
dictionary_weight: u32,
calldata_fuzz_config: CalldataFuzzDictionary,
) -> BoxedStrategy<BasicTxDetails> {
let random_contract = select_random_contract(contracts);
let senders = Rc::new(senders);
Expand All @@ -72,10 +82,19 @@ fn generate_call(
let func = select_random_function(abi, functions);
let senders = senders.clone();
let fuzz_state = fuzz_state.clone();
let calldata_fuzz_config = calldata_fuzz_config.clone();
func.prop_flat_map(move |func| {
let sender =
select_random_sender(fuzz_state.clone(), senders.clone(), dictionary_weight);
(sender, fuzz_contract_with_calldata(fuzz_state.clone(), contract, func))
(
sender,
fuzz_contract_with_calldata(
fuzz_state.clone(),
calldata_fuzz_config.clone(),
contract,
func,
),
)
})
})
.boxed()
Expand All @@ -93,7 +112,7 @@ fn select_random_sender(
let fuzz_strategy = proptest::strategy::Union::new_weighted(vec![
(
100 - dictionary_weight,
fuzz_param(&alloy_dyn_abi::DynSolType::Address)
fuzz_param(&alloy_dyn_abi::DynSolType::Address, None)
.prop_map(move |addr| addr.as_address().unwrap())
.boxed(),
),
Expand Down Expand Up @@ -165,6 +184,7 @@ fn select_random_function(
/// for that function's input types.
pub fn fuzz_contract_with_calldata(
fuzz_state: EvmFuzzState,
calldata_fuzz_config: CalldataFuzzDictionary,
contract: Address,
func: Function,
) -> impl Strategy<Value = (Address, Bytes)> {
Expand All @@ -173,7 +193,7 @@ pub fn fuzz_contract_with_calldata(
// `prop_oneof!` / `TupleUnion` `Arc`s for cheap cloning
#[allow(clippy::arc_with_non_send_sync)]
let strats = prop_oneof![
60 => fuzz_calldata(func.clone()),
60 => fuzz_calldata_with_config(func.clone(), Some(calldata_fuzz_config)),
40 => fuzz_calldata_from_state(func, fuzz_state),
];
strats.prop_map(move |calldata| {
Expand Down
4 changes: 3 additions & 1 deletion crates/evm/fuzz/src/strategies/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,9 @@ mod param;
pub use param::{fuzz_param, fuzz_param_from_state};

mod calldata;
pub use calldata::fuzz_calldata;
pub use calldata::{
fuzz_calldata, fuzz_calldata_with_config, CalldataFuzzDictionary, CalldataFuzzDictionaryConfig,
};

mod state;
pub use state::{
Expand Down
50 changes: 38 additions & 12 deletions crates/evm/fuzz/src/strategies/param.rs
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
use super::state::EvmFuzzState;
use crate::strategies::calldata::CalldataFuzzDictionary;
use alloy_dyn_abi::{DynSolType, DynSolValue};
use alloy_primitives::{Address, FixedBytes, I256, U256};
use arbitrary::Unstructured;
Expand All @@ -10,12 +11,30 @@ const MAX_ARRAY_LEN: usize = 256;
/// Given a parameter type, returns a strategy for generating values for that type.
///
/// Works with ABI Encoder v2 tuples.
pub fn fuzz_param(param: &DynSolType) -> BoxedStrategy<DynSolValue> {
pub fn fuzz_param(
param: &DynSolType,
config: Option<CalldataFuzzDictionary>,
) -> BoxedStrategy<DynSolValue> {
let param = param.to_owned();
match param {
DynSolType::Address => any::<[u8; 32]>()
.prop_map(|x| DynSolValue::Address(Address::from_word(x.into())))
.boxed(),
DynSolType::Address => {
let cfg = config.clone();
if cfg.is_some() && !cfg.unwrap().addresses.is_empty() {
let dict_len = config.clone().unwrap().addresses.len();
any::<prop::sample::Index>()
.prop_map(move |index| index.index(dict_len))
.prop_map(move |index| {
DynSolValue::Address(
config.clone().unwrap().addresses.get(index).cloned().unwrap(),
)
})
.boxed()
} else {
any::<[u8; 32]>()
.prop_map(|x| DynSolValue::Address(Address::from_word(x.into())))
.boxed()
}
}
DynSolType::Int(n) => {
let strat = super::IntStrategy::new(n, vec![]);
let strat = strat.prop_map(move |x| DynSolValue::Int(x, n));
Expand Down Expand Up @@ -48,15 +67,22 @@ pub fn fuzz_param(param: &DynSolType) -> BoxedStrategy<DynSolValue> {
)
})
.boxed(),
DynSolType::Tuple(params) => {
params.iter().map(fuzz_param).collect::<Vec<_>>().prop_map(DynSolValue::Tuple).boxed()
}
DynSolType::FixedArray(param, size) => proptest::collection::vec(fuzz_param(&param), size)
.prop_map(DynSolValue::FixedArray)
.boxed(),
DynSolType::Array(param) => proptest::collection::vec(fuzz_param(&param), 0..MAX_ARRAY_LEN)
.prop_map(DynSolValue::Array)
DynSolType::Tuple(params) => params
.iter()
.map(|p| fuzz_param(p, config.clone()))
.collect::<Vec<_>>()
.prop_map(DynSolValue::Tuple)
.boxed(),
DynSolType::FixedArray(param, size) => {
proptest::collection::vec(fuzz_param(&param, config), size)
.prop_map(DynSolValue::FixedArray)
.boxed()
}
DynSolType::Array(param) => {
proptest::collection::vec(fuzz_param(&param, config), 0..MAX_ARRAY_LEN)
.prop_map(DynSolValue::Array)
.boxed()
}
DynSolType::CustomStruct { .. } => panic!("unsupported type"),
}
}
Expand Down
2 changes: 2 additions & 0 deletions crates/forge/tests/it/config.rs
Original file line number Diff line number Diff line change
Expand Up @@ -123,6 +123,7 @@ pub fn test_opts() -> TestOptions {
dictionary_weight: 40,
max_fuzz_dictionary_addresses: 10_000,
max_fuzz_dictionary_values: 10_000,
max_calldata_fuzz_dictionary_addresses: None,
},
})
.invariant(InvariantConfig {
Expand All @@ -136,6 +137,7 @@ pub fn test_opts() -> TestOptions {
include_push_bytes: true,
max_fuzz_dictionary_addresses: 10_000,
max_fuzz_dictionary_values: 10_000,
max_calldata_fuzz_dictionary_addresses: None,
},
shrink_sequence: true,
shrink_run_limit: 2usize.pow(18u32),
Expand Down

0 comments on commit 4991196

Please sign in to comment.