Fix claimFor

contracts/OperatorRewardsCollector.sol

@@ -5,6 +5,7 @@ import './library/UtilLib.sol';
 
 import './interfaces/IOperatorRewardsCollector.sol';
 import './interfaces/IStaderConfig.sol';
+import './interfaces/ISDUtilityPool.sol';
 
 import '@openzeppelin/contracts-upgradeable/access/AccessControlUpgradeable.sol';
 
@@ -13,9 +14,6 @@ contract OperatorRewardsCollector is IOperatorRewardsCollector, AccessControlUpg
 
     mapping(address => uint256) public balances;
 
-    mapping(address => uint256[]) public owedAmounts;
-    mapping(address => uint256[]) public claimableAmounts;
-
     /// @custom:oz-upgrades-unsafe-allow constructor
     constructor() {
         _disableInitializers();
@@ -40,36 +38,33 @@ contract OperatorRewardsCollector is IOperatorRewardsCollector, AccessControlUpg
     }
 
     function claim() external {
-        address operator = msg.sender;
-        uint256 amount = balances[operator];
-        balances[operator] -= amount;
-
-        address operatorRewardsAddr = UtilLib.getOperatorRewardAddress(msg.sender, staderConfig);
-        UtilLib.sendValue(operatorRewardsAddr, amount);
-        emit Claimed(operatorRewardsAddr, amount);
+        return claimFor(msg.sender);
     }
 
-    function claimFor(address operator) external {
-        uint256 toSendAmount;
-        for (uint256 i = 0; i < owedAmounts[operator].length; i++) {
-            if (balances[operator] >= owedAmounts[operator][i]) {
-                toSendAmount = owedAmounts[operator][i];
-                balances[operator] -= owedAmounts[operator][i];
-                claimableAmounts[operator][i] = owedAmounts[operator][i];
-                owedAmounts[operator][i] = 0;
-            } else {
-                toSendAmount = balances[operator];
-                owedAmounts[operator][i] -= balances[operator];
-                claimableAmounts[operator][i] = balances[operator];
-                balances[operator] = 0;
-                break;
-            }
+    function claimFor(address operator) public override {
+        // Retrieve operator liquidation details
+        ISDUtilityPool sdUtilityPool = ISDUtilityPool(staderConfig.getSDUtilityPool());
+        OperatorLiquidaton memory operatorLiquidation = sdUtilityPool.getOperatorLiquidation(operator);
+
+        // If the liquidation is not repaid, check balance and then proceed with repayment
+        if (!operatorLiquidation.isRepaid) {
+            // Ensure that the balance is sufficient
+            require(balances[operator] >= operatorLiquidation.amount, 'Insufficient balance');
+
+            // Repay the liquidation and update the operator's balance
+            sdUtilityPool.repayLiquidation(operator);
+            balances[operator] -= operatorLiquidation.amount;
         }
 
-        if (balances[operator] > 0) {
-            address operatorRewardsAddr = UtilLib.getOperatorRewardAddress(operator, staderConfig);
-            UtilLib.sendValue(operatorRewardsAddr, balances[operator]);
-            emit Claimed(operatorRewardsAddr, balances[operator]);
+        // Calculate payout amount
+        uint256 payoutAmount = balances[operator];
+        balances[operator] = 0;
+
+        // If there's an amount to send, transfer it to the operator's rewards address
+        if (payoutAmount > 0) {
+            address rewardsAddress = UtilLib.getOperatorRewardAddress(operator, staderConfig);
+            UtilLib.sendValue(rewardsAddress, payoutAmount);
+            emit Claimed(rewardsAddress, payoutAmount);
         }
     }
 

contracts/SDUtilityPool.sol

@@ -80,6 +80,7 @@ contract SDUtilityPool is ISDUtilityPool, AccessControlUpgradeable, PausableUpgr
     mapping(address => uint256) public override delegatorCTokenBalance;
     mapping(uint256 => DelegatorWithdrawInfo) public override delegatorWithdrawRequests;
     mapping(address => uint256[]) public override requestIdsByDelegatorAddress;
+    mapping(address => OperatorLiquidaton) private operatorLiquidation;
 
     /// @custom:oz-upgrades-unsafe-allow constructor
     constructor() {
@@ -375,6 +376,12 @@ contract SDUtilityPool is ISDUtilityPool, AccessControlUpgradeable, PausableUpgr
         return _utilizerBalanceStoredInternal(account);
     }
 
+    function repayLiquidation(address account) external override {
+        UtilLib.onlyStaderContract(msg.sender, staderConfig, staderConfig.OPERATOR_REWARD_COLLECTOR());
+
+        operatorLiquidation[account].isRepaid = true;
+    }
+
     /**
      * @notice Accrue fee then return the up-to-date exchange rate
      * @return Calculated exchange rate scaled by 1e18
@@ -768,4 +775,8 @@ contract SDUtilityPool is ISDUtilityPool, AccessControlUpgradeable, PausableUpgr
         uint256 totalEth = totalValidators * 2 ether;
         return totalEth;
     }
+
+    function getOperatorLiquidation(address account) external view override returns (OperatorLiquidaton memory) {
+        return operatorLiquidation[account];
+    }
 }

contracts/interfaces/ISDUtilityPool.sol

@@ -15,6 +15,11 @@ struct Config {
     uint256 ltv;
 }
 
+struct OperatorLiquidaton {
+    uint256 amount;
+    bool isRepaid;
+}
+
 interface ISDUtilityPool {
     error SDTransferFailed();
     error CannotFindRequestId();
@@ -93,6 +98,8 @@ interface ISDUtilityPool {
 
     function repayViaSDCollateral(address utilizer, uint256 repayAmount) external;
 
+    function repayLiquidation(address account) external;
+
     function withdrawProtocolFee(uint256 _amount) external;
 
     function accrueFee() external;
@@ -149,6 +156,8 @@ interface ISDUtilityPool {
 
     function utilizerData(address) external view returns (uint256 principal, uint256 utilizeIndex);
 
+    function getOperatorLiquidation(address) external view returns (OperatorLiquidaton memory);
+
     function delegatorWithdrawRequests(uint256)
         external
         view