From 4404500a6192e4d3c29aaba5996c456cb3cd5a75 Mon Sep 17 00:00:00 2001 From: Ben DiFrancesco Date: Mon, 11 Dec 2023 15:55:30 -0500 Subject: [PATCH] Implement internal accounting for stake accounting --- src/UniStaker.sol | 31 +++++++- test/UniStaker.t.sol | 185 ++++++++++++++++++++++++++++++++++++++++++- 2 files changed, 210 insertions(+), 6 deletions(-) diff --git a/src/UniStaker.sol b/src/UniStaker.sol index 5c8ebfd..8f33bf0 100644 --- a/src/UniStaker.sol +++ b/src/UniStaker.sol @@ -8,9 +8,25 @@ import {SafeERC20} from "openzeppelin/token/ERC20/utils/SafeERC20.sol"; import {ReentrancyGuard} from "openzeppelin/utils/ReentrancyGuard.sol"; contract UniStaker is ReentrancyGuard { + type DepositIdentifier is uint256; + + struct Deposit { + uint256 balance; + address owner; + address delegatee; + } + IERC20 public immutable REWARDS_TOKEN; IERC20Delegates public immutable STAKE_TOKEN; + DepositIdentifier private nextDepositId; + + uint256 public totalSupply; + + mapping(address depositor => uint256 amount) public totalDeposits; + + mapping(DepositIdentifier depositId => Deposit deposit) public deposits; + mapping(address delegatee => DelegationSurrogate surrogate) public surrogates; constructor(IERC20 _rewardsToken, IERC20Delegates _stakeToken) { @@ -19,13 +35,17 @@ contract UniStaker is ReentrancyGuard { } function stake(uint256 _amount, address _delegatee) - public + external nonReentrant - returns (uint256 _depositId) + returns (DepositIdentifier _depositId) { DelegationSurrogate _surrogate = _fetchOrDeploySurrogate(_delegatee); _stakeTokenSafeTransferFrom(msg.sender, address(_surrogate), _amount); - _depositId = 1; + _depositId = _useDepositId(); + + totalSupply += _amount; + totalDeposits[msg.sender] += _amount; + deposits[_depositId] = Deposit({balance: _amount, owner: msg.sender, delegatee: _delegatee}); } function _fetchOrDeploySurrogate(address _delegatee) @@ -43,4 +63,9 @@ contract UniStaker is ReentrancyGuard { function _stakeTokenSafeTransferFrom(address _from, address _to, uint256 _value) internal { SafeERC20.safeTransferFrom(IERC20(address(STAKE_TOKEN)), _from, _to, _value); } + + function _useDepositId() internal returns (DepositIdentifier _depositId) { + _depositId = nextDepositId; + nextDepositId = DepositIdentifier.wrap(DepositIdentifier.unwrap(_depositId) + 1); + } } diff --git a/test/UniStaker.t.sol b/test/UniStaker.t.sol index 213e337..b4ce1a6 100644 --- a/test/UniStaker.t.sol +++ b/test/UniStaker.t.sol @@ -33,13 +33,22 @@ contract UniStakerTest is Test { function _stake(address _depositor, uint256 _amount, address _delegatee) internal - returns (uint256 _depositId) + returns (UniStaker.DepositIdentifier _depositId) { vm.startPrank(_depositor); govToken.approve(address(uniStaker), _amount); _depositId = uniStaker.stake(_amount, _delegatee); vm.stopPrank(); } + + function _fetchDeposit(UniStaker.DepositIdentifier _depositId) + internal + view + returns (UniStaker.Deposit memory) + { + (uint256 _balance, address _owner, address _delegatee) = uniStaker.deposits(_depositId); + return UniStaker.Deposit({balance: _balance, owner: _owner, delegatee: _delegatee}); + } } contract Constructor is UniStakerTest { @@ -59,7 +68,7 @@ contract Constructor is UniStakerTest { } contract Stake is UniStakerTest { - function testFuzz_DeploysAndTransfersTokensToANewSurrogateWhenAUserStakes( + function testFuzz_DeploysAndTransfersTokensToANewSurrogateWhenAnAccountStakes( address _depositor, uint256 _amount, address _delegatee @@ -103,7 +112,7 @@ contract Stake is UniStakerTest { assertEq(govToken.balanceOf(_depositor2), 0); } - function testFuzz_DeploysAndTransferTokenToTwoSurrogatesWhenUsersStakesToDifferentDelegatees( + function testFuzz_DeploysAndTransferTokenToTwoSurrogatesWhenAccountsStakesToDifferentDelegatees( address _depositor1, uint256 _amount1, address _depositor2, @@ -136,4 +145,174 @@ contract Stake is UniStakerTest { assertEq(govToken.balanceOf(_depositor1), 0); assertEq(govToken.balanceOf(_depositor2), 0); } + + function testFuzz_UpdatesTheTotalSupplyWhenAnAccountStakes( + address _depositor, + uint256 _amount, + address _delegatee + ) public { + _amount = _boundMintAmount(_amount); + _mintGovToken(_depositor, _amount); + + _stake(_depositor, _amount, _delegatee); + + assertEq(uniStaker.totalSupply(), _amount); + } + + function testFuzz_UpdatesTheTotalSupplyWhenTwoAccountsStake( + address _depositor1, + uint256 _amount1, + address _depositor2, + uint256 _amount2, + address _delegatee1, + address _delegatee2 + ) public { + _amount1 = _boundMintAmount(_amount1); + _amount2 = _boundMintAmount(_amount2); + _mintGovToken(_depositor1, _amount1); + _mintGovToken(_depositor2, _amount2); + + _stake(_depositor1, _amount1, _delegatee1); + assertEq(uniStaker.totalSupply(), _amount1); + + _stake(_depositor2, _amount2, _delegatee2); + assertEq(uniStaker.totalSupply(), _amount1 + _amount2); + } + + function testFuzz_UpdatesAnAccountsTotalDepositsWhenItStakes( + address _depositor, + uint256 _amount1, + uint256 _amount2, + address _delegatee1, + address _delegatee2 + ) public { + _amount1 = _boundMintAmount(_amount1); + _amount2 = _boundMintAmount(_amount2); + _mintGovToken(_depositor, _amount1 + _amount2); + + // First stake + check total + _stake(_depositor, _amount1, _delegatee1); + assertEq(uniStaker.totalDeposits(_depositor), _amount1); + + // Second stake + check total + _stake(_depositor, _amount2, _delegatee2); + assertEq(uniStaker.totalDeposits(_depositor), _amount1 + _amount2); + } + + function testFuzz_UpdatesDifferentAccountsTotalDepositsIndependently( + address _depositor1, + uint256 _amount1, + address _depositor2, + uint256 _amount2, + address _delegatee1, + address _delegatee2 + ) public { + vm.assume(_depositor1 != _depositor2); + _amount1 = _boundMintAmount(_amount1); + _amount2 = _boundMintAmount(_amount2); + _mintGovToken(_depositor1, _amount1); + _mintGovToken(_depositor2, _amount2); + + _stake(_depositor1, _amount1, _delegatee1); + assertEq(uniStaker.totalDeposits(_depositor1), _amount1); + + _stake(_depositor2, _amount2, _delegatee2); + assertEq(uniStaker.totalDeposits(_depositor2), _amount2); + } + + function testFuzz_TracksTheBalanceForASpecificDeposit( + address _depositor, + uint256 _amount, + address _delegatee + ) public { + _amount = _boundMintAmount(_amount); + _mintGovToken(_depositor, _amount); + + UniStaker.DepositIdentifier _depositId = _stake(_depositor, _amount, _delegatee); + UniStaker.Deposit memory _deposit = _fetchDeposit(_depositId); + assertEq(_deposit.balance, _amount); + assertEq(_deposit.owner, _depositor); + assertEq(_deposit.delegatee, _delegatee); + } + + function testFuzz_TracksTheBalanceForDifferentDepositsFromTheSameAccountIndependently( + address _depositor, + uint256 _amount1, + uint256 _amount2, + address _delegatee1, + address _delegatee2 + ) public { + _amount1 = _boundMintAmount(_amount1); + _amount2 = _boundMintAmount(_amount2); + _mintGovToken(_depositor, _amount1 + _amount2); + + // Perform both deposits and track their identifiers separately + UniStaker.DepositIdentifier _depositId1 = _stake(_depositor, _amount1, _delegatee1); + UniStaker.DepositIdentifier _depositId2 = _stake(_depositor, _amount2, _delegatee2); + UniStaker.Deposit memory _deposit1 = _fetchDeposit(_depositId1); + UniStaker.Deposit memory _deposit2 = _fetchDeposit(_depositId2); + + // Check that the deposits have been recorded independently + assertEq(_deposit1.balance, _amount1); + assertEq(_deposit1.owner, _depositor); + assertEq(_deposit1.delegatee, _delegatee1); + assertEq(_deposit2.balance, _amount2); + assertEq(_deposit2.owner, _depositor); + assertEq(_deposit2.delegatee, _delegatee2); + } + + function testFuzz_TracksTheBalanceForDepositsFromDifferentAccountsIndependently( + address _depositor1, + address _depositor2, + uint256 _amount1, + uint256 _amount2, + address _delegatee1, + address _delegatee2 + ) public { + _amount1 = _boundMintAmount(_amount1); + _amount2 = _boundMintAmount(_amount2); + _mintGovToken(_depositor1, _amount1); + _mintGovToken(_depositor2, _amount2); + + // Perform both deposits and track their identifiers separately + UniStaker.DepositIdentifier _depositId1 = _stake(_depositor1, _amount1, _delegatee1); + UniStaker.DepositIdentifier _depositId2 = _stake(_depositor2, _amount2, _delegatee2); + UniStaker.Deposit memory _deposit1 = _fetchDeposit(_depositId1); + UniStaker.Deposit memory _deposit2 = _fetchDeposit(_depositId2); + + // Check that the deposits have been recorded independently + assertEq(_deposit1.balance, _amount1); + assertEq(_deposit1.owner, _depositor1); + assertEq(_deposit1.delegatee, _delegatee1); + assertEq(_deposit2.balance, _amount2); + assertEq(_deposit2.owner, _depositor2); + assertEq(_deposit2.delegatee, _delegatee2); + } + + mapping(UniStaker.DepositIdentifier depositId => bool isUsed) isIdUsed; + + function test_NeverReusesADepositIdentifier() public { + address _depositor = address(0xdeadbeef); + uint256 _amount = 116; + address _delegatee = address(0xaceface); + + UniStaker.DepositIdentifier _depositId; + + for (uint256 _i; _i < 10_000; _i++) { + // Perform the stake and save the deposit identifier + _amount = _bound(_amount, 0, 100_000_000_000e18); + _mintGovToken(_depositor, _amount); + _depositId = _stake(_depositor, _amount, _delegatee); + + // Ensure the identifier hasn't yet been used + assertFalse(isIdUsed[_depositId]); + // Record the fact this deposit Id has been used + isIdUsed[_depositId] = true; + + // Reset all the inputs for the next deposit by hashing the last inputs + _depositor = address(uint160(uint256(keccak256(abi.encode(_depositor))))); + _amount = uint256(keccak256(abi.encode(_amount))); + _delegatee = address(uint160(uint256(keccak256(abi.encode(_delegatee))))); + } + } }