foundry_cheatcodes/
utils.rs

1//! Implementations of [`Utilities`](spec::Group::Utilities) cheatcodes.
2
3use crate::{Cheatcode, Cheatcodes, CheatcodesExecutor, CheatsCtxt, Result, Vm::*};
4use alloy_dyn_abi::{DynSolType, DynSolValue, Resolver, TypedData, eip712_parser::EncodeType};
5use alloy_ens::namehash;
6use alloy_primitives::{B64, Bytes, I256, U256, aliases::B32, keccak256, map::HashMap};
7use alloy_sol_types::SolValue;
8use foundry_common::{TYPE_BINDING_PREFIX, fs};
9use foundry_config::fs_permissions::FsAccessKind;
10use foundry_evm_core::constants::DEFAULT_CREATE2_DEPLOYER;
11use foundry_evm_fuzz::strategies::BoundMutator;
12use proptest::prelude::Strategy;
13use rand::{Rng, RngCore, seq::SliceRandom};
14use revm::context::JournalTr;
15use std::path::PathBuf;
16
17/// Contains locations of traces ignored via cheatcodes.
18///
19/// The way we identify location in traces is by (node_idx, item_idx) tuple where node_idx is an
20/// index of a call trace node, and item_idx is a value between 0 and `node.ordering.len()` where i
21/// represents point after ith item, and 0 represents the beginning of the node trace.
22#[derive(Debug, Default, Clone)]
23pub struct IgnoredTraces {
24    /// Mapping from (start_node_idx, start_item_idx) to (end_node_idx, end_item_idx) representing
25    /// ranges of trace nodes to ignore.
26    pub ignored: HashMap<(usize, usize), (usize, usize)>,
27    /// Keeps track of (start_node_idx, start_item_idx) of the last `vm.pauseTracing` call.
28    pub last_pause_call: Option<(usize, usize)>,
29}
30
31impl Cheatcode for labelCall {
32    fn apply(&self, state: &mut Cheatcodes) -> Result {
33        let Self { account, newLabel } = self;
34        state.labels.insert(*account, newLabel.clone());
35        Ok(Default::default())
36    }
37}
38
39impl Cheatcode for getLabelCall {
40    fn apply(&self, state: &mut Cheatcodes) -> Result {
41        let Self { account } = self;
42        Ok(match state.labels.get(account) {
43            Some(label) => label.abi_encode(),
44            None => format!("unlabeled:{account}").abi_encode(),
45        })
46    }
47}
48
49impl Cheatcode for computeCreateAddressCall {
50    fn apply(&self, _state: &mut Cheatcodes) -> Result {
51        let Self { nonce, deployer } = self;
52        ensure!(*nonce <= U256::from(u64::MAX), "nonce must be less than 2^64");
53        Ok(deployer.create(nonce.to()).abi_encode())
54    }
55}
56
57impl Cheatcode for computeCreate2Address_0Call {
58    fn apply(&self, _state: &mut Cheatcodes) -> Result {
59        let Self { salt, initCodeHash, deployer } = self;
60        Ok(deployer.create2(salt, initCodeHash).abi_encode())
61    }
62}
63
64impl Cheatcode for computeCreate2Address_1Call {
65    fn apply(&self, _state: &mut Cheatcodes) -> Result {
66        let Self { salt, initCodeHash } = self;
67        Ok(DEFAULT_CREATE2_DEPLOYER.create2(salt, initCodeHash).abi_encode())
68    }
69}
70
71impl Cheatcode for ensNamehashCall {
72    fn apply(&self, _state: &mut Cheatcodes) -> Result {
73        let Self { name } = self;
74        Ok(namehash(name).abi_encode())
75    }
76}
77
78impl Cheatcode for bound_0Call {
79    fn apply(&self, state: &mut Cheatcodes) -> Result {
80        let Self { current, min, max } = *self;
81        let Some(mutated) = U256::bound(current, min, max, state.test_runner()) else {
82            bail!("cannot bound {current} in [{min}, {max}] range")
83        };
84        Ok(mutated.abi_encode())
85    }
86}
87
88impl Cheatcode for bound_1Call {
89    fn apply(&self, state: &mut Cheatcodes) -> Result {
90        let Self { current, min, max } = *self;
91        let Some(mutated) = I256::bound(current, min, max, state.test_runner()) else {
92            bail!("cannot bound {current} in [{min}, {max}] range")
93        };
94        Ok(mutated.abi_encode())
95    }
96}
97
98impl Cheatcode for randomUint_0Call {
99    fn apply(&self, state: &mut Cheatcodes) -> Result {
100        random_uint(state, None, None)
101    }
102}
103
104impl Cheatcode for randomUint_1Call {
105    fn apply(&self, state: &mut Cheatcodes) -> Result {
106        let Self { min, max } = *self;
107        random_uint(state, None, Some((min, max)))
108    }
109}
110
111impl Cheatcode for randomUint_2Call {
112    fn apply(&self, state: &mut Cheatcodes) -> Result {
113        let Self { bits } = *self;
114        random_uint(state, Some(bits), None)
115    }
116}
117
118impl Cheatcode for randomAddressCall {
119    fn apply(&self, state: &mut Cheatcodes) -> Result {
120        Ok(DynSolValue::type_strategy(&DynSolType::Address)
121            .new_tree(state.test_runner())
122            .unwrap()
123            .current()
124            .abi_encode())
125    }
126}
127
128impl Cheatcode for randomInt_0Call {
129    fn apply(&self, state: &mut Cheatcodes) -> Result {
130        random_int(state, None)
131    }
132}
133
134impl Cheatcode for randomInt_1Call {
135    fn apply(&self, state: &mut Cheatcodes) -> Result {
136        let Self { bits } = *self;
137        random_int(state, Some(bits))
138    }
139}
140
141impl Cheatcode for randomBoolCall {
142    fn apply(&self, state: &mut Cheatcodes) -> Result {
143        let rand_bool: bool = state.rng().random();
144        Ok(rand_bool.abi_encode())
145    }
146}
147
148impl Cheatcode for randomBytesCall {
149    fn apply(&self, state: &mut Cheatcodes) -> Result {
150        let Self { len } = *self;
151        ensure!(
152            len <= U256::from(usize::MAX),
153            format!("bytes length cannot exceed {}", usize::MAX)
154        );
155        let mut bytes = vec![0u8; len.to::<usize>()];
156        state.rng().fill_bytes(&mut bytes);
157        Ok(bytes.abi_encode())
158    }
159}
160
161impl Cheatcode for randomBytes4Call {
162    fn apply(&self, state: &mut Cheatcodes) -> Result {
163        let rand_u32 = state.rng().next_u32();
164        Ok(B32::from(rand_u32).abi_encode())
165    }
166}
167
168impl Cheatcode for randomBytes8Call {
169    fn apply(&self, state: &mut Cheatcodes) -> Result {
170        let rand_u64 = state.rng().next_u64();
171        Ok(B64::from(rand_u64).abi_encode())
172    }
173}
174
175impl Cheatcode for pauseTracingCall {
176    fn apply_full(
177        &self,
178        ccx: &mut crate::CheatsCtxt,
179        executor: &mut dyn CheatcodesExecutor,
180    ) -> Result {
181        let Some(tracer) = executor.tracing_inspector() else {
182            // No tracer -> nothing to pause
183            return Ok(Default::default());
184        };
185
186        // If paused earlier, ignore the call
187        if ccx.state.ignored_traces.last_pause_call.is_some() {
188            return Ok(Default::default());
189        }
190
191        let cur_node = &tracer.traces().nodes().last().expect("no trace nodes");
192        ccx.state.ignored_traces.last_pause_call = Some((cur_node.idx, cur_node.ordering.len()));
193
194        Ok(Default::default())
195    }
196}
197
198impl Cheatcode for resumeTracingCall {
199    fn apply_full(
200        &self,
201        ccx: &mut crate::CheatsCtxt,
202        executor: &mut dyn CheatcodesExecutor,
203    ) -> Result {
204        let Some(tracer) = executor.tracing_inspector() else {
205            // No tracer -> nothing to unpause
206            return Ok(Default::default());
207        };
208
209        let Some(start) = ccx.state.ignored_traces.last_pause_call.take() else {
210            // Nothing to unpause
211            return Ok(Default::default());
212        };
213
214        let node = &tracer.traces().nodes().last().expect("no trace nodes");
215        ccx.state.ignored_traces.ignored.insert(start, (node.idx, node.ordering.len()));
216
217        Ok(Default::default())
218    }
219}
220
221impl Cheatcode for interceptInitcodeCall {
222    fn apply(&self, state: &mut Cheatcodes) -> Result {
223        let Self {} = self;
224        if !state.intercept_next_create_call {
225            state.intercept_next_create_call = true;
226        } else {
227            bail!("vm.interceptInitcode() has already been called")
228        }
229        Ok(Default::default())
230    }
231}
232
233impl Cheatcode for setArbitraryStorage_0Call {
234    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
235        let Self { target } = self;
236        ccx.state.arbitrary_storage().mark_arbitrary(target, false);
237
238        Ok(Default::default())
239    }
240}
241
242impl Cheatcode for setArbitraryStorage_1Call {
243    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
244        let Self { target, overwrite } = self;
245        ccx.state.arbitrary_storage().mark_arbitrary(target, *overwrite);
246
247        Ok(Default::default())
248    }
249}
250
251impl Cheatcode for copyStorageCall {
252    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
253        let Self { from, to } = self;
254
255        ensure!(
256            !ccx.state.has_arbitrary_storage(to),
257            "target address cannot have arbitrary storage"
258        );
259
260        if let Ok(from_account) = ccx.ecx.journaled_state.load_account(*from) {
261            let from_storage = from_account.storage.clone();
262            if let Ok(mut to_account) = ccx.ecx.journaled_state.load_account(*to) {
263                to_account.storage = from_storage;
264                if let Some(arbitrary_storage) = &mut ccx.state.arbitrary_storage {
265                    arbitrary_storage.mark_copy(from, to);
266                }
267            }
268        }
269
270        Ok(Default::default())
271    }
272}
273
274impl Cheatcode for sortCall {
275    fn apply(&self, _state: &mut Cheatcodes) -> Result {
276        let Self { array } = self;
277
278        let mut sorted_values = array.clone();
279        sorted_values.sort();
280
281        Ok(sorted_values.abi_encode())
282    }
283}
284
285impl Cheatcode for shuffleCall {
286    fn apply(&self, state: &mut Cheatcodes) -> Result {
287        let Self { array } = self;
288
289        let mut shuffled_values = array.clone();
290        let rng = state.rng();
291        shuffled_values.shuffle(rng);
292
293        Ok(shuffled_values.abi_encode())
294    }
295}
296
297impl Cheatcode for setSeedCall {
298    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
299        let Self { seed } = self;
300        ccx.state.set_seed(*seed);
301        Ok(Default::default())
302    }
303}
304
305/// Helper to generate a random `uint` value (with given bits or bounded if specified)
306/// from type strategy.
307fn random_uint(state: &mut Cheatcodes, bits: Option<U256>, bounds: Option<(U256, U256)>) -> Result {
308    if let Some(bits) = bits {
309        // Generate random with specified bits.
310        ensure!(bits <= U256::from(256), "number of bits cannot exceed 256");
311        return Ok(DynSolValue::type_strategy(&DynSolType::Uint(bits.to::<usize>()))
312            .new_tree(state.test_runner())
313            .unwrap()
314            .current()
315            .abi_encode());
316    }
317
318    if let Some((min, max)) = bounds {
319        ensure!(min <= max, "min must be less than or equal to max");
320        // Generate random between range min..=max
321        let exclusive_modulo = max - min;
322        let mut random_number: U256 = state.rng().random();
323        if exclusive_modulo != U256::MAX {
324            let inclusive_modulo = exclusive_modulo + U256::from(1);
325            random_number %= inclusive_modulo;
326        }
327        random_number += min;
328        return Ok(random_number.abi_encode());
329    }
330
331    // Generate random `uint256` value.
332    Ok(DynSolValue::type_strategy(&DynSolType::Uint(256))
333        .new_tree(state.test_runner())
334        .unwrap()
335        .current()
336        .abi_encode())
337}
338
339/// Helper to generate a random `int` value (with given bits if specified) from type strategy.
340fn random_int(state: &mut Cheatcodes, bits: Option<U256>) -> Result {
341    let no_bits = bits.unwrap_or(U256::from(256));
342    ensure!(no_bits <= U256::from(256), "number of bits cannot exceed 256");
343    Ok(DynSolValue::type_strategy(&DynSolType::Int(no_bits.to::<usize>()))
344        .new_tree(state.test_runner())
345        .unwrap()
346        .current()
347        .abi_encode())
348}
349
350impl Cheatcode for eip712HashType_0Call {
351    fn apply(&self, state: &mut Cheatcodes) -> Result {
352        let Self { typeNameOrDefinition } = self;
353
354        let type_def = get_canonical_type_def(typeNameOrDefinition, state, None)?;
355
356        Ok(keccak256(type_def.as_bytes()).to_vec())
357    }
358}
359
360impl Cheatcode for eip712HashType_1Call {
361    fn apply(&self, state: &mut Cheatcodes) -> Result {
362        let Self { bindingsPath, typeName } = self;
363
364        let path = state.config.ensure_path_allowed(bindingsPath, FsAccessKind::Read)?;
365        let type_def = get_type_def_from_bindings(typeName, path, &state.config.root)?;
366
367        Ok(keccak256(type_def.as_bytes()).to_vec())
368    }
369}
370
371impl Cheatcode for eip712HashStruct_0Call {
372    fn apply(&self, state: &mut Cheatcodes) -> Result {
373        let Self { typeNameOrDefinition, abiEncodedData } = self;
374
375        let type_def = get_canonical_type_def(typeNameOrDefinition, state, None)?;
376        let primary = &type_def[..type_def.find('(').unwrap_or(type_def.len())];
377
378        get_struct_hash(primary, &type_def, abiEncodedData)
379    }
380}
381
382impl Cheatcode for eip712HashStruct_1Call {
383    fn apply(&self, state: &mut Cheatcodes) -> Result {
384        let Self { bindingsPath, typeName, abiEncodedData } = self;
385
386        let path = state.config.ensure_path_allowed(bindingsPath, FsAccessKind::Read)?;
387        let type_def = get_type_def_from_bindings(typeName, path, &state.config.root)?;
388
389        get_struct_hash(typeName, &type_def, abiEncodedData)
390    }
391}
392
393impl Cheatcode for eip712HashTypedDataCall {
394    fn apply(&self, _state: &mut Cheatcodes) -> Result {
395        let Self { jsonData } = self;
396        let typed_data: TypedData = serde_json::from_str(jsonData)?;
397        let digest = typed_data.eip712_signing_hash()?;
398
399        Ok(digest.to_vec())
400    }
401}
402
403/// Returns EIP-712 canonical type definition from the provided string type representation or type
404/// name. If type name provided, then it looks up bindings from file generated by `forge bind-json`.
405fn get_canonical_type_def(
406    name_or_def: &String,
407    state: &mut Cheatcodes,
408    path: Option<PathBuf>,
409) -> Result<String> {
410    let type_def = if name_or_def.contains('(') {
411        // If the input contains '(', it must be the type definition.
412        EncodeType::parse(name_or_def).and_then(|parsed| parsed.canonicalize())?
413    } else {
414        // Otherwise, it must be the type name.
415        let path = path.as_ref().unwrap_or(&state.config.bind_json_path);
416        let path = state.config.ensure_path_allowed(path, FsAccessKind::Read)?;
417        get_type_def_from_bindings(name_or_def, path, &state.config.root)?
418    };
419
420    Ok(type_def)
421}
422
423/// Returns the EIP-712 type definition from the bindings in the provided path.
424/// Assumes that read validation for the path has already been checked.
425fn get_type_def_from_bindings(name: &String, path: PathBuf, root: &PathBuf) -> Result<String> {
426    let content = fs::read_to_string(&path)?;
427
428    let type_defs: HashMap<&str, &str> = content
429        .lines()
430        .filter_map(|line| {
431            let relevant = line.trim().strip_prefix(TYPE_BINDING_PREFIX)?;
432            let (name, def) = relevant.split_once('=')?;
433            Some((name.trim(), def.trim().strip_prefix('"')?.strip_suffix("\";")?))
434        })
435        .collect();
436
437    match type_defs.get(name.as_str()) {
438        Some(value) => Ok(value.to_string()),
439        None => {
440            let bindings =
441                type_defs.keys().map(|k| format!(" - {k}")).collect::<Vec<String>>().join("\n");
442
443            bail!(
444                "'{}' not found in '{}'.{}",
445                name,
446                path.strip_prefix(root).unwrap_or(&path).to_string_lossy(),
447                if bindings.is_empty() {
448                    String::new()
449                } else {
450                    format!("\nAvailable bindings:\n{bindings}\n")
451                }
452            );
453        }
454    }
455}
456
457/// Returns the EIP-712 struct hash for provided name, definition and ABI encoded data.
458fn get_struct_hash(primary: &str, type_def: &String, abi_encoded_data: &Bytes) -> Result {
459    let mut resolver = Resolver::default();
460
461    // Populate the resolver by ingesting the canonical type definition, and then get the
462    // corresponding `DynSolType` of the primary type.
463    resolver
464        .ingest_string(type_def)
465        .map_err(|e| fmt_err!("Resolver failed to ingest type definition: {e}"))?;
466
467    let resolved_sol_type = resolver
468        .resolve(primary)
469        .map_err(|e| fmt_err!("Failed to resolve EIP-712 primary type '{primary}': {e}"))?;
470
471    // ABI-decode the bytes into `DynSolValue::CustomStruct`.
472    let sol_value = resolved_sol_type.abi_decode(abi_encoded_data.as_ref()).map_err(|e| {
473        fmt_err!("Failed to ABI decode using resolved_sol_type directly for '{primary}': {e}.")
474    })?;
475
476    // Use the resolver to properly encode the data.
477    let encoded_data: Vec<u8> = resolver
478        .encode_data(&sol_value)
479        .map_err(|e| fmt_err!("Failed to EIP-712 encode data for struct '{primary}': {e}"))?
480        .ok_or_else(|| fmt_err!("EIP-712 data encoding returned 'None' for struct '{primary}'"))?;
481
482    // Compute the type hash of the primary type.
483    let type_hash = resolver
484        .type_hash(primary)
485        .map_err(|e| fmt_err!("Failed to compute typeHash for EIP712 type '{primary}': {e}"))?;
486
487    // Compute the struct hash of the concatenated type hash and encoded data.
488    let mut bytes_to_hash = Vec::with_capacity(32 + encoded_data.len());
489    bytes_to_hash.extend_from_slice(type_hash.as_slice());
490    bytes_to_hash.extend_from_slice(&encoded_data);
491
492    Ok(keccak256(&bytes_to_hash).to_vec())
493}
foundry_cheatcodes/utils.rs

foundry_cheatcodes/
utils.rs
