add reward handling param + use latest redemption rate only

src/StakingNode.sol

@@ -171,7 +171,7 @@ contract StakingNode is IStakingNode, StakingNodeEvents, ReentrancyGuardUpgradea
     }
 
     //--------------------------------------------------------------------------------------
-    //----------------------------------  EXPEDITED WITHDRAWAL   ---------------------------
+    //----------------------------------  SURPLUS WITHDRAWAL   ---------------------------
     //--------------------------------------------------------------------------------------
 
     /**

src/StakingNodesManager.sol

@@ -27,7 +27,7 @@ interface StakingNodesManagerEvents {
     event RegisteredStakingNodeImplementationContract(address upgradeableBeaconAddress, address implementationContract);
     event UpgradedStakingNodeImplementationContract(address implementationContract, uint256 nodesCount);
     event NodeInitialized(address nodeAddress, uint64 initializedVersion);
-    event PrincipalWithdrawalProcessed(uint256 nodeId, uint256 amountToReinvest, uint256 amountToQueue);
+    event PrincipalWithdrawalProcessed(uint256 nodeId, uint256 amountToReinvest, uint256 amountToQueue, uint256 rewardsAmount);
     event ETHReceived(address sender, uint256 amount);
 }
 
@@ -55,6 +55,8 @@ contract StakingNodesManager is
     error NoValidatorsProvided();
     error ValidatorRegistrationPaused();
     error InvalidRewardsType(RewardsType rewardsType);
+    error ValidatorUnused(bytes publicKey);
+    error ValidatorNotWithdrawn(bytes publicKey, IEigenPod.VALIDATOR_STATUS status);
 
     //--------------------------------------------------------------------------------------
     //----------------------------------  ROLES  -------------------------------------------
@@ -507,8 +509,10 @@ contract StakingNodesManager is
         if (address(nodes[nodeId]) != msg.sender) {
             revert NotStakingNode(msg.sender, nodeId);
         }
+        _processRewards(nodeId, rewardsType, msg.value);
+    }
 
-        uint256 rewards = msg.value;
+    function _processRewards(uint256 nodeId, RewardsType rewardsType, uint256 rewards) internal {
         IRewardsReceiver receiver;
 
         if (rewardsType == RewardsType.ConsensusLayer) {
@@ -524,7 +528,7 @@ contract StakingNodesManager is
             revert TransferFailed();
         }
 
-        emit WithdrawnETHRewardsProcessed(nodeId, rewardsType, msg.value);
+        emit WithdrawnETHRewardsProcessed(nodeId, rewardsType, rewards);
     }
 
     /**
@@ -548,15 +552,32 @@ contract StakingNodesManager is
         uint256 amountToReinvest = action.amountToReinvest;
         uint256 amountToQueue = action.amountToQueue;
 
-        // Calculate the total amount to be processed by summing reinvestment and queuing amounts
-        uint256 totalAmount = amountToReinvest + amountToQueue;
+        // The rewardsAmount is trusted off-chain input provided in the WithdrawalAction struct.
+        // It represents the portion of the withdrawn amount that is considered as rewards.
+        // This value is determined off-chain by analyzing the difference between
+        // the initial stake and the total withdrawn amount.
+        //
+        // This design trade-off is a result of how Eigenlayer M3 pepe no long providees
+        // clear separation between principal and rewards amount and they both exit through the 
+        // Queued Withdrawals mechanism.
+        // 
+        // SECURITY NOTE:
+        // The accuracy and integrity of this value relies on the off-chain process
+        // that calculates it. There's an implicit trust that the WITHDRAWAL_MANAGER_ROLE
+        // will provide correct and verified data and that principal is not counted as Rewards
+        // and applied a fee.
+        uint256 rewardsAmount = action.rewardsAmount;
+
+        // Calculate the total amount to be processed by summing reinvestment, rewards and queuing amounts
+        uint256 totalAmount = amountToReinvest + amountToQueue + rewardsAmount;
 
         // Retrieve the staking node object using the nodeId
         IStakingNode node = nodes[nodeId];
 
         // Deallocate the specified total amount of ETH from the staking node
         node.deallocateStakedETH(totalAmount);
 
+
         // If there is an amount specified to reinvest, process it through ynETH
         if (amountToReinvest > 0) {
             ynETH.processWithdrawnETH{value: amountToReinvest}();
@@ -569,11 +590,29 @@ contract StakingNodesManager is
                 revert TransferFailed();
             }
         }
+
+        // If there is an amount of rewards specified, handle that
+        if (rewardsAmount > 0) {
+
+            // IMPORTANT: Impact on totalAssets()
+            // After charging the rewards fee, the totalAssets() of the system may decrease.
+            // Steps:
+            // 1. The full rewardsAmount is removed from the staking node's balance (which is part of totalAssets).
+            // 2. Only the remainingRewards (after fees) are reinvested back to the system.
+            // 3. The fees are sent to a separate fee receiver and are no longer part of the system's totalAssets.
+
+            (bool sent, ) = address(rewardsDistributor.consensusLayerReceiver()).call{value: rewardsAmount}("");
+            if (!sent) {
+                revert TransferFailed();
+            }
+            // process rewards immediately to avoid large totalAssets() fluctuations
+            rewardsDistributor.processRewards();
+        }
+
         // Emit an event to log the processed principal withdrawal details
-        emit PrincipalWithdrawalProcessed(nodeId, amountToReinvest, amountToQueue);
+        emit PrincipalWithdrawalProcessed(nodeId, amountToReinvest, amountToQueue, rewardsAmount);
     }
 
-
     //--------------------------------------------------------------------------------------
     //----------------------------------  VIEWS  -------------------------------------------
     //--------------------------------------------------------------------------------------

src/WithdrawalQueueManager.sol

@@ -153,14 +153,26 @@ contract WithdrawalQueueManager is IWithdrawalQueueManager, ERC721EnumerableUpgr
     //----------------------------------  WITHDRAWAL REQUESTS  -----------------------------
     //--------------------------------------------------------------------------------------
 
+
+    /**
+     * @notice Requests a withdrawal of a specified amount of redeemable assets without additional data.
+     * @dev This is a convenience function that calls the main requestWithdrawal function with empty data.
+     * @param amount The amount of redeemable assets to withdraw.
+     * @return tokenId The token ID associated with the withdrawal request.
+     */
+    function requestWithdrawal(uint256 amount) external returns (uint256 tokenId) {
+        return requestWithdrawal(amount, "");
+    }
+
     /**
      * @notice Requests a withdrawal of a specified amount of redeemable assets.
      * @dev Transfers the specified amount of redeemable assets from the sender to this contract, creates a withdrawal request,
      *      and mints a token representing this request. Emits a WithdrawalRequested event upon success.
      * @param amount The amount of redeemable assets to withdraw.
+     * @param data Extra data payload associated with the request
      * @return tokenId The token ID associated with the withdrawal request.
      */
