Raunak/upgradeable uch

contracts/core/UniversalChannelHandler.sol

@@ -11,10 +11,13 @@ import {
     IbcMwPacketReceiver,
     IbcMwEventsEmitter
 } from "../interfaces/IbcMiddleware.sol";
-import {IbcReceiver, IbcReceiverBase} from "../interfaces/IbcReceiver.sol";
+import {IbcReceiver} from "../interfaces/IbcReceiver.sol";
+import {IbcReceiverBaseUpgradeable} from "../interfaces/IbcReceiverUpgradeable.sol";
 import {ChannelOrder, ChannelEnd, IbcPacket, AckPacket, UniversalPacket, IbcUtils} from "../libs/Ibc.sol";
 
-contract UniversalChannelHandler is IbcReceiverBase, IbcUniversalChannelMW {
+contract UniversalChannelHandler is IbcReceiverBaseUpgradeable, IbcUniversalChannelMW {
+    uint256[49] private __gap;
+
     bytes32[] public connectedChannels;
     string public constant VERSION = "1.0";
     uint256 public constant MW_ID = 1;
@@ -24,7 +27,13 @@ contract UniversalChannelHandler is IbcReceiverBase, IbcUniversalChannelMW {
 
     event UCHPacketSent(address source, bytes32 destination);
 
-    constructor(IbcDispatcher _dispatcher) IbcReceiverBase(_dispatcher) {}
+    constructor() {
+        _disableInitializers();
+    }
+
+    function initialize(IbcDispatcher _dispatcher) public initializer {
+        __IbcReceiverBase_init(_dispatcher);
+    }
     /**
      * @dev Close a universal channel.
      * Cannot send or receive packets after the channel is closed.
@@ -182,6 +191,10 @@ contract UniversalChannelHandler is IbcReceiverBase, IbcUniversalChannelMW {
 
     function onChanOpenConfirm(bytes32 channelId) external onlyIbcDispatcher {}
 
+    function setDispatcher(IbcDispatcher _dispatcher) external onlyOwner {
+        dispatcher = _dispatcher;
+    }
+
     function _connectChannel(bytes32 channelId, string calldata version)
         internal
         returns (string memory checkedVersion)

contracts/interfaces/IUniversalChannelHandler.sol

@@ -0,0 +1,33 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.0;
+
+import {IbcDispatcher, IbcEventsEmitter} from "./IbcDispatcher.sol";
+
+import {L1Header, OpL2StateProof, Ics23Proof} from "./ProofVerifier.sol";
+import {IbcUniversalChannelMW} from "./IbcMiddleware.sol";
+import {
+    Channel,
+    ChannelEnd,
+    ChannelOrder,
+    IbcPacket,
+    ChannelState,
+    AckPacket,
+    IBCErrors,
+    IbcUtils,
+    Ibc
+} from "../libs/Ibc.sol";
+
+interface IUniversalChannelHandler is IbcUniversalChannelMW {
+    function registerMwStack(uint256 mwBitmap, address[] calldata mwAddrs) external;
+    function openChannel(
+        string calldata version,
+        ChannelOrder ordering,
+        bool feeEnabled,
+        string[] calldata connectionHops,
+        string calldata counterpartyPortIdentifier
+    ) external;
+    function closeChannel(bytes32 channelId) external;
+    function setDispatcher(IbcDispatcher dispatcher) external;
+    function dispatcher() external returns (IbcDispatcher dispatcher);
+    function connectedChannels(uint256) external view returns (bytes32 channel);
+}

contracts/interfaces/IbcReceiverUpgradeable.sol

@@ -0,0 +1,45 @@
+//SPDX-License-Identifier: UNLICENSED
+
+pragma solidity ^0.8.9;
+
+import {OwnableUpgradeable} from "@openzeppelin-upgradeable/contracts/access/OwnableUpgradeable.sol";
+import {IbcDispatcher} from "./IbcDispatcher.sol";
+import {ChannelOrder, ChannelEnd, IbcPacket, AckPacket} from "../libs/Ibc.sol";
+
+contract IbcReceiverBaseUpgradeable is OwnableUpgradeable {
+    IbcDispatcher public dispatcher;
+
+    error notIbcDispatcher();
+    error UnsupportedVersion();
+    error ChannelNotFound();
+
+    /**
+     * @dev Modifier to restrict access to only the IBC dispatcher.
+     * Only the address with the IBC_ROLE can execute the function.
+     * Should add this modifier to all IBC-related callback functions.
+     */
+    modifier onlyIbcDispatcher() {
+        if (msg.sender != address(dispatcher)) {
+            revert notIbcDispatcher();
+        }
+        _;
+    }
+
+    constructor() {
+        _disableInitializers();
+    }
+
+    /// This function is called for plain Ether transfers, i.e. for every call with empty calldata.
+    // An empty function body is sufficient to receive packet fee refunds.
+    receive() external payable {}
+
+    /**
+     * @dev initializer function that takes an IbcDispatcher address and grants the IBC_ROLE to the Polymer IBC
+     * Dispatcher.
+     * @param _dispatcher The address of the IbcDispatcher contract.
+     */
+    function __IbcReceiverBase_init(IbcDispatcher _dispatcher) internal onlyInitializing {
+        __Ownable_init();
+        dispatcher = _dispatcher;
+    }
+}

script/Deploy.s.sol

@@ -7,6 +7,7 @@ import "../contracts/utils/DummyLightClient.sol";
 import "../contracts/core/Dispatcher.sol";
 import {Mars} from "../contracts/examples/Mars.sol";
 import {IDispatcher} from "../contracts/core/Dispatcher.sol";
+import {IUniversalChannelHandler} from "../contracts/interfaces/IUniversalChannelHandler.sol";
 import {ERC1967Proxy} from "@openzeppelin/contracts/proxy/ERC1967/ERC1967Proxy.sol";
 import "../contracts/core/OpProofVerifier.sol";
 import "../contracts/core/OpLightClient.sol";
@@ -106,9 +107,18 @@ contract Deploy is Script {
     }
 
     function deployUniversalChannelHandler(address dispatcher) public broadcast returns (address addr_) {
-        UniversalChannelHandler handler = new UniversalChannelHandler{salt: _implSalt()}(IbcDispatcher(dispatcher));
-        console.log("UniversalChannelHandler deployed at %s", address(handler));
-        return address(handler);
+        UniversalChannelHandler uchImplementation = new UniversalChannelHandler();
+        IUniversalChannelHandler proxy = IUniversalChannelHandler(
+            address(
+                new ERC1967Proxy(
+                    address(uchImplementation),
+                    abi.encodeWithSelector(UniversalChannelHandler.initialize.selector, dispatcher)
+                )
+            )
+        );
+        console.log("UniversalChannelHandler implementation deployed at %s", address(uchImplementation));
+        console.log("UniversalChannelHandler proxy deployed at %s", address(proxy));
+        return address(proxy);
     }
 
     function deployEarth(address middleware) public broadcast returns (address addr_) {

test/TestUtils.t.sol

@@ -1,11 +1,12 @@
 import {IDispatcher} from "../contracts/interfaces/IDispatcher.sol";
+import {IUniversalChannelHandler} from "../contracts/interfaces/IUniversalChannelHandler.sol";
 import {LightClient} from "../contracts/interfaces/LightClient.sol";
 import {ERC1967Proxy} from "@openzeppelin/contracts/proxy/ERC1967/ERC1967Proxy.sol";
 import {Dispatcher} from "../contracts/core/Dispatcher.sol";
-// import {StdCheats} from "../forge-std/StdCheats.sol";
+import {UniversalChannelHandler} from "../contracts/core/UniversalChannelHandler.sol";
+
 import {Test} from "forge-std/Test.sol";
 import {Vm} from "forge-std/Vm.sol";
-// import {StdCheats} from "forge-std/StdCheats.sol";
 
 pragma solidity ^0.8.0;
 
@@ -29,6 +30,21 @@ abstract contract TestUtilsTest {
         );
     }
 
+    function deployUCHProxyAndImpl(address dispatcherProxy)
+        public
+        returns (IUniversalChannelHandler proxy, UniversalChannelHandler uchImplementation)
+    {
+        uchImplementation = new UniversalChannelHandler();
+        proxy = IUniversalChannelHandler(
+            address(
+                new ERC1967Proxy(
+                    address(uchImplementation),
+                    abi.encodeWithSelector(UniversalChannelHandler.initialize.selector, dispatcherProxy)
+                )
+            )
+        );
+    }
+
     function getProxyImplementation(address proxy, Vm vm) public view returns (address dispatcherImplementation) {
         dispatcherImplementation = address(uint160(uint256(vm.load(address(proxy), _IMPLEMENTATION_SLOT))));
     }

test/VirtualChain.sol

@@ -4,6 +4,7 @@ pragma solidity ^0.8.13;
 import "forge-std/Test.sol";
 import "@openzeppelin/contracts/utils/Strings.sol";
 import {IbcDispatcher, IbcEventsEmitter} from "../contracts/interfaces/IbcDispatcher.sol";
+import {IUniversalChannelHandler} from "../contracts/interfaces/IUniversalChannelHandler.sol";
 import {IDispatcher} from "../contracts/interfaces/IDispatcher.sol";
 import "../contracts/libs/Ibc.sol";
 import "../contracts/core/Dispatcher.sol";
@@ -27,7 +28,7 @@ struct ChannelSetting {
 
 struct VirtualChainData {
     IDispatcher dispatcherProxy;
-    UniversalChannelHandler ucHandler;
+    IUniversalChannelHandler ucHandlerProxy;
     Mars mars;
     Earth earth;
     GeneralMiddleware mw1;
@@ -39,7 +40,7 @@ struct VirtualChainData {
 contract VirtualChain is Test, IbcEventsEmitter, TestUtilsTest {
     IDispatcher public dispatcherProxy;
     Dispatcher public dispatcherImplementation;
-    UniversalChannelHandler public ucHandler;
+    IUniversalChannelHandler public ucHandlerProxy;
     GeneralMiddleware public mw1;
     GeneralMiddleware public mw2;
 
@@ -60,26 +61,26 @@ contract VirtualChain is Test, IbcEventsEmitter, TestUtilsTest {
         _seed = seed;
 
         (dispatcherProxy, dispatcherImplementation) = deployDispatcherProxyAndImpl(portPrefix, new DummyLightClient());
-        ucHandler = new UniversalChannelHandler(dispatcherProxy);
+        (ucHandlerProxy,) = deployUCHProxyAndImpl(address(dispatcherProxy));
 
         mars = new Mars(dispatcherProxy);
-        earth = new Earth(address(ucHandler));
+        earth = new Earth(address(ucHandlerProxy));
         // initialize portIds for counterparty chains
-        address[3] memory portContracts = [address(ucHandler), address(mars), address(earth)];
+        address[3] memory portContracts = [address(ucHandlerProxy), address(mars), address(earth)];
         for (uint256 i = 0; i < portContracts.length; i++) {
             portIds[portContracts[i]] = string(abi.encodePacked(portPrefix, toHexStr(portContracts[i])));
         }
         connectionHops = new string[](2);
         connectionHops[0] = newConnectionId();
         connectionHops[1] = newConnectionId();
 
-        mw1 = new GeneralMiddleware(1 << 1, address(ucHandler));
-        mw2 = new GeneralMiddleware(1 << 2, address(ucHandler));
+        mw1 = new GeneralMiddleware(1 << 1, address(ucHandlerProxy));
+        mw2 = new GeneralMiddleware(1 << 2, address(ucHandlerProxy));
     }
 
     // return virtualChainData
     function getVirtualChainData() external view returns (VirtualChainData memory) {
-        return VirtualChainData(dispatcherProxy, ucHandler, mars, earth, mw1, mw2, connectionHops);
+        return VirtualChainData(dispatcherProxy, ucHandlerProxy, mars, earth, mw1, mw2, connectionHops);
     }
 
     // expectedChannel returns a Channel struct with expected values
@@ -167,7 +168,7 @@ contract VirtualChain is Test, IbcEventsEmitter, TestUtilsTest {
                 remoteChain.portIds(address(remoteEnd))
             );
         }
-        vm.prank(address(ucHandler));
+        vm.prank(address(ucHandlerProxy));
         dispatcherProxy.channelOpenInit(setting.version, setting.ordering, setting.feeEnabled, connectionHops, cpPortId);
     }