From 58965e3116157aae9b526b4f4bdee5770771dd8a Mon Sep 17 00:00:00 2001 From: Vincent Geddes <117534+vgeddes@users.noreply.github.com> Date: Sun, 14 Apr 2024 15:19:07 +0200 Subject: [PATCH] Add upgradable `Shell` contract (#1170) * Add temporary rescue mechanism to gateway contract * Simplified the implementation * Add more comments --- contracts/src/DeployGatewayLogic.sol | 2 +- contracts/src/DeployScript.sol | 3 - contracts/src/FundAgent.sol | 1 - contracts/src/Gateway.sol | 36 ++------- contracts/src/Shell.sol | 29 +++++++ contracts/src/Upgrade.sol | 37 +++++++++ contracts/src/interfaces/IGateway.sol | 3 - contracts/src/interfaces/IShell.sol | 8 ++ contracts/src/interfaces/IUpgradable.sol | 13 +++ contracts/src/upgrades/rococo/GatewayV2.sol | 5 +- contracts/test/Gateway.t.sol | 88 ++------------------- contracts/test/Shell.t.sol | 48 +++++++++++ contracts/test/mocks/GatewayMock.sol | 4 +- contracts/test/mocks/GatewayUpgradeMock.sol | 64 --------------- 14 files changed, 152 insertions(+), 189 deletions(-) create mode 100644 contracts/src/Shell.sol create mode 100644 contracts/src/Upgrade.sol create mode 100644 contracts/src/interfaces/IShell.sol create mode 100644 contracts/src/interfaces/IUpgradable.sol create mode 100644 contracts/test/Shell.t.sol delete mode 100644 contracts/test/mocks/GatewayUpgradeMock.sol diff --git a/contracts/src/DeployGatewayLogic.sol b/contracts/src/DeployGatewayLogic.sol index 8c5cc70282..b9242d20b9 100644 --- a/contracts/src/DeployGatewayLogic.sol +++ b/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(); } diff --git a/contracts/src/DeployScript.sol b/contracts/src/DeployScript.sol index 8290d9ff8e..53d74fac2a 100644 --- a/contracts/src/DeployScript.sol +++ b/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(); } } diff --git a/contracts/src/FundAgent.sol b/contracts/src/FundAgent.sol index 95192299c0..d4c3201d14 100644 --- a/contracts/src/FundAgent.sol +++ b/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"; diff --git a/contracts/src/Gateway.sol b/contracts/src/Gateway.sol index 29900f438f..0bb2c098af 100644 --- a/contracts/src/Gateway.sol +++ b/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; diff --git a/contracts/src/Shell.sol b/contracts/src/Shell.sol new file mode 100644 index 0000000000..2654a8685b --- /dev/null +++ b/contracts/src/Shell.sol @@ -0,0 +1,29 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: 2023 Snowfork +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 {} +} diff --git a/contracts/src/Upgrade.sol b/contracts/src/Upgrade.sol new file mode 100644 index 0000000000..bee96074a0 --- /dev/null +++ b/contracts/src/Upgrade.sol @@ -0,0 +1,37 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: 2023 Snowfork +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); + } +} diff --git a/contracts/src/interfaces/IGateway.sol b/contracts/src/interfaces/IGateway.sol index e0d4d81da2..5bcb2fd0dd 100644 --- a/contracts/src/interfaces/IGateway.sol +++ b/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); diff --git a/contracts/src/interfaces/IShell.sol b/contracts/src/interfaces/IShell.sol new file mode 100644 index 0000000000..6811953848 --- /dev/null +++ b/contracts/src/interfaces/IShell.sol @@ -0,0 +1,8 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: 2023 Snowfork +pragma solidity 0.8.23; + +interface IShell { + // Upgrade gateway shell to a new implementation + function upgrade(address impl, bytes32 implCodeHash, bytes calldata initializerParams) external; +} diff --git a/contracts/src/interfaces/IUpgradable.sol b/contracts/src/interfaces/IUpgradable.sol new file mode 100644 index 0000000000..3574e5c5e3 --- /dev/null +++ b/contracts/src/interfaces/IUpgradable.sol @@ -0,0 +1,13 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: 2023 Snowfork +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); +} diff --git a/contracts/src/upgrades/rococo/GatewayV2.sol b/contracts/src/upgrades/rococo/GatewayV2.sol index 3918e3d790..e671e19422 100644 --- a/contracts/src/upgrades/rococo/GatewayV2.sol +++ b/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)); } } diff --git a/contracts/test/Gateway.t.sol b/contracts/test/Gateway.t.sol index 287f059a65..00f1682a0b 100644 --- a/contracts/test/Gateway.t.sol +++ b/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)); } diff --git a/contracts/test/Shell.t.sol b/contracts/test/Shell.t.sol new file mode 100644 index 0000000000..b399ae8acc --- /dev/null +++ b/contracts/test/Shell.t.sol @@ -0,0 +1,48 @@ +// SPDX-License-Identifier: Apache-2.0 +pragma solidity 0.8.23; + +import {Test} from "forge-std/Test.sol"; +import {Strings} from "openzeppelin/utils/Strings.sol"; +import {console} from "forge-std/console.sol"; + +import {IGateway} from "../src/interfaces/IGateway.sol"; +import {IInitializable} from "../src/interfaces/IInitializable.sol"; +import {IUpgradable} from "../src/interfaces/IUpgradable.sol"; +import {IShell} from "../src/interfaces/IShell.sol"; +import {Gateway} from "../src/Gateway.sol"; +import {GatewayProxy} from "../src/GatewayProxy.sol"; +import {Shell} from "../src/Shell.sol"; +import {Upgrade} from "../src/Upgrade.sol"; +import {GatewayMock, GatewayV2} from "./mocks/GatewayMock.sol"; + +contract ShellTest is Test { + GatewayProxy public gateway; + Shell public shell; + + function setUp() public { + shell = new Shell(address(this)); + gateway = new GatewayProxy(address(shell), bytes("")); + } + + function testUpgradeShell() public { + // Upgrade to this new logic contract + address newLogic = address(new GatewayV2()); + bytes memory initParams = abi.encode(42); + + // Expect the gateway to emit `Upgrade.Upgraded` + vm.expectEmit(true, false, false, true); + emit IUpgradable.Upgraded(newLogic); + + // Perform the upgrade + IShell(address(gateway)).upgrade(newLogic, newLogic.codehash, initParams); + + // Verify that the upgrade occured + + // Execute code only available in the new impl + assertEq(GatewayV2(address(gateway)).getValue(), 42); + + // Should no longer be able to upgrade via trusted operator + vm.expectRevert(); + IShell(address(gateway)).upgrade(newLogic, newLogic.codehash, initParams); + } +} diff --git a/contracts/test/mocks/GatewayMock.sol b/contracts/test/mocks/GatewayMock.sol index b1014811be..6f3f28b22b 100644 --- a/contracts/test/mocks/GatewayMock.sol +++ b/contracts/test/mocks/GatewayMock.sol @@ -89,8 +89,10 @@ library AdditionalStorage { } } +import {IInitializable} from "../../src/interfaces/IInitializable.sol"; + // Used to test upgrades. -contract GatewayV2 { +contract GatewayV2 is IInitializable { // Reinitialize gateway with some additional storage fields function initialize(bytes memory params) external { AdditionalStorage.Layout storage $ = AdditionalStorage.layout(); diff --git a/contracts/test/mocks/GatewayUpgradeMock.sol b/contracts/test/mocks/GatewayUpgradeMock.sol deleted file mode 100644 index 130ba07841..0000000000 --- a/contracts/test/mocks/GatewayUpgradeMock.sol +++ /dev/null @@ -1,64 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: 2023 Snowfork -pragma solidity 0.8.23; - -import {Channel, InboundMessage, OperatingMode, ParaID, Command, ChannelID, MultiAddress} from "../../src/Types.sol"; -import {IGateway} from "../../src/interfaces/IGateway.sol"; -import {IInitializable} from "../../src/interfaces/IInitializable.sol"; -import {Verification} from "../../src/Verification.sol"; -import {UD60x18, convert} from "prb/math/src/UD60x18.sol"; - -contract GatewayUpgradeMock is IGateway, IInitializable { - /** - * Getters - */ - function operatingMode() external pure returns (OperatingMode) { - return OperatingMode.Normal; - } - - function channelOperatingModeOf(ChannelID) external pure returns (OperatingMode) { - return OperatingMode.Normal; - } - - function channelNoncesOf(ChannelID) external pure returns (uint64, uint64) { - return (0, 0); - } - - function agentOf(bytes32) external pure returns (address) { - return address(0); - } - - function implementation() external pure returns (address) { - return address(0); - } - - function isTokenRegistered(address) external pure returns (bool) { - return true; - } - - function submitV1(InboundMessage calldata, bytes32[] calldata, Verification.Proof calldata) external {} - - function quoteRegisterTokenFee() external pure returns (uint256) { - return 1; - } - - function registerToken(address) external payable {} - - function quoteSendTokenFee(address, ParaID, uint128) external pure returns (uint256) { - return 1; - } - - function sendToken(address, ParaID, MultiAddress calldata, uint128, uint128) external payable {} - - event Initialized(uint256 d0, uint256 d1); - - function initialize(bytes memory data) external { - // Just decode and exit - (uint256 d0, uint256 d1) = abi.decode(data, (uint256, uint256)); - emit Initialized(d0, d1); - } - - function pricingParameters() external pure returns (UD60x18, uint128) { - return (convert(0), uint128(0)); - } -}