-    function requestWithdrawal(uint256 amount) external nonReentrant returns (uint256 tokenId) {
+    function requestWithdrawal(uint256 amount, bytes memory data) public nonReentrant returns (uint256 tokenId) {
         if (amount == 0) {
             revert AmountMustBeGreaterThanZero();
         }
@@ -174,7 +186,8 @@ contract WithdrawalQueueManager is IWithdrawalQueueManager, ERC721EnumerableUpgr
             feeAtRequestTime: withdrawalFee,
             redemptionRateAtRequestTime: currentRate,
             creationTimestamp: block.timestamp,
-            processed: false
+            processed: false,
+            data: data
         });
 
         pendingRequestedRedemptionAmount += calculateRedemptionAmount(amount, currentRate);
@@ -214,7 +227,14 @@ contract WithdrawalQueueManager is IWithdrawalQueueManager, ERC721EnumerableUpgr
         }
 
         withdrawalRequests[tokenId].processed = true;
-        uint256 unitOfAccountAmount = calculateRedemptionAmount(request.amount, request.redemptionRateAtRequestTime);
+
+        // Redemption rate at claim time is the minimum between
+        // the redemption rate at request time and the current redemption Rate
+        uint256 currentRate = redemptionAssetsVault.redemptionRate();
+        uint256 redemptionRate = request.redemptionRateAtRequestTime < currentRate ? request.redemptionRateAtRequestTime : currentRate;
+
+        uint256 unitOfAccountAmount = calculateRedemptionAmount(request.amount, redemptionRate);
+
         pendingRequestedRedemptionAmount -= unitOfAccountAmount;
 
         _burn(tokenId);
@@ -228,10 +248,10 @@ contract WithdrawalQueueManager is IWithdrawalQueueManager, ERC721EnumerableUpgr
         }
 
         // Transfer net amount (unitOfAccountAmount - feeAmount) to the receiver
