Add upgradable Shell contract (#1170)

* Add temporary rescue mechanism to gateway contract

* Simplified the implementation

* Add more comments

contracts/src/DeployGatewayLogic.sol

@@ -26,7 +26,7 @@ contract DeployGatewayLogic is Script {
         uint8 foreignTokenDecimals = uint8(vm.envUint("FOREIGN_TOKEN_DECIMALS"));
 
         AgentExecutor executor = new AgentExecutor();
-        new Gateway(address(beefyClient), address(executor), bridgeHubParaID, bridgeHubAgentID, foreignTokenDecimals);
+        new Gateway(beefyClient, address(executor), bridgeHubParaID, bridgeHubAgentID, foreignTokenDecimals);
 
         vm.stopBroadcast();
     }

contracts/src/DeployScript.sol

@@ -9,7 +9,6 @@ import {BeefyClient} from "./BeefyClient.sol";
 import {IGateway} from "./interfaces/IGateway.sol";
 import {GatewayProxy} from "./GatewayProxy.sol";
 import {Gateway} from "./Gateway.sol";
-import {GatewayUpgradeMock} from "../test/mocks/GatewayUpgradeMock.sol";
 import {Agent} from "./Agent.sol";
 import {AgentExecutor} from "./AgentExecutor.sol";
 import {ChannelID, ParaID, OperatingMode} from "./Types.sol";
@@ -99,8 +98,6 @@ contract DeployScript is Script {
         payable(bridgeHubAgent).safeNativeTransfer(initialDeposit);
         payable(assetHubAgent).safeNativeTransfer(initialDeposit);
 
-        new GatewayUpgradeMock();
-
         vm.stopBroadcast();
     }
 }

contracts/src/FundAgent.sol

@@ -9,7 +9,6 @@ import {BeefyClient} from "./BeefyClient.sol";
 import {IGateway} from "./interfaces/IGateway.sol";
 import {GatewayProxy} from "./GatewayProxy.sol";
 import {Gateway} from "./Gateway.sol";
-import {GatewayUpgradeMock} from "../test/mocks/GatewayUpgradeMock.sol";
 import {Agent} from "./Agent.sol";
 import {AgentExecutor} from "./AgentExecutor.sol";
 import {ParaID} from "./Types.sol";

contracts/src/Gateway.sol

@@ -19,8 +19,10 @@ import {
     Ticket,
     Costs
 } from "./Types.sol";
+import {Upgrade} from "./Upgrade.sol";
 import {IGateway} from "./interfaces/IGateway.sol";
 import {IInitializable} from "./interfaces/IInitializable.sol";
+import {IUpgradable} from "./interfaces/IUpgradable.sol";
 import {ERC1967} from "./utils/ERC1967.sol";
 import {Address} from "./utils/Address.sol";
 import {SafeNativeTransfer} from "./utils/SafeTransfer.sol";
@@ -46,7 +48,7 @@ import {AssetsStorage} from "./storage/AssetsStorage.sol";
 
 import {UD60x18, ud60x18, convert} from "prb/math/src/UD60x18.sol";
 
