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::{DynSolType, DynSolValue, Resolver, TypedData, eip712_parser::EncodeType};
5use alloy_ens::namehash;
6use alloy_primitives::{B64, Bytes, 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 proptest::prelude::Strategy;
12use rand::{Rng, RngCore, seq::SliceRandom};
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(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
276impl Cheatcode for setSeedCall {
277    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
278        let Self { seed } = self;
279        ccx.state.set_seed(*seed);
280        Ok(Default::default())
281    }
282}
283
284/// Helper to generate a random `uint` value (with given bits or bounded if specified)
285/// from type strategy.
286fn random_uint(state: &mut Cheatcodes, bits: Option<U256>, bounds: Option<(U256, U256)>) -> Result {
287    if let Some(bits) = bits {
288        // Generate random with specified bits.
289        ensure!(bits <= U256::from(256), "number of bits cannot exceed 256");
290        return Ok(DynSolValue::type_strategy(&DynSolType::Uint(bits.to::<usize>()))
291            .new_tree(state.test_runner())
292            .unwrap()
293            .current()
294            .abi_encode());
295    }
296
297    if let Some((min, max)) = bounds {
298        ensure!(min <= max, "min must be less than or equal to max");
299        // Generate random between range min..=max
300        let exclusive_modulo = max - min;
301        let mut random_number: U256 = state.rng().random();
302        if exclusive_modulo != U256::MAX {
303            let inclusive_modulo = exclusive_modulo + U256::from(1);
304            random_number %= inclusive_modulo;
305        }
306        random_number += min;
307        return Ok(random_number.abi_encode());
308    }
309
310    // Generate random `uint256` value.
311    Ok(DynSolValue::type_strategy(&DynSolType::Uint(256))
312        .new_tree(state.test_runner())
313        .unwrap()
314        .current()
315        .abi_encode())
316}
317
318/// Helper to generate a random `int` value (with given bits if specified) from type strategy.
319fn random_int(state: &mut Cheatcodes, bits: Option<U256>) -> Result {
320    let no_bits = bits.unwrap_or(U256::from(256));
321    ensure!(no_bits <= U256::from(256), "number of bits cannot exceed 256");
322    Ok(DynSolValue::type_strategy(&DynSolType::Int(no_bits.to::<usize>()))
323        .new_tree(state.test_runner())
324        .unwrap()
325        .current()
326        .abi_encode())
327}
328
329impl Cheatcode for eip712HashType_0Call {
330    fn apply(&self, state: &mut Cheatcodes) -> Result {
331        let Self { typeNameOrDefinition } = self;
332
333        let type_def = get_canonical_type_def(typeNameOrDefinition, state, None)?;
334
335        Ok(keccak256(type_def.as_bytes()).to_vec())
336    }
337}
338
339impl Cheatcode for eip712HashType_1Call {
340    fn apply(&self, state: &mut Cheatcodes) -> Result {
341        let Self { bindingsPath, typeName } = self;
342
343        let path = state.config.ensure_path_allowed(bindingsPath, FsAccessKind::Read)?;
344        let type_def = get_type_def_from_bindings(typeName, path, &state.config.root)?;
345
346        Ok(keccak256(type_def.as_bytes()).to_vec())
347    }
348}
349
350impl Cheatcode for eip712HashStruct_0Call {
351    fn apply(&self, state: &mut Cheatcodes) -> Result {
352        let Self { typeNameOrDefinition, abiEncodedData } = self;
353
354        let type_def = get_canonical_type_def(typeNameOrDefinition, state, None)?;
355        let primary = &type_def[..type_def.find('(').unwrap_or(type_def.len())];
356
357        get_struct_hash(primary, &type_def, abiEncodedData)
358    }
359}
360
361impl Cheatcode for eip712HashStruct_1Call {
362    fn apply(&self, state: &mut Cheatcodes) -> Result {
363        let Self { bindingsPath, typeName, abiEncodedData } = self;
364
365        let path = state.config.ensure_path_allowed(bindingsPath, FsAccessKind::Read)?;
366        let type_def = get_type_def_from_bindings(typeName, path, &state.config.root)?;
367
368        get_struct_hash(typeName, &type_def, abiEncodedData)
369    }
370}
371
372impl Cheatcode for eip712HashTypedDataCall {
373    fn apply(&self, _state: &mut Cheatcodes) -> Result {
374        let Self { jsonData } = self;
375        let typed_data: TypedData = serde_json::from_str(jsonData)?;
376        let digest = typed_data.eip712_signing_hash()?;
377
378        Ok(digest.to_vec())
379    }
380}
381
382/// Returns EIP-712 canonical type definition from the provided string type representation or type
383/// name. If type name provided, then it looks up bindings from file generated by `forge bind-json`.
384fn get_canonical_type_def(
385    name_or_def: &String,
386    state: &mut Cheatcodes,
387    path: Option<PathBuf>,
388) -> Result<String> {
389    let type_def = if name_or_def.contains('(') {
390        // If the input contains '(', it must be the type definition.
391        EncodeType::parse(name_or_def).and_then(|parsed| parsed.canonicalize())?
392    } else {
393        // Otherwise, it must be the type name.
394        let path = path.as_ref().unwrap_or(&state.config.bind_json_path);
395        let path = state.config.ensure_path_allowed(path, FsAccessKind::Read)?;
396        get_type_def_from_bindings(name_or_def, path, &state.config.root)?
397    };
398
399    Ok(type_def)
400}
401
402/// Returns the EIP-712 type definition from the bindings in the provided path.
403/// Assumes that read validation for the path has already been checked.
404fn get_type_def_from_bindings(name: &String, path: PathBuf, root: &PathBuf) -> Result<String> {
405    let content = fs::read_to_string(&path)?;
406
407    let type_defs: HashMap<&str, &str> = content
408        .lines()
409        .filter_map(|line| {
410            let relevant = line.trim().strip_prefix(TYPE_BINDING_PREFIX)?;
411            let (name, def) = relevant.split_once('=')?;
412            Some((name.trim(), def.trim().strip_prefix('"')?.strip_suffix("\";")?))
413        })
414        .collect();
415
416    match type_defs.get(name.as_str()) {
417        Some(value) => Ok(value.to_string()),
418        None => {
419            let bindings =
420                type_defs.keys().map(|k| format!(" - {k}")).collect::<Vec<String>>().join("\n");
421
422            bail!(
423                "'{}' not found in '{}'.{}",
424                name,
425                path.strip_prefix(root).unwrap_or(&path).to_string_lossy(),
426                if bindings.is_empty() {
427                    String::new()
428                } else {
429                    format!("\nAvailable bindings:\n{bindings}\n")
430                }
431            );
432        }
433    }
434}
435
436/// Returns the EIP-712 struct hash for provided name, definition and ABI encoded data.
437fn get_struct_hash(primary: &str, type_def: &String, abi_encoded_data: &Bytes) -> Result {
438    let mut resolver = Resolver::default();
439
440    // Populate the resolver by ingesting the canonical type definition, and then get the
441    // corresponding `DynSolType` of the primary type.
442    resolver
443        .ingest_string(type_def)
444        .map_err(|e| fmt_err!("Resolver failed to ingest type definition: {e}"))?;
445
446    let resolved_sol_type = resolver
447        .resolve(primary)
448        .map_err(|e| fmt_err!("Failed to resolve EIP-712 primary type '{primary}': {e}"))?;
449
450    // ABI-decode the bytes into `DynSolValue::CustomStruct`.
451    let sol_value = resolved_sol_type.abi_decode(abi_encoded_data.as_ref()).map_err(|e| {
452        fmt_err!("Failed to ABI decode using resolved_sol_type directly for '{primary}': {e}.")
453    })?;
454
455    // Use the resolver to properly encode the data.
456    let encoded_data: Vec<u8> = resolver
457        .encode_data(&sol_value)
458        .map_err(|e| fmt_err!("Failed to EIP-712 encode data for struct '{primary}': {e}"))?
459        .ok_or_else(|| fmt_err!("EIP-712 data encoding returned 'None' for struct '{primary}'"))?;
460
461    // Compute the type hash of the primary type.
462    let type_hash = resolver
463        .type_hash(primary)
464        .map_err(|e| fmt_err!("Failed to compute typeHash for EIP712 type '{primary}': {e}"))?;
465
466    // Compute the struct hash of the concatenated type hash and encoded data.
467    let mut bytes_to_hash = Vec::with_capacity(32 + encoded_data.len());
468    bytes_to_hash.extend_from_slice(type_hash.as_slice());
469    bytes_to_hash.extend_from_slice(&encoded_data);
470
471    Ok(keccak256(&bytes_to_hash).to_vec())
472}
foundry_cheatcodes/utils.rs

foundry_cheatcodes/
utils.rs
