From 9d2125b013cbcb61dce2546379a79a4d99ba2f78 Mon Sep 17 00:00:00 2001
From: grandizzy <38490174+grandizzy@users.noreply.github.com>
Date: Fri, 22 Mar 2024 15:29:29 +0200
Subject: [PATCH] fix(invariant): call override strategy panic (#7469)

* fix(invariant): override call strat panic

* Add test
---
 crates/evm/fuzz/src/strategies/invariants.rs         | 12 +++++++++++-
 crates/forge/tests/it/invariant.rs                   |  1 +
 .../fuzz/invariant/common/InvariantReentrancy.t.sol  | 12 +++++++++++-
 3 files changed, 23 insertions(+), 2 deletions(-)

diff --git a/crates/evm/fuzz/src/strategies/invariants.rs b/crates/evm/fuzz/src/strategies/invariants.rs
index 5fd766bb4f79..1d226b003478 100644
--- a/crates/evm/fuzz/src/strategies/invariants.rs
+++ b/crates/evm/fuzz/src/strategies/invariants.rs
@@ -25,7 +25,17 @@ pub fn override_call_strat(
     .prop_flat_map(move |target_address| {
         let fuzz_state = fuzz_state.clone();
         let calldata_fuzz_config = calldata_fuzz_config.clone();
-        let (_, abi, functions) = &contracts.lock()[&target_address];
+
+        let contracts = &contracts.lock();
+        let (_, abi, functions) = contracts.get(&target_address).unwrap_or({
+            // Choose a random contract if target selected by lazy strategy is not in fuzz run
+            // identified contracts. This can happen when contract is created in `setUp` call
+            // but is not included in targetContracts.
+            let rand_index = rand::thread_rng().gen_range(0..contracts.iter().len());
+            let (_, contract_specs) = contracts.iter().nth(rand_index).unwrap();
+            contract_specs
+        });
+
         let func = select_random_function(abi, functions);
         func.prop_flat_map(move |func| {
             fuzz_contract_with_calldata(&fuzz_state, &calldata_fuzz_config, target_address, func)
diff --git a/crates/forge/tests/it/invariant.rs b/crates/forge/tests/it/invariant.rs
index e902b6aa9f18..82cf4909e790 100644
--- a/crates/forge/tests/it/invariant.rs
+++ b/crates/forge/tests/it/invariant.rs
@@ -153,6 +153,7 @@ async fn test_invariant() {
 async fn test_invariant_override() {
     let filter = Filter::new(".*", ".*", ".*fuzz/invariant/common/InvariantReentrancy.t.sol");
     let mut runner = TEST_DATA_DEFAULT.runner();
+    runner.test_options.invariant.fail_on_revert = false;
     runner.test_options.invariant.call_override = true;
     let results = runner.test_collect(&filter);
 
diff --git a/testdata/default/fuzz/invariant/common/InvariantReentrancy.t.sol b/testdata/default/fuzz/invariant/common/InvariantReentrancy.t.sol
index 086d5f99ab9e..06b4b21d761f 100644
--- a/testdata/default/fuzz/invariant/common/InvariantReentrancy.t.sol
+++ b/testdata/default/fuzz/invariant/common/InvariantReentrancy.t.sol
@@ -5,7 +5,9 @@ import "ds-test/test.sol";
 
 contract Malicious {
     function world() public {
-        // Does not matter, since it will get overridden.
+        // add code so contract is accounted as valid sender
+        // see https://github.com/foundry-rs/foundry/issues/4245
+        payable(msg.sender).transfer(1);
     }
 }
 
@@ -39,6 +41,14 @@ contract InvariantReentrancy is DSTest {
         vuln = new Vulnerable(address(mal));
     }
 
+    // do not include `mal` in identified contracts
+    // see https://github.com/foundry-rs/foundry/issues/4245
+    function targetContracts() public view returns (address[] memory) {
+        address[] memory targets = new address[](1);
+        targets[0] = address(vuln);
+        return targets;
+    }
+
     function invariantNotStolen() public {
         require(vuln.stolen() == false, "stolen");
     }