From 5731020bc2f12b849240654f2c9619343608f966 Mon Sep 17 00:00:00 2001
From: RnkSngh <raunaksingh2.718@gmail.com>
Date: Wed, 17 Apr 2024 13:26:53 -0400
Subject: [PATCH] add upgrade logic ot uch

---
 contracts/core/UniversalChannelHandler.sol |  5 ++++-
 test/universal.channel.t.sol               | 25 ++++++++++++++++++++++
 2 files changed, 29 insertions(+), 1 deletion(-)

diff --git a/contracts/core/UniversalChannelHandler.sol b/contracts/core/UniversalChannelHandler.sol
index ef5a8be9..245283d9 100644
--- 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
index 37a0e571..ffd37155 100644
--- 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.