feat: add storage-safe bank V2

2026-08-21 14:58:24 -06:00
parent eb2fd5245a
commit ca6a99b913
6 changed files with 419 additions and 1 deletions
@@ -0,0 +1,186 @@
+// 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);
+    }
+}