feat invariant (#5868) - configure calldata fuzzed addresses dictionary

crates/config/src/fuzz.rs

@@ -86,6 +86,10 @@ 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.
+    #[serde(deserialize_with = "crate::deserialize_usize_or_max")]
+    pub max_calldata_fuzz_dictionary_addresses: usize,
 }
 
 impl Default for FuzzDictionaryConfig {
@@ -98,6 +102,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: 0,
         }
     }
 }

crates/evm/evm/src/executors/invariant/mod.rs

@@ -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;
 
 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.
 ///
@@ -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());
@@ -245,6 +251,7 @@ impl<'a> InvariantExecutor<'a> {
             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::<Vec<_>>());
 
         let (reverts, error) = failures.into_inner().into_inner();
@@ -283,12 +290,16 @@ impl<'a> InvariantExecutor<'a> {
         let targeted_contracts: FuzzRunIdentifiedContracts =
             Arc::new(Mutex::new(targeted_contracts));
 
+        let calldata_fuzz_config =
+            CalldataFuzzDictionary::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();
@@ -306,6 +317,7 @@ impl<'a> InvariantExecutor<'a> {
                     fuzz_state.clone(),
                     targeted_contracts.clone(),
                     target_contract_ref.clone(),
+                    calldata_fuzz_config.clone(),
                 ),
                 target_contract_ref,
             ));
@@ -314,7 +326,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

crates/evm/fuzz/src/strategies/calldata.rs

@@ -1,18 +1,85 @@
-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};
+
+/// Clonable wrapper around [CalldataFuzzDictionary].
+#[derive(Debug, Clone)]
+pub struct CalldataFuzzDictionary {
+    pub inner: Arc<CalldataFuzzDictionaryConfig>,
+}
+
+impl CalldataFuzzDictionary {
+    pub fn new(config: &FuzzDictionaryConfig, state: EvmFuzzState) -> Self {
+        Self { inner: Arc::new(CalldataFuzzDictionaryConfig::new(config, state)) }
+    }
+}
+
+#[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()
+    }
+}
+
+/// Represents custom configuration for invariant fuzzed calldata strategies.
+///
+/// At the moment only the dictionary of addresses to be used for a fuzzed `function(address)` can
+/// be configured, but support for other types can be added.
+impl CalldataFuzzDictionaryConfig {
+    /// Creates config with the set of addresses that can be used for fuzzing invariant calldata (if
+    /// `max_calldata_fuzz_dictionary_addresses` configured).
+    /// The set of addresses contains a number of `max_calldata_fuzz_dictionary_addresses` random
+    /// addresses plus all addresses that already had their PUSH bytes collected (retrieved from
+    /// `EvmFuzzState`, if `include_push_bytes` config enabled).
+    pub fn new(config: &FuzzDictionaryConfig, state: EvmFuzzState) -> Self {
+        let mut addresses: HashSet<Address> = HashSet::new();
+        let dict_size = config.max_calldata_fuzz_dictionary_addresses;
+
+        if dict_size > 0 {
+            loop {
+                if addresses.len() == dict_size {
+                    break
+                }
+                addresses.insert(Address::random());
+            }
+
+            // Add all addresses that already had their PUSH bytes collected.
+            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)
+}
+
+/// Given a function, it returns a strategy which generates valid calldata
+/// for that function's input types, following custom configuration rules.
+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

crates/evm/fuzz/src/strategies/invariants.rs

@@ -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};
@@ -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();
 
@@ -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()
@@ -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
@@ -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);
@@ -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()
@@ -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(),
         ),
@@ -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)> {
@@ -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| {

crates/evm/fuzz/src/strategies/mod.rs

@@ -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::{

crates/evm/fuzz/src/strategies/param.rs

@@ -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;
@@ -10,12 +11,32 @@ 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 => {
+            if config.is_some() {
+                let fuzz_config = config.unwrap().inner;
+                let address_dict_len = fuzz_config.addresses.len();
+                if address_dict_len > 0 {
+                    // Create strategy to return random address from configured dictionary.
+                    return any::<prop::sample::Index>()
+                        .prop_map(move |index| index.index(address_dict_len))
+                        .prop_map(move |index| {
+                            DynSolValue::Address(fuzz_config.addresses.get(index).cloned().unwrap())
+                        })
+                        .boxed()
+                }
+            }
+
+            // If no config for addresses dictionary then create unbounded addresses strategy.
+            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));
@@ -48,15 +69,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"),
     }
 }