From 2a21766237a9c29970519f7fa00482faecd06fdd Mon Sep 17 00:00:00 2001 From: golem Date: Mon, 17 Aug 2026 17:10:35 -0600 Subject: [PATCH] feat: add V1 custody accounting --- src/BankV1.sol | 34 ++++ test/BankV1.t.sol | 326 ++++++++++++++++++++++++++++++ test/mocks/FeeOnTransferToken.sol | 26 +++ test/mocks/ReentrantToken.sol | 104 ++++++++++ 4 files changed, 490 insertions(+) create mode 100644 test/BankV1.t.sol create mode 100644 test/mocks/FeeOnTransferToken.sol create mode 100644 test/mocks/ReentrantToken.sol diff --git a/src/BankV1.sol b/src/BankV1.sol index 13ffca5..0f55426 100644 --- a/src/BankV1.sol +++ b/src/BankV1.sol @@ -2,6 +2,7 @@ pragma solidity 0.8.35; import {IERC20} from "@openzeppelin/contracts/token/ERC20/IERC20.sol"; +import {SafeERC20} from "@openzeppelin/contracts/token/ERC20/utils/SafeERC20.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"; @@ -9,7 +10,15 @@ import {OwnableUpgradeable} from "@openzeppelin/contracts-upgradeable/access/Own import {PausableUpgradeable} from "@openzeppelin/contracts-upgradeable/utils/PausableUpgradeable.sol"; contract BankV1 is Initializable, UUPSUpgradeable, OwnableUpgradeable, PausableUpgradeable, ReentrancyGuardTransient { + using SafeERC20 for IERC20; + error InvalidAsset(address asset); + error ZeroAmount(); + error InsufficientBalance(address account, uint256 available, uint256 requested); + error UnexpectedAssetDelta(uint256 expected, uint256 actual); + + event Deposited(address indexed account, uint256 amount); + event Withdrawn(address indexed account, uint256 amount); IERC20 internal _asset; mapping(address account => uint256 balance) internal _balances; @@ -40,6 +49,31 @@ contract BankV1 is Initializable, UUPSUpgradeable, OwnableUpgradeable, PausableU _unpause(); } + function deposit(uint256 amount) external whenNotPaused nonReentrant { + if (amount == 0) revert ZeroAmount(); + + uint256 reservesBefore = _asset.balanceOf(address(this)); + _asset.safeTransferFrom(msg.sender, address(this), amount); + uint256 reservesAfter = _asset.balanceOf(address(this)); + uint256 received = reservesAfter >= reservesBefore ? reservesAfter - reservesBefore : 0; + if (received != amount) revert UnexpectedAssetDelta(amount, received); + + _balances[msg.sender] += amount; + _totalLiabilities += amount; + emit Deposited(msg.sender, amount); + } + + function withdraw(uint256 amount) external whenNotPaused nonReentrant { + if (amount == 0) revert ZeroAmount(); + uint256 available = _balances[msg.sender]; + if (amount > available) revert InsufficientBalance(msg.sender, available, amount); + + _balances[msg.sender] = available - amount; + _totalLiabilities -= amount; + _asset.safeTransfer(msg.sender, amount); + emit Withdrawn(msg.sender, amount); + } + function asset() external view returns (IERC20) { return _asset; } diff --git a/test/BankV1.t.sol b/test/BankV1.t.sol new file mode 100644 index 0000000..e53edc1 --- /dev/null +++ b/test/BankV1.t.sol @@ -0,0 +1,326 @@ +// SPDX-License-Identifier: MIT +pragma solidity 0.8.35; + +import {IERC20Errors} from "@openzeppelin/contracts/interfaces/draft-IERC6093.sol"; +import {PausableUpgradeable} from "@openzeppelin/contracts-upgradeable/utils/PausableUpgradeable.sol"; +import {ReentrancyGuardTransient} from "@openzeppelin/contracts/utils/ReentrancyGuardTransient.sol"; +import {Upgrades} from "openzeppelin-foundry-upgrades/Upgrades.sol"; + +import {BankV1} from "../src/BankV1.sol"; +import {BankTestBase} from "./helpers/BankTestBase.sol"; +import {FeeOnTransferToken} from "./mocks/FeeOnTransferToken.sol"; +import {ReentrantToken} from "./mocks/ReentrantToken.sol"; + +event Deposited(address indexed account, uint256 amount); +event Withdrawn(address indexed account, uint256 amount); + +error ZeroAmount(); +error InsufficientBalance(address account, uint256 available, uint256 requested); +error UnexpectedAssetDelta(uint256 expected, uint256 actual); + +contract BankV1CustodyTest is BankTestBase { + uint256 private constant MAX_DEPOSIT = 1_000_000e6; + + function testDepositCreditsExactCustomerAndLiabilityAgainstReceivedReserves() public { + _mintAndApprove(alice, 1_000e6, 100e6); + + vm.prank(alice); + bank.deposit(100e6); + + assertEq(token.balanceOf(alice), 900e6); + assertEq(token.balanceOf(proxy), 100e6); + assertEq(bank.balanceOf(alice), 100e6); + assertEq(bank.totalLiabilities(), 100e6); + } + + function testDepositEmitsDepositedEvent() public { + _mintAndApprove(alice, 100e6, 100e6); + + vm.expectEmit(true, false, false, true, proxy); + emit Deposited(alice, 100e6); + vm.prank(alice); + bank.deposit(100e6); + } + + function testDepositRejectsZeroAmount() public { + vm.prank(alice); + vm.expectRevert(ZeroAmount.selector); + bank.deposit(0); + } + + function testDepositRejectsCallsWhilePaused() public { + _mintAndApprove(alice, 100e6, 100e6); + vm.prank(owner); + bank.pause(); + + vm.prank(alice); + vm.expectRevert(PausableUpgradeable.EnforcedPause.selector); + bank.deposit(100e6); + } + + function testDepositRollsBackWhenAllowanceIsInadequate() public { + _mintAndApprove(alice, 100e6, 99e6); + + vm.prank(alice); + vm.expectRevert(abi.encodeWithSelector(IERC20Errors.ERC20InsufficientAllowance.selector, proxy, 99e6, 100e6)); + bank.deposit(100e6); + + _assertEmptyAccounting(alice); + assertEq(token.balanceOf(alice), 100e6); + } + + function testDepositRollsBackWhenWalletBalanceIsInadequate() public { + _mintAndApprove(alice, 99e6, 100e6); + + vm.prank(alice); + vm.expectRevert(abi.encodeWithSelector(IERC20Errors.ERC20InsufficientBalance.selector, alice, 99e6, 100e6)); + bank.deposit(100e6); + + _assertEmptyAccounting(alice); + assertEq(token.balanceOf(alice), 99e6); + } + + function testDepositsKeepTwoCustomersAccountingIndependent() public { + _mintAndApprove(alice, 1_000e6, 300e6); + _mintAndApprove(bob, 1_000e6, 700e6); + + vm.prank(alice); + bank.deposit(300e6); + vm.prank(bob); + bank.deposit(700e6); + + assertEq(bank.balanceOf(alice), 300e6); + assertEq(bank.balanceOf(bob), 700e6); + assertEq(bank.totalLiabilities(), 1_000e6); + assertEq(token.balanceOf(proxy), 1_000e6); + } + + function testFeeOnTransferDepositRevertsAndRollsBackTokenAndAccounting() public { + FeeOnTransferToken feeToken = new FeeOnTransferToken(); + BankV1 feeBank = _deployBank(address(feeToken)); + feeToken.mint(alice, 100e6); + vm.prank(alice); + feeToken.approve(address(feeBank), 100e6); + + vm.prank(alice); + vm.expectRevert(abi.encodeWithSelector(UnexpectedAssetDelta.selector, 100e6, 99e6)); + feeBank.deposit(100e6); + + assertEq(feeToken.balanceOf(alice), 100e6); + assertEq(feeToken.balanceOf(address(feeBank)), 0); + assertEq(feeBank.balanceOf(alice), 0); + assertEq(feeBank.totalLiabilities(), 0); + } + + function testDepositSwallowsNestedRevertAndCreditsOnlyOnce() public { + ReentrantToken reentrantToken = new ReentrantToken(); + BankV1 reentrantBank = _deployBank(address(reentrantToken)); + reentrantToken.mint(alice, 100e6); + vm.prank(alice); + reentrantToken.approve(address(reentrantBank), 100e6); + reentrantToken.configureDepositCallback(address(reentrantBank), false); + + vm.prank(alice); + reentrantBank.deposit(100e6); + + assertTrue(reentrantToken.nestedCallAttempted()); + assertFalse(reentrantToken.nestedCallSucceeded()); + assertEq(reentrantToken.nestedRevertSelector(), ReentrancyGuardTransient.ReentrancyGuardReentrantCall.selector); + assertEq(reentrantToken.observedAccountBalance(), 0); + assertEq(reentrantToken.observedLiabilities(), 0); + assertEq(reentrantToken.balanceOf(alice), 0); + assertEq(reentrantToken.balanceOf(address(reentrantBank)), 100e6); + assertEq(reentrantBank.balanceOf(alice), 100e6); + assertEq(reentrantBank.totalLiabilities(), 100e6); + } + + function testDepositPropagatesNestedRevertAtomicallyWhenConfigured() public { + ReentrantToken reentrantToken = new ReentrantToken(); + BankV1 reentrantBank = _deployBank(address(reentrantToken)); + reentrantToken.mint(alice, 100e6); + vm.prank(alice); + reentrantToken.approve(address(reentrantBank), 100e6); + reentrantToken.configureDepositCallback(address(reentrantBank), true); + + vm.prank(alice); + vm.expectRevert(ReentrancyGuardTransient.ReentrancyGuardReentrantCall.selector); + reentrantBank.deposit(100e6); + + assertEq(reentrantToken.balanceOf(alice), 100e6); + assertEq(reentrantToken.balanceOf(address(reentrantBank)), 0); + assertEq(reentrantBank.balanceOf(alice), 0); + assertEq(reentrantBank.totalLiabilities(), 0); + } + + function testWithdrawDebitsExactCustomerLiabilityAndReserves() public { + _deposit(alice, 1_000e6); + + vm.prank(alice); + bank.withdraw(400e6); + + assertEq(token.balanceOf(alice), 400e6); + assertEq(token.balanceOf(proxy), 600e6); + assertEq(bank.balanceOf(alice), 600e6); + assertEq(bank.totalLiabilities(), 600e6); + } + + function testWithdrawEmitsWithdrawnEvent() public { + _deposit(alice, 100e6); + + vm.expectEmit(true, false, false, true, proxy); + emit Withdrawn(alice, 40e6); + vm.prank(alice); + bank.withdraw(40e6); + } + + function testWithdrawRejectsZeroAmount() public { + vm.prank(alice); + vm.expectRevert(ZeroAmount.selector); + bank.withdraw(0); + } + + function testWithdrawRejectsCallsWhilePaused() public { + _deposit(alice, 100e6); + vm.prank(owner); + bank.pause(); + + vm.prank(alice); + vm.expectRevert(PausableUpgradeable.EnforcedPause.selector); + bank.withdraw(100e6); + } + + function testWithdrawReportsAvailableAndRequestedOnInsufficientInternalBalance() public { + _deposit(alice, 40e6); + + vm.prank(alice); + vm.expectRevert(abi.encodeWithSelector(InsufficientBalance.selector, alice, 40e6, 41e6)); + bank.withdraw(41e6); + + assertEq(bank.balanceOf(alice), 40e6); + assertEq(bank.totalLiabilities(), 40e6); + assertEq(token.balanceOf(proxy), 40e6); + } + + function testWithdrawUpdatesAccountingBeforeTransferCallbackAndCannotDoubleDebit() public { + ReentrantToken reentrantToken = new ReentrantToken(); + BankV1 reentrantBank = _deployBank(address(reentrantToken)); + reentrantToken.mint(alice, 100e6); + vm.prank(alice); + reentrantToken.approve(address(reentrantBank), 100e6); + vm.prank(alice); + reentrantBank.deposit(100e6); + reentrantToken.configureWithdrawalCallback(address(reentrantBank), false); + + vm.prank(alice); + reentrantBank.withdraw(40e6); + + assertEq(reentrantToken.observedAccountBalance(), 60e6); + assertEq(reentrantToken.observedLiabilities(), 60e6); + assertTrue(reentrantToken.nestedCallAttempted()); + assertFalse(reentrantToken.nestedCallSucceeded()); + assertEq(reentrantToken.nestedRevertSelector(), ReentrancyGuardTransient.ReentrancyGuardReentrantCall.selector); + assertEq(reentrantToken.balanceOf(alice), 40e6); + assertEq(reentrantToken.balanceOf(address(reentrantBank)), 60e6); + assertEq(reentrantBank.balanceOf(alice), 60e6); + assertEq(reentrantBank.totalLiabilities(), 60e6); + } + + function testOneCustomersWithdrawalLeavesOtherCustomerUnchanged() public { + _deposit(alice, 100e6); + _deposit(bob, 200e6); + + vm.prank(alice); + bank.withdraw(40e6); + + assertEq(bank.balanceOf(alice), 60e6); + assertEq(bank.balanceOf(bob), 200e6); + assertEq(bank.totalLiabilities(), 260e6); + assertEq(token.balanceOf(proxy), 260e6); + } + + function testDirectTransferCreatesSurplusThatRemainsAfterFullWithdrawal() public { + _deposit(alice, 100e6); + vm.prank(owner); + token.mint(bob, 25e6); + vm.prank(bob); + token.transfer(proxy, 25e6); + + assertEq(token.balanceOf(proxy), 125e6); + assertEq(bank.totalLiabilities(), 100e6); + + vm.prank(alice); + bank.withdraw(100e6); + + assertEq(token.balanceOf(proxy), 25e6); + assertEq(bank.balanceOf(alice), 0); + assertEq(bank.totalLiabilities(), 0); + } + + function testFuzzDepositPreservesExactAccounting(uint256 amountSeed) public { + uint256 amount = bound(amountSeed, 1, MAX_DEPOSIT); + _mintAndApprove(alice, amount, amount); + + vm.prank(alice); + bank.deposit(amount); + + assertEq(token.balanceOf(alice), 0); + assertEq(token.balanceOf(proxy), amount); + assertEq(bank.balanceOf(alice), amount); + assertEq(bank.totalLiabilities(), amount); + } + + function testFuzzWithdrawPreservesExactAccounting(uint256 depositSeed, uint256 withdrawalSeed) public { + uint256 deposited = bound(depositSeed, 1, MAX_DEPOSIT); + uint256 withdrawn = bound(withdrawalSeed, 1, deposited); + _deposit(alice, deposited); + + vm.prank(alice); + bank.withdraw(withdrawn); + + uint256 remaining = deposited - withdrawn; + assertEq(token.balanceOf(alice), withdrawn); + assertEq(token.balanceOf(proxy), remaining); + assertEq(bank.balanceOf(alice), remaining); + assertEq(bank.totalLiabilities(), remaining); + } + + function testFuzzOverWithdrawAlwaysReverts(uint256 depositSeed, uint256 excessSeed) public { + uint256 deposited = bound(depositSeed, 1, MAX_DEPOSIT); + uint256 excess = bound(excessSeed, 1, MAX_DEPOSIT); + uint256 requested = deposited + excess; + _deposit(alice, deposited); + + vm.prank(alice); + vm.expectRevert(abi.encodeWithSelector(InsufficientBalance.selector, alice, deposited, requested)); + bank.withdraw(requested); + + assertEq(bank.balanceOf(alice), deposited); + assertEq(bank.totalLiabilities(), deposited); + assertEq(token.balanceOf(proxy), deposited); + } + + function _mintAndApprove(address account, uint256 mintAmount, uint256 approveAmount) private { + vm.prank(owner); + token.mint(account, mintAmount); + vm.prank(account); + token.approve(proxy, approveAmount); + } + + function _deposit(address account, uint256 amount) private { + _mintAndApprove(account, amount, amount); + vm.prank(account); + bank.deposit(amount); + } + + function _deployBank(address asset_) private returns (BankV1 deployedBank) { + address deployedProxy = + Upgrades.deployUUPSProxy("BankV1.sol:BankV1", abi.encodeCall(BankV1.initialize, (asset_, owner))); + deployedBank = BankV1(deployedProxy); + } + + function _assertEmptyAccounting(address account) private view { + assertEq(token.balanceOf(proxy), 0); + assertEq(bank.balanceOf(account), 0); + assertEq(bank.totalLiabilities(), 0); + } +} diff --git a/test/mocks/FeeOnTransferToken.sol b/test/mocks/FeeOnTransferToken.sol new file mode 100644 index 0000000..80f5280 --- /dev/null +++ b/test/mocks/FeeOnTransferToken.sol @@ -0,0 +1,26 @@ +// SPDX-License-Identifier: MIT +pragma solidity 0.8.35; + +import {ERC20} from "@openzeppelin/contracts/token/ERC20/ERC20.sol"; + +contract FeeOnTransferToken is ERC20 { + constructor() ERC20("Fee-on-Transfer Token", "FOT") {} + + function decimals() public pure override returns (uint8) { + return 6; + } + + function mint(address to, uint256 amount) external { + _mint(to, amount); + } + + function transferFrom(address from, address to, uint256 amount) public override returns (bool) { + address spender = _msgSender(); + _spendAllowance(from, spender, amount); + + uint256 received = amount * 99 / 100; + _transfer(from, to, received); + _burn(from, amount - received); + return true; + } +} diff --git a/test/mocks/ReentrantToken.sol b/test/mocks/ReentrantToken.sol new file mode 100644 index 0000000..98dabf3 --- /dev/null +++ b/test/mocks/ReentrantToken.sol @@ -0,0 +1,104 @@ +// SPDX-License-Identifier: MIT +pragma solidity 0.8.35; + +import {ERC20} from "@openzeppelin/contracts/token/ERC20/ERC20.sol"; + +interface IReentrantBankTarget { + function deposit(uint256 amount) external; + function withdraw(uint256 amount) external; + function balanceOf(address account) external view returns (uint256); + function totalLiabilities() external view returns (uint256); +} + +contract ReentrantToken is ERC20 { + enum Callback { + None, + Deposit, + Withdraw + } + + IReentrantBankTarget public callbackTarget; + Callback public callback; + bool public propagateRevert; + bool public nestedCallAttempted; + bool public nestedCallSucceeded; + bytes4 public nestedRevertSelector; + uint256 public observedAccountBalance; + uint256 public observedLiabilities; + + constructor() ERC20("Reentrant Token", "REENT") {} + + function decimals() public pure override returns (uint8) { + return 6; + } + + function mint(address to, uint256 amount) external { + _mint(to, amount); + } + + function configureDepositCallback(address bank, bool propagate) external { + callbackTarget = IReentrantBankTarget(bank); + callback = Callback.Deposit; + propagateRevert = propagate; + _resetObservations(); + } + + function configureWithdrawalCallback(address bank, bool propagate) external { + callbackTarget = IReentrantBankTarget(bank); + callback = Callback.Withdraw; + propagateRevert = propagate; + _resetObservations(); + } + + function clearCallback() external { + callback = Callback.None; + propagateRevert = false; + _resetObservations(); + } + + function transferFrom(address from, address to, uint256 amount) public override returns (bool) { + if (callback == Callback.Deposit && _msgSender() == address(callbackTarget)) { + observedAccountBalance = callbackTarget.balanceOf(from); + observedLiabilities = callbackTarget.totalLiabilities(); + _attemptNestedCall(abi.encodeCall(IReentrantBankTarget.deposit, (1))); + } + return super.transferFrom(from, to, amount); + } + + function transfer(address to, uint256 amount) public override returns (bool) { + if (callback == Callback.Withdraw && _msgSender() == address(callbackTarget)) { + observedAccountBalance = callbackTarget.balanceOf(to); + observedLiabilities = callbackTarget.totalLiabilities(); + _attemptNestedCall(abi.encodeCall(IReentrantBankTarget.withdraw, (1))); + } + return super.transfer(to, amount); + } + + function _attemptNestedCall(bytes memory callData) private { + nestedCallAttempted = true; + bytes memory revertData; + (nestedCallSucceeded, revertData) = address(callbackTarget).call(callData); + + if (!nestedCallSucceeded && revertData.length >= 4) { + bytes4 selector; + assembly ("memory-safe") { + selector := mload(add(revertData, 0x20)) + } + nestedRevertSelector = selector; + } + + if (!nestedCallSucceeded && propagateRevert) { + assembly ("memory-safe") { + revert(add(revertData, 0x20), mload(revertData)) + } + } + } + + function _resetObservations() private { + nestedCallAttempted = false; + nestedCallSucceeded = false; + nestedRevertSelector = bytes4(0); + observedAccountBalance = 0; + observedLiabilities = 0; + } +}