foundry_cheatcodes/
utils.rs

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

foundry_cheatcodes/
utils.rs
