feat: restricted withdrawal functions to withdrawer or admin

kilnfi · Apr 14, 2023 · ed1b36b · ed1b36b
1 parent aad23d5
commit ed1b36b
Show file tree

Hide file tree

Showing 2 changed files with 372 additions and 24 deletions.
diff --git a/src/contracts/StakingContract.sol b/src/contracts/StakingContract.sol
@@ -33,6 +33,7 @@ contract StakingContract {
     error UnsortedIndexes();
     error InvalidPublicKeys();
     error InvalidSignatures();
+    error InvalidWithdrawer();
     error AlreadyInitialized();
     error InvalidDepositValue();
     error NotEnoughValidators();
@@ -599,6 +600,7 @@ contract StakingContract {
         }
         for (uint256 i = 0; i < _publicKeys.length; ) {
             bytes memory publicKey = BytesLib.slice(_publicKeys, i, PUBLIC_KEY_LENGTH);
+            _onlyWithdrawerOrAdmin(publicKey);
             _deployAndWithdraw(publicKey, EXECUTION_LAYER_SALT_PREFIX, StakingContractStorageLib.getELDispatcher());
             unchecked {
                 i += PUBLIC_KEY_LENGTH;
@@ -616,6 +618,7 @@ contract StakingContract {
         }
         for (uint256 i = 0; i < _publicKeys.length; ) {
             bytes memory publicKey = BytesLib.slice(_publicKeys, i, PUBLIC_KEY_LENGTH);
+            _onlyWithdrawerOrAdmin(publicKey);
             _deployAndWithdraw(publicKey, CONSENSUS_LAYER_SALT_PREFIX, StakingContractStorageLib.getCLDispatcher());
             unchecked {
                 i += PUBLIC_KEY_LENGTH;
@@ -625,14 +628,14 @@ contract StakingContract {
 
     /// @notice Withdraw both Consensus and Execution Layer Fees for given validators public keys
     /// @dev Funds are sent to the withdrawer account
-    /// @dev This method is public on purpose
     /// @param _publicKeys Validators to withdraw fees from
     function batchWithdraw(bytes calldata _publicKeys) external {
         if (_publicKeys.length % PUBLIC_KEY_LENGTH != 0) {
             revert InvalidPublicKeys();
         }
         for (uint256 i = 0; i < _publicKeys.length; ) {
             bytes memory publicKey = BytesLib.slice(_publicKeys, i, PUBLIC_KEY_LENGTH);
+            _onlyWithdrawerOrAdmin(publicKey);
             _deployAndWithdraw(publicKey, EXECUTION_LAYER_SALT_PREFIX, StakingContractStorageLib.getELDispatcher());
             _deployAndWithdraw(publicKey, CONSENSUS_LAYER_SALT_PREFIX, StakingContractStorageLib.getCLDispatcher());
             unchecked {
@@ -643,24 +646,25 @@ contract StakingContract {
 
     /// @notice Withdraw the Execution Layer Fee for a given validator public key
     /// @dev Funds are sent to the withdrawer account
-    /// @dev This method is public on purpose
     /// @param _publicKey Validator to withdraw Execution Layer Fees from
     function withdrawELFee(bytes calldata _publicKey) external {
+        _onlyWithdrawerOrAdmin(_publicKey);
         _deployAndWithdraw(_publicKey, EXECUTION_LAYER_SALT_PREFIX, StakingContractStorageLib.getELDispatcher());
     }
 
     /// @notice Withdraw the Consensus Layer Fee for a given validator public key
     /// @dev Funds are sent to the withdrawer account
-    /// @dev This method is public on purpose
     /// @param _publicKey Validator to withdraw Consensus Layer Fees from
     function withdrawCLFee(bytes calldata _publicKey) external {
+        _onlyWithdrawerOrAdmin(_publicKey);
         _deployAndWithdraw(_publicKey, CONSENSUS_LAYER_SALT_PREFIX, StakingContractStorageLib.getCLDispatcher());
     }
 
     /// @notice Withdraw both Consensus and Execution Layer Fee for a given validator public key
     /// @dev Reverts if any is null
     /// @param _publicKey Validator to withdraw Execution and Consensus Layer Fees from
     function withdraw(bytes calldata _publicKey) external {
+        _onlyWithdrawerOrAdmin(_publicKey);
         _deployAndWithdraw(_publicKey, EXECUTION_LAYER_SALT_PREFIX, StakingContractStorageLib.getELDispatcher());
         _deployAndWithdraw(_publicKey, CONSENSUS_LAYER_SALT_PREFIX, StakingContractStorageLib.getCLDispatcher());
     }
@@ -696,6 +700,15 @@ contract StakingContract {
     /// ██ ██  ██ ██    ██    ██      ██   ██ ██  ██ ██ ██   ██ ██
     /// ██ ██   ████    ██    ███████ ██   ██ ██   ████ ██   ██ ███████
 
+    function _onlyWithdrawerOrAdmin(bytes memory _publicKey) internal view {
+        if (
+            msg.sender != _getWithdrawer(_getPubKeyRoot(_publicKey)) &&
+            StakingContractStorageLib.getAdmin() != msg.sender
+        ) {
+            revert InvalidWithdrawer();
+        }
+    }
+
     function _onlyActiveOperator(uint256 _operatorIndex) internal view {
         StakingContractStorageLib.OperatorInfo storage operatorInfo = StakingContractStorageLib.getOperators().value[
             _operatorIndex
@@ -1148,12 +1161,10 @@ contract StakingContract {
         bytes32 feeRecipientSalt = sha256(abi.encodePacked(_prefix, publicKeyRoot));
         address implementation = StakingContractStorageLib.getFeeRecipientImplementation();
         address feeRecipientAddress = Clones.predictDeterministicAddress(implementation, feeRecipientSalt);
-        if (feeRecipientAddress.balance > 0) {
-            if (feeRecipientAddress.code.length == 0) {
-                Clones.cloneDeterministic(implementation, feeRecipientSalt);
-                IFeeRecipient(feeRecipientAddress).init(_dispatcher, publicKeyRoot);
-            }
-            IFeeRecipient(feeRecipientAddress).withdraw();
+        if (feeRecipientAddress.code.length == 0) {
+            Clones.cloneDeterministic(implementation, feeRecipientSalt);
+            IFeeRecipient(feeRecipientAddress).init(_dispatcher, publicKeyRoot);
         }
+        IFeeRecipient(feeRecipientAddress).withdraw();
     }
 }