Merge pull request #10 from symbioticfi/statemind-fixes

Statemind fixes

README.md

@@ -14,7 +14,7 @@ This repository provides a framework for developing middleware in a modular and
 
   - **Operators**: Manages operator registration and operator's vault.
 
-  - **KeyManager**: Manages operator keys. Variants include `KeyManager256`, `KeyManagerBytes`, and `NoKeyManager`.
+  - **KeyManager**: Manages operator keys. Variants include `KeyManagerAddress`, `KeyManager256`, `KeyManagerBytes`, and `NoKeyManager`.
 
   - **AccessManager**: Controls access to restricted functions. Implementations include `OwnableAccessManager`, `OzAccessControl`, `OzAccessManaged`, and `NoAccessManager`.
 
@@ -63,7 +63,7 @@ Features:
 #### SelfRegisterMiddleware
 
 ```solidity
-contract SelfRegisterMiddleware is SharedVaults, SelfRegisterOperators, KeyManager256, ECDSASig, NoAccessManager, TimestampCapture, EqualStakePower {
+contract SelfRegisterMiddleware is SharedVaults, SelfRegisterOperators, KeyManagerAddress, ECDSASig, NoAccessManager, TimestampCapture, EqualStakePower {
     // Implementation details...
 }
 ```
@@ -89,7 +89,7 @@ Features:
 #### SelfRegisterSqrtTaskMiddleware
 
 ```solidity
-contract SelfRegisterSqrtTaskMiddleware is SharedVaults, SelfRegisterOperators, KeyManager256, ECDSASig, OwnableAccessManager, TimestampCapture, EqualStakePower {
+contract SelfRegisterSqrtTaskMiddleware is SharedVaults, SelfRegisterOperators, KeyManagerAddress, ECDSASig, OwnableAccessManager, TimestampCapture, EqualStakePower {
     // Implementation details...
 }
 ```
@@ -209,10 +209,9 @@ contract MyCustomMiddleware is BaseMiddleware, Operators, KeyStorage256, Ownable
     1. Granting roles to addresses using `grantRole(bytes32 role, address account)`
     2. Setting role admins with `_setRoleAdmin(bytes32 role, bytes32 adminRole)` 
     3. Assigning roles to function selectors via `_setSelectorRole(bytes4 selector, bytes32 role)`
-  - `OzAccessManaged`: Wraps OpenZeppelin's AccessManaged contract to integrate with external access control systems
   - `OzAccessManaged`: Wraps OpenZeppelin's AccessManaged contract to integrate with external access control systems. This allows for more complex access control scenarios where permissions are managed externally, providing flexibility and scalability in managing roles and permissions.
 
-- **Key Manager**: Choose a `KeyManager` implementation that suits your key management needs. Use `KeyManager256` for managing 256-bit keys, `KeyManagerBytes` for handling arbitrary-length keys, or `NoKeyManager` if key management is not required.
+- **Key Manager**: Choose a `KeyManager` implementation that suits your key management needs. Use `KeyManagerAddress` for managing address keys, `KeyManager256` for managing 256-bit keys, `KeyManagerBytes` for handling arbitrary-length keys, or `NoKeyManager` if key management is not required.
 
 This framework provides flexibility in building middleware by allowing you to mix and match various extensions based on your requirements. By following the modular approach and best practices outlined, you can develop robust middleware solutions that integrate seamlessly with the network.
 

src/examples/self-register-network/SelfRegisterEd25519Middleware.sol

@@ -5,7 +5,7 @@ import {BaseMiddleware} from "../../middleware/BaseMiddleware.sol";
 import {SharedVaults} from "../../extensions/SharedVaults.sol";
 import {SelfRegisterOperators} from "../../extensions/operators/SelfRegisterOperators.sol";
 
-import {NoAccessManager} from "../../extensions/managers/access/NoAccessManager.sol";
+import {OwnableAccessManager} from "../../extensions/managers/access/OwnableAccessManager.sol";
 import {TimestampCapture} from "../../extensions/managers/capture-timestamps/TimestampCapture.sol";
 import {EqualStakePower} from "../../extensions/managers/stake-powers/EqualStakePower.sol";
 import {KeyManager256} from "../../extensions/managers/keys/KeyManager256.sol";
