foundry_evm_fuzz/strategies/
state.rs

1use crate::{BasicTxDetails, invariant::FuzzRunIdentifiedContracts};
2use alloy_dyn_abi::{DynSolType, DynSolValue, EventExt, FunctionExt};
3use alloy_json_abi::{Function, JsonAbi};
4use alloy_primitives::{
5    Address, B256, Bytes, Log, U256,
6    map::{AddressIndexSet, AddressMap, B256IndexSet, HashMap},
7};
8use foundry_common::{
9    ignore_metadata_hash, mapping_slots::MappingSlots, slot_identifier::SlotIdentifier,
10};
11use foundry_compilers::artifacts::StorageLayout;
12use foundry_config::FuzzDictionaryConfig;
13use foundry_evm_core::utils::StateChangeset;
14use parking_lot::{RawRwLock, RwLock, lock_api::RwLockReadGuard};
15use revm::{
16    bytecode::opcode,
17    database::{CacheDB, DatabaseRef, DbAccount},
18    state::AccountInfo,
19};
20use std::{collections::BTreeMap, fmt, sync::Arc};
21
22/// The maximum number of bytes we will look at in bytecodes to find push bytes (24 KiB).
23///
24/// This is to limit the performance impact of fuzz tests that might deploy arbitrarily sized
25/// bytecode (as is the case with Solmate).
26const PUSH_BYTE_ANALYSIS_LIMIT: usize = 24 * 1024;
27
28/// A set of arbitrary 32 byte data from the VM used to generate values for the strategy.
29///
30/// Wrapped in a shareable container.
31#[derive(Clone, Debug)]
32pub struct EvmFuzzState {
33    inner: Arc<RwLock<FuzzDictionary>>,
34    /// Addresses of external libraries deployed in test setup, excluded from fuzz test inputs.
35    pub deployed_libs: Vec<Address>,
36    /// Records mapping accesses. Used to identify storage slots belonging to mappings and sampling
37    /// the values in the [`FuzzDictionary`].
38    ///
39    /// Only needed when [`StorageLayout`] is available.
40    pub(crate) mapping_slots: Option<AddressMap<MappingSlots>>,
41}
42
43impl EvmFuzzState {
44    pub fn new<DB: DatabaseRef>(
45        db: &CacheDB<DB>,
46        config: FuzzDictionaryConfig,
47        deployed_libs: &[Address],
48    ) -> Self {
49        // Sort accounts to ensure deterministic dictionary generation from the same setUp state.
50        let mut accs = db.cache.accounts.iter().collect::<Vec<_>>();
51        accs.sort_by_key(|(address, _)| *address);
52
53        // Create fuzz dictionary and insert values from db state.
54        let mut dictionary = FuzzDictionary::new(config);
55        dictionary.insert_db_values(accs);
56        Self {
57            inner: Arc::new(RwLock::new(dictionary)),
58            deployed_libs: deployed_libs.to_vec(),
59            mapping_slots: None,
60        }
61    }
62
63    pub fn with_mapping_slots(mut self, mapping_slots: AddressMap<MappingSlots>) -> Self {
64        self.mapping_slots = Some(mapping_slots);
65        self
66    }
67
68    pub fn collect_values(&self, values: impl IntoIterator<Item = B256>) {
69        let mut dict = self.inner.write();
70        for value in values {
71            dict.insert_value(value);
72        }
73    }
74
75    /// Collects state changes from a [StateChangeset] and logs into an [EvmFuzzState] according to
76    /// the given [FuzzDictionaryConfig].
77    pub fn collect_values_from_call(
78        &self,
79        fuzzed_contracts: &FuzzRunIdentifiedContracts,
80        tx: &BasicTxDetails,
81        result: &Bytes,
82        logs: &[Log],
83        state_changeset: &StateChangeset,
84        run_depth: u32,
85    ) {
86        let mut dict = self.inner.write();
87        {
88            let targets = fuzzed_contracts.targets.lock();
89            let (target_abi, target_function) = targets.fuzzed_artifacts(tx);
90            dict.insert_logs_values(target_abi, logs, run_depth);
91            dict.insert_result_values(target_function, result, run_depth);
92            // Get storage layouts for contracts in the state changeset
93            let storage_layouts = targets.get_storage_layouts();
94            dict.insert_new_state_values(
95                state_changeset,
96                &storage_layouts,
97                self.mapping_slots.as_ref(),
98            );
99        }
100    }
101
102    /// Removes all newly added entries from the dictionary.
103    ///
104    /// Should be called between fuzz/invariant runs to avoid accumulating data derived from fuzz
105    /// inputs.
106    pub fn revert(&self) {
107        self.inner.write().revert();
108    }
109
110    pub fn dictionary_read(&self) -> RwLockReadGuard<'_, RawRwLock, FuzzDictionary> {
111        self.inner.read()
112    }
113
114    /// Logs stats about the current state.
115    pub fn log_stats(&self) {
116        self.inner.read().log_stats();
117    }
118}
119
120// We're using `IndexSet` to have a stable element order when restoring persisted state, as well as
121// for performance when iterating over the sets.
122#[derive(Default)]
123pub struct FuzzDictionary {
124    /// Collected state values.
125    state_values: B256IndexSet,
126    /// Addresses that already had their PUSH bytes collected.
127    addresses: AddressIndexSet,
128    /// Configuration for the dictionary.
129    config: FuzzDictionaryConfig,
130    /// Number of state values initially collected from db.
131    /// Used to revert new collected values at the end of each run.
132    db_state_values: usize,
133    /// Number of address values initially collected from db.
134    /// Used to revert new collected addresses at the end of each run.
135    db_addresses: usize,
136    /// Sample typed values that are collected from call result and used across invariant runs.
137    sample_values: HashMap<DynSolType, B256IndexSet>,
138
139    misses: usize,
140    hits: usize,
141}
142
143impl fmt::Debug for FuzzDictionary {
144    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
145        f.debug_struct("FuzzDictionary")
146            .field("state_values", &self.state_values.len())
147            .field("addresses", &self.addresses)
148            .finish()
149    }
150}
151
152impl FuzzDictionary {
153    pub fn new(config: FuzzDictionaryConfig) -> Self {
154        let mut dictionary = Self { config, ..Default::default() };
155        dictionary.prefill();
156        dictionary
157    }
158
159    /// Insert common values into the dictionary at initialization.
160    fn prefill(&mut self) {
161        self.insert_value(B256::ZERO);
162    }
163
164    /// Insert values from initial db state into fuzz dictionary.
165    /// These values are persisted across invariant runs.
166    fn insert_db_values(&mut self, db_state: Vec<(&Address, &DbAccount)>) {
167        for (address, account) in db_state {
168            // Insert basic account information
169            self.insert_value(address.into_word());
170            // Insert push bytes
171            self.insert_push_bytes_values(address, &account.info);
172            // Insert storage values.
173            if self.config.include_storage {
174                // Sort storage values before inserting to ensure deterministic dictionary.
175                let values = account.storage.iter().collect::<BTreeMap<_, _>>();
176                for (slot, value) in values {
177                    self.insert_storage_value(slot, value, None, None);
178                }
179            }
180        }
181
182        // We need at least some state data if DB is empty,
183        // otherwise we can't select random data for state fuzzing.
184        if self.values().is_empty() {
185            // Prefill with a random address.
186            self.insert_value(Address::random().into_word());
187        }
188
189        // Record number of values and addresses inserted from db to be used for reverting at the
190        // end of each run.
191        self.db_state_values = self.state_values.len();
192        self.db_addresses = self.addresses.len();
193    }
194
195    /// Insert values collected from call result into fuzz dictionary.
196    fn insert_result_values(
197        &mut self,
198        function: Option<&Function>,
199        result: &Bytes,
200        run_depth: u32,
201    ) {
202        if let Some(function) = function
203            && !function.outputs.is_empty()
204        {
205            // Decode result and collect samples to be used in subsequent fuzz runs.
206            if let Ok(decoded_result) = function.abi_decode_output(result) {
207                self.insert_sample_values(decoded_result, run_depth);
208            }
209        }
210    }
211
212    /// Insert values from call log topics and data into fuzz dictionary.
213    fn insert_logs_values(&mut self, abi: Option<&JsonAbi>, logs: &[Log], run_depth: u32) {
214        let mut samples = Vec::new();
215        // Decode logs with known events and collect samples from indexed fields and event body.
216        for log in logs {
217            let mut log_decoded = false;
218            // Try to decode log with events from contract abi.
219            if let Some(abi) = abi {
220                for event in abi.events() {
221                    if let Ok(decoded_event) = event.decode_log(log) {
222                        samples.extend(decoded_event.indexed);
223                        samples.extend(decoded_event.body);
224                        log_decoded = true;
225                        break;
226                    }
227                }
228            }
229
230            // If we weren't able to decode event then we insert raw data in fuzz dictionary.
231            if !log_decoded {
232                for &topic in log.topics() {
233                    self.insert_value(topic);
234                }
235                let chunks = log.data.data.chunks_exact(32);
236                let rem = chunks.remainder();
237                for chunk in chunks {
238                    self.insert_value(chunk.try_into().unwrap());
239                }
240                if !rem.is_empty() {
241                    self.insert_value(B256::right_padding_from(rem));
242                }
243            }
244        }
245
246        // Insert samples collected from current call in fuzz dictionary.
247        self.insert_sample_values(samples, run_depth);
248    }
249
250    /// Insert values from call state changeset into fuzz dictionary.
251    /// These values are removed at the end of current run.
252    fn insert_new_state_values(
253        &mut self,
254        state_changeset: &StateChangeset,
255        storage_layouts: &HashMap<Address, Arc<StorageLayout>>,
256        mapping_slots: Option<&AddressMap<MappingSlots>>,
257    ) {
258        for (address, account) in state_changeset {
259            // Insert basic account information.
260            self.insert_value(address.into_word());
261            // Insert push bytes.
262            self.insert_push_bytes_values(address, &account.info);
263            // Insert storage values.
264            if self.config.include_storage {
265                let storage_layout = storage_layouts.get(address).cloned();
266                trace!(
267                    "{address:?} has mapping_slots {}",
268                    mapping_slots.is_some_and(|m| m.contains_key(address))
269                );
270                let mapping_slots = mapping_slots.and_then(|m| m.get(address));
271                for (slot, value) in &account.storage {
272                    self.insert_storage_value(
273                        slot,
274                        &value.present_value,
275                        storage_layout.as_deref(),
276                        mapping_slots,
277                    );
278                }
279            }
280        }
281    }
282
283    /// Insert values from push bytes into fuzz dictionary.
284    /// Values are collected only once for a given address.
285    /// If values are newly collected then they are removed at the end of current run.
286    fn insert_push_bytes_values(&mut self, address: &Address, account_info: &AccountInfo) {
287        if self.config.include_push_bytes && !self.addresses.contains(address) {
288            // Insert push bytes
289            if let Some(code) = &account_info.code {
290                self.insert_address(*address);
291                self.collect_push_bytes(ignore_metadata_hash(code.original_byte_slice()));
292            }
293        }
294    }
295
296    fn collect_push_bytes(&mut self, code: &[u8]) {
297        let mut i = 0;
298        let len = code.len().min(PUSH_BYTE_ANALYSIS_LIMIT);
299        while i < len {
300            let op = code[i];
301            if (opcode::PUSH1..=opcode::PUSH32).contains(&op) {
302                let push_size = (op - opcode::PUSH1 + 1) as usize;
303                let push_start = i + 1;
304                let push_end = push_start + push_size;
305                // As a precaution, if a fuzz test deploys malformed bytecode (such as using
306                // `CREATE2`) this will terminate the loop early.
307                if push_start > code.len() || push_end > code.len() {
308                    break;
309                }
310
311                let push_value = U256::try_from_be_slice(&code[push_start..push_end]).unwrap();
312                if push_value != U256::ZERO {
313                    // Never add 0 to the dictionary as it's always present.
314                    self.insert_value(push_value.into());
315
316                    // Also add the value below and above the push value to the dictionary.
317                    self.insert_value((push_value - U256::from(1)).into());
318
319                    if push_value != U256::MAX {
320                        self.insert_value((push_value + U256::from(1)).into());
321                    }
322                }
323
324                i += push_size;
325            }
326            i += 1;
327        }
328    }
329
330    /// Insert values from single storage slot and storage value into fuzz dictionary.
331    /// Uses [`SlotIdentifier`] to identify storage slots types.
332    fn insert_storage_value(
333        &mut self,
334        slot: &U256,
335        value: &U256,
336        layout: Option<&StorageLayout>,
337        mapping_slots: Option<&MappingSlots>,
338    ) {
339        let slot = B256::from(*slot);
340        let value = B256::from(*value);
341
342        // Always insert the slot itself
343        self.insert_value(slot);
344
345        // If we have a storage layout, use SlotIdentifier for better type identification
346        if let Some(slot_identifier) =
347            layout.map(|l| SlotIdentifier::new(l.clone().into()))
348            // Identify Slot Type
349            && let Some(slot_info) = slot_identifier.identify(&slot, mapping_slots) && slot_info.decode(value).is_some()
350        {
351            trace!(?slot_info, "inserting typed storage value");
352            self.sample_values.entry(slot_info.slot_type.dyn_sol_type).or_default().insert(value);
353        } else {
354            self.insert_value(value);
355            let value = U256::from_be_bytes(value.0);
356
357            // Also add the value below and above the storage value to the dictionary
358            if value != U256::ZERO {
359                let below_value = value - U256::from(1);
360                self.insert_value(B256::from(below_value));
361            }
362            if value != U256::MAX {
363                let above_value = value + U256::from(1);
364                self.insert_value(B256::from(above_value));
365            }
366        }
367    }
368
369    /// Insert address into fuzz dictionary.
370    /// If address is newly collected then it is removed by index at the end of current run.
371    fn insert_address(&mut self, address: Address) {
372        if self.addresses.len() < self.config.max_fuzz_dictionary_addresses {
373            self.addresses.insert(address);
374        }
375    }
376
377    /// Insert raw value into fuzz dictionary.
378    /// If value is newly collected then it is removed by index at the end of current run.
379    fn insert_value(&mut self, value: B256) {
380        if self.state_values.len() < self.config.max_fuzz_dictionary_values {
381            let new_value = self.state_values.insert(value);
382            let counter = if new_value { &mut self.misses } else { &mut self.hits };
383            *counter += 1;
384        }
385    }
386
387    /// Insert sample values that are reused across multiple runs.
388    /// The number of samples is limited to invariant run depth.
389    /// If collected samples limit is reached then values are inserted as regular values.
390    pub fn insert_sample_values(
391        &mut self,
392        sample_values: impl IntoIterator<Item = DynSolValue>,
393        limit: u32,
394    ) {
395        for sample in sample_values {
396            if let (Some(sample_type), Some(sample_value)) = (sample.as_type(), sample.as_word()) {
397                if let Some(values) = self.sample_values.get_mut(&sample_type) {
398                    if values.len() < limit as usize {
399                        values.insert(sample_value);
400                    } else {
401                        // Insert as state value (will be removed at the end of the run).
402                        self.insert_value(sample_value);
403                    }
404                } else {
405                    self.sample_values.entry(sample_type).or_default().insert(sample_value);
406                }
407            }
408        }
409    }
410
411    pub fn values(&self) -> &B256IndexSet {
412        &self.state_values
413    }
414
415    pub fn len(&self) -> usize {
416        self.state_values.len()
417    }
418
419    pub fn is_empty(&self) -> bool {
420        self.state_values.is_empty()
421    }
422
423    #[inline]
424    pub fn samples(&self, param_type: &DynSolType) -> Option<&B256IndexSet> {
425        self.sample_values.get(param_type)
426    }
427
428    #[inline]
429    pub fn addresses(&self) -> &AddressIndexSet {
430        &self.addresses
431    }
432
433    /// Revert values and addresses collected during the run by truncating to initial db len.
434    pub fn revert(&mut self) {
435        self.state_values.truncate(self.db_state_values);
436        self.addresses.truncate(self.db_addresses);
437    }
438
439    pub fn log_stats(&self) {
440        trace!(
441            addresses.len = self.addresses.len(),
442            sample.len = self.sample_values.len(),
443            state.len = self.state_values.len(),
444            state.misses = self.misses,
445            state.hits = self.hits,
446            "FuzzDictionary stats",
447        );
448    }
449}