refactor: new upgrade flow

contracts-abi/abi/VanillaRegistry.abi

@@ -172,7 +172,7 @@
       {
         "name": "",
         "type": "tuple",
-        "internalType": "struct IVanillaRegistry.StakedValidator",
+        "internalType": "struct IVanillaRegistryV2.StakedValidator",
         "components": [
           {
             "name": "exists",

contracts-abi/script.sh

@@ -34,7 +34,7 @@ extract_and_save_abi "$BASE_DIR/out/SettlementGateway.sol/SettlementGateway.json
 # Extract ABI for L1Gateway.json
 extract_and_save_abi "$BASE_DIR/out/L1Gateway.sol/L1Gateway.json" "$ABI_DIR/L1Gateway.abi"
 
-extract_and_save_abi "$BASE_DIR/out/VanillaRegistry.sol/VanillaRegistry.json" "$ABI_DIR/VanillaRegistry.abi"
+extract_and_save_abi "$BASE_DIR/out/VanillaRegistryV2.sol/VanillaRegistryV2.json" "$ABI_DIR/VanillaRegistry.abi"
 
 extract_and_save_abi "$BASE_DIR/out/BlockTracker.sol/BlockTracker.json" "$ABI_DIR/BlockTracker.abi"
 

contracts/.solhintignore

@@ -1,3 +1,4 @@
 lib/
 test/
 node_modules/
+contracts/upgrades/

contracts/contracts/interfaces/IValidatorOptInRouter.sol

@@ -1,7 +1,7 @@
 // SPDX-License-Identifier: BSL 1.1
 pragma solidity 0.8.26;
 
-import {IVanillaRegistry} from "./IVanillaRegistry.sol";
+import {IVanillaRegistry} from "../upgrades/IVanillaRegistry.sol";
 import {IMevCommitAVS} from "./IMevCommitAVS.sol";
 import {IMevCommitMiddleware} from "./IMevCommitMiddleware.sol";
 

contracts/contracts/interfaces/IVanillaRegistryV2.sol

@@ -0,0 +1,162 @@
+// SPDX-License-Identifier: BSL 1.1
+pragma solidity 0.8.26;
+
+import { BlockHeightOccurrence } from "../utils/Occurrence.sol";
+
+/// @title IVanillaRegistryV2
+/// @notice Interface for the VanillaRegistry contract for validators.
+interface IVanillaRegistryV2 {
+
+    /// @dev Struct representing a validator staked with the registry.
+    struct StakedValidator {
+        bool exists;
+        address withdrawalAddress;
+        uint256 balance;
+        BlockHeightOccurrence.Occurrence unstakeOccurrence;
+    }
+
+    /// @dev Event emitted when a validator is staked.
+    event Staked(address indexed msgSender, address indexed withdrawalAddress, bytes valBLSPubKey, uint256 amount);
+
+    /// @dev Event emitted when ETH is added to the staked balance a validator. 
+    event StakeAdded(address indexed msgSender, address indexed withdrawalAddress, bytes valBLSPubKey, uint256 amount, uint256 newBalance);
+
+    /// @dev Event emitted when a validator is unstaked.
+    event Unstaked(address indexed msgSender, address indexed withdrawalAddress, bytes valBLSPubKey, uint256 amount);
+
+    /// @dev Event emitted when a validator's stake is withdrawn.
+    event StakeWithdrawn(address indexed msgSender, address indexed withdrawalAddress, bytes valBLSPubKey, uint256 amount);
+
+    /// @dev Event emitted when total stake is withdrawn.
+    event TotalStakeWithdrawn(address indexed msgSender, address indexed withdrawalAddress, uint256 totalAmount);
+
+    /// @dev Event emitted when a validator is slashed.
+    event Slashed(address indexed msgSender, address indexed slashReceiver, address indexed withdrawalAddress, bytes valBLSPubKey, uint256 amount);
+
+    /// @dev Event emitted when the min stake parameter is set.
+    event MinStakeSet(address indexed msgSender, uint256 newMinStake);
+
+    /// @dev Event emitted when the slash oracle parameter is set.
+    event SlashOracleSet(address indexed msgSender, address newSlashOracle);
+
+    /// @dev Event emitted when the slash receiver parameter is set.
+    event SlashReceiverSet(address indexed msgSender, address newSlashReceiver);
+
+    /// @dev Event emitted when the unstake period blocks parameter is set.
+    event UnstakePeriodBlocksSet(address indexed msgSender, uint256 newUnstakePeriodBlocks);
+
+    /// @dev Event emitted when the slashing payout period blocks parameter is set.
+    event SlashingPayoutPeriodBlocksSet(address indexed msgSender, uint256 newSlashingPayoutPeriodBlocks);
+
+    /// @dev Event emitted when a staker is whitelisted.
+    event StakerWhitelisted(address indexed msgSender, address staker);
+
+    /// @dev Event emitted when a staker is removed from the whitelist.
+    event StakerRemovedFromWhitelist(address indexed msgSender, address staker);
+
+    error ValidatorRecordMustExist(bytes valBLSPubKey);
+    error ValidatorRecordMustNotExist(bytes valBLSPubKey);
+    error ValidatorCannotBeUnstaking(bytes valBLSPubKey);
+    error SenderIsNotWithdrawalAddress(address sender, address withdrawalAddress);
+    error InvalidBLSPubKeyLength(uint256 expected, uint256 actual);
+    error SenderIsNotSlashOracle(address sender, address slashOracle);
+    error WithdrawalAddressMustBeSet();
+    error MustUnstakeToWithdraw();
+    error AtLeastOneRecipientRequired();
+    error StakeTooLowForNumberOfKeys(uint256 msgValue, uint256 required);
+    error WithdrawingTooSoon();
+    error WithdrawalAddressMismatch(address actualWithdrawalAddress, address expectedWithdrawalAddress);
+    error WithdrawalFailed();
+    error NoFundsToWithdraw();
+    error SlashingTransferFailed();
+    error SlashAmountMustBePositive();
+    error SlashAmountMustBeLessThanMinStake();
+    error SlashOracleMustBeSet();
+    error SlashReceiverMustBeSet();
+    error UnstakePeriodMustBePositive();
+    error SlashingPayoutPeriodMustBePositive();
+    error SenderIsNotWhitelistedStaker(address sender);
+    error StakerAlreadyWhitelisted(address staker);
+    error StakerNotWhitelisted(address staker);
+
+    /// @dev Initializes the contract with the provided parameters.
+    function initialize(
+        uint256 _minStake, 
+        address _slashOracle,
+        address _slashReceiver,
+        uint256 _unstakePeriodBlocks, 
+        uint256 _slashingPayoutPeriodBlocks,
+        address _owner
+    ) external;
+
+    /* 
+     * @dev Stakes ETH on behalf of one or multiple validators via their BLS pubkey.
+     * @param blsPubKeys The validator BLS public keys to stake.
+     */
+    function stake(bytes[] calldata blsPubKeys) external payable;
+
+    /* 
+     * @dev Stakes ETH on behalf of one or multiple validators via their BLS pubkey,
+     * and specifies an address other than msg.sender to be the withdrawal address.
+     * @param blsPubKeys The validator BLS public keys to stake.
+     * @param withdrawalAddress The address to receive the staked ETH.
+     */
+    function delegateStake(bytes[] calldata blsPubKeys, address withdrawalAddress) external payable;
+
+    /* 
+     * @dev Adds ETH to the staked balance of one or multiple validators via their BLS pubkey.
+     * @dev A staking entry must already exist for each provided BLS pubkey.
+     * @param blsPubKeys The BLS public keys to add stake to.
+     */
+    function addStake(bytes[] calldata blsPubKeys) external payable;
+
+    /* 
+     * @dev Unstakes ETH on behalf of one or multiple validators via their BLS pubkey.
+     * @param blsPubKeys The BLS public keys to unstake.
+     */
+    function unstake(bytes[] calldata blsPubKeys) external;
+
+    /* 
+     * @dev Withdraws ETH on behalf of one or multiple validators via their BLS pubkey.
+     * @param blsPubKeys The BLS public keys to withdraw.
+     */
+    function withdraw(bytes[] calldata blsPubKeys) external;
+
+    /// @dev Allows oracle to slash some portion of stake for one or multiple validators via their BLS pubkey.
+    /// @param blsPubKeys The BLS public keys to slash.
+    /// @param payoutIfDue Whether to payout slashed funds to receiver if the payout period is due.
+    function slash(bytes[] calldata blsPubKeys, bool payoutIfDue) external;
+
+    /// @dev Enables the owner to pause the contract.
+    function pause() external;
+
+    /// @dev Enables the owner to unpause the contract.
+    function unpause() external;
+
+    /// @dev Enables the owner to set the minimum stake parameter.
+    function setMinStake(uint256 newMinStake) external;
+
+    /// @dev Enables the owner to set the slash oracle parameter.
+    function setSlashOracle(address newSlashOracle) external;
+
+    /// @dev Enables the owner to set the slash receiver parameter.
+    function setSlashReceiver(address newSlashReceiver) external;
+
+    /// @dev Enables the owner to set the unstake period parameter.
+    function setUnstakePeriodBlocks(uint256 newUnstakePeriodBlocks) external;
+
+    /// @dev Returns true if a validator is considered "opted-in" to mev-commit via this registry.
+    function isValidatorOptedIn(bytes calldata valBLSPubKey) external view returns (bool);
+
+    /// @dev Returns stored staked validator struct for a given BLS pubkey.
+    function getStakedValidator(bytes calldata valBLSPubKey) external view returns (StakedValidator memory);
+
+    /// @dev Returns the staked amount for a given BLS pubkey.
+    function getStakedAmount(bytes calldata valBLSPubKey) external view returns (uint256);
+
+    /// @dev Returns true if a validator is currently unstaking.
+    function isUnstaking(bytes calldata valBLSPubKey) external view returns (bool);
+
+    /// @dev Returns the number of blocks remaining until an unstaking validator can withdraw their staked ETH.
+    function getBlocksTillWithdrawAllowed(bytes calldata valBLSPubKey) external view returns (uint256);
+}

contracts/contracts/interfaces/IVanillaRegistry.sol → contracts/contracts/upgrades/IVanillaRegistry.sol

@@ -48,12 +48,6 @@ interface IVanillaRegistry {
     /// @dev Event emitted when the slashing payout period blocks parameter is set.
     event SlashingPayoutPeriodBlocksSet(address indexed msgSender, uint256 newSlashingPayoutPeriodBlocks);
 
-    /// @dev Event emitted when a staker is whitelisted.
-    event StakerWhitelisted(address indexed msgSender, address staker);
-
-    /// @dev Event emitted when a staker is removed from the whitelist.
-    event StakerRemovedFromWhitelist(address indexed msgSender, address staker);
-
     error ValidatorRecordMustExist(bytes valBLSPubKey);
     error ValidatorRecordMustNotExist(bytes valBLSPubKey);
     error ValidatorCannotBeUnstaking(bytes valBLSPubKey);
@@ -69,15 +63,13 @@ interface IVanillaRegistry {
     error WithdrawalFailed();
     error NoFundsToWithdraw();
     error SlashingTransferFailed();
+    error MinStakeMustBePositive();
     error SlashAmountMustBePositive();
     error SlashAmountMustBeLessThanMinStake();
     error SlashOracleMustBeSet();
     error SlashReceiverMustBeSet();
     error UnstakePeriodMustBePositive();
     error SlashingPayoutPeriodMustBePositive();
-    error SenderIsNotWhitelistedStaker(address sender);
-    error StakerAlreadyWhitelisted(address staker);
-    error StakerNotWhitelisted(address staker);
 
     /// @dev Initializes the contract with the provided parameters.
     function initialize(
@@ -159,4 +151,4 @@ interface IVanillaRegistry {
 
     /// @dev Returns the number of blocks remaining until an unstaking validator can withdraw their staked ETH.
     function getBlocksTillWithdrawAllowed(bytes calldata valBLSPubKey) external view returns (uint256);
-}
+}

contracts/contracts/validator-registry/VanillaRegistry.sol → contracts/contracts/upgrades/VanillaRegistry.sol

@@ -1,7 +1,7 @@
 // SPDX-License-Identifier: BSL 1.1
 pragma solidity 0.8.26;
 
-import {IVanillaRegistry} from "../interfaces/IVanillaRegistry.sol";
+import {IVanillaRegistry} from "./IVanillaRegistry.sol";
 import {VanillaRegistryStorage} from "./VanillaRegistryStorage.sol";
 import {BlockHeightOccurrence} from "../utils/Occurrence.sol";
 import {Ownable2StepUpgradeable} from "@openzeppelin/contracts-upgradeable/access/Ownable2StepUpgradeable.sol";
@@ -31,28 +31,12 @@ contract VanillaRegistry is IVanillaRegistry, VanillaRegistryStorage,
         _;
     }
 
-    /// @dev Modifier to confirm the sender is whitelisted.
-    modifier onlyWhitelistedStaker() {
-        require(whitelistedStakers[msg.sender], IVanillaRegistry.SenderIsNotWhitelistedStaker(msg.sender));
-        _;
-    }
-
     /// @dev See https://docs.openzeppelin.com/upgrades-plugins/1.x/writing-upgradeable#initializing_the_implementation_contract
     /// @custom:oz-upgrades-unsafe-allow constructor
     constructor() {
         _disableInitializers();
     }
 
-    /// @dev Receive function is disabled for this contract to prevent unintended interactions.
-    receive() external payable {
-        revert Errors.InvalidReceive();
-    }
-
-    /// @dev Fallback function to revert all calls, ensuring no unintended interactions.
-    fallback() external payable {
-        revert Errors.InvalidFallback();
-    }
-
     /// @dev Initializes the contract with the provided parameters.
     function initialize(
         uint256 _minStake, 
@@ -62,20 +46,29 @@ contract VanillaRegistry is IVanillaRegistry, VanillaRegistryStorage,
         uint256 _slashingPayoutPeriodBlocks,
         address _owner
     ) external initializer {
-        __Pausable_init();
         _setMinStake(_minStake);
         _setSlashOracle(_slashOracle);
         _setUnstakePeriodBlocks(_unstakePeriodBlocks);
         FeePayout.init(slashingFundsTracker, _slashReceiver, _slashingPayoutPeriodBlocks);
         __Ownable_init(_owner);
     }
 
+    /// @dev Receive function is disabled for this contract to prevent unintended interactions.
+    receive() external payable {
+        revert Errors.InvalidReceive();
+    }
+
+    /// @dev Fallback function to revert all calls, ensuring no unintended interactions.
+    fallback() external payable {
+        revert Errors.InvalidFallback();
+    }
+
     /* 
      * @dev Stakes ETH on behalf of one or multiple validators via their BLS pubkey.
      * @param blsPubKeys The validator BLS public keys to stake.
      */
     function stake(bytes[] calldata blsPubKeys) external payable
-        onlyValidBLSPubKeys(blsPubKeys) onlyWhitelistedStaker() whenNotPaused() {
+        onlyValidBLSPubKeys(blsPubKeys) whenNotPaused() {
         _stake(blsPubKeys, msg.sender);
     }
 
@@ -96,7 +89,7 @@ contract VanillaRegistry is IVanillaRegistry, VanillaRegistryStorage,
      * @dev A staking entry must already exist for each provided BLS pubkey.
      * @param blsPubKeys The BLS public keys to add stake to.
      */
-    function addStake(bytes[] calldata blsPubKeys) external payable onlyWhitelistedStaker() whenNotPaused() {
+    function addStake(bytes[] calldata blsPubKeys) external payable whenNotPaused() {
         _addStake(blsPubKeys);
     }
 
@@ -186,26 +179,6 @@ contract VanillaRegistry is IVanillaRegistry, VanillaRegistryStorage,
         FeePayout.transferToRecipient(slashingFundsTracker);
     }
 
-    /// @dev Enables the owner to whitelist stakers.
-    function whitelistStakers(address[] calldata stakers) external onlyOwner {
-        uint256 len = stakers.length;
-        for (uint256 i = 0; i < len; ++i) {
-            require(!whitelistedStakers[stakers[i]], IVanillaRegistry.StakerAlreadyWhitelisted(stakers[i]));
-            whitelistedStakers[stakers[i]] = true;
-            emit StakerWhitelisted(msg.sender, stakers[i]);
-        }
-    }
-
-    /// @dev Enables the owner to remove stakers from the whitelist.
-    function removeWhitelistedStakers(address[] calldata stakers) external onlyOwner {
-        uint256 len = stakers.length;
-        for (uint256 i = 0; i < len; ++i) {
-            require(whitelistedStakers[stakers[i]], IVanillaRegistry.StakerNotWhitelisted(stakers[i]));
-            whitelistedStakers[stakers[i]] = false;
-            emit StakerRemovedFromWhitelist(msg.sender, stakers[i]);
-        }
-    }
-
     /// @dev Returns true if a validator is considered "opted-in" to mev-commit via this registry.
     function isValidatorOptedIn(bytes calldata valBLSPubKey) external view returns (bool) {
         return _isValidatorOptedIn(valBLSPubKey);
@@ -392,6 +365,7 @@ contract VanillaRegistry is IVanillaRegistry, VanillaRegistryStorage,
 
     /// @dev Internal function to set the minimum stake parameter.
     function _setMinStake(uint256 newMinStake) internal {
+        require(newMinStake != 0, IVanillaRegistry.MinStakeMustBePositive());
         minStake = newMinStake;
         emit MinStakeSet(msg.sender, newMinStake);
     }
@@ -426,13 +400,11 @@ contract VanillaRegistry is IVanillaRegistry, VanillaRegistryStorage,
 
     /// @dev Internal function to check if a validator is considered "opted-in" to mev-commit via this registry.
     function _isValidatorOptedIn(bytes calldata valBLSPubKey) internal view returns (bool) {
-        return !_isUnstaking(valBLSPubKey) && 
-        stakedValidators[valBLSPubKey].balance >= minStake &&
-        whitelistedStakers[stakedValidators[valBLSPubKey].withdrawalAddress];
+        return !_isUnstaking(valBLSPubKey) && stakedValidators[valBLSPubKey].balance >= minStake;
     }
 
     /// @dev Internal function to check if a validator is currently unstaking.
     function _isUnstaking(bytes calldata valBLSPubKey) internal view returns (bool) {
         return stakedValidators[valBLSPubKey].unstakeOccurrence.exists;
     }
-}
+}

contracts/contracts/validator-registry/VanillaRegistryStorage.sol → contracts/contracts/upgrades/VanillaRegistryStorage.sol

@@ -1,7 +1,7 @@
 // SPDX-License-Identifier: BSL 1.1
 pragma solidity 0.8.26;
 
-import {IVanillaRegistry} from "../interfaces/IVanillaRegistry.sol";
+import {IVanillaRegistry} from "./IVanillaRegistry.sol";
 import {FeePayout} from "../utils/FeePayout.sol";
 
 /// @title VanillaRegistryStorage
@@ -26,9 +26,6 @@ contract VanillaRegistryStorage {
     /// @dev Mapping of withdrawal addresses to claimable ETH that was force withdrawn by the owner.
     mapping(address withdrawalAddress => uint256 amountToClaim) public forceWithdrawnFunds;
 
-    /// @dev Mapping of staker addresses to whether they are whitelisted.
-    mapping(address staker => bool whitelisted) public whitelistedStakers;
-
     /// @dev See https://docs.openzeppelin.com/upgrades-plugins/1.x/writing-upgradeable#storage-gaps
-    uint256[47] private __gap;
-}
+    uint256[48] private __gap;
+}

contracts/contracts/validator-registry/ValidatorOptInRouter.sol

@@ -3,7 +3,7 @@ pragma solidity 0.8.26;
 
 import {ValidatorOptInRouterStorage} from "./ValidatorOptInRouterStorage.sol";
 import {IValidatorOptInRouter} from "../interfaces/IValidatorOptInRouter.sol";
-import {IVanillaRegistry} from "../interfaces/IVanillaRegistry.sol";
+import {IVanillaRegistry} from "../upgrades/IVanillaRegistry.sol";
 import {IMevCommitAVS} from "../interfaces/IMevCommitAVS.sol";
 import {Ownable2StepUpgradeable} from "@openzeppelin/contracts-upgradeable/access/Ownable2StepUpgradeable.sol";
 import {UUPSUpgradeable} from "@openzeppelin/contracts-upgradeable/proxy/utils/UUPSUpgradeable.sol";