diff --git a/contracts/staking/Vester.sol b/contracts/staking/Vester.sol index 81ef8f2..487dc0b 100644 --- a/contracts/staking/Vester.sol +++ b/contracts/staking/Vester.sol @@ -15,7 +15,6 @@ contract Vester is IVester, IERC20, ReentrancyGuard, Governable { string public name; string public symbol; uint8 public decimals = 18; - // remove? uint256 public vestingDuration; address public esToken; address public pairToken; @@ -25,6 +24,7 @@ contract Vester is IVester, IERC20, ReentrancyGuard, Governable { uint256 public override totalSupply; uint256 public pairSupply; + bool public needCheckStake; mapping(address account => uint256 amount) public balances; mapping(address account => uint256 amount) public override pairAmounts; @@ -55,18 +55,17 @@ contract Vester is IVester, IERC20, ReentrancyGuard, Governable { address _esToken, address _pairToken, address _claimableToken, - address _rewardTracker + address _rewardTracker, + bool _needCheckStake ) { name = _name; symbol = _symbol; - vestingDuration = _vestingDuration; - esToken = _esToken; pairToken = _pairToken; claimableToken = _claimableToken; - rewardTracker = _rewardTracker; + needCheckStake = _needCheckStake; } function setHandler(address _handler, bool _isActive) external onlyGov { @@ -122,6 +121,10 @@ contract Vester is IVester, IERC20, ReentrancyGuard, Governable { emit Withdraw(account, claimedAmount, balance); } + function setRewardTracker(address _rewardTracker) external onlyGov { + rewardTracker = _rewardTracker; + } + function setCumulativeRewardDeductions(address _account, uint256 _amount) external override nonReentrant { _validateHandler(); cumulativeRewardDeductions[_account] = _amount; @@ -267,7 +270,6 @@ contract Vester is IVester, IERC20, ReentrancyGuard, Governable { */ function _deposit(address _account, uint256 _amount) private { require(_amount > 0, "Vester: invalid _amount"); - _updateVesting(_account); IERC20(esToken).safeTransferFrom(_account, address(this), _amount); @@ -283,7 +285,11 @@ contract Vester is IVester, IERC20, ReentrancyGuard, Governable { _mintPair(_account, pairAmountDiff); } } - + if (needCheckStake && hasRewardTracker()) { + // if u want to transfer 100 esCec to cec, u need to have 100 cec in stake + uint256 cecAmount = IRewardTracker(rewardTracker).depositBalances(_account, claimableToken); + require(balances[_account] <= cecAmount, "Vester: insufficient cec balance"); + } uint256 maxAmount = getMaxVestableAmount(_account); require(getTotalVested(_account) <= maxAmount, "Vester: max vestable amount exceeded"); diff --git a/test/testStaking.ts b/test/testStaking.ts index 5e26e29..8766598 100644 --- a/test/testStaking.ts +++ b/test/testStaking.ts @@ -59,6 +59,7 @@ describe('RewardRouter', function() { ZeroAddress, cec.target, ZeroAddress, + false ); await stakedCecTracker.setInPrivateTransferMode(true) diff --git a/test/testVester.ts b/test/testVester.ts index 8d488c3..210eefc 100644 --- a/test/testVester.ts +++ b/test/testVester.ts @@ -29,6 +29,7 @@ describe("Vester", function () { ZeroAddress, cec.target, ZeroAddress, + false ); await cec.mint(owner.address, expandDecimals(100000, 18)); await esCec.mint(owner.address, expandDecimals(100000, 18));