foundry_cheatcodes/test/

revert_handlers.rs

use crate::{Error, Result};
use alloy_primitives::{Address, Bytes, address, hex};
use alloy_sol_types::{SolError, SolValue};
use foundry_common::ContractsByArtifact;
use foundry_evm_core::decode::RevertDecoder;
use revm::interpreter::{InstructionResult, return_ok};
use spec::Vm;

use super::{
    assume::{AcceptableRevertParameters, AssumeNoRevert},
    expect::ExpectedRevert,
};

/// For some cheatcodes we may internally change the status of the call, i.e. in `expectRevert`.
/// Solidity will see a successful call and attempt to decode the return data. Therefore, we need
/// to populate the return with dummy bytes so the decode doesn't fail.
///
/// 8192 bytes was arbitrarily chosen because it is long enough for return values up to 256 words in
/// size.
static DUMMY_CALL_OUTPUT: Bytes = Bytes::from_static(&[0u8; 8192]);

/// Same reasoning as [DUMMY_CALL_OUTPUT], but for creates.
const DUMMY_CREATE_ADDRESS: Address = address!("0x0000000000000000000000000000000000000001");

/// Common parameters for expected or assumed reverts. Allows for code reuse.
pub(crate) trait RevertParameters {
    fn reverter(&self) -> Option<Address>;
    fn reason(&self) -> Option<&[u8]>;
    fn partial_match(&self) -> bool;
}

impl RevertParameters for AcceptableRevertParameters {
    fn reverter(&self) -> Option<Address> {
        self.reverter
    }

    fn reason(&self) -> Option<&[u8]> {
        Some(&self.reason)
    }

    fn partial_match(&self) -> bool {
        self.partial_match
    }
}

/// Core logic for handling reverts that may or may not be expected (or assumed).
fn handle_revert(
    is_cheatcode: bool,
    revert_params: &impl RevertParameters,
    status: InstructionResult,
    retdata: &Bytes,
    known_contracts: &Option<ContractsByArtifact>,
    reverter: Option<&Address>,
) -> Result<(), Error> {
    // If expected reverter address is set then check it matches the actual reverter.
    if let (Some(expected_reverter), Some(&actual_reverter)) = (revert_params.reverter(), reverter)
        && expected_reverter != actual_reverter
    {
        return Err(fmt_err!(
            "Reverter != expected reverter: {} != {}",
            actual_reverter,
            expected_reverter
        ));
    }

    let expected_reason = revert_params.reason();
    // If None, accept any revert.
    let Some(expected_reason) = expected_reason else {
        return Ok(());
    };

    let actual_revert = if retdata.is_empty() && !expected_reason.is_empty() {
        if status == InstructionResult::Revert {
            bail!("call reverted as expected, but without data");
        }
        RevertDecoder::new().decode(retdata, Some(status)).into_bytes()
    } else {
        retdata.to_vec()
    };

    // Compare only the first 4 bytes if partial match.
    if revert_params.partial_match()
        && let (Some(actual_prefix), Some(expected_prefix)) =
            (actual_revert.get(..4), expected_reason.get(..4))
        && actual_prefix == expected_prefix
    {
        return Ok(());
    }

    // Compare complete payloads before decoding, which may ignore trailing ABI data.
    if actual_revert == expected_reason {
        return Ok(());
    }

    // Unwrap string errors for legacy raw-message matching.
    let actual_reason = decode_revert(actual_revert);

    if actual_reason == expected_reason
        || (is_cheatcode && memchr::memmem::find(&actual_reason, expected_reason).is_some())
    {
        return Ok(());
    }

    let (actual, expected) = if let Some(contracts) = known_contracts {
        let decoder = RevertDecoder::new().with_abis(contracts.values().map(|c| &c.abi));
        (
            &decoder.decode(actual_reason.as_slice(), Some(status)),
            &decoder.decode(expected_reason, Some(status)),
        )
    } else {
        (&stringify(&actual_reason), &stringify(expected_reason))
    };

    // Lossy decoding can render distinct payloads identically; append the raw data to
    // disambiguate. Note that `retdata` is the original revert data, unlike `actual_reason`,
    // which may have been synthesized from the status or unwrapped by `decode_revert`.
    if expected == actual {
        return Err(fmt_err!(
            "Error != expected error: {actual} (raw {}) != {expected} (raw {})",
            hex::encode_prefixed(retdata),
            hex::encode_prefixed(expected_reason)
        ));
    }

    Err(fmt_err!("Error != expected error: {} != {}", actual, expected))
}

pub(crate) fn handle_assume_no_revert(
    assume_no_revert: &AssumeNoRevert,
    status: InstructionResult,
    retdata: &Bytes,
    known_contracts: &Option<ContractsByArtifact>,
) -> Result<()> {
    // if a generic AssumeNoRevert, return Ok(). Otherwise, iterate over acceptable reasons and try
    // to match against any, otherwise, return an Error with the revert data
    if assume_no_revert.reasons.is_empty() {
        Ok(())
    } else {
        assume_no_revert
            .reasons
            .iter()
            .find_map(|reason| {
                handle_revert(
                    false,
                    reason,
                    status,
                    retdata,
                    known_contracts,
                    assume_no_revert.reverted_by.as_ref(),
                )
                .ok()
            })
            .ok_or_else(|| retdata.clone().into())
    }
}

