chore: improve multisig builder transaction simulation building (#144)

* improve multisig builder transaction simulation building

* use encodeCall in tests

* inline safeOverrides internal function

* refactor transaction overrides and simplify tests

* remove unnecessary comment

Author: jackchuma

script/universal/DoubleNestedMultisigBuilder.sol

@@ -11,7 +11,7 @@ abstract contract DoubleNestedMultisigBuilder is MultisigScript {
     /*
      * @custom:deprecated Use `sign(address[] memory _safes)` instead.
      */
-    function sign(address _signerSafe, address _intermediateSafe) public {
+    function sign(address _signerSafe, address _intermediateSafe) external {
         sign(_toArray(_signerSafe, _intermediateSafe));
     }
 

script/universal/MultisigBuilder.sol

@@ -11,7 +11,7 @@ abstract contract MultisigBuilder is MultisigScript {
     /*
      * @custom:deprecated Use `sign(address[] memory _safes)` instead, with an empty array.
      */
-    function sign() public {
+    function sign() external {
         sign(new address[](0));
     }
 

script/universal/MultisigScript.sol

@@ -433,8 +433,10 @@ abstract contract MultisigScript is Script {
     {
         IMulticall3.Call3[] memory calls = _simulateForSignerCalls(_safes, _datas, _value);
 
+        bytes32 firstCallDataHash = _getTransactionHash(_safes[0], _datas[0], _value);
+
         // Now define the state overrides for the simulation.
-        Simulation.StateOverride[] memory overrides = _overrides(_safes);
+        Simulation.StateOverride[] memory overrides = _overrides(_safes, firstCallDataHash);
 
         bytes memory txData = abi.encodeCall(IMulticall3.aggregate3, (calls));
         console.log("---\nSimulation link:");
@@ -451,13 +453,12 @@ abstract contract MultisigScript is Script {
 
     function _simulateForSignerCalls(address[] memory _safes, bytes[] memory _datas, uint256 _value)
         private
-        pure
+        view
         returns (IMulticall3.Call3[] memory)
     {
         IMulticall3.Call3[] memory calls = new IMulticall3.Call3[](_safes.length);
         for (uint256 i; i < _safes.length; i++) {
-            address signer = i == 0 ? MULTICALL3_ADDRESS : _safes[i - 1];
-            // uint256 value = i == _safes.length - 1 ? _value : 0;
+            address signer = i == 0 ? msg.sender : _safes[i - 1];
 
             calls[i] = IMulticall3.Call3({
                 target: _safes[i],
@@ -471,35 +472,31 @@ abstract contract MultisigScript is Script {
         return calls;
     }
 
-    function _overrides(address[] memory _safes) internal view returns (Simulation.StateOverride[] memory) {
-        Simulation.StateOverride[] memory simOverrides = _simulationOverrides();
-        Simulation.StateOverride[] memory overrides =
-            new Simulation.StateOverride[](_safes.length + simOverrides.length);
-        for (uint256 i = 0; i < _safes.length; i++) {
-            address owner = i == 0 ? MULTICALL3_ADDRESS : address(0);
-            overrides[i] = _safeOverrides(_safes[i], owner);
-        }
-        for (uint256 i = 0; i < simOverrides.length; i++) {
-            overrides[i + _safes.length] = simOverrides[i];
-        }
-        return overrides;
-    }
-
     // The state change simulation can set the threshold, owner address and/or nonce.
     // This allows simulation of the final transaction by overriding the threshold to 1.
     // State changes reflected in the simulation as a result of these overrides will
     // not be reflected in the prod execution.
-    function _safeOverrides(address _safe, address _owner)
+    function _overrides(address[] memory _safes, bytes32 firstCallDataHash)
         internal
         view
-        virtual
-        returns (Simulation.StateOverride memory)
+        returns (Simulation.StateOverride[] memory)
     {
-        uint256 _nonce = _getNonce(_safe);
-        if (_owner == address(0)) {
-            return Simulation.overrideSafeThresholdAndNonce(_safe, _nonce);
+        Simulation.StateOverride[] memory simOverrides = _simulationOverrides();
+        Simulation.StateOverride[] memory overrides =
+            new Simulation.StateOverride[](_safes.length + simOverrides.length);
+
+        uint256 nonce = _getNonce(_safes[0]);
+        overrides[0] = Simulation.overrideSafeThresholdApprovalAndNonce(_safes[0], nonce, msg.sender, firstCallDataHash);
+
+        for (uint256 i = 1; i < _safes.length; i++) {
+            overrides[i] = Simulation.overrideSafeThresholdAndNonce(_safes[i], _getNonce(_safes[i]));
         }
-        return Simulation.overrideSafeThresholdOwnerAndNonce(_safe, _owner, _nonce);
+
+        for (uint256 i = 0; i < simOverrides.length; i++) {
+            overrides[i + _safes.length] = simOverrides[i];
+        }
+
+        return overrides;
     }
 
     // Get the nonce to use for the given safe, for signing and simulations.

script/universal/NestedMultisigBuilder.sol

@@ -11,7 +11,7 @@ abstract contract NestedMultisigBuilder is MultisigScript {
     /*
      * @custom:deprecated Use `sign(address[] memory _safes)` instead.
      */
-    function sign(address _signerSafe) public {
+    function sign(address _signerSafe) external {
         sign(_toArray(_signerSafe));
     }
 

script/universal/Simulation.sol

@@ -54,29 +54,44 @@ library Simulation {
         return accesses;
     }
 
-    function overrideSafeThresholdAndNonce(address _safe, uint256 _nonce)
+    function overrideSafeThresholdApprovalAndNonce(address _safe, uint256 _nonce, address owner, bytes32 dataHash)
         internal
         view
         returns (StateOverride memory)
     {
+        // solhint-disable-next-line max-line-length
         StateOverride memory state = StateOverride({contractAddress: _safe, overrides: new StorageOverride[](0)});
         state = addThresholdOverride(_safe, state);
         state = addNonceOverride(_safe, state, _nonce);
+        state = addApprovalOverride(state, owner, dataHash);
         return state;
     }
 
-    function overrideSafeThresholdOwnerAndNonce(address _safe, address _owner, uint256 _nonce)
+    function overrideSafeThresholdAndNonce(address _safe, uint256 _nonce)
         internal
         view
         returns (StateOverride memory)
     {
         StateOverride memory state = StateOverride({contractAddress: _safe, overrides: new StorageOverride[](0)});
         state = addThresholdOverride(_safe, state);
-        state = addOwnerOverride(_safe, state, _owner);
         state = addNonceOverride(_safe, state, _nonce);
         return state;
     }
 
+    function addApprovalOverride(StateOverride memory _state, address owner, bytes32 dataHash)
+        internal
+        pure
+        returns (StateOverride memory)
+    {
+        return addOverride(
+            _state,
+            StorageOverride({
+                key: keccak256(abi.encode(dataHash, keccak256(abi.encode(owner, uint256(8))))),
+                value: bytes32(uint256(0x1))
+            })
+        );
+    }
+
     function addThresholdOverride(address _safe, StateOverride memory _state)
         internal
         view

test/universal/DoubleNestedMultisigBuilder.t.sol

@@ -79,14 +79,20 @@ contract DoubleNestedMultisigBuilderTest is Test, DoubleNestedMultisigBuilder {
 
     function test_sign_double_nested_safe1() external {
         vm.recordLogs();
-        sign(safe1, safe3);
+        bytes memory txData = abi.encodeCall(DoubleNestedMultisigBuilder.sign, (safe1, safe3));
+        vm.prank(wallet1.addr);
+        (bool success,) = address(this).call(txData);
+        vm.assertTrue(success);
         Vm.Log[] memory logs = vm.getRecordedLogs();
         assertEq(keccak256(logs[logs.length - 1].data), keccak256(abi.encode(dataToSign1)));
     }
 
     function test_sign_double_nested_safe2() external {
         vm.recordLogs();
-        sign(safe2, safe3);
+        bytes memory txData = abi.encodeCall(DoubleNestedMultisigBuilder.sign, (safe2, safe3));
+        vm.prank(wallet2.addr);
+        (bool success,) = address(this).call(txData);
+        vm.assertTrue(success);
         Vm.Log[] memory logs = vm.getRecordedLogs();
         assertEq(keccak256(logs[logs.length - 1].data), keccak256(abi.encode(dataToSign2)));
     }

test/universal/MultisigBuilder.t.sol

@@ -58,7 +58,10 @@ contract MultisigBuilderTest is Test, MultisigBuilder {
         buildCallsInternal = _buildCallsNoValue;
 
         vm.recordLogs();
-        sign();
+        bytes memory txData = abi.encodeCall(MultisigBuilder.sign, ());
+        vm.prank(wallet1.addr);
+        (bool success,) = address(this).call(txData);
+        vm.assertTrue(success);
         Vm.Log[] memory logs = vm.getRecordedLogs();
         assertEq(keccak256(logs[logs.length - 1].data), keccak256(abi.encode(dataToSignNoValue)));
     }
@@ -67,7 +70,10 @@ contract MultisigBuilderTest is Test, MultisigBuilder {
         buildCallsInternal = _buildCallsWithValue;
 
         vm.recordLogs();
-        sign();
+        bytes memory txData = abi.encodeCall(MultisigBuilder.sign, ());
+        vm.prank(wallet1.addr);
+        (bool success,) = address(this).call(txData);
+        vm.assertTrue(success);
         Vm.Log[] memory logs = vm.getRecordedLogs();
         assertEq(keccak256(logs[logs.length - 1].data), keccak256(abi.encode(dataToSignWithValue)));
     }

test/universal/NestedMultisigBuilder.t.sol

@@ -73,14 +73,20 @@ contract NestedMultisigBuilderTest is Test, NestedMultisigBuilder {
 
     function test_sign_safe1() external {
         vm.recordLogs();
-        sign(safe1);
+        bytes memory txData = abi.encodeCall(NestedMultisigBuilder.sign, (safe1));
+        vm.prank(wallet1.addr);
+        (bool success,) = address(this).call(txData);
+        vm.assertTrue(success);
         Vm.Log[] memory logs = vm.getRecordedLogs();
         assertEq(keccak256(logs[logs.length - 1].data), keccak256(abi.encode(dataToSign1)));
     }
 
     function test_sign_safe2() external {
         vm.recordLogs();
-        sign(safe2);
+        bytes memory txData = abi.encodeCall(NestedMultisigBuilder.sign, (safe2));
+        vm.prank(wallet2.addr);
+        (bool success,) = address(this).call(txData);
+        vm.assertTrue(success);
         Vm.Log[] memory logs = vm.getRecordedLogs();
         assertEq(keccak256(logs[logs.length - 1].data), keccak256(abi.encode(dataToSign2)));
     }