feat: track active validator count in pods (#474)

Layr-Labs · Mar 12, 2024 · 9ac0ef8 · 9ac0ef8
1 parent e12b03f
commit 9ac0ef8
Show file tree

Hide file tree

Showing 3 changed files with 36 additions and 4 deletions.
diff --git a/src/contracts/pods/EigenPod.sol b/src/contracts/pods/EigenPod.sol
@@ -91,9 +91,12 @@ contract EigenPod is IEigenPod, Initializable, ReentrancyGuardUpgradeable, Eigen
     /// @notice This variable tracks any ETH deposited into this contract via the `receive` fallback function
     uint256 public nonBeaconChainETHBalanceWei;
 
-     /// @notice This variable tracks the total amount of partial withdrawals claimed via merkle proofs prior to a switch to ZK proofs for claiming partial withdrawals
+    /// @notice This variable tracks the total amount of partial withdrawals claimed via merkle proofs prior to a switch to ZK proofs for claiming partial withdrawals
     uint64 public sumOfPartialWithdrawalsClaimedGwei;
 
+    /// @notice Number of validators with proven withdrawal credentials, who do not have proven full withdrawals
+    uint256 activeValidatorCount;
+
     modifier onlyEigenPodManager() {
         require(msg.sender == address(eigenPodManager), "EigenPod.onlyEigenPodManager: not eigenPodManager");
         _;
@@ -479,6 +482,7 @@ contract EigenPod is IEigenPod, Initializable, ReentrancyGuardUpgradeable, Eigen
         });
 
         // Proofs complete - update this validator's status, record its proven balance, and save in state:
+        activeValidatorCount++;
         validatorInfo.status = VALIDATOR_STATUS.ACTIVE;
         validatorInfo.validatorIndex = validatorIndex;
         validatorInfo.mostRecentBalanceUpdateTimestamp = oracleTimestamp;
@@ -697,9 +701,12 @@ contract EigenPod is IEigenPod, Initializable, ReentrancyGuardUpgradeable, Eigen
          * Finally, the validator is fully withdrawn. Update their status and place in state:
          */
 
-        validatorInfo.restakedBalanceGwei = 0;
-        validatorInfo.status = VALIDATOR_STATUS.WITHDRAWN;
+        if (validatorInfo.status != VALIDATOR_STATUS.WITHDRAWN) {
+            activeValidatorCount--;
+            validatorInfo.status = VALIDATOR_STATUS.WITHDRAWN;
+        }
 
+        validatorInfo.restakedBalanceGwei = 0;        
         _validatorPubkeyHashToInfo[validatorPubkeyHash] = validatorInfo;
 
         emit FullWithdrawalRedeemed(validatorIndex, withdrawalTimestamp, recipient, withdrawalAmountGwei);

diff --git a/src/test/harnesses/EigenPodHarness.sol b/src/test/harnesses/EigenPodHarness.sol
@@ -20,6 +20,14 @@ contract EPInternalFunctions is EigenPod, Test {
         _GENESIS_TIME
     ) {}
 
+    function getActiveValidatorCount() public view returns (uint256) {
+        return activeValidatorCount;
+    }
+
+    function setActiveValidatorCount(uint _count) public {
+        activeValidatorCount = _count;
+    }
+
     function verifyWithdrawalCredentials(
         uint64 oracleTimestamp,
         bytes32 beaconStateRoot,

diff --git a/src/test/unit/EigenPodUnit.t.sol b/src/test/unit/EigenPodUnit.t.sol
@@ -453,6 +453,8 @@ contract EigenPodUnitTests_VerifyWithdrawalCredentialsTests is EigenPodHarnessSe
         uint64 effectiveBalanceGwei = validatorFields.getEffectiveBalanceGwei();
         assertGt(effectiveBalanceGwei, MAX_RESTAKED_BALANCE_GWEI_PER_VALIDATOR, "Proof file has an effective balance less than 32 ETH");
 
+        uint activeValidatorCountBefore = eigenPodHarness.getActiveValidatorCount();
+
         // Verify withdrawal credentials
         vm.expectEmit(true, true, true, true);
         emit ValidatorRestaked(validatorIndex);
@@ -467,6 +469,8 @@ contract EigenPodUnitTests_VerifyWithdrawalCredentialsTests is EigenPodHarnessSe
         );
 
         // Checks
+        uint activeValidatorCountAfter = eigenPodHarness.getActiveValidatorCount();
+        assertEq(activeValidatorCountAfter, activeValidatorCountBefore + 1, "active validator count should increase when proving withdrawal credentials");
         assertEq(restakedBalanceWei, uint256(MAX_RESTAKED_BALANCE_GWEI_PER_VALIDATOR) * uint256(1e9), "Returned restaked balance gwei should be max");
         _assertWithdrawalCredentialsSet(MAX_RESTAKED_BALANCE_GWEI_PER_VALIDATOR);
     }
@@ -480,6 +484,8 @@ contract EigenPodUnitTests_VerifyWithdrawalCredentialsTests is EigenPodHarnessSe
         uint64 effectiveBalanceGwei = validatorFields.getEffectiveBalanceGwei();
         assertLt(effectiveBalanceGwei, MAX_RESTAKED_BALANCE_GWEI_PER_VALIDATOR, "Proof file has an effective balance greater than 32 ETH");
 
+        uint activeValidatorCountBefore = eigenPodHarness.getActiveValidatorCount();
+
         // Verify withdrawal credentials
         vm.expectEmit(true, true, true, true);
         emit ValidatorRestaked(validatorIndex);
@@ -494,6 +500,8 @@ contract EigenPodUnitTests_VerifyWithdrawalCredentialsTests is EigenPodHarnessSe
         );
 
         // Checks
+        uint activeValidatorCountAfter = eigenPodHarness.getActiveValidatorCount();
+        assertEq(activeValidatorCountAfter, activeValidatorCountBefore + 1, "active validator count should increase when proving withdrawal credentials");
         assertEq(restakedBalanceWei, uint256(effectiveBalanceGwei) * uint256(1e9), "Returned restaked balance gwei incorrect");
         _assertWithdrawalCredentialsSet(effectiveBalanceGwei);
     }
@@ -920,10 +928,19 @@ contract EigenPodUnitTests_WithdrawalTests is EigenPodHarnessSetup, ProofParsing
             mostRecentBalanceUpdateTimestamp: 0,
             status: IEigenPod.VALIDATOR_STATUS.ACTIVE
         });
+
+        // Since we're withdrawing using an ACTIVE validator, ensure we have
+        // a validator count to decrement
+        uint activeValidatorCountBefore = 1 + eigenPodHarness.getActiveValidatorCount();
+        eigenPodHarness.setActiveValidatorCount(activeValidatorCountBefore);
 
-        // Process full withdrawal
+        // Process full withdrawal.
         IEigenPod.VerifiedWithdrawal memory vw = eigenPodHarness.processFullWithdrawal(0, pubkeyHash, 0, podOwner, withdrawalAmount, validatorInfo);
 
+        // Validate that our activeValidatorCount decreased
+        uint activeValidatorCountAfter = eigenPodHarness.getActiveValidatorCount();
+        assertEq(activeValidatorCountAfter, activeValidatorCountBefore - 1, "active validator count should decrease when withdrawing active validator");
+
         // Get expected amounts based on withdrawalAmount
         uint64 amountETHToQueue;
         uint64 amountETHToSend;