fix(invariant): decode custom error with target contract abis (#7559)

* fix(invariant): decode custom error with target contract abis

* Changes after review: don't collect

crates/evm/evm/src/executors/invariant/error.rs

@@ -5,7 +5,9 @@ use alloy_primitives::{Address, Bytes, Log};
 use eyre::Result;
 use foundry_common::contracts::{ContractsByAddress, ContractsByArtifact};
 use foundry_evm_core::{constants::CALLER, decode::RevertDecoder};
-use foundry_evm_fuzz::{BaseCounterExample, CounterExample, FuzzedCases, Reason};
+use foundry_evm_fuzz::{
+    invariant::FuzzRunIdentifiedContracts, BaseCounterExample, CounterExample, FuzzedCases, Reason,
+};
 use foundry_evm_traces::{load_contracts, CallTraceArena, TraceKind, Traces};
 use itertools::Itertools;
 use parking_lot::RwLock;
@@ -98,8 +100,10 @@ pub struct FailedInvariantCaseData {
 }
 
 impl FailedInvariantCaseData {
+    #[allow(clippy::too_many_arguments)]
     pub fn new(
         invariant_contract: &InvariantContract<'_>,
+        targeted_contracts: &FuzzRunIdentifiedContracts,
         error_func: Option<&Function>,
         calldata: &[BasicTxDetails],
         call_result: RawCallResult,
@@ -112,8 +116,16 @@ impl FailedInvariantCaseData {
         } else {
             (None, "Revert")
         };
+
+        // Collect abis of fuzzed and invariant contracts to decode custom error.
+        let targets = targeted_contracts.lock();
+        let abis = targets
+            .iter()
+            .map(|contract| &contract.1 .1)
+            .chain(std::iter::once(invariant_contract.abi));
+
         let revert_reason = RevertDecoder::new()
-            .with_abi(invariant_contract.abi)
+            .with_abis(abis)
             .decode(call_result.result.as_ref(), Some(call_result.exit_reason));
 
         Self {

crates/evm/evm/src/executors/invariant/funcs.rs

@@ -6,7 +6,7 @@ use alloy_primitives::Log;
 use foundry_common::{ContractsByAddress, ContractsByArtifact};
 use foundry_evm_core::constants::CALLER;
 use foundry_evm_coverage::HitMaps;
-use foundry_evm_fuzz::invariant::{BasicTxDetails, InvariantContract};
+use foundry_evm_fuzz::invariant::{BasicTxDetails, FuzzRunIdentifiedContracts, InvariantContract};
 use foundry_evm_traces::{load_contracts, TraceKind, Traces};
 use revm::primitives::U256;
 use std::borrow::Cow;
@@ -16,6 +16,7 @@ use std::borrow::Cow;
 /// Either returns the call result if successful, or nothing if there was an error.
 pub fn assert_invariants(
     invariant_contract: &InvariantContract<'_>,
+    targeted_contracts: &FuzzRunIdentifiedContracts,
     executor: &Executor,
     calldata: &[BasicTxDetails],
     invariant_failures: &mut InvariantFailures,
@@ -49,6 +50,7 @@ pub fn assert_invariants(
         if invariant_failures.error.is_none() {
             let case_data = FailedInvariantCaseData::new(
                 invariant_contract,
+                targeted_contracts,
                 Some(func),
                 calldata,
                 call_result,

crates/evm/evm/src/executors/invariant/mod.rs

@@ -179,6 +179,7 @@ impl<'a> InvariantExecutor<'a> {
         // This does not count as a fuzz run. It will just register the revert.
         let last_call_results = RefCell::new(assert_invariants(
             &invariant_contract,
+            &targeted_contracts,
             &self.executor,
             &[],
             &mut failures.borrow_mut(),
@@ -711,14 +712,15 @@ fn can_continue(
         !executor.is_success(*contract.0, false, Cow::Borrowed(state_changeset), false)
     });
 
-    // Assert invariants IFF the call did not revert and the handlers did not fail.
+    // Assert invariants IF the call did not revert and the handlers did not fail.
     if !call_result.reverted && !handlers_failed {
         if let Some(traces) = call_result.traces {
             run_traces.push(traces);
         }
 
         call_results = assert_invariants(
             invariant_contract,
+            targeted_contracts,
             executor,
             calldata,
             failures,
@@ -735,6 +737,7 @@ fn can_continue(
         if fail_on_revert {
             let case_data = FailedInvariantCaseData::new(
                 invariant_contract,
+                targeted_contracts,
                 None,
                 calldata,
                 call_result,

crates/forge/tests/it/invariant.rs

@@ -145,6 +145,10 @@ async fn test_invariant() {
                 "default/fuzz/invariant/common/InvariantAssume.t.sol:InvariantAssume",
                 vec![("invariant_dummy()", true, None, None, None)],
             ),
+            (
+                "default/fuzz/invariant/common/InvariantCustomError.t.sol:InvariantCustomError",
+                vec![("invariant_decode_error()", true, None, None, None)],
+            ),
         ]),
     );
 }
@@ -411,3 +415,24 @@ async fn test_invariant_assume_respects_restrictions() {
         )]),
     );
 }
+
+#[tokio::test(flavor = "multi_thread")]
+async fn test_invariant_decode_custom_error() {
+    let filter = Filter::new(".*", ".*", ".*fuzz/invariant/common/InvariantCustomError.t.sol");
+    let mut runner = TEST_DATA_DEFAULT.runner();
+    runner.test_options.invariant.fail_on_revert = true;
+    let results = runner.test_collect(&filter);
+    assert_multiple(
+        &results,
+        BTreeMap::from([(
+            "default/fuzz/invariant/common/InvariantCustomError.t.sol:InvariantCustomError",
+            vec![(
+                "invariant_decode_error()",
+                false,
+                Some("InvariantCustomError(111, \"custom\")".into()),
+                None,
+                None,
+            )],
+        )]),
+    );
+}

testdata/default/fuzz/invariant/common/InvariantCustomError.t.sol

@@ -0,0 +1,35 @@
+// SPDX-License-Identifier: MIT OR Apache-2.0
+pragma solidity ^0.8.0;
+
+import "ds-test/test.sol";
+import "cheats/Vm.sol";
+
+contract ContractWithCustomError {
+    error InvariantCustomError(uint256, string);
+
+    function revertWithInvariantCustomError() external {
+        revert InvariantCustomError(111, "custom");
+    }
+}
+
+contract Handler is DSTest {
+    ContractWithCustomError target;
+
+    constructor() {
+        target = new ContractWithCustomError();
+    }
+
+    function revertTarget() external {
+        target.revertWithInvariantCustomError();
+    }
+}
+
+contract InvariantCustomError is DSTest {
+    Handler handler;
+
+    function setUp() external {
+        handler = new Handler();
+    }
+
+    function invariant_decode_error() public {}
+}