add upgrade logic ot uch

open-ibc · Apr 22, 2024 · 5731020 · 5731020
1 parent a72eecd
commit 5731020
Show file tree

Hide file tree

Showing 2 changed files with 29 additions and 1 deletion.
diff --git a/contracts/core/UniversalChannelHandler.sol b/contracts/core/UniversalChannelHandler.sol
@@ -14,8 +14,9 @@ import {
 import {IbcReceiver} from "../interfaces/IbcReceiver.sol";
 import {IbcReceiverBaseUpgradeable} from "../interfaces/IbcReceiverUpgradeable.sol";
 import {ChannelOrder, ChannelEnd, IbcPacket, AckPacket, UniversalPacket, IbcUtils} from "../libs/Ibc.sol";
+import {UUPSUpgradeable} from "@openzeppelin/contracts/proxy/utils/UUPSUpgradeable.sol";
 
-contract UniversalChannelHandler is IbcReceiverBaseUpgradeable, IbcUniversalChannelMW {
+contract UniversalChannelHandler is IbcReceiverBaseUpgradeable, UUPSUpgradeable, IbcUniversalChannelMW {
     uint256[49] private __gap;
 
     bytes32[] public connectedChannels;
@@ -177,6 +178,8 @@ contract UniversalChannelHandler is IbcReceiverBaseUpgradeable, IbcUniversalChan
         dispatcher = _dispatcher;
     }
 
+    function _authorizeUpgrade(address newImplementation) internal override onlyOwner {}
+
     function _connectChannel(bytes32 channelId, string calldata version)
         internal
         returns (string memory checkedVersion)

diff --git a/test/universal.channel.t.sol b/test/universal.channel.t.sol
@@ -245,6 +245,31 @@ contract UniversalChannelPacketTest is Base, IbcMwEventsEmitter {
         vm.stopPrank();
     }
 
+    function test_uch_uch_upgrade__ok() public {
+        IUniversalChannelHandler uch = eth1.ucHandlerProxy();
+        UniversalChannelHandler newUCHImplementation = new UniversalChannelHandler();
+        vm.startPrank(address(eth1)); // Prank eth1 since that address is the owner
+        assertFalse(address(getProxyImplementation(address(uch), vm)) == address(newUCHImplementation));
+        UUPSUpgradeable(address(uch)).upgradeTo(address(address(newUCHImplementation)));
+
+        assertEq(
+            address(getProxyImplementation(address(uch), vm)),
+            address(newUCHImplementation),
+            "new uch implementation not set correctly in uch proxy"
+        );
+        vm.stopPrank();
+    }
+
+    function test_nonOwner_cannot_upgrade_uch() public {
+        IUniversalChannelHandler uch = eth1.ucHandlerProxy();
+        UniversalChannelHandler newUCHImplementation = new UniversalChannelHandler();
+        address notOwner = vm.addr(1);
+        vm.startPrank(notOwner);
+        vm.expectRevert("Ownable: caller is not the owner");
+        UUPSUpgradeable(address(uch)).upgradeTo(address(address(newUCHImplementation)));
+        vm.stopPrank();
+    }
+
     /**
      * Test packet flow from chain A to chain B via UniversalChannel MW and optionally other MW that sits on top of
      * UniversalChannel MW.