feat: stake cap

src/CuratedModule.sol

@@ -6,14 +6,15 @@ pragma solidity 0.8.33;
 import { ICuratedModule } from "./interfaces/ICuratedModule.sol";
 import { IMetaRegistry } from "./interfaces/IMetaRegistry.sol";
 import { IStakingModule, IStakingModuleV2 } from "./interfaces/IStakingModule.sol";
-import { NodeOperator } from "./interfaces/IBaseModule.sol";
+import { IBaseModule, NodeOperator } from "./interfaces/IBaseModule.sol";
 
 import { BaseModule } from "./abstract/BaseModule.sol";
 
 import { SigningKeys } from "./lib/SigningKeys.sol";
 import { CuratedDepositAllocator } from "./lib/allocator/CuratedDepositAllocator.sol";
 import { NodeOperatorOps } from "./lib/NodeOperatorOps.sol";
 import { StakeTracker } from "./lib/StakeTracker.sol";
+import { ValidatorBalanceLimits } from "./lib/ValidatorBalanceLimits.sol";
 
 contract CuratedModule is ICuratedModule, BaseModule {
     IMetaRegistry public immutable META_REGISTRY;
@@ -93,6 +94,17 @@ contract CuratedModule is ICuratedModule, BaseModule {
         _incrementModuleNonce();
     }
 
+    /// @inheritdoc IBaseModule
+    function reportValidatorBalance(
+        uint256 nodeOperatorId,
+        uint256 keyIndex,
+        uint256 currentBalanceWei
+    ) public override(BaseModule, IBaseModule) {
+        super.reportValidatorBalance(nodeOperatorId, keyIndex, currentBalanceWei);
+        // Balance reports may change the allocated stake and, consequently, the stake cap headroom.
+        _updateDepositableValidatorsCount({ nodeOperatorId: nodeOperatorId, incrementNonceIfUpdated: true });
+    }
+
     /// @inheritdoc IStakingModuleV2
     function allocateDeposits(
         uint256 maxDepositAmount,
@@ -219,8 +231,19 @@ contract CuratedModule is ICuratedModule, BaseModule {
         bool incrementNonceIfUpdated
     ) internal override returns (bool depositableChanged) {
         if (newCount > 0) {
-            uint256 weight = _metaRegistry().getNodeOperatorWeight(nodeOperatorId);
+            IMetaRegistry metaRegistry = _metaRegistry();
+            uint256 weight = metaRegistry.getNodeOperatorWeight(nodeOperatorId);
             if (weight == 0) newCount = 0;
+
+            if (newCount > 0) {
+                uint256 cap = metaRegistry.maximumStakeCapPerNodeOperator();
+                uint256 currentStake = StakeTracker.getOperatorBalance(_baseStorage(), nodeOperatorId);
+                uint256 capCapacity;
+                if (currentStake < cap) {
+                    capCapacity = (cap - currentStake) / ValidatorBalanceLimits.MIN_ACTIVATION_BALANCE;
+                }
+                if (newCount > capCapacity) newCount = capCapacity;
+            }
         }
 
         depositableChanged = super._applyDepositableValidatorsCount({
@@ -245,7 +268,19 @@ contract CuratedModule is ICuratedModule, BaseModule {
             topUpLimits: topUpLimits
         });
 
-        StakeTracker.increaseKeyBalances($, operatorIds, keyIndices, allocations);
+        uint256[] memory allocatedOperatorIds = StakeTracker.increaseKeyBalances(
+            $,
+            operatorIds,
+            keyIndices,
+            allocations
+        );
+
+        for (uint256 i; i < allocatedOperatorIds.length; ++i) {
+            _updateDepositableValidatorsCount({
+                nodeOperatorId: allocatedOperatorIds[i],
+                incrementNonceIfUpdated: false
+            });
+        }
     }
 
     function _validateTopUpPublicKeys(

src/MetaRegistry.sol

@@ -24,6 +24,8 @@ import { MAX_BP } from "./lib/Constants.sol";
 contract MetaRegistry is IMetaRegistry, Initializable, AccessControlEnumerableUpgradeable {
     using ExternalOperatorLib for ExternalOperator;
 
+    uint64 internal constant INITIALIZED_VERSION = 2;
+
     struct CachedOperatorGroup {
         string name;
         uint64[] subNodeOperatorIds;
@@ -54,11 +56,13 @@ contract MetaRegistry is IMetaRegistry, Initializable, AccessControlEnumerableUp
         mapping(uint256 providerId => WeightBoostProviderEntry entry) weightBoostProviders;
         mapping(address provider => uint256 providerId) weightBoostProviderIdByAddress;
         uint256 weightBoostProvidersCount;
+        uint256 maximumStakeCapPerNodeOperator;
     }
 
     bytes32 public constant MANAGE_OPERATOR_GROUPS_ROLE = keccak256("MANAGE_OPERATOR_GROUPS_ROLE");
     bytes32 public constant SET_OPERATOR_INFO_ROLE = keccak256("SET_OPERATOR_INFO_ROLE");
     bytes32 public constant SET_BOND_CURVE_WEIGHT_ROLE = keccak256("SET_BOND_CURVE_WEIGHT_ROLE");
+    bytes32 public constant MANAGE_STAKE_CAP_ROLE = keccak256("MANAGE_STAKE_CAP_ROLE");
 
     // ID of the stub node operator group that means "not in any group". This value is used for all node operators that are not assigned to any group, so it
     // can't be used as a real group ID.
@@ -91,17 +95,33 @@ contract MetaRegistry is IMetaRegistry, Initializable, AccessControlEnumerableUp
     }
 
     /// @inheritdoc IMetaRegistry
-    function initialize(address admin) external initializer {
+    function initialize(address admin, uint256 initialCap) external reinitializer(INITIALIZED_VERSION) {
         if (admin == address(0)) revert ZeroAdminAddress();
 
         _grantRole(DEFAULT_ADMIN_ROLE, admin);
+        _setMaximumStakeCapPerNodeOperator(initialCap);
+    }
+
+    /// @inheritdoc IMetaRegistry
+    function finalizeUpgradeV2(uint256 initialCap) external reinitializer(INITIALIZED_VERSION) {
+        _setMaximumStakeCapPerNodeOperator(initialCap);
     }
 
     /// @inheritdoc IMetaRegistry
     function getInitializedVersion() external view returns (uint64) {
         return _getInitializedVersion();
     }
 
+    /// @inheritdoc IMetaRegistry
+    function maximumStakeCapPerNodeOperator() external view returns (uint256 cap) {
+        return _storage().maximumStakeCapPerNodeOperator;
+    }
+
+    /// @inheritdoc IMetaRegistry
+    function setMaximumStakeCapPerNodeOperator(uint256 newCap) external onlyRole(MANAGE_STAKE_CAP_ROLE) {
+        _setMaximumStakeCapPerNodeOperator(newCap);
+    }
+
     /// @inheritdoc IMetaRegistry
     function setOperatorMetadataAsAdmin(
         uint256 nodeOperatorId,
@@ -530,6 +550,18 @@ contract MetaRegistry is IMetaRegistry, Initializable, AccessControlEnumerableUp
         MODULE.requestFullDepositInfoUpdate();
     }
 
+    function _setMaximumStakeCapPerNodeOperator(uint256 newCap) internal {
+        if (newCap == 0 || newCap % 1 ether != 0) revert InvalidStakeCap();
+
+        MetaRegistryStorage storage $ = _storage();
+        uint256 previousCap = $.maximumStakeCapPerNodeOperator;
+        if (newCap == previousCap) revert SameMaximumStakeCap();
+
+        $.maximumStakeCapPerNodeOperator = newCap;
+        emit MaximumStakeCapPerNodeOperatorSet(previousCap, newCap);
+        _requestFullDepositInfoUpdate();
+    }
+
     function _storeOperatorMetadata(uint256 nodeOperatorId, OperatorMetadata memory metadata) internal {
         if (bytes(metadata.name).length > MAX_NAME_LENGTH) revert OperatorNameTooLong();
         if (bytes(metadata.description).length > MAX_DESCRIPTION_LENGTH) revert OperatorDescriptionTooLong();

src/interfaces/IMetaRegistry.sol

@@ -53,6 +53,7 @@ interface IMetaRegistry {
     event GroupWeightsRefreshed(uint256 indexed groupId);
     event OperatorMetadataSet(uint256 indexed nodeOperatorId, OperatorMetadata metadata);
     event NodeOperatorEffectiveWeightChanged(uint256 indexed nodeOperatorId, uint256 oldWeight, uint256 newWeight);
+    event MaximumStakeCapPerNodeOperatorSet(uint256 previousCap, uint256 newCap);
 
     error ZeroModuleAddress();
     error ZeroAdminAddress();
@@ -75,6 +76,8 @@ interface IMetaRegistry {
     error ModuleAddressNotCached();
     error OperatorNameTooLong();
     error OperatorDescriptionTooLong();
+    error InvalidStakeCap();
+    error SameMaximumStakeCap();
 
     /// @notice Role allowed to manage operator groups.
     function MANAGE_OPERATOR_GROUPS_ROLE() external view returns (bytes32);
@@ -88,6 +91,9 @@ interface IMetaRegistry {
     /// @notice Role allowed to set bond curve weights.
     function SET_BOND_CURVE_WEIGHT_ROLE() external view returns (bytes32);
 
+    /// @notice Role allowed to manage the maximum allocated stake cap per Node Operator.
+    function MANAGE_STAKE_CAP_ROLE() external view returns (bytes32);
+
     /// @notice Curated module allowed to call module-only hooks.
     function MODULE() external view returns (ICuratedModule);
 
@@ -120,11 +126,27 @@ interface IMetaRegistry {
 
     /// @notice Initialize the registry.
     /// @param admin Address to receive DEFAULT_ADMIN_ROLE.
-    function initialize(address admin) external;
+    /// @param initialCap Initial maximum stake cap in wei.
+    function initialize(address admin, uint256 initialCap) external;
+
+    /// @notice Finalize the upgrade that introduces the maximum stake cap.
+    /// @dev Must be called atomically with the implementation upgrade.
+    /// @param initialCap Initial maximum stake cap in wei.
+    function finalizeUpgradeV2(uint256 initialCap) external;
 
     /// @notice Returns the initialized version of the contract.
     function getInitializedVersion() external view returns (uint64);
 
+    /// @notice Returns the maximum CuratedModule stake allocated to a single Node Operator.
+    /// @return cap Maximum stake cap in wei.
+    function maximumStakeCapPerNodeOperator() external view returns (uint256 cap);
+
+    /// @notice Sets the maximum CuratedModule stake allocated to a single Node Operator.
+    /// @dev The cap limits new initial deposits and top-ups. It does not initiate exits and may be exceeded
+    ///      by allocated balance changes reported after allocation. The value must be positive and divisible by 1 ETH.
+    /// @param newCap New maximum stake cap in wei.
+    function setMaximumStakeCapPerNodeOperator(uint256 newCap) external;
+
     /// @notice Set or update metadata for a node operator (callable by SET_OPERATOR_INFO_ROLE).
     /// @param nodeOperatorId Node operator ID.
     /// @param metadata Metadata payload to persist.

src/lib/StakeTracker.sol

@@ -41,16 +41,17 @@ library StakeTracker {
     }
 
     /// @dev Applies per-key top-up allocations, updates key allocated balances, and aggregates stake deltas per operator.
+    /// @return allocatedOperatorIds Unique IDs of operators whose tracked stake increased.
     function increaseKeyBalances(
         ModuleLinearStorage.BaseModuleStorage storage $,
         uint256[] calldata operatorIds,
         uint256[] calldata keyIndices,
         uint256[] calldata allocations
-    ) external {
-        uint256[] memory allocatedOperatorIds = new uint256[](operatorIds.length);
+    ) external returns (uint256[] memory allocatedOperatorIds) {
+        allocatedOperatorIds = new uint256[](operatorIds.length);
         uint256[] memory increments = new uint256[](operatorIds.length);
         TransientUintUintMap operatorIndexes = TransientUintUintMapLib.create();
-        uint256 touchedOperatorsCount;
+        uint256 allocatedOperatorsCount;
 
         for (uint256 i; i < allocations.length; ++i) {
             uint256 allocationWei = allocations[i];
@@ -67,22 +68,26 @@ library StakeTracker {
 
             uint256 operatorIndex = operatorIndexes.get(operatorIds[i]);
             if (operatorIndex == 0) {
-                operatorIndex = touchedOperatorsCount;
+                operatorIndex = allocatedOperatorsCount;
                 allocatedOperatorIds[operatorIndex] = operatorIds[i];
                 increments[operatorIndex] = appliedIncrementWei;
                 unchecked {
-                    ++touchedOperatorsCount;
+                    ++allocatedOperatorsCount;
                 }
                 // Store index + 1 so zero can remain the "not seen yet" sentinel in the transient map.
-                operatorIndexes.set(operatorIds[i], touchedOperatorsCount);
+                operatorIndexes.set(operatorIds[i], allocatedOperatorsCount);
             } else {
                 unchecked {
                     increments[operatorIndex - 1] += appliedIncrementWei;
                 }
             }
         }
 
-        for (uint256 i; i < touchedOperatorsCount; ++i) {
+        assembly ("memory-safe") {
+            mstore(allocatedOperatorIds, allocatedOperatorsCount)
+        }
+
+        for (uint256 i; i < allocatedOperatorIds.length; ++i) {
             increaseOperatorBalance($, allocatedOperatorIds[i], increments[i]);
         }
     }

src/lib/allocator/CuratedDepositAllocator.sol

@@ -70,7 +70,7 @@ library CuratedDepositAllocator {
     ///        (non-zero weight, non-zero quantized top-up capacity),
     ///        so a subset cannot bias its share by omitting other eligible operators.
     ///      - Per-operator capacity is computed as:
-    ///        `(active_validators * 2048 ETH) - current_operator_balance`, floored at zero.
+    ///        `min(active_validators * 2048 ETH, stake_cap) - current_operator_balance`, floored at zero.
     ///      - `current_operator_balance` here is the module's tracked stake view, not a live decrementing oracle value:
     ///        active balance decreases are intentionally reflected later via withdrawal reporting.
     /// @param $ Base module storage pointer.
@@ -95,7 +95,7 @@ library CuratedDepositAllocator {
     ///        (non-zero weight, non-zero quantized top-up capacity),
     ///        so a subset cannot bias its share by omitting other eligible operators.
     ///      - Per-operator capacity is computed as:
-    ///        `(active_validators * 2048 ETH) - current_operator_balance`, floored at zero.
+    ///        `min(active_validators * 2048 ETH, stake_cap) - current_operator_balance`, floored at zero.
     ///      - `current_operator_balance` is intentionally based on tracked stake that preserves prior observed highs
     ///        until withdrawal settlement, so active slashing/leakage is accounted when penalties are finalized.
     ///      - Per-key top-up limits are not used as caps for operator-level allocation; they are
@@ -269,11 +269,12 @@ library CuratedDepositAllocator {
         currentStakeByOperatorId = new uint256[](operatorsCount);
 
         IMetaRegistry metaRegistry = ICuratedModule(address(this)).META_REGISTRY();
+        uint256 stakeCap = metaRegistry.maximumStakeCapPerNodeOperator();
 
         // Build global share baseline across all eligible operators (non-zero weight + usable capacity).
         for (uint256 i; i < operatorsCount; ++i) {
             uint256 balance = StakeTracker.getOperatorBalance($, i);
-            uint256 capacity = quantizeForTopUp(_topUpCapacity($.nodeOperators[i], balance));
+            uint256 capacity = quantizeForTopUp(_topUpCapacity($.nodeOperators[i], balance, stakeCap));
             if (capacity == 0) continue;
             capacitiesByOperatorId[i] = capacity;
 
@@ -293,11 +294,17 @@ library CuratedDepositAllocator {
     }
 
     /// @dev Maximum top-up capacity for an operator:
-    ///      (active validators * 2048 ETH) - current balance, floored at zero.
-    function _topUpCapacity(NodeOperator storage no, uint256 balanceWei) internal view returns (uint256 capacity) {
+    ///      min(active validators * 2048 ETH, stake cap) - current balance, floored at zero.
+    function _topUpCapacity(
+        NodeOperator storage no,
+        uint256 balanceWei,
+        uint256 stakeCap
+    ) internal view returns (uint256 capacity) {
         unchecked {
-            uint256 maxTotal = (no.totalDepositedKeys - no.totalWithdrawnKeys) *
-                ValidatorBalanceLimits.MAX_EFFECTIVE_BALANCE;
+            uint256 maxTotal = Math.min(
+                (no.totalDepositedKeys - no.totalWithdrawnKeys) * ValidatorBalanceLimits.MAX_EFFECTIVE_BALANCE,
+                stakeCap
+            );
             if (maxTotal > balanceWei) capacity = maxTotal - balanceWei;
         }
     }