From f8ea1e579db7c9167aa26b1e6b0a7b2c8e95088a Mon Sep 17 00:00:00 2001
From: LayneHaber <layne.haber@gmail.com>
Date: Sun, 1 Jan 2023 19:04:32 -0700
Subject: [PATCH] fix: persistent message nonce tracking

---
 .../contracts/messaging/MerkleTreeManager.sol    | 16 +++++++++++++++-
 .../messaging/connectors/SpokeConnector.sol      |  7 +------
 .../messaging/MerkleTreeManager.t.sol            | 15 +++++++++++++++
 3 files changed, 31 insertions(+), 7 deletions(-)

diff --git a/packages/deployments/contracts/contracts/messaging/MerkleTreeManager.sol b/packages/deployments/contracts/contracts/messaging/MerkleTreeManager.sol
index e124b35abf..da5a7bdcdc 100644
--- 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
index c943e63f74..0054325213 100644
--- 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
index 20fe2d0f50..34337bcc9a 100644
--- a/packages/deployments/contracts/contracts_forge/messaging/MerkleTreeManager.t.sol
+++ b/packages/deployments/contracts/contracts_forge/messaging/MerkleTreeManager.t.sol
@@ -146,6 +146,8 @@ 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));
@@ -153,6 +155,17 @@ contract MerkleTreeManagerTest is ForgeHelper {
     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");