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

foundry_cheatcodes/
utils.rs
