foundry_cheatcodes/evm/
mock.rs

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

foundry_cheatcodes/evm/
mock.rs
