foundry_cheatcodes/evm/
mock.rs

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

foundry_cheatcodes/evm/
mock.rs
