foundry_evm_fuzz/invariant/
mod.rs

Help
1use alloy_json_abi::{Function, JsonAbi};
2use alloy_primitives::{Address, Bytes, Selector};
3use itertools::Either;
4use parking_lot::Mutex;
5use std::{collections::BTreeMap, sync::Arc};
6
7mod call_override;
8pub use call_override::RandomCallGenerator;
9
10mod filters;
11pub use filters::{ArtifactFilters, SenderFilters};
12use foundry_common::{ContractsByAddress, ContractsByArtifact};
13use foundry_evm_core::utils::{get_function, StateChangeset};
14
15/// Contracts identified as targets during a fuzz run.
16///
17/// During execution, any newly created contract is added as target and used through the rest of
18/// the fuzz run if the collection is updatable (no `targetContract` specified in `setUp`).
19#[derive(Clone, Debug)]
20pub struct FuzzRunIdentifiedContracts {
21    /// Contracts identified as targets during a fuzz run.
22    pub targets: Arc<Mutex<TargetedContracts>>,
23    /// Whether target contracts are updatable or not.
24    pub is_updatable: bool,
25}
26
27impl FuzzRunIdentifiedContracts {
28    /// Creates a new `FuzzRunIdentifiedContracts` instance.
29    pub fn new(targets: TargetedContracts, is_updatable: bool) -> Self {
30        Self { targets: Arc::new(Mutex::new(targets)), is_updatable }
31    }
32
33    /// If targets are updatable, collect all contracts created during an invariant run (which
34    /// haven't been discovered yet).
35    pub fn collect_created_contracts(
36        &self,
37        state_changeset: &StateChangeset,
38        project_contracts: &ContractsByArtifact,
39        setup_contracts: &ContractsByAddress,
40        artifact_filters: &ArtifactFilters,
41        created_contracts: &mut Vec<Address>,
42    ) -> eyre::Result<()> {
43        if !self.is_updatable {
44            return Ok(());
45        }
46
47        let mut targets = self.targets.lock();
48        for (address, account) in state_changeset {
49            if setup_contracts.contains_key(address) {
50                continue;
51            }
52            if !account.is_touched() {
53                continue;
54            }
55            let Some(code) = &account.info.code else {
56                continue;
57            };
58            if code.is_empty() {
59                continue;
60            }
61            let Some((artifact, contract)) =
62                project_contracts.find_by_deployed_code(code.original_byte_slice())
63            else {
64                continue;
65            };
66            let Some(functions) =
67                artifact_filters.get_targeted_functions(artifact, &contract.abi)?
68            else {
69                continue;
70            };
71            created_contracts.push(*address);
72            let contract = TargetedContract {
73                identifier: artifact.name.clone(),
74                abi: contract.abi.clone(),
75                targeted_functions: functions,
76                excluded_functions: Vec::new(),
77            };
78            targets.insert(*address, contract);
79        }
80        Ok(())
81    }
82
83    /// Clears targeted contracts created during an invariant run.
84    pub fn clear_created_contracts(&self, created_contracts: Vec<Address>) {
85        if !created_contracts.is_empty() {
86            let mut targets = self.targets.lock();
87            for addr in &created_contracts {
88                targets.remove(addr);
89            }
90        }
91    }
92}
93
94/// A collection of contracts identified as targets for invariant testing.
95#[derive(Clone, Debug, Default)]
96pub struct TargetedContracts {
97    /// The inner map of targeted contracts.
98    pub inner: BTreeMap<Address, TargetedContract>,
99}
100
101impl TargetedContracts {
102    /// Returns a new `TargetedContracts` instance.
103    pub fn new() -> Self {
104        Self::default()
105    }
106
107    /// Returns fuzzed contract abi and fuzzed function from address and provided calldata.
108    ///
109    /// Used to decode return values and logs in order to add values into fuzz dictionary.
110    pub fn fuzzed_artifacts(&self, tx: &BasicTxDetails) -> (Option<&JsonAbi>, Option<&Function>) {
111        match self.inner.get(&tx.call_details.target) {
112            Some(c) => (
113                Some(&c.abi),
114                c.abi.functions().find(|f| f.selector() == tx.call_details.calldata[..4]),
115            ),
116            None => (None, None),
117        }
118    }
119
120    /// Returns flatten target contract address and functions to be fuzzed.
121    /// Includes contract targeted functions if specified, else all mutable contract functions.
122    pub fn fuzzed_functions(&self) -> impl Iterator<Item = (&Address, &Function)> {
123        self.inner
124            .iter()
125            .filter(|(_, c)| !c.abi.functions.is_empty())
126            .flat_map(|(contract, c)| c.abi_fuzzed_functions().map(move |f| (contract, f)))
127    }
128
129    /// Identifies fuzzed contract and function based on given tx details and returns unique metric
130    /// key composed from contract identifier and function name.
131    pub fn fuzzed_metric_key(&self, tx: &BasicTxDetails) -> Option<String> {
132        self.inner.get(&tx.call_details.target).and_then(|contract| {
133            contract
134                .abi
135                .functions()
136                .find(|f| f.selector() == tx.call_details.calldata[..4])
137                .map(|function| format!("{}.{}", contract.identifier.clone(), function.name))
138        })
139    }
140}
141
142impl std::ops::Deref for TargetedContracts {
143    type Target = BTreeMap<Address, TargetedContract>;
144
145    fn deref(&self) -> &Self::Target {
146        &self.inner
147    }
148}
149
150impl std::ops::DerefMut for TargetedContracts {
151    fn deref_mut(&mut self) -> &mut Self::Target {
152        &mut self.inner
153    }
154}
155
156/// A contract identified as target for invariant testing.
157#[derive(Clone, Debug)]
158pub struct TargetedContract {
159    /// The contract identifier. This is only used in error messages.
160    pub identifier: String,
161    /// The contract's ABI.
162    pub abi: JsonAbi,
163    /// The targeted functions of the contract.
164    pub targeted_functions: Vec<Function>,
165    /// The excluded functions of the contract.
166    pub excluded_functions: Vec<Function>,
167}
168
169impl TargetedContract {
170    /// Returns a new `TargetedContract` instance.
171    pub fn new(identifier: String, abi: JsonAbi) -> Self {
172        Self { identifier, abi, targeted_functions: Vec::new(), excluded_functions: Vec::new() }
173    }
174
175    /// Helper to retrieve functions to fuzz for specified abi.
176    /// Returns specified targeted functions if any, else mutable abi functions that are not
177    /// marked as excluded.
178    pub fn abi_fuzzed_functions(&self) -> impl Iterator<Item = &Function> {
179        if !self.targeted_functions.is_empty() {
180            Either::Left(self.targeted_functions.iter())
181        } else {
182            Either::Right(self.abi.functions().filter(|&func| {
183                !matches!(
184                    func.state_mutability,
185                    alloy_json_abi::StateMutability::Pure | alloy_json_abi::StateMutability::View
186                ) && !self.excluded_functions.contains(func)
187            }))
188        }
189    }
190
191    /// Returns the function for the given selector.
192    pub fn get_function(&self, selector: Selector) -> eyre::Result<&Function> {
193        get_function(&self.identifier, selector, &self.abi)
194    }
195
196    /// Adds the specified selectors to the targeted functions.
197    pub fn add_selectors(
198        &mut self,
199        selectors: impl IntoIterator<Item = Selector>,
200        should_exclude: bool,
201    ) -> eyre::Result<()> {
202        for selector in selectors {
203            if should_exclude {
204                self.excluded_functions.push(self.get_function(selector)?.clone());
205            } else {
206                self.targeted_functions.push(self.get_function(selector)?.clone());
207            }
208        }
209        Ok(())
210    }
211}
212
213/// Details of a transaction generated by invariant strategy for fuzzing a target.
214#[derive(Clone, Debug)]
215pub struct BasicTxDetails {
216    // Transaction sender address.
217    pub sender: Address,
218    // Transaction call details.
219    pub call_details: CallDetails,
220}
221
222/// Call details of a transaction generated to fuzz invariant target.
223#[derive(Clone, Debug)]
224pub struct CallDetails {
225    // Address of target contract.
226    pub target: Address,
227    // The data of the transaction.
228    pub calldata: Bytes,
229}
230
231/// Test contract which is testing its invariants.
232#[derive(Clone, Debug)]
233pub struct InvariantContract<'a> {
234    /// Address of the test contract.
235    pub address: Address,
236    /// Invariant function present in the test contract.
237    pub invariant_function: &'a Function,
238    /// If true, `afterInvariant` function is called after each invariant run.
239    pub call_after_invariant: bool,
240    /// ABI of the test contract.
241    pub abi: &'a JsonAbi,
242}
foundry_evm_fuzz/invariant/mod.rs

foundry_evm_fuzz/invariant/
mod.rs
