From 20f1922c7326fe33e4e14dcbe017c976079e94ec Mon Sep 17 00:00:00 2001 From: Ben DiFrancesco Date: Tue, 12 Dec 2023 07:43:25 -0500 Subject: [PATCH] Implement and test the basic accounting mechanics for stake withdrawals --- src/UniStaker.sol | 12 ++++ test/UniStaker.t.sol | 160 ++++++++++++++++++++++++++++++++++++++++++- 2 files changed, 171 insertions(+), 1 deletion(-) diff --git a/src/UniStaker.sol b/src/UniStaker.sol index 8f33bf0..0564c34 100644 --- a/src/UniStaker.sol +++ b/src/UniStaker.sol @@ -10,6 +10,8 @@ import {ReentrancyGuard} from "openzeppelin/utils/ReentrancyGuard.sol"; contract UniStaker is ReentrancyGuard { type DepositIdentifier is uint256; + error UniStaker__Unauthorized(bytes32 reason, address caller); + struct Deposit { uint256 balance; address owner; @@ -48,6 +50,16 @@ contract UniStaker is ReentrancyGuard { deposits[_depositId] = Deposit({balance: _amount, owner: msg.sender, delegatee: _delegatee}); } + function withdraw(DepositIdentifier _depositId, uint256 _amount) external nonReentrant { + Deposit storage deposit = deposits[_depositId]; + if (msg.sender != deposit.owner) revert UniStaker__Unauthorized("not owner", msg.sender); + + deposit.balance -= _amount; // overflow prevents withdrawing more than balance + totalSupply -= _amount; + totalDeposits[msg.sender] -= _amount; + _stakeTokenSafeTransferFrom(address(surrogates[deposit.delegatee]), deposit.owner, _amount); + } + function _fetchOrDeploySurrogate(address _delegatee) internal returns (DelegationSurrogate _surrogate) diff --git a/test/UniStaker.t.sol b/test/UniStaker.t.sol index 7a6c7a2..b4eca50 100644 --- a/test/UniStaker.t.sol +++ b/test/UniStaker.t.sol @@ -49,6 +49,15 @@ contract UniStakerTest is Test { (uint256 _balance, address _owner, address _delegatee) = uniStaker.deposits(_depositId); return UniStaker.Deposit({balance: _balance, owner: _owner, delegatee: _delegatee}); } + + function _boundMintAndStake(address _depositor, uint256 _amount, address _delegatee) + internal + returns (uint256 _boundedAmount, UniStaker.DepositIdentifier _depositId) + { + _boundedAmount = _boundMintAmount(_amount); + _mintGovToken(_depositor, _boundedAmount); + _depositId = _stake(_depositor, _boundedAmount, _delegatee); + } } contract Constructor is UniStakerTest { @@ -112,7 +121,7 @@ contract Stake is UniStakerTest { assertEq(govToken.balanceOf(_depositor2), 0); } - function testFuzz_DeploysAndTransferTokenToTwoSurrogatesWhenAccountsStakesToDifferentDelegatees( + function testFuzz_DeploysAndTransfersTokenToTwoSurrogatesWhenAccountsStakesToDifferentDelegatees( address _depositor1, uint256 _amount1, address _depositor2, @@ -330,3 +339,152 @@ contract Stake is UniStakerTest { } } } + +contract Withdraw is UniStakerTest { + function testFuzz_AllowsDepositorToWithdrawFullStake( + address _depositor, + uint256 _amount, + address _delegatee + ) public { + UniStaker.DepositIdentifier _depositId; + (_amount, _depositId) = _boundMintAndStake(_depositor, _amount, _delegatee); + + vm.prank(_depositor); + uniStaker.withdraw(_depositId, _amount); + + UniStaker.Deposit memory _deposit = _fetchDeposit(_depositId); + address _surrogate = address(uniStaker.surrogates(_deposit.delegatee)); + + assertEq(govToken.balanceOf(_depositor), _amount); + assertEq(_deposit.balance, 0); + assertEq(govToken.balanceOf(_surrogate), 0); + } + + function testFuzz_AllowsDepositorToWithdrawPartialStake( + address _depositor, + uint256 _depositAmount, + address _delegatee, + uint256 _withdrawalAmount + ) public { + UniStaker.DepositIdentifier _depositId; + (_depositAmount, _depositId) = _boundMintAndStake(_depositor, _depositAmount, _delegatee); + _withdrawalAmount = bound(_withdrawalAmount, 0, _depositAmount); + + vm.prank(_depositor); + uniStaker.withdraw(_depositId, _withdrawalAmount); + + UniStaker.Deposit memory _deposit = _fetchDeposit(_depositId); + address _surrogate = address(uniStaker.surrogates(_deposit.delegatee)); + + assertEq(govToken.balanceOf(_depositor), _withdrawalAmount); + assertEq(_deposit.balance, _depositAmount - _withdrawalAmount); + assertEq(govToken.balanceOf(_surrogate), _depositAmount - _withdrawalAmount); + } + + function testFuzz_UpdatesTheTotalSupplyWhenAnAccountWithdraws( + address _depositor, + uint256 _depositAmount, + address _delegatee, + uint256 _withdrawalAmount + ) public { + UniStaker.DepositIdentifier _depositId; + (_depositAmount, _depositId) = _boundMintAndStake(_depositor, _depositAmount, _delegatee); + _withdrawalAmount = bound(_withdrawalAmount, 0, _depositAmount); + + vm.prank(_depositor); + uniStaker.withdraw(_depositId, _withdrawalAmount); + + assertEq(uniStaker.totalSupply(), _depositAmount - _withdrawalAmount); + } + + function testFuzz_UpdatesTheTotalSupplyWhenTwoAccountsWithdraw( + address _depositor1, + uint256 _depositAmount1, + address _delegatee1, + address _depositor2, + uint256 _depositAmount2, + address _delegatee2, + uint256 _withdrawalAmount1, + uint256 _withdrawalAmount2 + ) public { + // Make two separate deposits + UniStaker.DepositIdentifier _depositId1; + (_depositAmount1, _depositId1) = _boundMintAndStake(_depositor1, _depositAmount1, _delegatee1); + UniStaker.DepositIdentifier _depositId2; + (_depositAmount2, _depositId2) = _boundMintAndStake(_depositor2, _depositAmount2, _delegatee2); + + // Calculate withdrawal amounts + _withdrawalAmount1 = bound(_withdrawalAmount1, 0, _depositAmount1); + _withdrawalAmount2 = bound(_withdrawalAmount2, 0, _depositAmount2); + + // Execute both withdrawals + vm.prank(_depositor1); + uniStaker.withdraw(_depositId1, _withdrawalAmount1); + vm.prank(_depositor2); + uniStaker.withdraw(_depositId2, _withdrawalAmount2); + + uint256 _remainingDeposits = + _depositAmount1 + _depositAmount2 - _withdrawalAmount1 - _withdrawalAmount2; + assertEq(uniStaker.totalSupply(), _remainingDeposits); + } + + function testFuzz_UpdatesAnAccountsTotalDepositsWhenItWithdrawals( + address _depositor, + uint256 _depositAmount1, + uint256 _depositAmount2, + address _delegatee1, + address _delegatee2, + uint256 _withdrawalAmount + ) public { + // Make two separate deposits + UniStaker.DepositIdentifier _depositId1; + (_depositAmount1, _depositId1) = _boundMintAndStake(_depositor, _depositAmount1, _delegatee1); + UniStaker.DepositIdentifier _depositId2; + (_depositAmount2, _depositId2) = _boundMintAndStake(_depositor, _depositAmount2, _delegatee2); + + // Withdraw part of the first deposit + _withdrawalAmount = bound(_withdrawalAmount, 0, _depositAmount1); + vm.prank(_depositor); + uniStaker.withdraw(_depositId1, _withdrawalAmount); + + // Ensure the account's total balance + global balance accounting have been updated + assertEq( + uniStaker.totalDeposits(_depositor), _depositAmount1 + _depositAmount2 - _withdrawalAmount + ); + assertEq(uniStaker.totalSupply(), _depositAmount1 + _depositAmount2 - _withdrawalAmount); + } + + function testFuzz_RevertIf_TheWithdrawerIsNotTheDepositor( + address _depositor, + uint256 _amount, + address _delegatee, + address _notDepositor + ) public { + UniStaker.DepositIdentifier _depositId; + (_amount, _depositId) = _boundMintAndStake(_depositor, _amount, _delegatee); + vm.assume(_depositor != _notDepositor); + + vm.prank(_notDepositor); + vm.expectRevert( + abi.encodeWithSelector( + UniStaker.UniStaker__Unauthorized.selector, bytes32("not owner"), _notDepositor + ) + ); + uniStaker.withdraw(_depositId, _amount); + } + + function testFuzz_RevertIf_TheWithdrawalAmountIsGreaterThanTheBalance( + address _depositor, + uint256 _amount, + uint256 _amountOver, + address _delegatee + ) public { + UniStaker.DepositIdentifier _depositId; + (_amount, _depositId) = _boundMintAndStake(_depositor, _amount, _delegatee); + _amountOver = bound(_amountOver, 1, type(uint128).max); + + vm.prank(_depositor); + vm.expectRevert(); + uniStaker.withdraw(_depositId, _amount + _amountOver); + } +}