-contract Gateway is IGateway, IInitializable {
+contract Gateway is IGateway, IInitializable, IUpgradable {
     using Address for address;
     using SafeNativeTransfer for address payable;
 
@@ -84,11 +86,9 @@ contract Gateway is IGateway, IInitializable {
     error InvalidChannelUpdate();
     error AgentExecutionFailed(bytes returndata);
     error InvalidAgentExecutionPayload();
-    error InvalidCodeHash();
     error InvalidConstructorParams();
-    error AlreadyInitialized();
 
-    // handler functions are privileged
+    // Message handlers can only be dispatched by the gateway itself
     modifier onlySelf() {
         if (msg.sender != address(this)) {
             revert Unauthorized();
@@ -331,29 +331,7 @@ contract Gateway is IGateway, IInitializable {
     /// @dev Perform an upgrade of the gateway
     function upgrade(bytes calldata data) external onlySelf {
         UpgradeParams memory params = abi.decode(data, (UpgradeParams));
-
-        // Verify that the implementation is actually a contract
-        if (!params.impl.isContract()) {
-            revert InvalidCodeHash();
-        }
-
-        // As a sanity check, ensure that the codehash of implementation contract
-        // matches the codehash in the upgrade proposal
-        if (params.impl.codehash != params.implCodeHash) {
-            revert InvalidCodeHash();
-        }
-
-        // Update the proxy with the address of the new implementation
-        ERC1967.store(params.impl);
-
-        // Apply the initialization function of the implementation only if params were provided
-        if (params.initParams.length > 0) {
-            (bool success, bytes memory returndata) =
-                params.impl.delegatecall(abi.encodeCall(IInitializable.initialize, params.initParams));
-            Call.verifyResult(success, returndata);
-        }
-
-        emit Upgraded(params.impl);
+        Upgrade.upgrade(params.impl, params.implCodeHash, params.initParams);
     }
 
     // @dev Set the operating mode of the gateway
@@ -587,10 +565,6 @@ contract Gateway is IGateway, IInitializable {
 
         CoreStorage.Layout storage core = CoreStorage.layout();
 
-        if (core.channels[PRIMARY_GOVERNANCE_CHANNEL_ID].agent != address(0)) {
-            revert AlreadyInitialized();
-        }
-
         Config memory config = abi.decode(data, (Config));
 
         core.mode = config.mode;

contracts/src/Shell.sol

@@ -0,0 +1,29 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: 2023 Snowfork <[email protected]>
+pragma solidity 0.8.23;
+
+import {Upgrade} from "./Upgrade.sol";
+import {IInitializable} from "./interfaces/IInitializable.sol";
+import {IUpgradable} from "./interfaces/IUpgradable.sol";
+import {IShell} from "./interfaces/IShell.sol";
+
+// address recoveryOperator = vm.envOr("RECOVERY_OPERATOR", address(0));
+
+contract Shell is IShell, IUpgradable, IInitializable {
+    address public immutable operator;
+
+    error Unauthorised();
+
+    constructor(address _operator) {
+        operator = _operator;
+    }
+
+    function upgrade(address impl, bytes32 implCodeHash, bytes calldata initializerParams) external {
+        if (msg.sender != operator) {
+            revert Unauthorised();
+        }
+        Upgrade.upgrade(impl, implCodeHash, initializerParams);
+    }
+
+    function initialize(bytes memory params) external {}
+}

contracts/src/Upgrade.sol

@@ -0,0 +1,37 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: 2023 Snowfork <[email protected]>
+pragma solidity 0.8.23;
+
+import {ERC1967} from "./utils/ERC1967.sol";
+import {Call} from "./utils/Call.sol";
+import {Address} from "./utils/Address.sol";
+import {IInitializable} from "./interfaces/IInitializable.sol";
+import {IUpgradable} from "./interfaces/IUpgradable.sol";
+
+/// @dev Upgrades implementation contract
+library Upgrade {
+    using Address for address;
+
+    function upgrade(address impl, bytes32 implCodeHash, bytes memory initializerParams) internal {
+        // Verify that the implementation is actually a contract
+        if (!impl.isContract()) {
+            revert IUpgradable.InvalidContract();
+        }
+
+        // As a sanity check, ensure that the codehash of implementation contract
+        // matches the codehash in the upgrade proposal
+        if (impl.codehash != implCodeHash) {
+            revert IUpgradable.InvalidCodeHash();
+        }
+
+        // Update the proxy with the address of the new implementation
+        ERC1967.store(impl);
+
+        // Call the initializer
+        (bool success, bytes memory returndata) =
+            impl.delegatecall(abi.encodeCall(IInitializable.initialize, initializerParams));
+        Call.verifyResult(success, returndata);
+
+        emit IUpgradable.Upgraded(impl);
+    }
+}

contracts/src/interfaces/IGateway.sol

@@ -26,9 +26,6 @@ interface IGateway {
     // Emitted when a channel has been updated
     event ChannelUpdated(ChannelID indexed channelID);
 
-    // Emitted when the gateway is upgraded
-    event Upgraded(address indexed implementation);
-
     // Emitted when the operating mode is changed
     event OperatingModeChanged(OperatingMode mode);
 

contracts/src/interfaces/IShell.sol

@@ -0,0 +1,8 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: 2023 Snowfork <[email protected]>
+pragma solidity 0.8.23;
+
+interface IShell {
+    // Upgrade gateway shell to a new implementation
+    function upgrade(address impl, bytes32 implCodeHash, bytes calldata initializerParams) external;
+}

contracts/src/interfaces/IUpgradable.sol

@@ -0,0 +1,13 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: 2023 Snowfork <[email protected]>
+pragma solidity 0.8.23;
+
+interface IUpgradable {
+    // The new implementation address is a not a contract
+    error InvalidContract();
+    // The supplied codehash does not match the new implementation codehash
+    error InvalidCodeHash();
+
+    // The implementation contract was upgraded
+    event Upgraded(address indexed implementation);
+}

contracts/src/upgrades/rococo/GatewayV2.sol

@@ -9,6 +9,7 @@ import {PricingStorage} from "../../storage/PricingStorage.sol";
 
 contract GatewayV2 is Gateway {
     constructor(
+        address recoveryOperator,
         address beefyClient,
         address agentExecutor,
         ParaID bridgeHubParaID,
@@ -24,10 +25,6 @@ contract GatewayV2 is Gateway {
 
         PricingStorage.Layout storage pricing = PricingStorage.layout();
 
-        if (pricing.multiplier != convert(0)) {
-            revert AlreadyInitialized();
-        }
-
         pricing.multiplier = abi.decode(data, (UD60x18));
     }
 }

contracts/test/Gateway.t.sol

@@ -9,6 +9,7 @@ import {BeefyClient} from "../src/BeefyClient.sol";
 
 import {IGateway} from "../src/interfaces/IGateway.sol";
 import {IInitializable} from "../src/interfaces/IInitializable.sol";
+import {IUpgradable} from "../src/interfaces/IUpgradable.sol";
 import {Gateway} from "../src/Gateway.sol";
 import {GatewayMock, GatewayV2} from "./mocks/GatewayMock.sol";
 
@@ -19,6 +20,9 @@ import {Agent} from "../src/Agent.sol";
 import {Verification} from "../src/Verification.sol";
 import {Assets} from "../src/Assets.sol";
 import {SubstrateTypes} from "./../src/SubstrateTypes.sol";
+import {MultiAddress} from "../src/MultiAddress.sol";
+import {Channel, InboundMessage, OperatingMode, ParaID, Command, ChannelID, MultiAddress} from "../src/Types.sol";
+
 
 import {NativeTransferFailed} from "../src/utils/SafeTransfer.sol";
 import {PricingStorage} from "../src/storage/PricingStorage.sol";
@@ -46,7 +50,6 @@ import {
 } from "../src/Types.sol";
 
 import {WETH9} from "canonical-weth/WETH9.sol";
-import "./mocks/GatewayUpgradeMock.sol";
 import {UD60x18, ud60x18, convert} from "prb/math/src/UD60x18.sol";
 
 contract GatewayTest is Test {
@@ -477,91 +480,14 @@ contract GatewayTest is Test {
 
         // Expect the gateway to emit `Upgraded`
         vm.expectEmit(true, false, false, false);
-        emit IGateway.Upgraded(address(newLogic));
+        emit IUpgradable.Upgraded(address(newLogic));
 
         GatewayMock(address(gateway)).upgradePublic(abi.encode(params));
 
-        // Verify that the GatewayV2.setup was called
+        // Verify that the GatewayV2.initialize was called
         assertEq(GatewayV2(address(gateway)).getValue(), 42);
     }
 
-    function testUpgradeInitializerRunsOnlyOnce() public {
-        // Upgrade to this current logic contract
-        AgentExecutor executor = new AgentExecutor();
-        GatewayMock currentLogic =
-            new GatewayMock(address(0), address(executor), bridgeHubParaID, bridgeHubAgentID, foreignTokenDecimals);
-
-        Gateway.Config memory config = Gateway.Config({
-            mode: OperatingMode.Normal,
-            deliveryCost: outboundFee,
-            registerTokenFee: registerTokenFee,
-            assetHubParaID: assetHubParaID,
-            assetHubAgentID: assetHubAgentID,
-            assetHubCreateAssetFee: createTokenFee,
-            assetHubReserveTransferFee: sendTokenFee,
-            exchangeRate: exchangeRate,
-            multiplier: multiplier
-        });
-
-        UpgradeParams memory params = UpgradeParams({
-            impl: address(currentLogic),
-            implCodeHash: address(currentLogic).codehash,
-            initParams: abi.encode(config)
-        });
-
-        vm.expectRevert(Gateway.AlreadyInitialized.selector);
-        // Expect the gateway to emit `Upgraded`
-        GatewayMock(address(gateway)).upgradePublic(abi.encode(params));
-    }
-
-    function testUpgradeSkipsInitializerIfNoneProvided() public {
-        bytes32 agentID = keccak256("123");
-
-        testSetPricingParameters();
-        uint256 fee = IGateway(address(gateway)).quoteRegisterTokenFee();
-        assertEq(fee, 20000000000000001);
-
-        testCreateAgent();
-        assertNotEq(GatewayMock(address(gateway)).agentOf(agentID), address(0));
-
-        // Upgrade to this current logic contract
-        AgentExecutor executor = new AgentExecutor();
-        GatewayMock currentLogic =
-            new GatewayMock(address(0), address(executor), bridgeHubParaID, bridgeHubAgentID, foreignTokenDecimals);
-
-        bytes memory initParams; // empty
-        UpgradeParams memory params = UpgradeParams({
-            impl: address(currentLogic),
-            implCodeHash: address(currentLogic).codehash,
-            initParams: initParams
-        });
-
-        // Expect the gateway to emit `Upgraded`
-        GatewayMock(address(gateway)).upgradePublic(abi.encode(params));
-
-        // Verify that storage was not overwritten
-        fee = IGateway(address(gateway)).quoteRegisterTokenFee();
-        assertEq(fee, 20000000000000001);
-        assertNotEq(GatewayMock(address(gateway)).agentOf(agentID), address(0));
-    }
-
-    function testUpgradeGatewayMock() public {
-        GatewayUpgradeMock newLogic = new GatewayUpgradeMock();
-        uint256 d0 = 99;
-        uint256 d1 = 66;
-        bytes memory initParams = abi.encode(d0, d1);
-        console.logBytes(initParams);
-
-        UpgradeParams memory params =
-            UpgradeParams({impl: address(newLogic), implCodeHash: address(newLogic).codehash, initParams: initParams});
-
-        // Expect the gateway to emit `Initialized`
-        vm.expectEmit(true, false, false, true);
-        emit GatewayUpgradeMock.Initialized(d0, d1);
-
-        GatewayMock(address(gateway)).upgradePublic(abi.encode(params));
-    }
-
     function testUpgradeFailOnInitializationFailure() public {
         GatewayV2 newLogic = new GatewayV2();
 
@@ -581,7 +507,7 @@ contract GatewayTest is Test {
         UpgradeParams memory params =
             UpgradeParams({impl: address(newLogic), implCodeHash: bytes32(0), initParams: abi.encode(42)});
 
-        vm.expectRevert(Gateway.InvalidCodeHash.selector);
+        vm.expectRevert(IUpgradable.InvalidCodeHash.selector);
         GatewayMock(address(gateway)).upgradePublic(abi.encode(params));
     }