feat: add storage-safe bank V2
This commit is contained in:
golem
@@ -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);
|
||||
}
|
||||
}
|
||||
@@ -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;
|
||||
}
|
||||
}
|
||||
@@ -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 {}
|
||||
}
|
||||
@@ -0,0 +1,8 @@
|
||||
// SPDX-License-Identifier: MIT
|
||||
pragma solidity 0.8.35;
|
||||
|
||||
contract NonUUPSImplementation {
|
||||
function contractVersion() external pure returns (uint256) {
|
||||
return 2;
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user