feat: add getBlocksTillWithdrawAllowed query + tests

primev · Apr 25, 2024 · a23816c · a23816c
1 parent f211aa0
commit a23816c
Show file tree

Hide file tree

Showing 6 changed files with 138 additions and 4 deletions.
diff --git a/contracts-abi/abi/ValidatorRegistry.abi b/contracts-abi/abi/ValidatorRegistry.abi
@@ -4,6 +4,25 @@
     "inputs": [],
     "stateMutability": "nonpayable"
   },
+  {
+    "type": "function",
+    "name": "getBlocksTillWithdrawAllowed",
+    "inputs": [
+      {
+        "name": "valBLSPubKey",
+        "type": "bytes",
+        "internalType": "bytes"
+      }
+    ],
+    "outputs": [
+      {
+        "name": "",
+        "type": "uint256",
+        "internalType": "uint256"
+      }
+    ],
+    "stateMutability": "view"
+  },
   {
     "type": "function",
     "name": "getNumberOfStakedValidators",

diff --git a/contracts-abi/clients/ValidatorRegistry/ValidatorRegistry.go b/contracts-abi/clients/ValidatorRegistry/ValidatorRegistry.go
diff --git a/contracts/contracts/ValidatorRegistry.sol b/contracts/contracts/ValidatorRegistry.sol
@@ -145,6 +145,12 @@ contract ValidatorRegistry is OwnableUpgradeable, ReentrancyGuardUpgradeable {
         return unstakingBalances[valBLSPubKey];
     }
 
+    function getBlocksTillWithdrawAllowed(bytes calldata valBLSPubKey) external view returns (uint256) {
+        require(unstakeBlockNums[valBLSPubKey] > 0, "Unstake must be initiated to check withdrawal eligibility");
+        uint256 blocksSinceUnstakeInitiated = block.number - unstakeBlockNums[valBLSPubKey];
+        return blocksSinceUnstakeInitiated > unstakePeriodBlocks ? 0 : unstakePeriodBlocks - blocksSinceUnstakeInitiated;
+    }
+
     /// @return numStakedValidators uint number of currently staked validators. 
     /// @return stakedValsetVersion uint version of the staked valset at the time of query.
     function getNumberOfStakedValidators() external view returns (uint256, uint256) {

diff --git a/contracts/scripts/DeployScripts.s.sol b/contracts/scripts/DeployScripts.s.sol
@@ -114,10 +114,14 @@ contract DeployValidatorRegistry is Script {
     function run() external {
         vm.startBroadcast();
 
+        // 7000 blocks @ 200ms per block = 23.3 min. This allows two L1 epochs (finalization time) + settlement buffer, 
+        // to pass between validator unstake initiation and withdrawal.
+        uint256 unstakePeriodBlocks = 7000;
+
         // Can later be upgraded with https://docs.openzeppelin.com/upgrades-plugins/1.x/api-foundry-upgrades#Upgrades-upgradeProxy-address-string-bytes-
         address proxy = Upgrades.deployUUPSProxy(
             "ValidatorRegistry.sol",
-            abi.encodeCall(ValidatorRegistry.initialize, (3 ether, 64, msg.sender))
+            abi.encodeCall(ValidatorRegistry.initialize, (3 ether, unstakePeriodBlocks, msg.sender))
         );
         console.log("ValidatorRegistry UUPS proxy deployed to:", address(proxy));
 

diff --git a/contracts/scripts/ValidatorExampleScript.s.sol b/contracts/scripts/ValidatorExampleScript.s.sol
@@ -29,16 +29,33 @@ abstract contract ExampleScript is Script {
             console.log("--------------------");
         }
 
-        uint256 numStakedValidators = _validatorRegistry.getNumberOfStakedValidators();
+        (uint256 numStakedValidators, uint256 stakedValsetVersion) = _validatorRegistry.getNumberOfStakedValidators();
         console.log("Num Staked Validators:", numStakedValidators);
+        console.log("Staked Valset Version from len query:", stakedValsetVersion);
         if (numStakedValidators == 0) {
             return;
         }
-        bytes[] memory vals = _validatorRegistry.getStakedValidators(0, numStakedValidators);
+        bytes[] memory vals;
+        (vals, stakedValsetVersion) = _validatorRegistry.getStakedValidators(0, numStakedValidators);
         for (uint i = 0; i < vals.length; i++) {
             console.log("Staked validator from batch query: ");
             console.logBytes(vals[i]);
         }
+        console.log("Staked Valset Version from getStakedValidators query:", stakedValsetVersion);
+    }
+
+    function checkWithdrawal(bytes[] memory blsKeys) public view {
+        console.log("--------------------");
+        console.log("Checking Withdrawal related state...");
+        console.log("--------------------");
+        for (uint i = 0; i < blsKeys.length; i++) {
+            uint256 blocksTillWithdrawAllowed = _validatorRegistry.getBlocksTillWithdrawAllowed(blsKeys[i]);
+            console.log("--------------------");
+            console.log("BLS Key: ");
+            console.logBytes(blsKeys[i]);
+            console.log("Blocks till Withdraw Allowed:", blocksTillWithdrawAllowed);
+            console.log("--------------------");
+        }
     }
 }
 
@@ -87,6 +104,9 @@ contract UnstakeExample is ExampleScript {
         checkStaking(validators);
 
         console.log("Balance of 0xf39Fd6e51aad88F6F4ce6aB8827279cffFb92266:", defaultEOA.balance);
+
+        checkWithdrawal(validators);
+
         vm.stopBroadcast();
     }
 }
@@ -105,6 +125,8 @@ contract WithdrawExample is ExampleScript {
 
         checkStaking(validators);
 
+        checkWithdrawal(validators);
+
         _validatorRegistry.withdraw(validators);
         console.log("Withdraw initiated");