foundry_cheatcodes/evm/
mock.rs

1use crate::{Cheatcode, Cheatcodes, CheatsCtxt, Result, Vm::*};
2use alloy_primitives::{Address, Bytes, U256};
3use foundry_evm_core::evm::FoundryEvmNetwork;
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<FEN: FoundryEvmNetwork>(&self, state: &mut Cheatcodes<FEN>) -> 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<FEN: FoundryEvmNetwork>(&self, ccx: &mut CheatsCtxt<'_, '_, FEN>) -> Result {
56        let Self { callee, data, returnData } = self;
57        let _ = make_acc_non_empty(callee, ccx)?;
58
59        mock_call(ccx.state, callee, data, None, returnData, InstructionResult::Return);
60        Ok(Default::default())
61    }
62}
63
64impl Cheatcode for mockCall_1Call {
65    fn apply_stateful<FEN: FoundryEvmNetwork>(&self, ccx: &mut CheatsCtxt<'_, '_, FEN>) -> Result {
66        let Self { callee, msgValue, data, returnData } = self;
67        let _ = make_acc_non_empty(callee, ccx)?;
68
69        mock_call(ccx.state, callee, data, Some(msgValue), returnData, InstructionResult::Return);
70        Ok(Default::default())
71    }
72}
73
74impl Cheatcode for mockCall_2Call {
75    fn apply_stateful<FEN: FoundryEvmNetwork>(&self, ccx: &mut CheatsCtxt<'_, '_, FEN>) -> Result {
76        let Self { callee, data, returnData } = self;
77        let _ = make_acc_non_empty(callee, ccx)?;
78
79        mock_call(
80            ccx.state,
81            callee,
82            &Bytes::from(*data),
83            None,
84            returnData,
85            InstructionResult::Return,
86        );
87        Ok(Default::default())
88    }
89}
90
91impl Cheatcode for mockCall_3Call {
92    fn apply_stateful<FEN: FoundryEvmNetwork>(&self, ccx: &mut CheatsCtxt<'_, '_, FEN>) -> Result {
93        let Self { callee, msgValue, data, returnData } = self;
94        let _ = make_acc_non_empty(callee, ccx)?;
95
96        mock_call(
97            ccx.state,
98            callee,
99            &Bytes::from(*data),
100            Some(msgValue),
101            returnData,
102            InstructionResult::Return,
103        );
104        Ok(Default::default())
105    }
106}
107
108impl Cheatcode for mockCalls_0Call {
109    fn apply_stateful<FEN: FoundryEvmNetwork>(&self, ccx: &mut CheatsCtxt<'_, '_, FEN>) -> Result {
110        let Self { callee, data, returnData } = self;
111        let _ = make_acc_non_empty(callee, ccx)?;
112
113        mock_calls(ccx.state, callee, data, None, returnData, InstructionResult::Return);
114        Ok(Default::default())
115    }
116}
117
118impl Cheatcode for mockCalls_1Call {
119    fn apply_stateful<FEN: FoundryEvmNetwork>(&self, ccx: &mut CheatsCtxt<'_, '_, FEN>) -> Result {
120        let Self { callee, msgValue, data, returnData } = self;
121        let _ = make_acc_non_empty(callee, ccx)?;
122
123        mock_calls(ccx.state, callee, data, Some(msgValue), returnData, InstructionResult::Return);
124        Ok(Default::default())
125    }
126}
127
128impl Cheatcode for mockCallRevert_0Call {
129    fn apply_stateful<FEN: FoundryEvmNetwork>(&self, ccx: &mut CheatsCtxt<'_, '_, FEN>) -> Result {
130        let Self { callee, data, revertData } = self;
131        let _ = make_acc_non_empty(callee, ccx)?;
132
133        mock_call(ccx.state, callee, data, None, revertData, InstructionResult::Revert);
134        Ok(Default::default())
135    }
136}
137
138impl Cheatcode for mockCallRevert_1Call {
139    fn apply_stateful<FEN: FoundryEvmNetwork>(&self, ccx: &mut CheatsCtxt<'_, '_, FEN>) -> Result {
140        let Self { callee, msgValue, data, revertData } = self;
141        let _ = make_acc_non_empty(callee, ccx)?;
142
143        mock_call(ccx.state, callee, data, Some(msgValue), revertData, InstructionResult::Revert);
144        Ok(Default::default())
145    }
146}
147
148impl Cheatcode for mockCallRevert_2Call {
149    fn apply_stateful<FEN: FoundryEvmNetwork>(&self, ccx: &mut CheatsCtxt<'_, '_, FEN>) -> Result {
150        let Self { callee, data, revertData } = self;
151        let _ = make_acc_non_empty(callee, ccx)?;
152
153        mock_call(
154            ccx.state,
155            callee,
156            &Bytes::from(*data),
157            None,
158            revertData,
159            InstructionResult::Revert,
160        );
161        Ok(Default::default())
162    }
163}
164
165impl Cheatcode for mockCallRevert_3Call {
166    fn apply_stateful<FEN: FoundryEvmNetwork>(&self, ccx: &mut CheatsCtxt<'_, '_, FEN>) -> Result {
167        let Self { callee, msgValue, data, revertData } = self;
168        let _ = make_acc_non_empty(callee, ccx)?;
169
170        mock_call(
171            ccx.state,
172            callee,
173            &Bytes::from(*data),
174            Some(msgValue),
175            revertData,
176            InstructionResult::Revert,
177        );
178        Ok(Default::default())
179    }
180}
181
182impl Cheatcode for mockFunctionCall {
183    fn apply<FEN: FoundryEvmNetwork>(&self, state: &mut Cheatcodes<FEN>) -> Result {
184        let Self { callee, target, data } = self;
185        state.mocked_functions.entry(*callee).or_default().insert(data.clone(), *target);
186
187        Ok(Default::default())
188    }
189}
190
191fn mock_call<FEN: FoundryEvmNetwork>(
192    state: &mut Cheatcodes<FEN>,
193    callee: &Address,
194    cdata: &Bytes,
195    value: Option<&U256>,
196    rdata: &Bytes,
197    ret_type: InstructionResult,
198) {
199    mock_calls(state, callee, cdata, value, std::slice::from_ref(rdata), ret_type)
200}
201
202fn mock_calls<FEN: FoundryEvmNetwork>(
203    state: &mut Cheatcodes<FEN>,
204    callee: &Address,
205    cdata: &Bytes,
206    value: Option<&U256>,
207    rdata_vec: &[Bytes],
208    ret_type: InstructionResult,
209) {
210    state.mocked_calls.entry(*callee).or_default().insert(
211        MockCallDataContext { calldata: cdata.clone(), value: value.copied() },
212        rdata_vec
213            .iter()
214            .map(|rdata| MockCallReturnData { ret_type, data: rdata.clone() })
215            .collect::<VecDeque<_>>(),
216    );
217}
218
219// Etches a single byte onto the account if it is empty to circumvent the `extcodesize`
220// check Solidity might perform.
221fn make_acc_non_empty<FEN: FoundryEvmNetwork>(
222    callee: &Address,
223    ccx: &mut CheatsCtxt<'_, '_, FEN>,
224) -> Result {
225    let empty_bytecode = {
226        let acc = ccx.ecx.journal_mut().load_account(*callee)?;
227        acc.info.code.as_ref().is_none_or(Bytecode::is_empty)
228    };
229    if empty_bytecode {
230        let code = Bytecode::new_raw(Bytes::from_static(&[0u8]));
231        ccx.ecx.journal_mut().set_code(*callee, code);
232    }
233
234    Ok(Default::default())
235}
foundry_cheatcodes/evm/mock.rs

foundry_cheatcodes/evm/
mock.rs
