2127 test rollback strategies for smart contracts (#2154)

* downgrade demonstrated via test

* test downgrade to invalid address reverts

* turned require statements into reverts in scripts

contracts/script/FeeContract.s.sol

@@ -8,6 +8,8 @@ import { FeeContract as FC } from "../src/FeeContract.sol";
 contract DeployFeeContractScript is Script {
     string internal contractName = vm.envString("FEE_CONTRACT_ORIGINAL_NAME");
 
+    error OwnerNotAsExpected(address expectedOwner, address currentOwner);
+
     /// @dev Deploys both the proxy and the implementation contract.
     /// The proxy admin is set as the owner of the contract upon deployment.
     /// The `owner` parameter should be the address of the multisig wallet to ensure proper
@@ -37,10 +39,9 @@ contract DeployFeeContractScript is Script {
         FC feeContractProxy = FC(payable(proxyAddress));
 
         // verify post deployment details
-        require(
-            feeContractProxy.owner() == owner,
-            "Post Deployment Verification: The contract owner is the one you specified"
-        );
+        if (feeContractProxy.owner() != owner) {
+            revert OwnerNotAsExpected(owner, feeContractProxy.owner());
+        }
 
         // Get the implementation address
         implementationAddress = Upgrades.getImplementationAddress(proxyAddress);

contracts/script/LightClient.s.sol

@@ -10,6 +10,11 @@ import { ERC1967Proxy } from "@openzeppelin/contracts/proxy/ERC1967/ERC1967Proxy
 contract DeployLightClientScript is Script {
     string public contractName = vm.envString("LIGHT_CLIENT_CONTRACT_ORIGINAL_NAME");
 
+    // Deployment Errors
+    error SetPermissionedProverFailed();
+    error OwnerTransferFailed();
+    error RetentionPeriodIsNotSetCorrectly();
+
     /// @dev Deploys both the proxy and the implementation contract.
     /// The proxy admin is set as the owner of the contract upon deployment.
     /// The `owner` parameter should be the address of the multisig wallet to ensure proper
@@ -18,6 +23,7 @@ contract DeployLightClientScript is Script {
     /// @param stateHistoryRetentionPeriod state history retention period in seconds
     /// @param owner The address that will be set as the owner of the proxy (typically a multisig
     /// wallet).
+
     function run(uint32 numInitValidators, uint32 stateHistoryRetentionPeriod, address owner)
         public
         returns (
@@ -65,17 +71,11 @@ contract DeployLightClientScript is Script {
         lightClientProxy.transferOwnership(owner);
 
         // verify post deployment details
-        require(
-            lightClientProxy.permissionedProver() == permissionedProver,
-            "Post Deployment Verification: Set permissioned prover failed"
-        );
-        require(
-            lightClientProxy.owner() == owner, "Post Deployment Verification: Owner transfer failed"
-        );
-        require(
-            lightClientProxy.stateHistoryRetentionPeriod() == stateHistoryRetentionPeriod,
-            "Post Deployment Verification: stateHistoryRetentionPeriod is not as expected"
-        );
+        if (lightClientProxy.permissionedProver() != owner) revert SetPermissionedProverFailed();
+        if (lightClientProxy.owner() != owner) revert OwnerTransferFailed();
+        if (lightClientProxy.stateHistoryRetentionPeriod() != stateHistoryRetentionPeriod) {
+            revert RetentionPeriodIsNotSetCorrectly();
+        }
 
         // Get the implementation address
         implementationAddress = Upgrades.getImplementationAddress(proxyAddress);

contracts/test/LightClientUpgradeToVx.t.sol

@@ -12,12 +12,17 @@ import { LightClientV3 as LCV3 } from "../test/LightClientV3.sol";
 import { DeployLightClientContractWithoutMultiSigScript as DeployScript } from
     "./script/LightClientTestScript.s.sol";
 import { UpgradeLightClientScript as UpgradeScript } from "./script/UpgradeLightClientToV2.s.sol";
+import { DowngradeLightClientScript as DowngradeScript } from
+    "./script/DowngradeLightClientV2ToV1.s.sol";
+
 import { UpgradeLightClientScript as ULCV3 } from "./script/UpgradeLightClientToV3.s.sol";
 import { OwnableUpgradeable } from
     "@openzeppelin/contracts-upgradeable/access/OwnableUpgradeable.sol";
 import { Initializable } from "@openzeppelin/contracts-upgradeable/proxy/utils/Initializable.sol";
 import { BN254 } from "bn254/BN254.sol";
 import { IPlonkVerifier as V } from "../src/interfaces/IPlonkVerifier.sol";
+import { Upgrades } from "openzeppelin-foundry-upgrades/Upgrades.sol";
+import { ERC1967Utils } from "@openzeppelin/contracts/proxy/ERC1967/ERC1967Utils.sol";
 
 contract LightClientUpgradeToVxTest is Test {
     LCV1 public lcV1Proxy;
@@ -26,20 +31,24 @@ contract LightClientUpgradeToVxTest is Test {
 
     DeployScript public deployer = new DeployScript();
     UpgradeScript public upgraderV2 = new UpgradeScript();
+    DowngradeScript public downgrader = new DowngradeScript();
     ULCV3 public upgraderV3 = new ULCV3();
 
     LCV1.LightClientState public stateV1;
     LCV1.StakeTableState public stakeStateV1;
 
     address public admin;
     address public proxy;
+    address public lcV1Impl;
 
     uint32 public constant MAX_HISTORY_SECONDS = 864000; //10 days
 
     // deploy the first implementation with its proxy
     function setUp() public {
         (proxy, admin, stateV1, stakeStateV1) = deployer.run(5, MAX_HISTORY_SECONDS);
         lcV1Proxy = LCV1(proxy);
+        lcV1Impl = Upgrades.getImplementationAddress(proxy);
+        assertNotEq(lcV1Impl, address(0));
     }
 
     function testCorrectInitialization() public view {
@@ -96,6 +105,99 @@ contract LightClientUpgradeToVxTest is Test {
         assertEq(extraFieldV2, expectedExtendedLightClientState.extraField);
     }
 
+    // test that the data remains the same after upgrading the implementation
+    function testRollbackV2toV1() public {
+        // Upgrade LightClient and check that the genesis state is not changed and that the new
+        // field
+        // of the upgraded contract is set to 0
+        uint256 myNewField = 123;
+        uint256 extraField = 2;
+        lcV2Proxy = LCV2(upgraderV2.run(proxy, myNewField, extraField, admin));
+
+        assertEq(lcV2Proxy.newField(), myNewField);
+
+        LCV1.LightClientState memory expectedLightClientState =
+            LCV1.LightClientState(stateV1.viewNum, stateV1.blockHeight, stateV1.blockCommRoot);
+
+        LCV2.ExtendedLightClientState memory expectedExtendedLightClientState = LCV2
+            .ExtendedLightClientState(
+            stateV1.viewNum, stateV1.blockHeight, stateV1.blockCommRoot, extraField
+        );
+
+        // compare with the current version of the light client state
+        (uint64 viewNum, uint64 blockHeight, BN254.ScalarField blockCommRoot) =
+            lcV2Proxy.finalizedState();
+        assertEq(viewNum, expectedLightClientState.viewNum);
+        assertEq(blockHeight, expectedLightClientState.blockHeight);
+        assertEq(abi.encode(blockCommRoot), abi.encode(expectedLightClientState.blockCommRoot));
+
+        // compare with the extended light client state
+        (
+            uint64 viewNumV2,
+            uint64 blockHeightV2,
+            BN254.ScalarField blockCommRootV2,
+            uint256 extraFieldV2
+        ) = lcV2Proxy.extendedFinalizedState();
+        assertEq(viewNumV2, expectedExtendedLightClientState.viewNum);
+        assertEq(blockHeightV2, expectedExtendedLightClientState.blockHeight);
+        assertEq(
+            abi.encode(blockCommRootV2), abi.encode(expectedExtendedLightClientState.blockCommRoot)
+        );
+        assertEq(extraFieldV2, expectedExtendedLightClientState.extraField);
+
+        // Now time to downgrade
+        // rollback to lightclient v1 by upgrading v2 to v1
+        lcV1Proxy = LCV1(downgrader.run(proxy, admin, lcV1Impl));
+
+        // re-confirm that the proxy address is the same for both versions
+        assertEq(address(lcV1Proxy), address(lcV2Proxy));
+
+        // confirm that the implementation address is the same as the first time we deployed, so we
+        // know the downgrade worked
+        assertEq(address(Upgrades.getImplementationAddress(address(lcV1Proxy))), address(lcV1Impl));
+
+        // ensure that the genesis states are still the same as the original contract
+        testCorrectInitialization();
+    }
+
+    // test that the data remains the same after upgrading the implementation
+    function testExpectRevertRollbackV2toInvalidAddress() public {
+        // Upgrade LightClient and check that the genesis state is not changed and that the new
+        // field
+        // of the upgraded contract is set to 0
+        uint256 myNewField = 123;
+        uint256 extraField = 2;
+        lcV2Proxy = LCV2(upgraderV2.run(proxy, myNewField, extraField, admin));
+
+        assertEq(lcV2Proxy.newField(), myNewField);
+
+        // get current light client implementation address
+        address lcV2Impl = Upgrades.getImplementationAddress(address(lcV2Proxy));
+
+        // Now time to downgrade, use invalid lcV1 impl address
+        // we expect a revert when you try to downgrade to an invalid address, the proxy address
+        vm.expectRevert(
+            abi.encodeWithSelector(ERC1967Utils.ERC1967InvalidImplementation.selector, proxy)
+        );
+        LCV1(downgrader.run(proxy, admin, proxy));
+
+        // re-confirm that the proxy address is the same for both versions
+        assertEq(address(lcV1Proxy), address(lcV2Proxy));
+
+        // confirm that the implementation address is the same as the recently deployed light client
+        // v2
+        // since the downgrade would have reverted
+        assertEq(address(Upgrades.getImplementationAddress(address(lcV2Proxy))), lcV2Impl);
+
+        // ensure that the genesis states are still the same as the original contract
+        testCorrectInitialization();
+    }
+
+    /**
+     * TODO:
+     * show that downgrading to the wrong light client impl should fail
+     */
+
     // test that the data remains the same after upgrading the implementation
     function testExpectRevertUpgradeSameDataV1ToV2ReinitializeTwice() public {
         // Upgrade LightClient and check that the genesis state is not changed and that the new

contracts/test/script/DowngradeLightClientV2ToV1.s.sol

@@ -0,0 +1,39 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.19;
+
+import { Script } from "forge-std/Script.sol";
+
+import { LightClientV2 as LCV2 } from "../LightClientV2.sol";
+
+contract DowngradeLightClientScript is Script {
+    /// @notice runs the downgrade
+    /// @param mostRecentlyDeployedProxy address of deployed proxy
+    /// @param lightClientImplAddr address of old light client implementation that we're downgrading
+    /// to
+    /// @return address of the proxy
+    function run(address mostRecentlyDeployedProxy, address admin, address lightClientImplAddr)
+        external
+        returns (address)
+    {
+        vm.startBroadcast(admin);
+        address proxy = downgradeLightClient(mostRecentlyDeployedProxy, lightClientImplAddr);
+        vm.stopBroadcast();
+        return proxy;
+    }
+
+    /// @notice downgrades the light client contract by calling the upgrade function in the
+    /// implementation contract via the proxy
+    /// @param proxyAddress address of proxy
+    /// @param oldLightClient address of old implementation
+    /// @return address of the proxy
+    function downgradeLightClient(address proxyAddress, address oldLightClient)
+        public
+        returns (address)
+    {
+        LCV2 proxy = LCV2(proxyAddress); //make the function call on the previous implementation
+
+        // "upgrade" the proxy address so that it now points to the old implementation
+        proxy.upgradeToAndCall(oldLightClient, "");
+        return address(proxy);
+    }
+}