diff --git a/src/BankV1.sol b/src/BankV1.sol index 0f55426..35965fa 100644 --- a/src/BankV1.sol +++ b/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_); } diff --git a/src/BankV2.sol b/src/BankV2.sol new file mode 100644 index 0000000..ec158bc --- /dev/null +++ b/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; + } +} diff --git a/test/BankUpgrade.t.sol b/test/BankUpgrade.t.sol new file mode 100644 index 0000000..8f15e65 --- /dev/null +++ b/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); + } +} diff --git a/test/BankV2.t.sol b/test/BankV2.t.sol new file mode 100644 index 0000000..a95b475 --- /dev/null +++ b/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; + } +} diff --git a/test/mocks/IncompatibleBank.sol b/test/mocks/IncompatibleBank.sol new file mode 100644 index 0000000..56b3d3a --- /dev/null +++ b/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 {} +} diff --git a/test/mocks/NonUUPSImplementation.sol b/test/mocks/NonUUPSImplementation.sol new file mode 100644 index 0000000..d377035 --- /dev/null +++ b/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; + } +}