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

foundry_cheatcodes/
utils.rs
