From 692adefee4d8a8d95c7fa5c2121710e95dd0cdec Mon Sep 17 00:00:00 2001
From: grandizzy <grandizzy.the.egg@gmail.com>
Date: Tue, 27 Feb 2024 07:46:20 +0200
Subject: [PATCH] Add test from issue 5868

---
 crates/config/src/fuzz.rs                     |  1 +
 crates/forge/tests/it/invariant.rs            | 52 +++++++++++
 .../common/InvariantCalldataDictionary.t.sol  | 89 +++++++++++++++++++
 3 files changed, 142 insertions(+)
 create mode 100644 testdata/fuzz/invariant/common/InvariantCalldataDictionary.t.sol

diff --git a/crates/config/src/fuzz.rs b/crates/config/src/fuzz.rs
index ff138dc8a6b25..11d214670488c 100644
--- a/crates/config/src/fuzz.rs
+++ b/crates/config/src/fuzz.rs
@@ -88,6 +88,7 @@ pub struct FuzzDictionaryConfig {
     pub max_fuzz_dictionary_values: usize,
     /// How many random addresses to use and to recycle when fuzzing calldata.
     /// If not specified then `max_fuzz_dictionary_addresses` value applies.
+    #[serde(deserialize_with = "crate::deserialize_usize_or_max")]
     pub max_calldata_fuzz_dictionary_addresses: usize,
 }
 
diff --git a/crates/forge/tests/it/invariant.rs b/crates/forge/tests/it/invariant.rs
index 333392b29106d..3f72b59713465 100644
--- a/crates/forge/tests/it/invariant.rs
+++ b/crates/forge/tests/it/invariant.rs
@@ -118,6 +118,10 @@ async fn test_invariant() {
                     None,
                 )],
             ),
+            (
+                "fuzz/invariant/common/InvariantCalldataDictionary.t.sol:InvariantCalldataDictionary",
+                vec![("invariant_owner_never_changes()", true, None, None, None)],
+            ),
         ]),
     );
 }
@@ -244,3 +248,51 @@ async fn test_invariant_shrink() {
         }
     };
 }
+
+#[tokio::test(flavor = "multi_thread")]
+async fn test_invariant_calldata_fuzz_dictionary_addresses() {
+    let mut runner = runner().await;
+
+    // should not fail with default options (address dict not finite)
+    let opts = test_opts();
+    runner.test_options = opts.clone();
+    let results = runner
+        .test_collect(
+            &Filter::new(".*", ".*", ".*fuzz/invariant/common/InvariantCalldataDictionary.t.sol"),
+            opts,
+        )
+        .await;
+    assert_multiple(
+        &results,
+        BTreeMap::from([(
+            "fuzz/invariant/common/InvariantCalldataDictionary.t.sol:InvariantCalldataDictionary",
+            vec![("invariant_owner_never_changes()", true, None, None, None)],
+        )]),
+    );
+
+    // same test should fail when calldata address dict is bounded
+    let mut opts = test_opts();
+    // set address dictionary to single entry to fail fast
+    opts.invariant.dictionary.max_calldata_fuzz_dictionary_addresses = 1;
+    runner.test_options = opts.clone();
+
+    let results = runner
+        .test_collect(
+            &Filter::new(".*", ".*", ".*fuzz/invariant/common/InvariantCalldataDictionary.t.sol"),
+            opts,
+        )
+        .await;
+    assert_multiple(
+        &results,
+        BTreeMap::from([(
+            "fuzz/invariant/common/InvariantCalldataDictionary.t.sol:InvariantCalldataDictionary",
+            vec![(
+                "invariant_owner_never_changes()",
+                false,
+                Some("<empty revert data>".into()),
+                None,
+                None,
+            )],
+        )]),
+    );
+}
diff --git a/testdata/fuzz/invariant/common/InvariantCalldataDictionary.t.sol b/testdata/fuzz/invariant/common/InvariantCalldataDictionary.t.sol
new file mode 100644
index 0000000000000..de7449ab98410
--- /dev/null
+++ b/testdata/fuzz/invariant/common/InvariantCalldataDictionary.t.sol
@@ -0,0 +1,89 @@
+// SPDX-License-Identifier: MIT OR Apache-2.0
+pragma solidity 0.8.18;
+
+import "ds-test/test.sol";
+import "../../../cheats/Vm.sol";
+
+struct FuzzSelector {
+    address addr;
+    bytes4[] selectors;
+}
+
+// https://github.com/foundry-rs/foundry/issues/5868
+contract Owned {
+    address public owner;
+    address private ownerCandidate;
+
+    constructor() {
+        owner = msg.sender;
+    }
+
+    modifier onlyOwner() {
+        require(msg.sender == owner);
+        _;
+    }
+
+    modifier onlyOwnerCandidate() {
+        require(msg.sender == ownerCandidate);
+        _;
+    }
+
+    function transferOwnership(address candidate) external onlyOwner {
+        ownerCandidate = candidate;
+    }
+
+    function acceptOwnership() external onlyOwnerCandidate {
+        // console2.log("in contract %s accepta", msg.sender);
+        owner = ownerCandidate;
+    }
+
+    function curCandidate() external view returns (address) {
+        return ownerCandidate;
+    }
+}
+
+contract Handler is DSTest {
+    Vm constant vm = Vm(HEVM_ADDRESS);
+    Owned owned;
+
+    constructor(Owned _owned) {
+        owned = _owned;
+    }
+
+    function transferOwnership(address sender, address candidate) external {
+        vm.assume(sender != address(0));
+        vm.prank(sender);
+        owned.transferOwnership(candidate);
+    }
+
+    function acceptOwnership(address sender) external {
+        vm.assume(sender != address(0));
+        vm.prank(sender);
+        owned.acceptOwnership();
+    }
+}
+
+contract InvariantCalldataDictionary is DSTest {
+    address owner;
+    Owned owned;
+    Handler handler;
+
+    function setUp() public {
+        owner = address(this);
+        owned = new Owned();
+        handler = new Handler(owned);
+    }
+
+    function targetSelectors() public returns (FuzzSelector[] memory) {
+        FuzzSelector[] memory targets = new FuzzSelector[](1);
+        bytes4[] memory selectors = new bytes4[](2);
+        selectors[0] = handler.transferOwnership.selector;
+        selectors[1] = handler.acceptOwnership.selector;
+        targets[0] = FuzzSelector(address(handler), selectors);
+        return targets;
+    }
+
+    function invariant_owner_never_changes() public {
+        assertEq(owned.owner(), owner);
+    }
+}