fix(script): read the Safe nonce from the real chain state after simulating

script/MultiSigBatchBase.sol

@@ -1,5 +1,4 @@
 // SPDX-License-Identifier: UNLICENSED
-
 pragma solidity >=0.8.20 <0.9.0;
 
 import { SafeNonce } from "./SafeNonce.sol";
@@ -23,36 +22,54 @@ abstract contract MultiSigBatchBase is Script {
         _data.push(data_);
     }
 
-    /// @dev Proposes the batch at the next free Safe nonce. See {SafeNonce-next}.
+    /// @dev    Proposes the batch at the next free Safe nonce. See {SafeNonce-next}.
+    /// @param  safe_   The Safe to propose to.
+    /// @param  sender_ The owner signing the proposal.
     function _proposeBatch(address safe_, address sender_) internal {
-        _propose(sender_, SafeNonce.next(_safeMultiSig, safe_));
+        _safeMultiSig.initialize(safe_);
+        _propose(sender_, SafeNonce.next(_safeMultiSig));
     }
 
-    /// @dev Proposes the batch at an explicit nonce. The Safe's on-chain nonce only advances on execution, so
-    ///      proposing at it can collide with already queued proposals instead of queueing behind them.
+    /// @dev    Proposes the batch at an explicit nonce, for when the Safe transaction service cannot be queried or
+    ///         the batch must queue at a chosen position.
+    /// @param  safe_   The Safe to propose to.
+    /// @param  sender_ The owner signing the proposal.
+    /// @param  nonce_  The Safe nonce to propose at.
     function _proposeBatch(address safe_, address sender_, uint256 nonce_) internal {
-        console.log("Safe nonce:", nonce_);
-
         _safeMultiSig.initialize(safe_);
         _propose(sender_, nonce_);
     }
 
-    /// @dev Simulates the batch through the Safe itself, using synthetic owner approvals, so that the MultiSend
-    ///      encoding, the threshold check and any guard or fallback handler are exercised too.
+    /// @dev    Simulates the batch through the Safe itself, using synthetic owner approvals, so that the MultiSend
+    ///         encoding, the threshold check and any guard or fallback handler are exercised too.
+    /// @param  safe_ The Safe to simulate through.
     function _simulateBatch(address safe_) internal {
         _safeMultiSig.initialize(safe_);
 
         address[] memory owners_ = OwnerManager(safe_).getOwners();
 
+        uint256 snapshot_ = vm.snapshotState();
+
         // NOTE: `isolate` mode runs each top-level call as its own transaction, requiring the signer to pay for gas.
         for (uint256 i = 0; i < owners_.length; i++) {
             vm.deal(owners_[i], owners_[i].balance + 1 ether);
         }
 
-        require(_safeMultiSig.simulateTransactionsMultiSigNoSign(_targets, _data, owners_), "Simulation failed");
+        bool success_ = _safeMultiSig.simulateTransactionsMultiSigNoSign(_targets, _data, owners_);
+
+        // NOTE: The simulation executes the batch for real on the local fork, which advances the Safe nonce and
+        //       applies the batch. Restoring the state keeps the nonce the proposal is later signed at correct.
+        require(vm.revertToStateAndDelete(snapshot_), "State restore failed");
+
+        require(success_, "Simulation failed");
     }
 
