fix: keyregistrar and crosschainregistry checks

Author: 0xrajath

src/contracts/cloud/ComputeRegistry.sol

@@ -51,11 +51,13 @@ contract ComputeRegistry is Initializable, ComputeRegistryStorage, PermissionCon
     constructor(
         IReleaseManager _releaseManager,
         IAllocationManager _allocationManager,
+        IKeyRegistrar _keyRegistrar,
+        ICrossChainRegistry _crossChainRegistry,
         IPermissionController _permissionController,
         bytes32 _tosHash,
         string memory _version
     )
-        ComputeRegistryStorage(_releaseManager, _allocationManager, _tosHash)
+        ComputeRegistryStorage(_releaseManager, _allocationManager, _keyRegistrar, _crossChainRegistry, _tosHash)
         PermissionControllerMixin(_permissionController)
         SignatureUtilsMixin(_version)
     {
@@ -79,6 +81,13 @@ contract ComputeRegistry is Initializable, ComputeRegistryStorage, PermissionCon
         bytes32 operatorSetKey = operatorSet.key();
         require(!isOperatorSetRegistered[operatorSetKey], OperatorSetAlreadyRegistered());
 
+        // Check if the curve type has been set for the operator set
+        IKeyRegistrarTypes.CurveType curveType = KEY_REGISTRAR.getOperatorSetCurveType(operatorSet);
+        require(curveType != IKeyRegistrarTypes.CurveType.NONE, CurveTypeNotSet());
+
+        // Check if the operator set has an active generation reservation
+        require(CROSS_CHAIN_REGISTRY.hasActiveGenerationReservation(operatorSet), NoActiveGenerationReservation());
+
         // Check if there is at least one release for the operator set
         // The ReleaseManager will revert with `NoReleases()` if there are no releases for the operator set
         RELEASE_MANAGER.getLatestRelease(operatorSet);

src/contracts/cloud/ComputeRegistryStorage.sol

@@ -4,6 +4,8 @@ pragma solidity ^0.8.27;
 import "../interfaces/IComputeRegistry.sol";
 import "../interfaces/IReleaseManager.sol";
 import "../interfaces/IAllocationManager.sol";
+import "../interfaces/IKeyRegistrar.sol";
+import "../interfaces/ICrossChainRegistry.sol";
 import "../libraries/OperatorSetLib.sol";
 
 abstract contract ComputeRegistryStorage is IComputeRegistry {
@@ -22,6 +24,12 @@ abstract contract ComputeRegistryStorage is IComputeRegistry {
     /// @notice The AllocationManager contract
     IAllocationManager public immutable ALLOCATION_MANAGER;
 
+    /// @notice The KeyRegistrar contract
+    IKeyRegistrar public immutable KEY_REGISTRAR;
+
+    /// @notice The CrossChainRegistry contract
+    ICrossChainRegistry public immutable CROSS_CHAIN_REGISTRY;
+
     /// @notice The hash of the Terms of Service that AVS operators must sign
     bytes32 public immutable TOS_HASH;
 
@@ -42,9 +50,17 @@ abstract contract ComputeRegistryStorage is IComputeRegistry {
      */
     uint256[48] private __gap;
 
-    constructor(IReleaseManager _releaseManager, IAllocationManager _allocationManager, bytes32 _tosHash) {
+    constructor(
+        IReleaseManager _releaseManager,
+        IAllocationManager _allocationManager,
+        IKeyRegistrar _keyRegistrar,
+        ICrossChainRegistry _crossChainRegistry,
+        bytes32 _tosHash
+    ) {
         RELEASE_MANAGER = _releaseManager;
         ALLOCATION_MANAGER = _allocationManager;
+        KEY_REGISTRAR = _keyRegistrar;
+        CROSS_CHAIN_REGISTRY = _crossChainRegistry;
         TOS_HASH = _tosHash;
     }
 }

src/contracts/interfaces/IComputeRegistry.sol

@@ -26,6 +26,12 @@ interface IComputeRegistryErrors {
 
     /// @dev Thrown when an invalid operator set is provided
     error InvalidOperatorSet();
+
+    /// @dev Thrown when the curve type for an operator set has not been set
+    error CurveTypeNotSet();
+
+    /// @dev Thrown when the operator set does not have an active generation reservation
+    error NoActiveGenerationReservation();
 }
 
 interface IComputeRegistryEvents {

src/test/mocks/CrossChainRegistryMock.sol

@@ -1,13 +1,22 @@
 // SPDX-License-Identifier: BUSL-1.1
 pragma solidity ^0.8.27;
 
-import "forge-std/Test.sol";
+import "src/contracts/interfaces/ICrossChainRegistry.sol";
+import "src/contracts/libraries/OperatorSetLib.sol";
 
-import "src/contracts/multichain/CrossChainRegistry.sol";
-
-contract CrossChainRegistryMock is Test, ICrossChainRegistryTypes {
+contract CrossChainRegistryMock {
     using OperatorSetLib for OperatorSet;
 
-    receive() external payable {}
-    fallback() external payable {}
+    mapping(bytes32 => bool) public generationReservations;
+
+    function hasActiveGenerationReservation(OperatorSet memory operatorSet) external view returns (bool) {
+        bytes32 key = keccak256(abi.encode(operatorSet.avs, operatorSet.id));
+        return generationReservations[key];
+    }
+
+    // Helper function for testing
+    function setHasActiveGenerationReservation(OperatorSet memory operatorSet, bool hasReservation) external {
+        bytes32 key = keccak256(abi.encode(operatorSet.avs, operatorSet.id));
+        generationReservations[key] = hasReservation;
+    }
 }

src/test/mocks/KeyRegistrarMock.sol

@@ -0,0 +1,22 @@
+// SPDX-License-Identifier: BUSL-1.1
+pragma solidity ^0.8.27;
+
+import "src/contracts/interfaces/IKeyRegistrar.sol";
+import "src/contracts/libraries/OperatorSetLib.sol";
+
+contract KeyRegistrarMock {
+    using OperatorSetLib for OperatorSet;
+
+    mapping(bytes32 => IKeyRegistrarTypes.CurveType) public operatorSetCurveTypes;
+
+    function getOperatorSetCurveType(OperatorSet memory operatorSet) external view returns (IKeyRegistrarTypes.CurveType) {
+        bytes32 key = keccak256(abi.encode(operatorSet.avs, operatorSet.id));
+        return operatorSetCurveTypes[key];
+    }
+
+    // Helper function for testing
+    function setOperatorSetCurveType(OperatorSet memory operatorSet, IKeyRegistrarTypes.CurveType curveType) external {
+        bytes32 key = keccak256(abi.encode(operatorSet.avs, operatorSet.id));
+        operatorSetCurveTypes[key] = curveType;
+    }
+}

src/test/unit/ComputeRegistryUnit.t.sol

@@ -5,6 +5,8 @@ import "src/contracts/cloud/ComputeRegistry.sol";
 import "src/test/utils/EigenLayerUnitTestSetup.sol";
 import "src/test/mocks/MockAVSRegistrar.sol";
 import "src/test/mocks/ReleaseManagerMock.sol";
+import "src/test/mocks/KeyRegistrarMock.sol";
+import "src/test/mocks/CrossChainRegistryMock.sol";
 
 contract ComputeRegistryUnitTests is EigenLayerUnitTestSetup, IComputeRegistryErrors, IComputeRegistryEvents {
     using StdStyle for *;
@@ -19,6 +21,8 @@ contract ComputeRegistryUnitTests is EigenLayerUnitTestSetup, IComputeRegistryEr
     // Contracts
     ComputeRegistry computeRegistry;
     ReleaseManagerMock releaseManagerMock;
+    KeyRegistrarMock keyRegistrarMock;
+    CrossChainRegistryMock crossChainRegistryMock;
 
     // Test variables
     address defaultAVS;
@@ -31,6 +35,8 @@ contract ComputeRegistryUnitTests is EigenLayerUnitTestSetup, IComputeRegistryEr
 
         // Setup mock contracts
         releaseManagerMock = new ReleaseManagerMock();
+        keyRegistrarMock = new KeyRegistrarMock();
+        crossChainRegistryMock = new CrossChainRegistryMock();
 
         // Setup default test accounts
         defaultSignerPrivateKey = 0x1234;
@@ -41,6 +47,8 @@ contract ComputeRegistryUnitTests is EigenLayerUnitTestSetup, IComputeRegistryEr
         computeRegistry = new ComputeRegistry(
             IReleaseManager(address(releaseManagerMock)),
             IAllocationManager(address(allocationManagerMock)),
+            IKeyRegistrar(address(keyRegistrarMock)),
+            ICrossChainRegistry(address(crossChainRegistryMock)),
             IPermissionController(address(permissionController)),
             TOS_HASH,
             VERSION
@@ -52,6 +60,8 @@ contract ComputeRegistryUnitTests is EigenLayerUnitTestSetup, IComputeRegistryEr
         // Configure mocks
         allocationManagerMock.setIsOperatorSet(defaultOperatorSet, true);
         releaseManagerMock.setHasRelease(defaultOperatorSet, true);
+        keyRegistrarMock.setOperatorSetCurveType(defaultOperatorSet, IKeyRegistrarTypes.CurveType.ECDSA);
+        crossChainRegistryMock.setHasActiveGenerationReservation(defaultOperatorSet, true);
 
         // Setup permissions for default signer to call registerForCompute and deregisterFromCompute
         vm.startPrank(defaultAVS);
@@ -82,6 +92,8 @@ contract ComputeRegistryUnitTests_Initialization is ComputeRegistryUnitTests {
     function test_initialization() public view {
         assertEq(address(computeRegistry.RELEASE_MANAGER()), address(releaseManagerMock));
         assertEq(address(computeRegistry.ALLOCATION_MANAGER()), address(allocationManagerMock));
+        assertEq(address(computeRegistry.KEY_REGISTRAR()), address(keyRegistrarMock));
+        assertEq(address(computeRegistry.CROSS_CHAIN_REGISTRY()), address(crossChainRegistryMock));
         assertEq(address(computeRegistry.permissionController()), address(permissionController));
         assertEq(computeRegistry.TOS_HASH(), TOS_HASH);
         assertEq(computeRegistry.MAX_EXPIRY(), MAX_EXPIRY);
@@ -164,6 +176,9 @@ contract ComputeRegistryUnitTests_RegisterForCompute is ComputeRegistryUnitTests
         OperatorSet memory operatorSet = OperatorSet(defaultAVS, 1);
         allocationManagerMock.setIsOperatorSet(operatorSet, true);
         releaseManagerMock.setHasRelease(operatorSet, false);
+        // Set curve type and generation reservation so we get to the release check
+        keyRegistrarMock.setOperatorSetCurveType(operatorSet, IKeyRegistrarTypes.CurveType.ECDSA);
+        crossChainRegistryMock.setHasActiveGenerationReservation(operatorSet, true);
 
         bytes memory signature = _generateTOSSignature(operatorSet, defaultSigner, defaultSignerPrivateKey);
 
@@ -189,6 +204,60 @@ contract ComputeRegistryUnitTests_RegisterForCompute is ComputeRegistryUnitTests
         vm.expectRevert(ISignatureUtilsMixinErrors.InvalidSignature.selector);
         computeRegistry.registerForCompute(defaultOperatorSet, signature);
     }
+
+    function test_registerForCompute_revert_curveTypeNotSet() public {
+        OperatorSet memory operatorSet = OperatorSet(defaultAVS, 2);
+        allocationManagerMock.setIsOperatorSet(operatorSet, true);
+        // Curve type defaults to NONE in the mock
+        crossChainRegistryMock.setHasActiveGenerationReservation(operatorSet, true);
+        releaseManagerMock.setHasRelease(operatorSet, true);
+
+        bytes memory signature = _generateTOSSignature(operatorSet, defaultSigner, defaultSignerPrivateKey);
+
+        vm.prank(defaultSigner);
+        vm.expectRevert(CurveTypeNotSet.selector);
+        computeRegistry.registerForCompute(operatorSet, signature);
+    }
+
+    function test_registerForCompute_revert_noActiveGenerationReservation() public {
+        OperatorSet memory operatorSet = OperatorSet(defaultAVS, 3);
+        allocationManagerMock.setIsOperatorSet(operatorSet, true);
+        keyRegistrarMock.setOperatorSetCurveType(operatorSet, IKeyRegistrarTypes.CurveType.BN254);
+        // Generation reservation defaults to false in the mock
+        releaseManagerMock.setHasRelease(operatorSet, true);
+
+        bytes memory signature = _generateTOSSignature(operatorSet, defaultSigner, defaultSignerPrivateKey);
+
+        vm.prank(defaultSigner);
+        vm.expectRevert(NoActiveGenerationReservation.selector);
+        computeRegistry.registerForCompute(operatorSet, signature);
+    }
+
+    function test_registerForCompute_withDifferentCurveTypes() public {
+        // Test with ECDSA curve type
+        OperatorSet memory ecdsaOperatorSet = OperatorSet(defaultAVS, 4);
+        allocationManagerMock.setIsOperatorSet(ecdsaOperatorSet, true);
+        keyRegistrarMock.setOperatorSetCurveType(ecdsaOperatorSet, IKeyRegistrarTypes.CurveType.ECDSA);
+        crossChainRegistryMock.setHasActiveGenerationReservation(ecdsaOperatorSet, true);
+        releaseManagerMock.setHasRelease(ecdsaOperatorSet, true);
+
+        bytes memory signature1 = _generateTOSSignature(ecdsaOperatorSet, defaultSigner, defaultSignerPrivateKey);
+        vm.prank(defaultSigner);
+        computeRegistry.registerForCompute(ecdsaOperatorSet, signature1);
+        assertTrue(computeRegistry.isOperatorSetRegistered(ecdsaOperatorSet.key()));
+
+        // Test with BN254 curve type
+        OperatorSet memory bn254OperatorSet = OperatorSet(defaultAVS, 5);
+        allocationManagerMock.setIsOperatorSet(bn254OperatorSet, true);
+        keyRegistrarMock.setOperatorSetCurveType(bn254OperatorSet, IKeyRegistrarTypes.CurveType.BN254);
+        crossChainRegistryMock.setHasActiveGenerationReservation(bn254OperatorSet, true);
+        releaseManagerMock.setHasRelease(bn254OperatorSet, true);
+
+        bytes memory signature2 = _generateTOSSignature(bn254OperatorSet, defaultSigner, defaultSignerPrivateKey);
+        vm.prank(defaultSigner);
+        computeRegistry.registerForCompute(bn254OperatorSet, signature2);
+        assertTrue(computeRegistry.isOperatorSetRegistered(bn254OperatorSet.key()));
+    }
 }
 
 contract ComputeRegistryUnitTests_DeregisterFromCompute is ComputeRegistryUnitTests {