From 97075a18d8823b18d417701e1386e204c7fd411f Mon Sep 17 00:00:00 2001 From: Sergey <2901744+evercoinx@users.noreply.github.com.> Date: Tue, 3 Sep 2024 15:36:40 +0200 Subject: [PATCH] Implement method to claim withdrawal POL --- contracts/MaticX.sol | 99 ++++++++++++++---------- contracts/interfaces/IMaticX.sol | 2 + contracts/interfaces/IValidatorShare.sol | 2 + contracts/mocks/ValidatorShareMock.sol | 26 ++++--- 4 files changed, 79 insertions(+), 50 deletions(-) diff --git a/contracts/MaticX.sol b/contracts/MaticX.sol index 49bb3a59..d08ad9b7 100644 --- a/contracts/MaticX.sol +++ b/contracts/MaticX.sol @@ -359,12 +359,63 @@ contract MaticX is } /** - * @dev Claims tokens from validator share and sends them to the + * @dev Claims MATIC tokens from a validator share and sends them to the * address if the request is in the userWithdrawalRequests * @param _idx - User withdrawal request array index */ function claimWithdrawal(uint256 _idx) external override whenNotPaused { - _claimWithdrawal(msg.sender, _idx); + _claimWithdrawal(msg.sender, _idx, false); + } + + /** + * @dev Claims POL tokens from a validator share and sends them to the + * address if the request is in the userWithdrawalRequests + * @param _idx - User withdrawal request array index + */ + function claimWithdrawalPOL(uint256 _idx) external override whenNotPaused { + _claimWithdrawal(msg.sender, _idx, true); + } + + /** + * @dev Claims tokens from validator share and sends them to the + * address if the request is in the userWithdrawalRequests + * @param _to - Address of the withdrawal request owner + * @param _idx - User withdrawal request array index + */ + function _claimWithdrawal(address _to, uint256 _idx, bool pol) + internal + returns (uint256) + { + address token = _getToken(pol); + uint256 balanceBeforeClaim = IERC20Upgradeable(token).balanceOf( + address(this) + ); + + WithdrawalRequest[] storage userRequests = userWithdrawalRequests[_to]; + WithdrawalRequest memory userRequest = userRequests[_idx]; + require( + IStakeManager(stakeManager).epoch() >= userRequest.requestEpoch, + "Not able to claim yet" + ); + + pol + ? IValidatorShare(userRequest.validatorAddress) + .unstakeClaimTokens_newPOL(userRequest.validatorNonce) + : IValidatorShare(userRequest.validatorAddress) + .unstakeClaimTokens_new(userRequest.validatorNonce); + + // swap with the last item and pop it. + userRequests[_idx] = userRequests[userRequests.length - 1]; + userRequests.pop(); + + uint256 amountToClaim = + IERC20Upgradeable(token).balanceOf(address(this)) - + balanceBeforeClaim; + + IERC20Upgradeable(token).safeTransfer(_to, amountToClaim); + + emit ClaimWithdrawal(_to, _idx, amountToClaim); + return amountToClaim; } /** @@ -554,45 +605,6 @@ contract MaticX is return amountToMint; } - /** - * @dev Claims tokens from validator share and sends them to the - * address if the request is in the userWithdrawalRequests - * @param _to - Address of the withdrawal request owner - * @param _idx - User withdrawal request array index - */ - function _claimWithdrawal(address _to, uint256 _idx) - internal - returns (uint256) - { - uint256 amountToClaim = 0; - uint256 balanceBeforeClaim = IERC20Upgradeable(maticToken).balanceOf( - address(this) - ); - WithdrawalRequest[] storage userRequests = userWithdrawalRequests[_to]; - WithdrawalRequest memory userRequest = userRequests[_idx]; - require( - IStakeManager(stakeManager).epoch() >= userRequest.requestEpoch, - "Not able to claim yet" - ); - - IValidatorShare(userRequest.validatorAddress).unstakeClaimTokens_new( - userRequest.validatorNonce - ); - - // swap with the last item and pop it. - userRequests[_idx] = userRequests[userRequests.length - 1]; - userRequests.pop(); - - amountToClaim = - IERC20Upgradeable(maticToken).balanceOf(address(this)) - - balanceBeforeClaim; - - IERC20Upgradeable(maticToken).safeTransfer(_to, amountToClaim); - - emit ClaimWithdrawal(_to, _idx, amountToClaim); - return amountToClaim; - } - /** * @dev Function that converts arbitrary maticX to Matic * @param _balance - Balance in maticX @@ -849,4 +861,9 @@ contract MaticX is _validatorRegistry = validatorRegistry; _polToken = polToken; } + + /// @notice Returns the POL or MATIC token depending on the used flow. + function _getToken(bool pol) private view returns (address) { + return pol ? polToken : maticToken; + } } diff --git a/contracts/interfaces/IMaticX.sol b/contracts/interfaces/IMaticX.sol index 9e5d0d8d..a67082e5 100644 --- a/contracts/interfaces/IMaticX.sol +++ b/contracts/interfaces/IMaticX.sol @@ -61,6 +61,8 @@ interface IMaticX is IERC20Upgradeable { function claimWithdrawal(uint256 _idx) external; + function claimWithdrawalPOL(uint256 _idx) external; + function withdrawRewards(uint256 _validatorId) external returns (uint256); function withdrawValidatorsReward(uint256[] calldata _validatorIds) diff --git a/contracts/interfaces/IValidatorShare.sol b/contracts/interfaces/IValidatorShare.sol index 66416d56..c9b33bb8 100644 --- a/contracts/interfaces/IValidatorShare.sol +++ b/contracts/interfaces/IValidatorShare.sol @@ -31,6 +31,8 @@ interface IValidatorShare { function unstakeClaimTokens_new(uint256 _unbondNonce) external; + function unstakeClaimTokens_newPOL(uint256 _unbondNonce) external; + function restake() external returns (uint256, uint256); function withdrawRewards() external; diff --git a/contracts/mocks/ValidatorShareMock.sol b/contracts/mocks/ValidatorShareMock.sol index 4f9d5927..30b516a2 100644 --- a/contracts/mocks/ValidatorShareMock.sol +++ b/contracts/mocks/ValidatorShareMock.sol @@ -64,16 +64,11 @@ contract ValidatorShareMock is IValidatorShare { } function unstakeClaimTokens_new(uint256 _unbondNonce) external override { - uint256 withdrawPoolShare = user2WithdrawPoolShare[msg.sender][ - _unbondNonce - ]; - uint256 amount2Transfer = (withdrawPoolShare * withdrawPool) / - totalWithdrawPoolShares; + _unstakeClaimTokens_new(_unbondNonce); + } - withdrawPool -= amount2Transfer; - totalShares -= withdrawPoolShare; - totalWithdrawPoolShares -= withdrawPoolShare; - IERC20(token).transfer(msg.sender, amount2Transfer); + function unstakeClaimTokens_newPOL(uint256 _unbondNonce) external override { + _unstakeClaimTokens_new(_unbondNonce); } function restake() external override returns (uint256, uint256) { @@ -159,4 +154,17 @@ contract ValidatorShareMock is IValidatorShare { unbondNonces[msg.sender] = unbondNonce; user2WithdrawPoolShare[msg.sender][unbondNonce] = _claimAmount; } + + function _unstakeClaimTokens_new(uint256 _unbondNonce) private { + uint256 withdrawPoolShare = user2WithdrawPoolShare[msg.sender][ + _unbondNonce + ]; + uint256 amount2Transfer = (withdrawPoolShare * withdrawPool) / + totalWithdrawPoolShares; + + withdrawPool -= amount2Transfer; + totalShares -= withdrawPoolShare; + totalWithdrawPoolShares -= withdrawPoolShare; + IERC20(token).transfer(msg.sender, amount2Transfer); + } }