@@ -16,7 +16,7 @@ contract SelfRegisterEd25519Middleware is
     SelfRegisterOperators,
     KeyManager256,
     EdDSASig,
-    NoAccessManager,
+    OwnableAccessManager,
     TimestampCapture,
     EqualStakePower
 {
@@ -28,16 +28,18 @@ contract SelfRegisterEd25519Middleware is
      * @param operatorRegistry The address of the operator registry
      * @param operatorNetOptin The address of the operator network opt-in service
      * @param reader The address of the reader contract used for delegatecall
+     * @param owner The address of the owner
      */
     constructor(
         address network,
         uint48 slashingWindow,
         address vaultRegistry,
         address operatorRegistry,
         address operatorNetOptin,
-        address reader
+        address reader,
+        address owner
     ) {
-        initialize(network, slashingWindow, vaultRegistry, operatorRegistry, operatorNetOptin, reader);
+        initialize(network, slashingWindow, vaultRegistry, operatorRegistry, operatorNetOptin, reader, owner);
     }
 
     function initialize(
@@ -46,9 +48,11 @@ contract SelfRegisterEd25519Middleware is
         address vaultRegistry,
         address operatorRegistry,
         address operatorNetOptin,
-        address reader
+        address reader,
+        address owner
     ) internal initializer {
         __BaseMiddleware_init(network, slashingWindow, vaultRegistry, operatorRegistry, operatorNetOptin, reader);
         __SelfRegisterOperators_init("SelfRegisterEd25519Middleware");
+        __OwnableAccessManager_init(owner);
     }
 }

src/examples/self-register-network/SelfRegisterMiddleware.sol

@@ -6,17 +6,17 @@ import {SharedVaults} from "../../extensions/SharedVaults.sol";
 import {SelfRegisterOperators} from "../../extensions/operators/SelfRegisterOperators.sol";
 
 import {ECDSASig} from "../../extensions/managers/sigs/ECDSASig.sol";
-import {NoAccessManager} from "../../extensions/managers/access/NoAccessManager.sol";
+import {OwnableAccessManager} from "../../extensions/managers/access/OwnableAccessManager.sol";
 import {TimestampCapture} from "../../extensions/managers/capture-timestamps/TimestampCapture.sol";
 import {EqualStakePower} from "../../extensions/managers/stake-powers/EqualStakePower.sol";
-import {KeyManager256} from "../../extensions/managers/keys/KeyManager256.sol";
+import {KeyManagerAddress} from "../../extensions/managers/keys/KeyManagerAddress.sol";
 
 contract SelfRegisterMiddleware is
     SharedVaults,
     SelfRegisterOperators,
-    KeyManager256,
+    KeyManagerAddress,
     ECDSASig,
-    NoAccessManager,
+    OwnableAccessManager,
     TimestampCapture,
     EqualStakePower
 {
@@ -28,16 +28,18 @@ contract SelfRegisterMiddleware is
      * @param operatorRegistry The address of the operator registry
      * @param operatorNetOptin The address of the operator network opt-in service
      * @param reader The address of the reader contract used for delegatecall
+     * @param owner The address of the owner
      */
     constructor(
         address network,
         uint48 slashingWindow,
         address vaultRegistry,
         address operatorRegistry,
         address operatorNetOptin,
-        address reader
+        address reader,
+        address owner
     ) {
-        initialize(network, slashingWindow, vaultRegistry, operatorRegistry, operatorNetOptin, reader);
+        initialize(network, slashingWindow, vaultRegistry, operatorRegistry, operatorNetOptin, reader, owner);
     }
 
     function initialize(
@@ -46,9 +48,11 @@ contract SelfRegisterMiddleware is
         address vaultRegistry,
         address operatorRegistry,
         address operatorNetOptIn,
-        address reader
+        address reader,
+        address owner
     ) internal initializer {
         __BaseMiddleware_init(network, slashingWindow, vaultRegistry, operatorRegistry, operatorNetOptIn, reader);
         __SelfRegisterOperators_init("SelfRegisterMiddleware");
+        __OwnableAccessManager_init(owner);
     }
 }