-        redemptionAssetsVault.transferRedemptionAssets(receiver, unitOfAccountAmount - feeAmount);
+        redemptionAssetsVault.transferRedemptionAssets(receiver, unitOfAccountAmount - feeAmount, request.data);
 
         if (feeAmount > 0) {
-            redemptionAssetsVault.transferRedemptionAssets(feeReceiver, feeAmount);
+            redemptionAssetsVault.transferRedemptionAssets(feeReceiver, feeAmount, request.data);
         }
 
         emit WithdrawalClaimed(tokenId, msg.sender, receiver, request);
@@ -286,14 +306,14 @@ contract WithdrawalQueueManager is IWithdrawalQueueManager, ERC721EnumerableUpgr
     /**
      * @notice Calculates the redemption amount based on the provided amount and the redemption rate at the time of request.
      * @param amount The amount of the redeemable asset.
-     * @param redemptionRateAtRequestTime The redemption rate at the time the request was made, expressed in the same unit of decimals as the redeemable asset.
+     * @param redemptionRate The redemption rate expressed in the same unit of decimals as the redeemable asset.
      * @return The calculated redemption amount, adjusted for the decimal places of the redeemable asset.
      */
     function calculateRedemptionAmount(
         uint256 amount,
-        uint256 redemptionRateAtRequestTime
+        uint256 redemptionRate
     ) public view returns (uint256) {
-        return amount * redemptionRateAtRequestTime / (10 ** redeemableAsset.decimals());
+        return amount * redemptionRate / (10 ** redeemableAsset.decimals());
     }
 
     /**

src/interfaces/IRedemptionAssetsVault.sol

@@ -14,7 +14,8 @@ interface IRedemptionAssetsVault {
     /// @dev This is only for INTERNAL USE
     /// @param to The address to which the assets will be transferred.
     /// @param amount The amount in unit of account
-    function transferRedemptionAssets(address to, uint256 amount) external;
+    /// @param data Extra data payload for redemption request
+    function transferRedemptionAssets(address to, uint256 amount, bytes calldata data) external;
 
     /// @notice Withdraws redemption assets from the queue's balance
     /// @param amount The amount in unit of account

src/interfaces/IStakingNodesManager.sol

@@ -28,6 +28,7 @@ interface IStakingNodesManager {
         uint256 nodeId;
         uint256 amountToReinvest;
         uint256 amountToQueue;
+        uint256 rewardsAmount;
     }
 
     function eigenPodManager() external view returns (IEigenPodManager);

src/interfaces/IWithdrawalQueueManager.sol

@@ -10,9 +10,11 @@ interface IWithdrawalQueueManager {
         uint256 redemptionRateAtRequestTime;
         uint256 creationTimestamp;
         bool processed;
+        bytes data;
     }
 
     function requestWithdrawal(uint256 amount) external returns (uint256);
+    function requestWithdrawal(uint256 amount, bytes calldata data) external returns (uint256);
     function claimWithdrawal(uint256 tokenId, address receiver) external;
     function finalizeRequestsUpToIndex(uint256 _lastFinalizedIndex) external;
 }

src/ynETHRedemptionAssetsVault.sol

@@ -92,7 +92,7 @@ contract ynETHRedemptionAssetsVault is IRedemptionAssetsVault, Initializable, Ac
      * @param amount The amount of assets to transfer.
      * @dev Requires the caller to be the redeemer and the contract to not be paused.
      */
