uupl-smart-contract/test/BankUpgrade.t.sol

// SPDX-License-Identifier: MIT
pragma solidity 0.8.35;

import {ERC1967Utils} from "@openzeppelin/contracts/proxy/ERC1967/ERC1967Utils.sol";
import {Initializable} from "@openzeppelin/contracts/proxy/utils/Initializable.sol";
import {OwnableUpgradeable} from "@openzeppelin/contracts-upgradeable/access/OwnableUpgradeable.sol";
import {Options, Upgrades} from "openzeppelin-foundry-upgrades/Upgrades.sol";

import {BankV1} from "../src/BankV1.sol";
import {BankV2} from "../src/BankV2.sol";
import {BankTestBase} from "./helpers/BankTestBase.sol";
import {IncompatibleBank} from "./mocks/IncompatibleBank.sol";
import {NonUUPSImplementation} from "./mocks/NonUUPSImplementation.sol";

contract UpgradeValidator {
    function validate(string memory contractName, Options memory opts) external {
        Upgrades.validateUpgrade(contractName, opts);
    }
}

contract BankUpgradeTest is BankTestBase {
    struct Snapshot {
        address proxyAddress;
        address implementationAddress;
        address ownerAddress;
        address assetAddress;
        bool pausedState;
        uint256 aliceBalance;
        uint256 bobBalance;
        uint256 liabilities;
        uint256 reserves;
        uint256 surplus;
        uint256 version;
    }

    Snapshot private beforeUpgrade;

    function setUp() public override {
        super.setUp();
        _deposit(alice, 1_000e6);
        _deposit(bob, 500e6);

        vm.prank(owner);
        token.mint(stranger, 75e6);
        vm.prank(stranger);
        assertTrue(token.transfer(proxy, 75e6));

        vm.prank(owner);
        bank.pause();

        uint256 liabilities = bank.totalLiabilities();
        uint256 reserves = token.balanceOf(proxy);
        beforeUpgrade = Snapshot({
            proxyAddress: proxy,
            implementationAddress: Upgrades.getImplementationAddress(proxy),
            ownerAddress: bank.owner(),
            assetAddress: address(bank.asset()),
            pausedState: bank.paused(),
            aliceBalance: bank.balanceOf(alice),
            bobBalance: bank.balanceOf(bob),
            liabilities: liabilities,
            reserves: reserves,
            surplus: reserves - liabilities,
            version: bank.contractVersion()
        });
    }

    function testValidatedOwnerUpgradeChangesOnlyImplementationAndVersion() public {
        BankV2 upgraded = _validatedOwnerUpgrade();
        address implementationAfter = Upgrades.getImplementationAddress(proxy);

        assertEq(address(upgraded), beforeUpgrade.proxyAddress);
        assertNotEq(implementationAfter, beforeUpgrade.implementationAddress);
        assertGt(implementationAfter.code.length, 0);
        assertEq(beforeUpgrade.version, 1);
        assertEq(upgraded.contractVersion(), 2);
        _assertSnapshotPreserved(upgraded);
    }

    function testEveryV1MutationStillWorksAfterUpgrade() public {
        BankV2 upgraded = _validatedOwnerUpgrade();

        vm.prank(owner);
        upgraded.unpause();
        assertFalse(upgraded.paused());

        _mintAndApprove(stranger, 100e6);
        vm.prank(stranger);
        upgraded.deposit(100e6);

        vm.prank(alice);
        upgraded.withdraw(200e6);

        assertEq(address(upgraded.asset()), address(token));
        assertEq(upgraded.owner(), owner);
        assertEq(upgraded.balanceOf(alice), 800e6);
        assertEq(upgraded.balanceOf(bob), 500e6);
        assertEq(upgraded.balanceOf(stranger), 100e6);
        assertEq(upgraded.totalLiabilities(), 1_400e6);
        assertEq(token.balanceOf(proxy), 1_475e6);
        assertEq(upgraded.contractVersion(), 2);

        vm.prank(owner);
        upgraded.pause();
        assertTrue(upgraded.paused());
        vm.prank(owner);
        upgraded.unpause();
        assertFalse(upgraded.paused());
    }

    function testNewImplementationCannotBeInitializedDirectly() public {
        _validatedOwnerUpgrade();
        address upgradedImplementation = Upgrades.getImplementationAddress(proxy);

        vm.expectRevert(Initializable.InvalidInitialization.selector);
        BankV2(upgradedImplementation).initialize(address(token), owner);
    }

    function testNonOwnerUpgradeToAndCallRejectsWithUnauthorizedAccount() public {
        address candidate = address(new BankV2());

        vm.prank(stranger);
        vm.expectRevert(abi.encodeWithSelector(OwnableUpgradeable.OwnableUnauthorizedAccount.selector, stranger));
        bank.upgradeToAndCall(candidate, "");

        assertEq(Upgrades.getImplementationAddress(proxy), beforeUpgrade.implementationAddress);
        assertEq(bank.contractVersion(), 1);
    }

    function testValidateUpgradeRejectsIncompatibleApplicationStorageLayout() public {
        Options memory opts;
        opts.referenceContract = "BankV1.sol:BankV1";
        UpgradeValidator validator = new UpgradeValidator();

        try validator.validate("IncompatibleBank.sol:IncompatibleBank", opts) {
            fail("incompatible storage layout was accepted");
        } catch Error(string memory reason) {
            assertTrue(vm.contains(reason, "Upgrade safety validation failed"));
            assertTrue(vm.contains(reason, "Deleted `_asset`"));
            assertTrue(vm.contains(reason, "Inserted `_asset`"));
        }
    }

    function testOwnerUpgradeToAndCallRejectsNonUUPSImplementationAtRuntime() public {
        NonUUPSImplementation candidate = new NonUUPSImplementation();

        vm.prank(owner);
        vm.expectRevert(abi.encodeWithSelector(ERC1967Utils.ERC1967InvalidImplementation.selector, address(candidate)));
        bank.upgradeToAndCall(address(candidate), "");

        assertEq(Upgrades.getImplementationAddress(proxy), beforeUpgrade.implementationAddress);
        assertEq(bank.contractVersion(), 1);
    }

    function _validatedOwnerUpgrade() private returns (BankV2 upgraded) {
        Options memory opts;
        opts.referenceContract = "BankV1.sol:BankV1";
        Upgrades.upgradeProxy(proxy, "BankV2.sol:BankV2", "", opts, owner);
        upgraded = BankV2(proxy);
    }

    function _assertSnapshotPreserved(BankV2 upgraded) private view {
        assertEq(address(upgraded), beforeUpgrade.proxyAddress);
        assertEq(upgraded.owner(), beforeUpgrade.ownerAddress);
        assertEq(address(upgraded.asset()), beforeUpgrade.assetAddress);
        assertEq(upgraded.paused(), beforeUpgrade.pausedState);
        assertEq(upgraded.balanceOf(alice), beforeUpgrade.aliceBalance);
        assertEq(upgraded.balanceOf(bob), beforeUpgrade.bobBalance);
        assertEq(upgraded.totalLiabilities(), beforeUpgrade.liabilities);
        assertEq(token.balanceOf(proxy), beforeUpgrade.reserves);
        assertEq(token.balanceOf(proxy) - upgraded.totalLiabilities(), beforeUpgrade.surplus);
    }

    function _deposit(address account, uint256 amount) private {
        _mintAndApprove(account, amount);
        vm.prank(account);
        bank.deposit(amount);
    }

    function _mintAndApprove(address account, uint256 amount) private {
        vm.prank(owner);
        token.mint(account, amount);
        vm.prank(account);
        token.approve(proxy, amount);
    }
}