+    /// @dev    Signs and proposes the batch at `nonce_` through the initialized Safe client.
+    /// @param  sender_ The owner signing the proposal.
+    /// @param  nonce_  The Safe nonce to propose at.
     function _propose(address sender_, uint256 nonce_) private {
+        console.log("[nonce] proposing at nonce", nonce_);
+
         (address to_, bytes memory data_) = _safeMultiSig.getProposeTransactionsTargetAndData(_targets, _data);
 
         // NOTE: Batches are executed via DelegateCall to preserve `msg.sender` across the sub-calls, and the signed

script/SafeNonce.sol

@@ -1,5 +1,4 @@
 // SPDX-License-Identifier: UNLICENSED
-
 pragma solidity >=0.8.20 <0.9.0;
 
 import { HTTP } from "../lib/safe-utils/lib/solidity-http/src/HTTP.sol";
@@ -8,72 +7,72 @@ import { Safe } from "../lib/safe-utils/src/Safe.sol";
 import { console } from "../lib/forge-std/src/console.sol";
 import { Vm } from "../lib/forge-std/src/Vm.sol";
 
-/// @notice Gets the next free Safe nonce from the Safe transaction service.
-/// @dev    The Safe on-chain nonce advances only on execution. Two proposals at one nonce compete,
-///         and only one of them can execute.
+/// @title  Next free Safe nonce, read from the Safe transaction service.
+/// @author M0 Labs
 library SafeNonce {
     using HTTP for *;
     using Safe for *;
 
     Vm private constant _vm = Vm(address(uint160(uint256(keccak256("hevm cheat code")))));
 
     /// @notice Thrown if the Safe transaction service does not answer the pending proposals query.
-    /// @param statusCode_ The HTTP status code of the response.
-    /// @param response_ The body of the response.
+    /// @param  statusCode_ The HTTP status code of the response.
+    /// @param  response_   The body of the response.
     error PendingProposalsQueryFailed(uint256 statusCode_, string response_);
 
-    /// @notice Returns the next free nonce of `safe_`: the on-chain nonce, or one above the highest pending proposal.
-    /// @dev    Initializes `client_` for `safe_`. Reverts if the Safe transaction service does not answer.
-    function next(Safe.Client storage client_, address safe_) internal returns (uint256 nonce_) {
-        client_.initialize(safe_);
+    /// @notice Returns the next free nonce of the Safe `client_` is initialized for.
+    /// @dev    Reverts if the Safe transaction service does not answer. The propose helpers accept an explicit
+    ///         nonce to bypass the service.
+    /// @param  client_ The Safe client, initialized for the Safe.
+    /// @return The on-chain nonce, or one above the highest pending proposal.
+    function next(Safe.Client storage client_) internal returns (uint256) {
         uint256 onChain_ = client_.getNonce();
 
         HTTP.Response memory response_ = client_
             .instance()
             .http
             .instance()
-            .GET(
-                string.concat(
-                    client_.getApiKitUrl(block.chainid),
-                    "/v1/safes/",
-                    _vm.toString(safe_),
-                    "/multisig-transactions/?executed=false&nonce__gte=",
-                    _vm.toString(onChain_),
-                    "&ordering=-nonce&limit=1"
-                )
-            )
+            .GET(_getPendingProposalsUrl(client_, onChain_))
             .request();
 
-        if (response_.status < 200 || response_.status >= 300) {
-            revert PendingProposalsQueryFailed(response_.status, response_.data);
-        }
+        uint256 nonce_ = fromResponse(onChain_, response_);
 
-        uint256 pendingCount_ = _vm.parseJsonUint(response_.data, ".count");
-        uint256 highestPending_;
+        console.log("[nonce] Safe on-chain nonce %d, next free nonce %d", onChain_, nonce_);
 
-        if (pendingCount_ == 0) {
-            console.log("[nonce] Safe nonce %d, no pending proposals", onChain_);
-        } else {
-            highestPending_ = _vm.parseJsonUint(response_.data, ".results[0].nonce");
-            console.log(
-                "[nonce] Safe nonce %d, %d pending proposal(s) up to nonce %d",
-                onChain_,
-                pendingCount_,
-                highestPending_
-            );
+        return nonce_;
+    }
+
+    /// @notice Returns the next free nonce from a page of pending proposals at or above `onChain_`.
+    /// @dev    Reverts if `response_` is not a 2xx answer.
+    /// @param  onChain_  The Safe's on-chain nonce.
+    /// @param  response_ The transaction service page of pending proposals, highest nonce first.
+    /// @return `onChain_` if no proposal is pending, else the highest pending nonce plus one.
+    function fromResponse(uint256 onChain_, HTTP.Response memory response_) internal pure returns (uint256) {
+        if (response_.status < 200 || response_.status >= 300) {
+            revert PendingProposalsQueryFailed(response_.status, response_.data);
         }
 
-        nonce_ = nextFrom(onChain_, pendingCount_, highestPending_);
+        if (_vm.parseJsonUint(response_.data, ".count") == 0) return onChain_;
 
-        console.log("[nonce] proposing at nonce", nonce_);
+        return _vm.parseJsonUint(response_.data, ".results[0].nonce") + 1;
     }
 
-    /// @notice Returns `onChain_` if no proposal is pending at or above it, else the highest pending nonce plus one.
-    function nextFrom(
-        uint256 onChain_,
-        uint256 pendingCount_,
-        uint256 highestPending_
-    ) internal pure returns (uint256) {
-        return pendingCount_ == 0 || highestPending_ < onChain_ ? onChain_ : highestPending_ + 1;
+    /// @dev    Builds the transaction service query for the pending proposals at or above `onChain_`.
+    /// @param  client_  The Safe client, initialized for the Safe.
+    /// @param  onChain_ The Safe's on-chain nonce.
+    /// @return The URL, ordered by nonce descending and limited to the first result.
+    function _getPendingProposalsUrl(
+        Safe.Client storage client_,
+        uint256 onChain_
+    ) private view returns (string memory) {
+        return
+            string.concat(
+                client_.getApiKitUrl(block.chainid),
+                "/v1/safes/",
+                _vm.toString(client_.instance().safe),
+                "/multisig-transactions/?executed=false&nonce__gte=",
+                _vm.toString(onChain_),
+                "&ordering=-nonce&limit=1"
+            );
     }
 }

script/SafeTimelockBatchBase.sol

@@ -8,53 +8,121 @@ import { Enum } from "../lib/safe-utils/lib/safe-smart-account/contracts/common/
 import { Safe } from "../lib/safe-utils/src/Safe.sol";
 import { TimelockController } from "../lib/openzeppelin-contracts-upgradeable/lib/openzeppelin-contracts/contracts/governance/TimelockController.sol";
 
+import { console } from "../lib/forge-std/src/console.sol";
+
 abstract contract SafeTimelockBatchBase is TimelockBatchBase {
     using Safe for *;
 
     Safe.Client internal _safeMultiSig;
 
     /// @notice Thrown in case a transaction that's supposed to be cancelled is not pending.
-    /// @param id_ The identifier of the transaction.
+    /// @param  id_ The identifier of the transaction.
     error OperationNotPending(bytes32 id_);
 
     /// @notice Proposes to schedule a batch of transactions to a timelock contract.
-    /// @param safe_ The address of the Safe multisig to propose to.
-    /// @param timelock_ The address of the timelock.
-    /// @param sender_ The sender's address.
-    /// @param predecessor_ The predecessor transaction, if any.
-    /// @param salt_ The salt to build the transaction with, if any.
+    /// @dev    Proposes at the next free Safe nonce. See {SafeNonce-next}.
+    /// @param  safe_        The address of the Safe multisig to propose to.
+    /// @param  timelock_    The address of the timelock.
+    /// @param  sender_      The sender's address.
+    /// @param  predecessor_ The predecessor transaction, if any.
+    /// @param  salt_        The salt to build the transaction with, if any.
     function _proposeScheduleBatch(
         address safe_,
         address timelock_,
         address sender_,
         bytes32 predecessor_,
         bytes32 salt_
     ) internal {
-        uint256 delay = TimelockController(payable(timelock_)).getMinDelay();
-        bytes memory batchData = _getScheduleBatchCallData(predecessor_, salt_, delay);
+        bytes memory data_ = _getScheduleBatchData(timelock_, predecessor_, salt_);
+
+        _safeMultiSig.initialize(safe_);
+        _proposeToTimelock(timelock_, data_, sender_, SafeNonce.next(_safeMultiSig));
+    }
+
+    /// @notice Proposes to schedule a batch of transactions to a timelock contract at an explicit Safe nonce.
+    /// @dev    For when the Safe transaction service cannot be queried or the proposal must queue at a chosen position.
+    /// @param  safe_        The address of the Safe multisig to propose to.
+    /// @param  timelock_    The address of the timelock.
+    /// @param  sender_      The sender's address.
+    /// @param  predecessor_ The predecessor transaction, if any.
+    /// @param  salt_        The salt to build the transaction with, if any.
+    /// @param  nonce_       The Safe nonce to propose at.
+    function _proposeScheduleBatch(
+        address safe_,
+        address timelock_,
+        address sender_,
+        bytes32 predecessor_,
+        bytes32 salt_,
+        uint256 nonce_
+    ) internal {
+        bytes memory data_ = _getScheduleBatchData(timelock_, predecessor_, salt_);
 
-        _proposeToTimelock(safe_, timelock_, batchData, sender_);
+        _safeMultiSig.initialize(safe_);
+        _proposeToTimelock(timelock_, data_, sender_, nonce_);
     }
 
     /// @notice Proposes to cancel the execution of a pending message that was originally scheduled through a timelock.
-    /// @param safe_ The address of the Safe multisig to propose the transaction to.
-    /// @param timelock_ The address of the timelock.
-    /// @param sender_ The sender's address.
-    /// @param id_ The id of the scheduled transaction to cancel.
+    /// @dev    Proposes at the next free Safe nonce. See {SafeNonce-next}.
+    /// @param  safe_     The address of the Safe multisig to propose the transaction to.
+    /// @param  timelock_ The address of the timelock.
+    /// @param  sender_   The sender's address.
+    /// @param  id_       The id of the scheduled transaction to cancel.
     function _proposeCancel(address safe_, address timelock_, address sender_, bytes32 id_) internal {
-        TimelockController timelock = TimelockController(payable(timelock_));
-        if (!timelock.isOperationPending(id_)) {
-            revert OperationNotPending(id_);
-        }
+        bytes memory data_ = _getCancelData(timelock_, id_);
 
-        _proposeToTimelock(safe_, timelock_, abi.encodeCall(TimelockController.cancel, id_), sender_);
+        _safeMultiSig.initialize(safe_);
+        _proposeToTimelock(timelock_, data_, sender_, SafeNonce.next(_safeMultiSig));
     }
 
-    /// @dev Proposes a call to the timelock at the next free Safe nonce. See {SafeNonce-next}.
-    function _proposeToTimelock(address safe_, address timelock_, bytes memory data_, address sender_) private {
-        uint256 nonce_ = SafeNonce.next(_safeMultiSig, safe_);
+    /// @notice Proposes to cancel a pending timelock operation at an explicit Safe nonce.
+    /// @dev    For when the Safe transaction service cannot be queried or the proposal must queue at a chosen position.
+    /// @param  safe_     The address of the Safe multisig to propose the transaction to.
+    /// @param  timelock_ The address of the timelock.
+    /// @param  sender_   The sender's address.
+    /// @param  id_       The id of the scheduled transaction to cancel.
+    /// @param  nonce_    The Safe nonce to propose at.
+    function _proposeCancel(address safe_, address timelock_, address sender_, bytes32 id_, uint256 nonce_) internal {
+        bytes memory data_ = _getCancelData(timelock_, id_);
+
+        _safeMultiSig.initialize(safe_);
+        _proposeToTimelock(timelock_, data_, sender_, nonce_);
+    }
+
+    /// @dev    Signs and proposes a call to the timelock at `nonce_` through the initialized Safe client.
+    /// @param  timelock_ The address of the timelock.
+    /// @param  data_     The call data for the timelock.
+    /// @param  sender_   The sender's address.
+    /// @param  nonce_    The Safe nonce to propose at.
+    function _proposeToTimelock(address timelock_, bytes memory data_, address sender_, uint256 nonce_) private {
+        console.log("[nonce] proposing at nonce", nonce_);
+
         bytes memory signature_ = _safeMultiSig.sign(timelock_, data_, Enum.Operation.Call, sender_, nonce_, "");
 
         _safeMultiSig.proposeTransactionWithSignature(timelock_, data_, sender_, signature_, nonce_);
     }
+
+    /// @dev    Builds the `scheduleBatch` call for the batch, at the timelock's minimum delay.
+    /// @param  timelock_    The address of the timelock.
+    /// @param  predecessor_ The predecessor transaction, if any.
+    /// @param  salt_        The salt to build the transaction with, if any.
+    /// @return The call data for the timelock.
+    function _getScheduleBatchData(
+        address timelock_,
+        bytes32 predecessor_,
+        bytes32 salt_
+    ) private view returns (bytes memory) {
+        uint256 delay_ = TimelockController(payable(timelock_)).getMinDelay();
+
+        return _getScheduleBatchCallData(predecessor_, salt_, delay_);
+    }
+
+    /// @dev    Builds the `cancel` call for a pending operation.
+    /// @param  timelock_ The address of the timelock.
+    /// @param  id_       The id of the scheduled transaction to cancel.
+    /// @return The call data for the timelock.
+    function _getCancelData(address timelock_, bytes32 id_) private view returns (bytes memory) {
+        if (!TimelockController(payable(timelock_)).isOperationPending(id_)) revert OperationNotPending(id_);
+
+        return abi.encodeCall(TimelockController.cancel, id_);
+    }
 }

script/TimelockBatchBase.sol

@@ -87,6 +87,8 @@ abstract contract TimelockBatchBase is Script {
     /// @notice Simulates the timelock execution based on the accumulated call stack.
     /// @param timelock_ The address of the timelock contract to execute from.
     function _simulateBatch(address timelock_) internal {
+        uint256 snapshot_ = vm.snapshotState();
+
         vm.startPrank(timelock_);
 
         for (uint256 i = 0; i < _timelockTargets.length; i++) {
@@ -95,5 +97,9 @@ abstract contract TimelockBatchBase is Script {
         }
 
         vm.stopPrank();
+
+        // NOTE: The simulation executes the batch for real on the local fork. Restoring the state keeps the
+        //       proposal built from the current chain state, e.g. the timelock's minimum delay.
+        require(vm.revertToStateAndDelete(snapshot_), "State restore failed");
     }
 }

test/MultiSigBatchBase.t.sol

@@ -0,0 +1,38 @@
+// SPDX-License-Identifier: UNLICENSED
+pragma solidity >=0.8.20 <0.9.0;
+
+import { Test } from "../lib/forge-std/src/Test.sol";
+
+import { MockSafe } from "./utils/MockSafe.sol";
+import { MultiSigBatchBaseHarness } from "./utils/MultiSigBatchBaseHarness.sol";
+
+contract MultiSigBatchBaseTests is Test {
+    MultiSigBatchBaseHarness public harness;
+    MockSafe public safe;
+
+    address public owner = makeAddr("owner");
+    address public target = makeAddr("target");
+
+    function setUp() external {
+        // NOTE: The MultiSend address is resolved per chain, and the mock Safe never calls it.
+        vm.chainId(1);
+
+        harness = new MultiSigBatchBaseHarness();
+
+        address[] memory owners_ = new address[](1);
+        owners_[0] = owner;
+
+        safe = new MockSafe(owners_);
+    }
+
+    /* ============ _simulateBatch ============ */
+
+    function test_simulateBatch_leavesStateUntouched() external {
+        harness.addToBatch(target, "");
+
+        harness.simulateBatch(address(safe));
+
+        assertEq(safe.nonce(), 0);
+        assertEq(owner.balance, 0);
+    }
+}