feat: add storage-safe bank V2

src/BankV1.sol

@@ -30,7 +30,7 @@ contract BankV1 is Initializable, UUPSUpgradeable, OwnableUpgradeable, PausableU
         _disableInitializers();
     }
 
-    function initialize(address asset_, address initialOwner) external initializer {
+    function initialize(address asset_, address initialOwner) public initializer {
         if (asset_ == address(0)) {
             revert InvalidAsset(asset_);
         }

src/BankV2.sol

@@ -0,0 +1,29 @@
+// SPDX-License-Identifier: MIT
+pragma solidity 0.8.35;
+
+import {BankV1} from "./BankV1.sol";
+
+/// @custom:oz-upgrades-from src/BankV1.sol:BankV1
+contract BankV2 is BankV1 {
+    event BalanceTransferred(address indexed from, address indexed to, uint256 amount);
+
+    error InvalidRecipient(address recipient);
+    error SelfTransfer();
+
+    function transferBalance(address recipient, uint256 amount) external whenNotPaused {
+        if (amount == 0) revert ZeroAmount();
+        if (recipient == address(0)) revert InvalidRecipient(recipient);
+        if (recipient == msg.sender) revert SelfTransfer();
+        uint256 available = _balances[msg.sender];
+        if (amount > available) {
+            revert InsufficientBalance(msg.sender, available, amount);
+        }
+        _balances[msg.sender] = available - amount;
+        _balances[recipient] += amount;
+        emit BalanceTransferred(msg.sender, recipient, amount);
+    }
+
+    function contractVersion() public pure override returns (uint256) {
+        return 2;
+    }
+}

test/BankUpgrade.t.sol

@@ -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);
+    }
+}

test/BankV2.t.sol

@@ -0,0 +1,165 @@
+// SPDX-License-Identifier: MIT
+pragma solidity 0.8.35;
+
+import {PausableUpgradeable} from "@openzeppelin/contracts-upgradeable/utils/PausableUpgradeable.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";
+
+event BalanceTransferred(address indexed from, address indexed to, uint256 amount);
+
+contract BankV2Test is BankTestBase {
+    uint256 private constant ALICE_DEPOSIT = 1_000e6;
+    uint256 private constant BOB_DEPOSIT = 500e6;
+
+    BankV2 internal bankV2;
+
+    function setUp() public override {
+        super.setUp();
+        _deposit(alice, ALICE_DEPOSIT);
+        _deposit(bob, BOB_DEPOSIT);
+
+        Options memory opts;
+        opts.referenceContract = "BankV1.sol:BankV1";
+        Upgrades.upgradeProxy(proxy, "BankV2.sol:BankV2", "", opts, owner);
+        bankV2 = BankV2(proxy);
+    }
+
+    function testTransferBalanceMoves250MillionUnitsAndEmitsExactEvent() public {
+        vm.expectEmit(true, true, false, true, proxy);
+        emit BalanceTransferred(alice, bob, 250e6);
+        vm.prank(alice);
+        bankV2.transferBalance(bob, 250e6);
+
+        assertEq(bankV2.balanceOf(alice), 750e6);
+        assertEq(bankV2.balanceOf(bob), 750e6);
+    }
+
+    function testTransferBalanceLeavesLiabilitiesAndTokenReservesUnchanged() public {
+        uint256 liabilitiesBefore = bankV2.totalLiabilities();
+        uint256 reservesBefore = token.balanceOf(proxy);
+
+        vm.prank(alice);
+        bankV2.transferBalance(bob, 250e6);
+
+        assertEq(bankV2.totalLiabilities(), liabilitiesBefore);
+        assertEq(token.balanceOf(proxy), reservesBefore);
+    }
+
+    function testTransferBalanceRejectsZeroAmount() public {
+        vm.prank(alice);
+        vm.expectRevert(BankV1.ZeroAmount.selector);
+        bankV2.transferBalance(bob, 0);
+    }
+
+    function testTransferBalanceRejectsZeroRecipient() public {
+        vm.prank(alice);
+        vm.expectRevert(abi.encodeWithSelector(BankV2.InvalidRecipient.selector, address(0)));
+        bankV2.transferBalance(address(0), 1);
+    }
+
+    function testTransferBalanceRejectsSenderAsRecipient() public {
+        vm.prank(alice);
+        vm.expectRevert(BankV2.SelfTransfer.selector);
+        bankV2.transferBalance(alice, 1);
+    }
+
+    function testTransferBalanceReportsAvailableAndRequestedWhenBalanceIsInsufficient() public {
+        vm.prank(alice);
+        vm.expectRevert(
+            abi.encodeWithSelector(BankV1.InsufficientBalance.selector, alice, ALICE_DEPOSIT, ALICE_DEPOSIT + 1)
+        );
+        bankV2.transferBalance(bob, ALICE_DEPOSIT + 1);
+
+        assertEq(bankV2.balanceOf(alice), ALICE_DEPOSIT);
+        assertEq(bankV2.balanceOf(bob), BOB_DEPOSIT);
+    }
+
+    function testTransferBalanceRejectsCallsWhilePaused() public {
+        vm.prank(owner);
+        bankV2.pause();
+
+        vm.prank(alice);
+        vm.expectRevert(PausableUpgradeable.EnforcedPause.selector);
+        bankV2.transferBalance(bob, 1);
+
+        assertEq(bankV2.balanceOf(alice), ALICE_DEPOSIT);
+        assertEq(bankV2.balanceOf(bob), BOB_DEPOSIT);
+    }
+
+    function testTransferBalanceCreditsRecipientWithNoPreviousBalance() public {
+        assertEq(bankV2.balanceOf(stranger), 0);
+
+        vm.prank(alice);
+        bankV2.transferBalance(stranger, 125e6);
+
+        assertEq(bankV2.balanceOf(alice), 875e6);
+        assertEq(bankV2.balanceOf(stranger), 125e6);
+    }
+
+    function testFuzzTransferBalanceAcrossTrackedRecipientsAndSenderBoundedAmounts(
+        uint256 recipientSeed,
+        uint256 amountSeed
+    ) public {
+        address recipient = _trackedRecipient(bound(recipientSeed, 0, 2));
+        uint256 amount = bound(amountSeed, 1, ALICE_DEPOSIT);
+        uint256 recipientBefore = bankV2.balanceOf(recipient);
+        uint256 liabilitiesBefore = bankV2.totalLiabilities();
+        uint256 reservesBefore = token.balanceOf(proxy);
+
+        vm.prank(alice);
+        bankV2.transferBalance(recipient, amount);
+
+        assertEq(bankV2.balanceOf(alice), ALICE_DEPOSIT - amount);
+        assertEq(bankV2.balanceOf(recipient), recipientBefore + amount);
+        assertEq(bankV2.totalLiabilities(), liabilitiesBefore);
+        assertEq(token.balanceOf(proxy), reservesBefore);
+    }
+
+    function testV1DepositWithdrawalViewsPauseAndUnpauseStillWorkThroughV2Proxy() public {
+        _mintAndApprove(stranger, 100e6);
+        vm.prank(stranger);
+        bankV2.deposit(100e6);
+
+        vm.prank(alice);
+        bankV2.withdraw(100e6);
+
+        assertEq(address(bankV2.asset()), address(token));
+        assertEq(bankV2.owner(), owner);
+        assertEq(bankV2.balanceOf(alice), 900e6);
+        assertEq(bankV2.balanceOf(bob), BOB_DEPOSIT);
+        assertEq(bankV2.balanceOf(stranger), 100e6);
+        assertEq(bankV2.totalLiabilities(), 1_500e6);
+        assertEq(token.balanceOf(proxy), 1_500e6);
+        assertEq(bankV2.contractVersion(), 2);
+
+        vm.prank(owner);
+        bankV2.pause();
+        assertTrue(bankV2.paused());
+
+        vm.prank(owner);
+        bankV2.unpause();
+        assertFalse(bankV2.paused());
+    }
+
+    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);
+    }
+
+    function _trackedRecipient(uint256 index) private view returns (address) {
+        if (index == 0) return bob;
+        if (index == 1) return stranger;
+        return owner;
+    }
+}