src/examples/sqrt-task-network/SelfRegisterSqrtTaskMiddleware.sol

@@ -17,15 +17,22 @@ import {SelfRegisterOperators} from "../../extensions/operators/SelfRegisterOper
 
 import {ECDSASig} from "../../extensions/managers/sigs/ECDSASig.sol";
 import {OwnableAccessManager} from "../../extensions/managers/access/OwnableAccessManager.sol";
-import {KeyManager256} from "../../extensions/managers/keys/KeyManager256.sol";
+import {KeyManagerAddress} from "../../extensions/managers/keys/KeyManagerAddress.sol";
 import {TimestampCapture} from "../../extensions/managers/capture-timestamps/TimestampCapture.sol";
 import {EqualStakePower} from "../../extensions/managers/stake-powers/EqualStakePower.sol";
 
+/**
+ * @title SelfRegisterSqrtTaskMiddleware
+ * @notice Middleware for managing sqrt computation tasks with self-registering operators
+ * @dev Uses SelfRegisterOperators for operator management because it allows permissionless registration.
+ * Task validation is done by a single validator chosen at task creation time because this avoids
+ * having to iterate over all operators, making it more gas efficient and avoiding potential DOS attacks.
+ */
 contract SelfRegisterSqrtTaskMiddleware is
     SharedVaults,
     SelfRegisterOperators,
     ECDSASig,
-    KeyManager256,
+    KeyManagerAddress,
     OwnableAccessManager,
     TimestampCapture,
     EqualStakePower
@@ -35,19 +42,23 @@ contract SelfRegisterSqrtTaskMiddleware is
 
     error InvalidHints();
     error TaskCompleted();
+    error TooManyOperatorVaults();
+    error InactiveValidator();
 
-    event CreateTask(uint256 indexed taskIndex);
+    event CreateTask(uint256 indexed taskIndex, address indexed validator);
     event CompleteTask(uint256 indexed taskIndex, bool isValidAnswer);
 
     struct Task {
         uint48 captureTimestamp;
         uint256 value;
-        address operator;
+        address validator;
         bool completed;
     }
 
     bytes32 private constant COMPLETE_TASK_TYPEHASH = keccak256("CompleteTask(uint256 taskIndex,uint256 answer)");
 
+    uint256 public constant MAX_OPERATOR_VAULTS = 20;
+
     Task[] public tasks;
 
     constructor(
@@ -77,22 +88,21 @@ contract SelfRegisterSqrtTaskMiddleware is
         __SelfRegisterOperators_init("SelfRegisterSqrtTaskMiddleware");
     }
 
