Track the Lido withdrawal request amounts (#53)

* Track the Lido withdrawal request amounts * Fixed fork tests with new cross price * Fix config of gas reports * Format code * fix: quick fix for invariant failing. (#54) --------- Co-authored-by: Clément <[email protected]>
OriginProtocol · Dec 16, 2024 · 48fd1ff · 48fd1ff
1 parent 0148a94
commit 48fd1ff
Show file tree

Hide file tree

Showing 3 changed files with 28 additions and 9 deletions.
diff --git a/src/contracts/LidoARM.sol b/src/contracts/LidoARM.sol
@@ -26,6 +26,9 @@ contract LidoARM is Initializable, AbstractARM {
     /// @notice The amount of stETH in the Lido Withdrawal Queue
     uint256 public lidoWithdrawalQueueAmount;
 
+    /// @notice stores the requested amount for each Lido withdrawal
+    mapping(uint256 id => uint256 amount) public lidoWithdrawalRequests;
+
     event RequestLidoWithdrawals(uint256[] amounts, uint256[] requestIds);
     event ClaimLidoWithdrawals(uint256[] requestIds);
 
@@ -83,6 +86,9 @@ contract LidoARM is Initializable, AbstractARM {
         uint256 totalAmountRequested = 0;
         for (uint256 i = 0; i < amounts.length; i++) {
             totalAmountRequested += amounts[i];
+
+            // Store the amount of each withdrawal request
+            lidoWithdrawalRequests[requestIds[i]] = amounts[i];
         }
 
         // Increase the Ether outstanding from the Lido Withdrawal Queue
@@ -96,20 +102,31 @@ contract LidoARM is Initializable, AbstractARM {
      * Before calling this method, caller should check on the request NFTs to ensure the withdrawal was processed.
      */
     function claimLidoWithdrawals(uint256[] calldata requestIds) external {
-        uint256 etherBefore = address(this).balance;
-
         // Claim the NFTs for ETH.
         uint256 lastIndex = lidoWithdrawalQueue.getLastCheckpointIndex();
         uint256[] memory hintIds = lidoWithdrawalQueue.findCheckpointHints(requestIds, 1, lastIndex);
         lidoWithdrawalQueue.claimWithdrawals(requestIds, hintIds);
 
-        uint256 etherAfter = address(this).balance;
+        // Reduce the amount outstanding from the Lido Withdrawal Queue.
+        // The amount of ETH claimed from the Lido Withdrawal Queue can be less than the requested amount
+        // in the event of a mass slashing event of Lido validators.
+        uint256 totalAmountRequested = 0;
+        for (uint256 i = 0; i < requestIds.length; i++) {
+            // Read the requested amount from storage
+            uint256 requestAmount = lidoWithdrawalRequests[requestIds[i]];
+
+            // Validate the request came from this Lido ARM contract and not
+            // transferred in from another account.
+            require(requestAmount > 0, "LidoARM: invalid request");
+
+            totalAmountRequested += requestAmount;
+        }
 
-        // Reduce the Ether outstanding from the Lido Withdrawal Queue
-        lidoWithdrawalQueueAmount -= etherAfter - etherBefore;
+        // Store the reduced amount outstanding from the Lido Withdrawal Queue
+        lidoWithdrawalQueueAmount -= totalAmountRequested;
 
         // Wrap all the received ETH to WETH.
-        weth.deposit{value: etherAfter}();
+        weth.deposit{value: address(this).balance}();
 
         emit ClaimLidoWithdrawals(requestIds);
     }

diff --git a/test/invariants/handlers/LLMHandler.sol b/test/invariants/handlers/LLMHandler.sol
@@ -52,11 +52,11 @@ contract LLMHandler is BaseHandler {
         uint256 totalAmount = _bound(_seed, 0, min(MAX_AMOUNT * 3, steth.balanceOf(address(arm))));
 
         // We can only request only 1k amount at a time
-        uint256 batch = (totalAmount / MAX_AMOUNT) + 1;
+        uint256 batch = ((totalAmount + MAX_AMOUNT - 1) / MAX_AMOUNT);
         uint256[] memory amounts = new uint256[](batch);
         uint256 totalAmount_ = totalAmount;
         for (uint256 i = 0; i < batch; i++) {
-            if (totalAmount_ >= MAX_AMOUNT) {
+            if (totalAmount_ > MAX_AMOUNT) {
                 amounts[i] = MAX_AMOUNT;
                 totalAmount_ -= MAX_AMOUNT;
             } else {

diff --git a/test/invariants/mocks/MockSTETH.sol b/test/invariants/mocks/MockSTETH.sol
@@ -11,6 +11,8 @@ contract MockSTETH is ERC20 {
     //////////////////////////////////////////////////////
     Vm private constant vm = Vm(address(uint160(uint256(keccak256("hevm cheat code")))));
 
+    bool public enableBrutalizeAmount = false;
+
     //////////////////////////////////////////////////////
     /// --- VARIABLES
     //////////////////////////////////////////////////////
@@ -34,7 +36,7 @@ contract MockSTETH is ERC20 {
 
     function brutalizeAmount(uint256 amount) public returns (uint256) {
         // Only brutalize the sender doesn't sent all of their balance
-        if (balanceOf[msg.sender] != amount && amount > 0) {
+        if (balanceOf[msg.sender] != amount && amount > 0 && enableBrutalizeAmount) {
             // Get a random number between 0 and 1
             uint256 randomUint = vm.randomUint(0, 1);
             // If the amount is greater than the random number, subtract the random number from the amount