test/mocks/IncompatibleBank.sol

@@ -0,0 +1,30 @@
+// SPDX-License-Identifier: MIT
+pragma solidity 0.8.35;
+
+import {IERC20} from "@openzeppelin/contracts/token/ERC20/IERC20.sol";
+import {Initializable} from "@openzeppelin/contracts/proxy/utils/Initializable.sol";
+import {UUPSUpgradeable} from "@openzeppelin/contracts/proxy/utils/UUPSUpgradeable.sol";
+import {ReentrancyGuardTransient} from "@openzeppelin/contracts/utils/ReentrancyGuardTransient.sol";
+import {OwnableUpgradeable} from "@openzeppelin/contracts-upgradeable/access/OwnableUpgradeable.sol";
+import {PausableUpgradeable} from "@openzeppelin/contracts-upgradeable/utils/PausableUpgradeable.sol";
+
+contract IncompatibleBank is
+    Initializable,
+    UUPSUpgradeable,
+    OwnableUpgradeable,
+    PausableUpgradeable,
+    ReentrancyGuardTransient
+{
+    mapping(address account => uint256 balance) internal _balances;
+    IERC20 internal _asset;
+    uint256 internal _totalLiabilities;
+    uint256[47] private __gap;
+
+    function initialize(address asset_, address initialOwner) public initializer {
+        __Ownable_init(initialOwner);
+        __Pausable_init();
+        _asset = IERC20(asset_);
+    }
+
+    function _authorizeUpgrade(address) internal override onlyOwner {}
+}

test/mocks/NonUUPSImplementation.sol

@@ -0,0 +1,8 @@
+// SPDX-License-Identifier: MIT
+pragma solidity 0.8.35;
+
+contract NonUUPSImplementation {
+    function contractVersion() external pure returns (uint256) {
+        return 2;
+    }
+}