-    // allow anyone to register shared vaults
-    function _checkAccess() internal view override(AccessManager, OwnableAccessManager) {
-        if (
-            msg.sig == this.registerSharedVault.selector || msg.sig == this.unregisterSharedVault.selector
-                || msg.sig == this.pauseSharedVault.selector
-        ) {
-            return;
+    function createTask(uint256 value, address validator) external returns (uint256 taskIndex) {
+        bytes memory key = abi.encode(validator);
+        address operator = operatorByKey(key);
+        if (!keyWasActiveAt(getCaptureTimestamp(), key)) {
+            revert InactiveValidator();
+        }
+        if (!_isOperatorRegistered(operator)) {
+            revert OperatorNotRegistered();
         }
-        OwnableAccessManager._checkAccess();
-    }
-
-    function createTask(uint256 value, address operator) external returns (uint256 taskIndex) {
         taskIndex = tasks.length;
-        tasks.push(Task({captureTimestamp: getCaptureTimestamp(), value: value, operator: operator, completed: false}));
+        tasks.push(
+            Task({captureTimestamp: getCaptureTimestamp(), value: value, validator: validator, completed: false})
+        );
 
-        emit CreateTask(taskIndex);
+        emit CreateTask(taskIndex, validator);
     }
 
     function completeTask(
@@ -126,7 +136,7 @@ contract SelfRegisterSqrtTaskMiddleware is
 
         bytes32 hash_ = _hashTypedDataV4(keccak256(abi.encode(COMPLETE_TASK_TYPEHASH, taskIndex, answer)));
 
-        if (!SignatureChecker.isValidSignatureNow(task.operator, hash_, signature)) {
+        if (!SignatureChecker.isValidSignatureNow(task.validator, hash_, signature)) {
             revert InvalidSignature();
         }
     }
@@ -159,7 +169,17 @@ contract SelfRegisterSqrtTaskMiddleware is
 
     function _slash(uint256 taskIndex, bytes[] calldata stakeHints, bytes[] calldata slashHints) private {
         Task storage task = tasks[taskIndex];
-        address[] memory vaults = _activeVaultsAt(task.captureTimestamp, task.operator);
+        address operator = operatorByKey(abi.encode(task.validator));
+        _slashOperator(task.captureTimestamp, operator, stakeHints, slashHints);
+    }
+
+    function _slashOperator(
+        uint48 captureTimestamp,
+        address operator,
+        bytes[] calldata stakeHints,
+        bytes[] calldata slashHints
+    ) private {
+        address[] memory vaults = _activeVaultsAt(captureTimestamp, operator);
         uint256 vaultsLength = vaults.length;
 
         if (stakeHints.length != slashHints.length || stakeHints.length != vaultsLength) {
@@ -169,15 +189,14 @@ contract SelfRegisterSqrtTaskMiddleware is
         bytes32 subnetwork = _NETWORK().subnetwork(0);
         for (uint256 i; i < vaultsLength; ++i) {
             address vault = vaults[i];
-            uint256 slashAmount = IBaseDelegator(IVault(vault).delegator()).stakeAt(
-                subnetwork, task.operator, task.captureTimestamp, stakeHints[i]
-            );
+            uint256 slashAmount =
+                IBaseDelegator(IVault(vault).delegator()).stakeAt(subnetwork, operator, captureTimestamp, stakeHints[i]);
 
             if (slashAmount == 0) {
                 continue;
             }
 
-            _slashVault(task.captureTimestamp, vault, subnetwork, task.operator, slashAmount, slashHints[i]);
+            _slashVault(captureTimestamp, vault, subnetwork, operator, slashAmount, slashHints[i]);
         }
     }
 
@@ -192,13 +211,11 @@ contract SelfRegisterSqrtTaskMiddleware is
         _slashVault(epochStart, vault, subnetwork, operator, amount, hints);
     }
 
-    function _beforeRegisterSharedVault(
-        address sharedVault
-    ) internal override {
-        IBaseDelegator(IVault(sharedVault).delegator()).setMaxNetworkLimit(DEFAULT_SUBNETWORK, type(uint256).max);
-    }
-
     function _beforeRegisterOperatorVault(address operator, address vault) internal override {
+        super._beforeRegisterOperatorVault(operator, vault);
+        if (_operatorVaultsLength(operator) >= MAX_OPERATOR_VAULTS) {
+            revert TooManyOperatorVaults();
+        }
         IBaseDelegator(IVault(vault).delegator()).setMaxNetworkLimit(DEFAULT_SUBNETWORK, type(uint256).max);
     }
 }

src/examples/sqrt-task-network/SqrtTaskMiddleware.sol

@@ -34,7 +34,7 @@ contract SqrtTaskMiddleware is
     error InvalidSignature();
     error TaskCompleted();
 
-    event CreateTask(uint256 indexed taskIndex);
+    event CreateTask(uint256 indexed taskIndex, address indexed operator);
     event CompleteTask(uint256 indexed taskIndex, bool isValidAnswer);
 
     struct Task {
@@ -77,7 +77,7 @@ contract SqrtTaskMiddleware is
         taskIndex = tasks.length;
         tasks.push(Task({captureTimestamp: getCaptureTimestamp(), value: value, operator: operator, completed: false}));
 
-        emit CreateTask(taskIndex);
+        emit CreateTask(taskIndex, operator);
     }
 
     function completeTask(