foundry_cheatcodes/evm/
mock.rs

1use crate::{Cheatcode, Cheatcodes, CheatsCtxt, Result, Vm::*};
2use alloy_primitives::{Address, Bytes, U256};
3use revm::{
4    bytecode::Bytecode,
5    context::{ContextTr, JournalTr},
6    interpreter::InstructionResult,
7};
8use std::{cmp::Ordering, collections::VecDeque};
9
10/// Mocked call data.
11#[derive(Clone, Debug, Default, PartialEq, Eq, Hash)]
12pub struct MockCallDataContext {
13    /// The partial calldata to match for mock
14    pub calldata: Bytes,
15    /// The value to match for mock
16    pub value: Option<U256>,
17}
18
19/// Mocked return data.
20#[derive(Clone, Debug)]
21pub struct MockCallReturnData {
22    /// The return type for the mocked call
23    pub ret_type: InstructionResult,
24    /// Return data or error
25    pub data: Bytes,
26}
27
28impl PartialOrd for MockCallDataContext {
29    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
30        Some(self.cmp(other))
31    }
32}
33
34impl Ord for MockCallDataContext {
35    fn cmp(&self, other: &Self) -> Ordering {
36        // Calldata matching is reversed to ensure that a tighter match is
37        // returned if an exact match is not found. In case, there is
38        // a partial match to calldata that is more specific than
39        // a match to a msg.value, then the more specific calldata takes
40        // precedence.
41        self.calldata.cmp(&other.calldata).reverse().then(self.value.cmp(&other.value).reverse())
42    }
43}
44
45impl Cheatcode for clearMockedCallsCall {
46    fn apply(&self, state: &mut Cheatcodes) -> Result {
47        let Self {} = self;
48        state.mocked_calls = Default::default();
49        Ok(Default::default())
50    }
51}
52
53impl Cheatcode for mockCall_0Call {
54    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
55        let Self { callee, data, returnData } = self;
56        let _ = make_acc_non_empty(callee, ccx)?;
57
58        mock_call(ccx.state, callee, data, None, returnData, InstructionResult::Return);
59        Ok(Default::default())
60    }
61}
62
63impl Cheatcode for mockCall_1Call {
64    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
65        let Self { callee, msgValue, data, returnData } = self;
66        let _ = make_acc_non_empty(callee, ccx)?;
67
68        mock_call(ccx.state, callee, data, Some(msgValue), returnData, InstructionResult::Return);
69        Ok(Default::default())
70    }
71}
72
73impl Cheatcode for mockCall_2Call {
74    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
75        let Self { callee, data, returnData } = self;
76        let _ = make_acc_non_empty(callee, ccx)?;
77
78        mock_call(
79            ccx.state,
80            callee,
81            &Bytes::from(*data),
82            None,
83            returnData,
84            InstructionResult::Return,
85        );
86        Ok(Default::default())
87    }
88}
89
90impl Cheatcode for mockCall_3Call {
91    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
92        let Self { callee, msgValue, data, returnData } = self;
93        let _ = make_acc_non_empty(callee, ccx)?;
94
95        mock_call(
96            ccx.state,
97            callee,
98            &Bytes::from(*data),
99            Some(msgValue),
100            returnData,
101            InstructionResult::Return,
102        );
103        Ok(Default::default())
104    }
105}
106
107impl Cheatcode for mockCalls_0Call {
108    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
109        let Self { callee, data, returnData } = self;
110        let _ = make_acc_non_empty(callee, ccx)?;
111
112        mock_calls(ccx.state, callee, data, None, returnData, InstructionResult::Return);
113        Ok(Default::default())
114    }
115}
116
117impl Cheatcode for mockCalls_1Call {
118    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
119        let Self { callee, msgValue, data, returnData } = self;
120        let _ = make_acc_non_empty(callee, ccx)?;
121
122        mock_calls(ccx.state, callee, data, Some(msgValue), returnData, InstructionResult::Return);
123        Ok(Default::default())
124    }
125}
126
127impl Cheatcode for mockCallRevert_0Call {
128    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
129        let Self { callee, data, revertData } = self;
130        let _ = make_acc_non_empty(callee, ccx)?;
131
132        mock_call(ccx.state, callee, data, None, revertData, InstructionResult::Revert);
133        Ok(Default::default())
134    }
135}
136
137impl Cheatcode for mockCallRevert_1Call {
138    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
139        let Self { callee, msgValue, data, revertData } = self;
140        let _ = make_acc_non_empty(callee, ccx)?;
141
142        mock_call(ccx.state, callee, data, Some(msgValue), revertData, InstructionResult::Revert);
143        Ok(Default::default())
144    }
145}
146
147impl Cheatcode for mockCallRevert_2Call {
148    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
149        let Self { callee, data, revertData } = self;
150        let _ = make_acc_non_empty(callee, ccx)?;
151
152        mock_call(
153            ccx.state,
154            callee,
155            &Bytes::from(*data),
156            None,
157            revertData,
158            InstructionResult::Revert,
159        );
160        Ok(Default::default())
161    }
162}
163
164impl Cheatcode for mockCallRevert_3Call {
165    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
166        let Self { callee, msgValue, data, revertData } = self;
167        let _ = make_acc_non_empty(callee, ccx)?;
168
169        mock_call(
170            ccx.state,
171            callee,
172            &Bytes::from(*data),
173            Some(msgValue),
174            revertData,
175            InstructionResult::Revert,
176        );
177        Ok(Default::default())
178    }
179}
180
181impl Cheatcode for mockFunctionCall {
182    fn apply(&self, state: &mut Cheatcodes) -> Result {
183        let Self { callee, target, data } = self;
184        state.mocked_functions.entry(*callee).or_default().insert(data.clone(), *target);
185
186        Ok(Default::default())
187    }
188}
189
190fn mock_call(
191    state: &mut Cheatcodes,
192    callee: &Address,
193    cdata: &Bytes,
194    value: Option<&U256>,
195    rdata: &Bytes,
196    ret_type: InstructionResult,
197) {
198    mock_calls(state, callee, cdata, value, std::slice::from_ref(rdata), ret_type)
199}
200
201fn mock_calls(
202    state: &mut Cheatcodes,
203    callee: &Address,
204    cdata: &Bytes,
205    value: Option<&U256>,
206    rdata_vec: &[Bytes],
207    ret_type: InstructionResult,
208) {
209    state.mocked_calls.entry(*callee).or_default().insert(
210        MockCallDataContext { calldata: cdata.clone(), value: value.copied() },
211        rdata_vec
212            .iter()
213            .map(|rdata| MockCallReturnData { ret_type, data: rdata.clone() })
214            .collect::<VecDeque<_>>(),
215    );
216}
217
218// Etches a single byte onto the account if it is empty to circumvent the `extcodesize`
219// check Solidity might perform.
220fn make_acc_non_empty(callee: &Address, ecx: &mut CheatsCtxt) -> Result {
221    let acc = ecx.journal_mut().load_account(*callee)?;
222
223    let empty_bytecode = acc.info.code.as_ref().is_none_or(Bytecode::is_empty);
224    if empty_bytecode {
225        let code = Bytecode::new_raw(Bytes::from_static(&[0u8]));
226        ecx.journal_mut().set_code(*callee, code);
227    }
228
229    Ok(Default::default())
230}
foundry_cheatcodes/evm/mock.rs

foundry_cheatcodes/evm/
mock.rs