pub(crate) fn handle_expect_revert(
    is_cheatcode: bool,
    is_create: bool,
    internal_expect_revert: bool,
    expected_revert: &ExpectedRevert,
    status: InstructionResult,
    retdata: Bytes,
    known_contracts: &Option<ContractsByArtifact>,
) -> Result<(Option<Address>, Bytes)> {
    let success_return = || {
        if is_create {
            (Some(DUMMY_CREATE_ADDRESS), Bytes::new())
        } else {
            (None, DUMMY_CALL_OUTPUT.clone())
        }
    };

    // Check depths if it's not an expect cheatcode call and if internal expect reverts not enabled.
    if !is_cheatcode && !internal_expect_revert {
        ensure!(
            expected_revert.max_depth > expected_revert.depth,
            "call didn't revert at a lower depth than cheatcode call depth"
        );
    }

    if expected_revert.count == 0 {
        // If no specific reason or reverter is expected, we just check if it reverted
        if expected_revert.reverter.is_none() && expected_revert.reason.is_none() {
            ensure!(
                matches!(status, return_ok!()),
                "call reverted when it was expected not to revert"
            );
            return Ok(success_return());
        }

        // Flags to track if the reason and reverter match.
        let mut reason_match = expected_revert.reason.as_ref().map(|_| false);
        let mut reverter_match = expected_revert.reverter.as_ref().map(|_| false);

        // If we expect no reverts with a specific reason/reverter, but got a revert,
        // we need to check if it matches our criteria
        if matches!(status, return_ok!()) {
            // No revert occurred, which is what we expected
            Ok(success_return())
        } else {
            // We got a revert, but we expected 0 reverts
            // We need to check if this revert matches our expected criteria

            // Reverter check
            if let (Some(expected_reverter), Some(actual_reverter)) =
                (expected_revert.reverter, expected_revert.reverted_by)
                && expected_reverter == actual_reverter
            {
                reverter_match = Some(true);
            }

            // Reason check
            let expected_reason = expected_revert.reason();
            if let Some(expected_reason) = expected_reason {
                let mut actual_revert: Vec<u8> = retdata.to_vec();
                actual_revert = decode_revert(actual_revert);

                if actual_revert == expected_reason {
                    reason_match = Some(true);
                }
            }

            match (reason_match, reverter_match) {
                (Some(true), Some(true)) => Err(fmt_err!(
                    "expected 0 reverts with reason: {}, from address: {}, but got one",
                    stringify(expected_reason.unwrap_or_default()),
                    expected_revert.reverter.unwrap()
                )),
                (Some(true), None) => Err(fmt_err!(
                    "expected 0 reverts with reason: {}, but got one",
                    stringify(expected_reason.unwrap_or_default())
                )),
                (None, Some(true)) => Err(fmt_err!(
                    "expected 0 reverts from address: {}, but got one",
                    expected_revert.reverter.unwrap()
                )),
                _ => {
                    // The revert doesn't match our criteria, which means it's a different revert
                    // For expectRevert with count=0, any revert should fail the test
                    let decoded_revert = decode_revert(retdata.to_vec());

                    // Provide more specific error messages based on what was expected
                    if let Some(reverter) = expected_revert.reverter {
                        if expected_revert.reason.is_some() {
                            Err(fmt_err!(
                                "call reverted with '{}' from {}, but expected 0 reverts with reason '{}' from {}",
                                stringify(&decoded_revert),
                                expected_revert.reverted_by.unwrap_or_default(),
                                stringify(expected_reason.unwrap_or_default()),
                                reverter
                            ))
                        } else {
                            Err(fmt_err!(
                                "call reverted with '{}' from {}, but expected 0 reverts from {}",
                                stringify(&decoded_revert),
                                expected_revert.reverted_by.unwrap_or_default(),
                                reverter
                            ))
                        }
                    } else {
                        Err(fmt_err!(
                            "call reverted with '{}' when it was expected not to revert",
                            stringify(&decoded_revert)
                        ))
                    }
                }
            }
        }
    } else {
        ensure!(!matches!(status, return_ok!()), "next call did not revert as expected");

        handle_revert(
            is_cheatcode,
            expected_revert,
            status,
            &retdata,
            known_contracts,
            expected_revert.reverted_by.as_ref(),
        )?;
        Ok(success_return())
    }
}

fn stringify(data: &[u8]) -> String {
    if let Ok(s) = String::abi_decode(data) {
        return s;
    }
    if data.is_ascii() {
        return std::str::from_utf8(data).unwrap().to_owned();
    }
    hex::encode_prefixed(data)
}

fn decode_revert(revert: Vec<u8>) -> Vec<u8> {
    if matches!(
        revert.get(..4).map(|s| s.try_into().unwrap()),
        Some(Vm::CheatcodeError::SELECTOR | alloy_sol_types::Revert::SELECTOR)
    ) && let Ok(decoded) = Vec::<u8>::abi_decode(&revert[4..])
    {
        return decoded;
    }
    revert
}