From 8ef84425e9e19e02c94d8ccf9814fe9190e36e68 Mon Sep 17 00:00:00 2001
From: NoiTaTuM <noit63@gmail.com>
Date: Thu, 9 Jan 2025 19:57:02 -0300
Subject: [PATCH] Fix setCanonicalPeer not setting the peer as canonical if it
 didn't exist

---
 evm/src/assets/TbrBase.sol |   2 +-
 evm/test/ConfigTest.t.sol  | 144 +++++++++++++++++++++----------------
 2 files changed, 82 insertions(+), 64 deletions(-)

diff --git a/evm/src/assets/TbrBase.sol b/evm/src/assets/TbrBase.sol
index 475f649b..c9ad00fe 100644
--- a/evm/src/assets/TbrBase.sol
+++ b/evm/src/assets/TbrBase.sol
@@ -154,7 +154,7 @@ abstract contract TbrBase is PriceOracleIntegration {
     bool isAlreadyPeer = state.peers[peer];
     if (!isAlreadyPeer)
       _addPeer(targetChain, peer);
-    else if (state.canonicalPeer != peer)
+    if (state.canonicalPeer != peer)
       state.canonicalPeer = peer;
   }
 
diff --git a/evm/test/ConfigTest.t.sol b/evm/test/ConfigTest.t.sol
index 66fcd16f..6b42a528 100644
--- a/evm/test/ConfigTest.t.sol
+++ b/evm/test/ConfigTest.t.sol
@@ -16,9 +16,8 @@ import "tbr/assets/TbrIds.sol";
 contract ConfigTest is TbrTestBase {
   using BytesParsing for bytes;
 
-  function addCanonicalPeer(uint16 peerChain) internal {
+  function addCanonicalPeer(uint16 peerChain, bytes32 peer) internal {
     uint8 commandCount = 1;
-    vm.prank(owner);
     invokeTbr(
       abi.encodePacked(
         tbr.exec768.selector, 
@@ -27,7 +26,7 @@ contract ConfigTest is TbrTestBase {
         commandCount, 
         UPDATE_CANONICAL_PEER_ID,
         peerChain,
-        makeBytes32('peer')
+        peer
       )
     );
   }
@@ -540,6 +539,75 @@ contract ConfigTest is TbrTestBase {
     assertEq(isPeer, true);
   }
 
+  function testAddCanonicalPeer() public {
+    uint16 peerChain = 1;
+    bytes32 newPeer = makeBytes32("peer");
+    uint8 commandCount = 1;
+
+    vm.expectRevert(NotAuthorized.selector);
+    addCanonicalPeer(peerChain, newPeer);
+
+    vm.startPrank(owner);
+    addCanonicalPeer(peerChain, newPeer);
+
+    (bool isPeer, ) = invokeStaticTbr(
+      abi.encodePacked(
+        tbr.get1959.selector, 
+        DISPATCHER_PROTOCOL_VERSION0, 
+        CONFIG_QUERIES_ID,
+        commandCount, 
+        IS_PEER_ID,
+        peerChain,
+        newPeer
+      )
+    ).asBoolUnchecked(0);
+
+    assertEq(isPeer, true);
+
+    (bytes32 canonicalPeer, ) = invokeStaticTbr(
+      abi.encodePacked(
+        tbr.get1959.selector, 
+        DISPATCHER_PROTOCOL_VERSION0, 
+        CONFIG_QUERIES_ID,
+        commandCount, 
+        CANONICAL_PEER_ID,
+        peerChain
+      )
+    ).asBytes32Unchecked(0);
+    
+    assertEq(canonicalPeer, newPeer);
+
+    bytes32 anotherPeer = makeBytes32("anotherPeer");
+    addCanonicalPeer(peerChain, anotherPeer);
+
+    (isPeer, ) = invokeStaticTbr(
+      abi.encodePacked(
+        tbr.get1959.selector, 
+        DISPATCHER_PROTOCOL_VERSION0, 
+        CONFIG_QUERIES_ID,
+        commandCount, 
+        IS_PEER_ID,
+        peerChain,
+        anotherPeer
+      )
+    ).asBoolUnchecked(0);
+
+    assertEq(isPeer, true);
+
+    (canonicalPeer, ) = invokeStaticTbr(
+      abi.encodePacked(
+        tbr.get1959.selector, 
+        DISPATCHER_PROTOCOL_VERSION0, 
+        CONFIG_QUERIES_ID,
+        commandCount, 
+        CANONICAL_PEER_ID,
+        peerChain
+      )
+    ).asBytes32Unchecked(0);
+
+    assertEq(canonicalPeer, anotherPeer);
+  }
+
   function testUpdateMaxGasDropoff() public {
     uint16 chainId = 1;
     uint32 maxGasDropoff = 1e6;
@@ -574,7 +642,9 @@ contract ConfigTest is TbrTestBase {
       )
     );
 
-    addCanonicalPeer(chainId);
+    vm.prank(owner);
+    addCanonicalPeer(chainId, makeBytes32("peer"));
+
     vm.prank(admin);
     invokeTbr(
       abi.encodePacked(
@@ -677,7 +747,9 @@ contract ConfigTest is TbrTestBase {
       )
     );
 
-    addCanonicalPeer(chainId);
+    vm.prank(owner);
+    addCanonicalPeer(chainId, makeBytes32("peer"));
+
     vm.prank(admin);
     invokeTbr(
       abi.encodePacked(
@@ -739,7 +811,9 @@ contract ConfigTest is TbrTestBase {
       )
     );
 
-    addCanonicalPeer(chainId);
+    vm.prank(owner);
+    addCanonicalPeer(chainId, makeBytes32("peer"));
+
     vm.prank(admin);
     invokeTbr(
       abi.encodePacked(
@@ -767,51 +841,6 @@ contract ConfigTest is TbrTestBase {
     assertEq(paused_, paused);
   }
 
-  function testUpdateCanonicalPeer() public {
-    bytes32 newCanonicalPeer = makeBytes32("canonicalPeer");
-    uint16 peerChain = 1;
-    uint8 commandCount = 1;
-
-    vm.expectRevert(NotAuthorized.selector);
-    invokeTbr(
-      abi.encodePacked(
-        tbr.exec768.selector, 
-        DISPATCHER_PROTOCOL_VERSION0, 
-        CONFIG_ID,
-        commandCount, 
-        UPDATE_CANONICAL_PEER_ID,
-        peerChain,
-        newCanonicalPeer
-      )
-    );
-
-    vm.prank(owner);
-    invokeTbr(
-      abi.encodePacked(
-        tbr.exec768.selector, 
-        DISPATCHER_PROTOCOL_VERSION0, 
-        CONFIG_ID,
-        commandCount, 
-        UPDATE_CANONICAL_PEER_ID,
-        peerChain,
-        newCanonicalPeer
-      )
-    );
-
-    (bytes32 newCanonicalPeer_, ) = invokeStaticTbr(
-      abi.encodePacked(
-        tbr.get1959.selector, 
-        DISPATCHER_PROTOCOL_VERSION0, 
-        CONFIG_QUERIES_ID,
-        commandCount, 
-        CANONICAL_PEER_ID,
-        peerChain
-      )
-    ).asBytes32Unchecked(0);
-
-    assertEq(newCanonicalPeer_, newCanonicalPeer);
-  }
-
   function testSweepTokens() public {
     uint usdtAmount = 1e6;
     uint ethAmount = 1 ether;
@@ -839,7 +868,6 @@ contract ConfigTest is TbrTestBase {
   function testIsChainSupported() public {
     uint16 chainId = 1;
     uint8 commandCount = 1;
-    bytes32 fakePeer = makeBytes32("peer");
 
     (bool isSupported, ) = invokeStaticTbr(
       abi.encodePacked(
@@ -854,17 +882,7 @@ contract ConfigTest is TbrTestBase {
     assertEq(isSupported, false);
 
     vm.prank(owner);
-    invokeTbr(
-      abi.encodePacked(
-        tbr.exec768.selector, 
-        DISPATCHER_PROTOCOL_VERSION0, 
-        CONFIG_ID,
-        commandCount, 
-        UPDATE_CANONICAL_PEER_ID,
-        chainId, 
-        fakePeer
-      )
-    );
+    addCanonicalPeer(chainId, makeBytes32("peer"));
 
     (isSupported, ) = invokeStaticTbr(
       abi.encodePacked(