Skip to main content

foundry_cheatcodes/test/
revert_handlers.rs

1use crate::{Error, Result};
2use alloy_primitives::{Address, Bytes, hex};
3use alloy_sol_types::{SolError, SolValue};
4use foundry_common::ContractsByArtifact;
5use foundry_evm_core::decode::RevertDecoder;
6use revm::interpreter::InstructionResult;
7use spec::Vm;
8
9use super::{
10    assume::{AcceptableRevertParameters, AssumeNoRevert},
11    expect::ExpectedRevert,
12};
13
14/// For some cheatcodes we may internally change the status of the call, i.e. in `expectRevert`.
15/// Solidity will see a successful call and attempt to decode the return data. Therefore, we need
16/// to populate the return with dummy bytes so the decode doesn't fail.
17///
18/// 8192 bytes was arbitrarily chosen because it is long enough for return values up to 256 words in
19/// size.
20static DUMMY_CALL_OUTPUT: Bytes = Bytes::from_static(&[0u8; 8192]);
21
22/// Same reasoning as [DUMMY_CALL_OUTPUT], but for creates.
23const DUMMY_CREATE_ADDRESS: Address = Address::with_last_byte(1);
24
25/// Common parameters for expected or assumed reverts. Allows for code reuse.
26pub(crate) trait RevertParameters {
27    fn reverter(&self) -> Option<Address>;
28    fn reason(&self) -> Option<&[u8]>;
29    fn partial_match(&self) -> bool;
30}
31
32impl RevertParameters for AcceptableRevertParameters {
33    fn reverter(&self) -> Option<Address> {
34        self.reverter
35    }
36
37    fn reason(&self) -> Option<&[u8]> {
38        Some(&self.reason)
39    }
40
41    fn partial_match(&self) -> bool {
42        self.partial_match
43    }
44}
45
46/// Core logic for handling reverts that may or may not be expected (or assumed).
47fn handle_revert(
48    is_cheatcode: bool,
49    revert_params: &impl RevertParameters,
50    status: InstructionResult,
51    retdata: &Bytes,
52    known_contracts: &Option<ContractsByArtifact>,
53    reverter: Option<&Address>,
54) -> Result<(), Error> {
55    // If expected reverter address is set then check it matches the actual reverter.
56    if let (Some(expected_reverter), Some(&actual_reverter)) = (revert_params.reverter(), reverter)
57        && expected_reverter != actual_reverter
58    {
59        return Err(fmt_err!(
60            "Reverter != expected reverter: {} != {}",
61            actual_reverter,
62            expected_reverter
63        ));
64    }
65
66    let expected_reason = revert_params.reason();
67    // If None, accept any revert.
68    let Some(expected_reason) = expected_reason else {
69        return Ok(());
70    };
71
72    let actual_revert = if retdata.is_empty() && !expected_reason.is_empty() {
73        if status == InstructionResult::Revert {
74            bail!("call reverted as expected, but without data");
75        }
76        RevertDecoder::new().decode(retdata, Some(status)).into_bytes()
77    } else {
78        retdata.to_vec()
79    };
80
81    // Compare only the first 4 bytes if partial match.
82    if revert_params.partial_match()
83        && let (Some(actual_prefix), Some(expected_prefix)) =
84            (actual_revert.get(..4), expected_reason.get(..4))
85        && actual_prefix == expected_prefix
86    {
87        return Ok(());
88    }
89
90    // Compare complete payloads before decoding, which may ignore trailing ABI data.
91    if actual_revert == expected_reason {
92        return Ok(());
93    }
94
95    // Unwrap string errors for legacy raw-message matching.
96    let actual_reason = decode_revert(actual_revert);
97
98    if actual_reason == expected_reason
99        || (is_cheatcode && memchr::memmem::find(&actual_reason, expected_reason).is_some())
100    {
101        return Ok(());
102    }
103
104    let (actual, expected) = if let Some(contracts) = known_contracts {
105        let decoder = RevertDecoder::new().with_abis(contracts.values().map(|c| &c.abi));
106        (
107            &decoder.decode(actual_reason.as_slice(), Some(status)),
108            &decoder.decode(expected_reason, Some(status)),
109        )
110    } else {
111        (&stringify(&actual_reason), &stringify(expected_reason))
112    };
113
114    // Lossy decoding can render distinct payloads identically; append the raw data to
115    // disambiguate. Note that `retdata` is the original revert data, unlike `actual_reason`,
116    // which may have been synthesized from the status or unwrapped by `decode_revert`.
117    if expected == actual {
118        return Err(fmt_err!(
119            "Error != expected error: {actual} (raw {}) != {expected} (raw {})",
120            hex::encode_prefixed(retdata),
121            hex::encode_prefixed(expected_reason)
122        ));
123    }
124
125    Err(fmt_err!("Error != expected error: {} != {}", actual, expected))
126}
127
128pub(crate) fn handle_assume_no_revert(
129    assume_no_revert: &AssumeNoRevert,
130    status: InstructionResult,
131    retdata: &Bytes,
132    known_contracts: &Option<ContractsByArtifact>,
133) -> Result<()> {
134    // if a generic AssumeNoRevert, return Ok(). Otherwise, iterate over acceptable reasons and try
135    // to match against any, otherwise, return an Error with the revert data
136    if assume_no_revert.reasons.is_empty() {
137        Ok(())
138    } else {
139        assume_no_revert
140            .reasons
141            .iter()
142            .find_map(|reason| {
143                handle_revert(
144                    false,
145                    reason,
146                    status,
147                    retdata,
148                    known_contracts,
149                    assume_no_revert.reverted_by.as_ref(),
150                )
151                .ok()
152            })
153            .ok_or_else(|| retdata.clone().into())
154    }
155}
156
157pub(crate) fn handle_expect_revert(
158    is_cheatcode: bool,
159    is_create: bool,
160    internal_expect_revert: bool,
161    expected_revert: &ExpectedRevert,
162    status: InstructionResult,
163    retdata: Bytes,
164    known_contracts: &Option<ContractsByArtifact>,
165) -> Result<(Option<Address>, Bytes)> {
166    let success_return = || {
167        if is_create {
168            (Some(DUMMY_CREATE_ADDRESS), Bytes::new())
169        } else {
170            (None, DUMMY_CALL_OUTPUT.clone())
171        }
172    };
173
174    // Check depths if it's not an expect cheatcode call and if internal expect reverts not enabled.
175    if !is_cheatcode && !internal_expect_revert {
176        ensure!(
177            expected_revert.max_depth > expected_revert.depth,
178            "call didn't revert at a lower depth than cheatcode call depth"
179        );
180    }
181
182    if expected_revert.count == 0 {
183        // If no specific reason or reverter is expected, we just check if it reverted
184        if expected_revert.reverter.is_none() && expected_revert.reason.is_none() {
185            ensure!(status.is_ok(), "call reverted when it was expected not to revert");
186            return Ok(success_return());
187        }
188
189        // Flags to track if the reason and reverter match.
190        let mut reason_match = expected_revert.reason.as_ref().map(|_| false);
191        let mut reverter_match = expected_revert.reverter.as_ref().map(|_| false);
192
193        // If we expect no reverts with a specific reason/reverter, but got a revert,
194        // we need to check if it matches our criteria
195        if status.is_ok() {
196            // No revert occurred, which is what we expected
197            Ok(success_return())
198        } else {
199            // We got a revert, but we expected 0 reverts
200            // We need to check if this revert matches our expected criteria
201
202            // Reverter check
203            if let (Some(expected_reverter), Some(actual_reverter)) =
204                (expected_revert.reverter, expected_revert.reverted_by)
205                && expected_reverter == actual_reverter
206            {
207                reverter_match = Some(true);
208            }
209
210            // Reason check
211            let expected_reason = expected_revert.reason();
212            if let Some(expected_reason) = expected_reason {
213                let mut actual_revert: Vec<u8> = retdata.to_vec();
214                actual_revert = decode_revert(actual_revert);
215
216                if actual_revert == expected_reason {
217                    reason_match = Some(true);
218                }
219            }
220
221            match (reason_match, reverter_match) {
222                (Some(true), Some(true)) => Err(fmt_err!(
223                    "expected 0 reverts with reason: {}, from address: {}, but got one",
224                    stringify(expected_reason.unwrap_or_default()),
225                    expected_revert.reverter.unwrap()
226                )),
227                (Some(true), None) => Err(fmt_err!(
228                    "expected 0 reverts with reason: {}, but got one",
229                    stringify(expected_reason.unwrap_or_default())
230                )),
231                (None, Some(true)) => Err(fmt_err!(
232                    "expected 0 reverts from address: {}, but got one",
233                    expected_revert.reverter.unwrap()
234                )),
235                _ => {
236                    // The revert doesn't match our criteria, which means it's a different revert
237                    // For expectRevert with count=0, any revert should fail the test
238                    let decoded_revert = decode_revert(retdata.to_vec());
239
240                    // Provide more specific error messages based on what was expected
241                    if let Some(reverter) = expected_revert.reverter {
242                        if expected_revert.reason.is_some() {
243                            Err(fmt_err!(
244                                "call reverted with '{}' from {}, but expected 0 reverts with reason '{}' from {}",
245                                stringify(&decoded_revert),
246                                expected_revert.reverted_by.unwrap_or_default(),
247                                stringify(expected_reason.unwrap_or_default()),
248                                reverter
249                            ))
250                        } else {
251                            Err(fmt_err!(
252                                "call reverted with '{}' from {}, but expected 0 reverts from {}",
253                                stringify(&decoded_revert),
254                                expected_revert.reverted_by.unwrap_or_default(),
255                                reverter
256                            ))
257                        }
258                    } else {
259                        Err(fmt_err!(
260                            "call reverted with '{}' when it was expected not to revert",
261                            stringify(&decoded_revert)
262                        ))
263                    }
264                }
265            }
266        }
267    } else {
268        ensure!(!status.is_ok(), "next call did not revert as expected");
269
270        handle_revert(
271            is_cheatcode,
272            expected_revert,
273            status,
274            &retdata,
275            known_contracts,
276            expected_revert.reverted_by.as_ref(),
277        )?;
278        Ok(success_return())
279    }
280}
281
282fn stringify(data: &[u8]) -> String {
283    if let Ok(s) = String::abi_decode(data) {
284        return s;
285    }
286    if data.is_ascii() {
287        return std::str::from_utf8(data).unwrap().to_owned();
288    }
289    hex::encode_prefixed(data)
290}
291
292fn decode_revert(revert: Vec<u8>) -> Vec<u8> {
293    if matches!(
294        revert.get(..4).map(|s| s.try_into().unwrap()),
295        Some(Vm::CheatcodeError::SELECTOR | alloy_sol_types::Revert::SELECTOR)
296    ) && let Ok(decoded) = Vec::<u8>::abi_decode(&revert[4..])
297    {
298        return decoded;
299    }
300    revert
301}