fix: persistent message nonce tracking

connext · Jan 2, 2023 · f8ea1e5 · f8ea1e5
1 parent cc06049
commit f8ea1e5
Show file tree

Hide file tree

Showing 3 changed files with 31 additions and 7 deletions.
diff --git a/packages/deployments/contracts/contracts/messaging/MerkleTreeManager.sol b/packages/deployments/contracts/contracts/messaging/MerkleTreeManager.sol
@@ -58,6 +58,11 @@ contract MerkleTreeManager is ProposedOwnableUpgradeable {
    */
   mapping(bytes32 => LeafStatus) public leaves;
 
+  /**
+   * @notice domain => next available nonce for the domain.
+   */
+  mapping(uint32 => uint32) public nonces;
+
   // ============ Modifiers ============
 
   modifier onlyArborist() {
@@ -139,6 +144,15 @@ contract MerkleTreeManager is ProposedOwnableUpgradeable {
 
   // ========= Public Functions =========
 
+  /**
+   * @notice Used to increment nonce
+   * @param _domain The domain the nonce will be used for
+   * @return _nonce The incremented nonce
+   */
+  function incrementNonce(uint32 _domain) public onlyArborist returns (uint32 _nonce) {
+    _nonce = nonces[_domain]++;
+  }
+
   /**
    * @notice Used to track proven leaves
    * @param _leaf The leaf to mark as proven
@@ -205,5 +219,5 @@ contract MerkleTreeManager is ProposedOwnableUpgradeable {
   }
 
   // ============ Upgrade Gap ============
-  uint256[47] private __GAP; // gap for upgrade safety
+  uint256[46] private __GAP; // gap for upgrade safety
 }
diff --git a/packages/deployments/contracts/contracts/messaging/connectors/SpokeConnector.sol b/packages/deployments/contracts/contracts/messaging/connectors/SpokeConnector.sol
@@ -141,11 +141,6 @@ abstract contract SpokeConnector is Connector, ConnectorManager, WatcherClient,
    */
   mapping(address => bool) public allowlistedSenders;
 
-  /**
-   * @notice domain => next available nonce for the domain.
-   */
-  mapping(uint32 => uint32) public nonces;
-
   // ============ Modifiers ============
 
   modifier onlyAllowlistedSender() {
@@ -317,7 +312,7 @@ abstract contract SpokeConnector is Connector, ConnectorManager, WatcherClient,
     bytes memory _messageBody
   ) external onlyAllowlistedSender returns (bytes32, bytes memory) {
     // Get the next nonce for the destination domain, then increment it.
-    uint32 _nonce = nonces[_destinationDomain]++;
+    uint32 _nonce = MERKLE.incrementNonce(_destinationDomain);
 
     // Format the message into packed bytes.
     bytes memory _message = Message.formatMessage(

diff --git a/packages/deployments/contracts/contracts_forge/messaging/MerkleTreeManager.t.sol b/packages/deployments/contracts/contracts_forge/messaging/MerkleTreeManager.t.sol
@@ -146,13 +146,26 @@ contract MerkleTreeManagerTest is ForgeHelper {
     assertEq(root, expectedRoot);
   }
 
+  // ========= setArborist =========
+
   function test_Merkle__setArborist_shouldWork() public {
     vm.expectEmit(true, true, true, true);
     emit ArboristUpdated(address(this), address(1));
 
     merkle.setArborist(address(1));
   }
 
+  // ========= incrementNonce =========
+
+  function test_Merkle__incrementNonce_works() public {
+    uint32 domain = uint32(1);
+    uint32 returned = merkle.incrementNonce(domain);
+    assertEq(returned, 0);
+    assertEq(merkle.nonces(domain), returned + 1);
+  }
+
+  // ========= markAsProven =========
+
   function test_Merkle__markAsProven_failsIfNotArborist() public {
     vm.expectRevert("!arborist");
     vm.prank(address(1231));
@@ -172,6 +185,8 @@ contract MerkleTreeManagerTest is ForgeHelper {
     assertTrue(merkle.leaves(leaf) == MerkleTreeManager.LeafStatus.Proven);
   }
 
+  // ========= markAsProcessed =========
+
   function test_Merkle__markAsProcessed_failsIfBadStatus() public {
     bytes32 leaf = bytes32(bytes("123123"));
     vm.expectRevert("!proven");