-    function transferRedemptionAssets(address to, uint256 amount) public onlyRedeemer whenNotPaused nonReentrant {
+    function transferRedemptionAssets(address to, uint256 amount, bytes calldata /* data */) public onlyRedeemer whenNotPaused nonReentrant {
         uint256 balance = availableRedemptionAssets();
         if (balance < amount) {
             revert InsufficientAssetBalance(ETH_ASSET, amount, balance);

test/integration/M3/Base.t.sol

@@ -27,12 +27,17 @@ import {RewardsDistributor} from "../../../src/RewardsDistributor.sol";
 import {StakingNode} from "../../../src/StakingNode.sol";
 import {WithdrawalQueueManager} from "../../../src/WithdrawalQueueManager.sol";
 import {ynETHRedemptionAssetsVault} from "../../../src/ynETHRedemptionAssetsVault.sol";
+import {IStakingNode} from "../../../src/interfaces/IStakingNodesManager.sol";
+
 
 import "forge-std/console.sol";
 import "forge-std/Test.sol";
 
 contract Base is Test, Utils {
 
+    bytes public constant ZERO_SIGNATURE = hex"000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000";
+    bytes constant ZERO_PUBLIC_KEY = hex"000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000"; 
+
     // Utils
     ContractAddresses public contractAddresses;
     ContractAddresses.ChainAddresses public chainAddresses;
@@ -214,4 +219,68 @@ contract Base is Test, Utils {
             vm.stopPrank();
         }
     }
+
+    function createValidators(uint256[] memory nodeIds, uint256 count) public returns (uint40[] memory) {
+        uint40[] memory validatorIndices = new uint40[](count * nodeIds.length);
+        uint256 index = 0;
+
+        for (uint256 j = 0; j < nodeIds.length; j++) {
+            bytes memory withdrawalCredentials = stakingNodesManager.getWithdrawalCredentials(nodeIds[j]);
+
+            for (uint256 i = 0; i < count; i++) {
+                validatorIndices[index] = beaconChain.newValidator{value: 32 ether}(withdrawalCredentials);
+                index++;
+            }
+        }
+        return validatorIndices;
+    }
+
+    function registerValidators(uint256[] memory validatorNodeIds) public {
+        IStakingNodesManager.ValidatorData[] memory validatorData = new IStakingNodesManager.ValidatorData[](validatorNodeIds.length);
+
+        for (uint256 i = 0; i < validatorNodeIds.length; i++) {
+            bytes memory publicKey = abi.encodePacked(uint256(i));
+            publicKey = bytes.concat(publicKey, new bytes(ZERO_PUBLIC_KEY.length - publicKey.length));
+            validatorData[i] = IStakingNodesManager.ValidatorData({
+                publicKey: publicKey,
+                signature: ZERO_SIGNATURE,
+                nodeId: validatorNodeIds[i],
+                depositDataRoot: bytes32(0)
+            });
+        }
+
+        for (uint256 i = 0; i < validatorData.length; i++) {
+            uint256 amount = 32 ether;
+            bytes memory withdrawalCredentials = stakingNodesManager.getWithdrawalCredentials(validatorData[i].nodeId);
+            bytes32 depositDataRoot = stakingNodesManager.generateDepositRoot(validatorData[i].publicKey, validatorData[i].signature, withdrawalCredentials, amount);
+            validatorData[i].depositDataRoot = depositDataRoot;
+        }
+
+        vm.prank(actors.ops.VALIDATOR_MANAGER);
+        stakingNodesManager.registerValidators(validatorData);
+    }
+
+    function runSystemStateInvariants(
+        uint256 previousTotalAssets,
+        uint256 previousTotalSupply,
+        uint256[] memory previousStakingNodeBalances
+    ) public {  
+        assertEq(yneth.totalAssets(), previousTotalAssets, "Total assets integrity check failed");
+        assertEq(yneth.totalSupply(), previousTotalSupply, "Share mint integrity check failed");
+        for (uint i = 0; i < previousStakingNodeBalances.length; i++) {
+            IStakingNode stakingNodeInstance = stakingNodesManager.nodes(i);
+            uint256 currentStakingNodeBalance = stakingNodeInstance.getETHBalance();
+            assertEq(currentStakingNodeBalance, previousStakingNodeBalances[i], "Staking node balance integrity check failed for node ID: ");
+        }
+	}
+
+    function getAllStakingNodeBalances() public view returns (uint256[] memory) {
+        uint256[] memory balances = new uint256[](stakingNodesManager.nodesLength());
+        for (uint256 i = 0; i < stakingNodesManager.nodesLength(); i++) {
+            IStakingNode stakingNode = stakingNodesManager.nodes(i);
+            balances[i] = stakingNode.getETHBalance();
+        }
+        return balances;
+    }
+
 }

test/integration/M3/Withdrawals.t.sol

@@ -27,7 +27,6 @@ contract M3WithdrawalsTest is Base {
     uint256 public nodeId;
 
     uint256 public constant AMOUNT = 32 ether;
-    bytes public constant ZERO_SIGNATURE = hex"000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000";
 
     //
     // setup
@@ -230,7 +229,8 @@ contract M3WithdrawalsTest is Base {
             _actions[0] = IStakingNodesManager.WithdrawalAction({
                 nodeId: nodeId,
                 amountToReinvest: AMOUNT / 2, // 16 ETH
-                amountToQueue: AMOUNT / 2 // 16 ETH
+                amountToQueue: AMOUNT / 2, // 16 ETH
+                rewardsAmount: 0
             });
             vm.prank(actors.ops.WITHDRAWAL_MANAGER);
             stakingNodesManager.processPrincipalWithdrawals({