Implement method to claim withdrawal POL

stader-labs · Sep 3, 2024 · 97075a1 · 97075a1
1 parent 0af909c
commit 97075a1
Show file tree

Hide file tree

Showing 4 changed files with 79 additions and 50 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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);
+	}
 }