From 07d8c1c20484ab03b41aaf77db69b0bf5cfc634b Mon Sep 17 00:00:00 2001 From: RuslanProgrammer Date: Fri, 12 Jan 2024 18:03:11 +0100 Subject: [PATCH] change implementation of StETH to almost real --- contracts/mock/tokens/StETHMock.sol | 91 ++++++++++++++++++++++++----- test/Distribution.test.ts | 15 +++-- 2 files changed, 85 insertions(+), 21 deletions(-) diff --git a/contracts/mock/tokens/StETHMock.sol b/contracts/mock/tokens/StETHMock.sol index c1092a7..cc4ae37 100644 --- a/contracts/mock/tokens/StETHMock.sol +++ b/contracts/mock/tokens/StETHMock.sol @@ -2,30 +2,95 @@ pragma solidity ^0.8.20; import {ERC20} from "@openzeppelin/contracts/token/ERC20/ERC20.sol"; -import {PRECISION} from "@solarity/solidity-lib/utils/Globals.sol"; import {Ownable} from "@openzeppelin/contracts/access/Ownable.sol"; contract StETHMock is ERC20, Ownable { - uint256 public balanceMultiplicator = PRECISION; + uint256 public totalShares; + uint256 public totalPooledEther; - constructor() ERC20("Staked Ether Mock", "stETHMock") {} + mapping(address => uint256) private shares; - function mint(address account_, uint256 amount_) external { - require(amount_ <= 1000 * (10 ** decimals()), "StETHMock: amount is too big"); + constructor() ERC20("Staked Ether Mock", "stETHMock") { + _mintShares(address(this), 10 ** decimals()); - _mint(account_, amount_); + totalPooledEther = 10 ** decimals(); } - function _transfer(address sender_, address recipient_, uint256 amount_) internal override { - amount_ = (amount_ * PRECISION) / balanceMultiplicator; - super._transfer(sender_, recipient_, amount_); + function mint(address _account, uint256 _amount) external { + require(_amount <= 1000 * (10 ** decimals()), "StETHMock: amount is too big"); + + uint256 sharesAmount = getSharesByPooledEth(_amount); + + _mintShares(_account, sharesAmount); + + totalPooledEther += _amount; + } + + function setTotalPooledEther(uint256 _totalPooledEther) external onlyOwner { + totalPooledEther = _totalPooledEther; + } + + function totalSupply() public view override returns (uint256) { + return totalPooledEther; + } + + function balanceOf(address _account) public view override returns (uint256) { + return getPooledEthByShares(_sharesOf(_account)); + } + + function sharesOf(address _account) external view returns (uint256) { + return _sharesOf(_account); + } + + function getSharesByPooledEth(uint256 _ethAmount) public view returns (uint256) { + return (_ethAmount * totalShares) / totalPooledEther; + } + + function getPooledEthByShares(uint256 _sharesAmount) public view returns (uint256) { + return (_sharesAmount * totalPooledEther) / totalShares; + } + + function transferShares(address _recipient, uint256 _sharesAmount) external returns (uint256) { + _transferShares(msg.sender, _recipient, _sharesAmount); + uint256 tokensAmount = getPooledEthByShares(_sharesAmount); + return tokensAmount; + } + + function transferSharesFrom(address _sender, address _recipient, uint256 _sharesAmount) external returns (uint256) { + uint256 tokensAmount = getPooledEthByShares(_sharesAmount); + _spendAllowance(_sender, msg.sender, tokensAmount); + _transferShares(_sender, _recipient, _sharesAmount); + return tokensAmount; + } + + function _transfer(address _sender, address _recipient, uint256 _amount) internal override { + uint256 _sharesToTransfer = getSharesByPooledEth(_amount); + _transferShares(_sender, _recipient, _sharesToTransfer); } - function setBalanceMultiplicator(uint256 balanceMultiplicator_) external onlyOwner { - balanceMultiplicator = balanceMultiplicator_; + function _sharesOf(address _account) internal view returns (uint256) { + return shares[_account]; } - function balanceOf(address account_) public view override returns (uint256) { - return (super.balanceOf(account_) * balanceMultiplicator) / PRECISION; + function _transferShares(address _sender, address _recipient, uint256 _sharesAmount) internal { + require(_sender != address(0), "TRANSFER_FROM_ZERO_ADDR"); + require(_recipient != address(0), "TRANSFER_TO_ZERO_ADDR"); + require(_recipient != address(this), "TRANSFER_TO_STETH_CONTRACT"); + + uint256 currentSenderShares = shares[_sender]; + require(_sharesAmount <= currentSenderShares, "BALANCE_EXCEEDED"); + + shares[_sender] = currentSenderShares - _sharesAmount; + shares[_recipient] += _sharesAmount; + } + + function _mintShares(address _recipient, uint256 _sharesAmount) internal returns (uint256 newTotalShares) { + require(_recipient != address(0), "MINT_TO_ZERO_ADDR"); + + totalShares += _sharesAmount; + + shares[_recipient] += _sharesAmount; + + return totalShares; } } diff --git a/test/Distribution.test.ts b/test/Distribution.test.ts index fcc0483..d063234 100644 --- a/test/Distribution.test.ts +++ b/test/Distribution.test.ts @@ -181,7 +181,6 @@ describe('Distribution', () => { depositToken.approve(distribution, wei(1000)), depositToken.connect(SECOND).approve(distribution, wei(1000)), ]); - await l1Sender.transferOwnership(distribution); await reverter.snapshot(); @@ -1312,7 +1311,7 @@ describe('Distribution', () => { expect(await depositToken.balanceOf(distribution)).to.eq(wei(20)); await setNextTime(oneDay + oneDay); - await depositToken.setBalanceMultiplicator(wei(0.8, 25)); + await depositToken.setTotalPooledEther(((await depositToken.totalPooledEther()) * 8n) / 10n); expect(await depositToken.balanceOf(distribution)).to.eq(wei(16)); let tx = await distribution.withdraw(poolId, wei(999)); @@ -1330,7 +1329,7 @@ describe('Distribution', () => { it('should revert if trying to withdraw zero', async () => { await distribution.stake(poolId, wei(10)); - await depositToken.setBalanceMultiplicator(wei(0.0001, 25)); + await depositToken.setTotalPooledEther(wei(0.0001, 25)); await distribution.withdraw(poolId, wei(10)); @@ -1582,7 +1581,7 @@ describe('Distribution', () => { await setTime(oneDay * 9999); - await depositToken.setBalanceMultiplicator(wei(0.5)); + await depositToken.setTotalPooledEther(wei(0.5)); const overplus = await distribution.overplus(); expect(overplus).to.eq(0); @@ -1590,7 +1589,7 @@ describe('Distribution', () => { it('should return overplus if deposited token increased', async () => { await distribution.stake(0, wei(1)); - await depositToken.setBalanceMultiplicator(wei(2, 25)); + await depositToken.setTotalPooledEther((await depositToken.totalPooledEther()) * 2n); let overplus = await distribution.overplus(); expect(overplus).to.eq(wei(1)); @@ -1600,12 +1599,12 @@ describe('Distribution', () => { overplus = await distribution.overplus(); expect(overplus).to.eq(wei(1)); - await depositToken.setBalanceMultiplicator(wei(1, 25)); + await depositToken.setTotalPooledEther((await depositToken.totalPooledEther()) / 2n); overplus = await distribution.overplus(); expect(overplus).to.eq(0); - await depositToken.setBalanceMultiplicator(wei(5, 25)); + await depositToken.setTotalPooledEther((await depositToken.totalPooledEther()) * 5n); overplus = await distribution.overplus(); expect(overplus).to.eq(wei(5.5)); @@ -1626,7 +1625,7 @@ describe('Distribution', () => { await distribution.stake(1, wei(1)); - await depositToken.setBalanceMultiplicator(wei(2, 25)); + await depositToken.setTotalPooledEther((await depositToken.totalPooledEther()) * 2n); const overplus = await distribution.overplus(); expect(overplus).to.eq(wei(1));