test: restricted strategy list

src/contracts/interfaces/IAVSDirectory.sol

@@ -329,7 +329,7 @@ interface IAVSDirectory is IAVSDirectoryEvents, IAVSDirectoryErrors, ISignatureU
     /// @notice Returns whether or not a list of strategies are all allocatable/slashable for an operator set.
     function isOperatorSetStrategyBatch(
         OperatorSet calldata operatorSet,
-        IStrategy[] calldata strategy
+        IStrategy[] calldata strategies
     ) external view returns (bool);
 
     /**

src/test/mocks/AVSDirectoryMock.sol

@@ -2,29 +2,47 @@
 pragma solidity ^0.8.9;
 
 import "forge-std/Test.sol";
+import "src/contracts/interfaces/IStrategy.sol";
 import "src/contracts/interfaces/IAVSDirectory.sol";
 
 contract AVSDirectoryMock is Test {
     receive() external payable {}
     fallback() external payable {}
 
-    mapping(address => mapping(bytes32 => bool)) public _isOperatorSlashable;
-    mapping(bytes32 => bool) public _isOperatorSetBatch;
+    mapping(bytes32 => bool) _isOperatorSlashable;
+    mapping(bytes32 => bool) _isOperatorSetBatch;
+    mapping(bytes32 => bool) _isOperatorSetStrategy;
+    mapping(bytes32 => bool) _isOperatorSetStrategyBatch;
+
+    // Getters
 
     function isOperatorSlashable(address operator, OperatorSet memory operatorSet) public virtual view returns (bool) {
-        return _isOperatorSlashable[operator][bytes32(abi.encode(operatorSet))];
+        return _isOperatorSlashable[keccak256(abi.encode(operator, operatorSet))];
     }
 
     function isOperatorSetBatch(OperatorSet[] memory operatorSets) public virtual view returns (bool) {
         return _isOperatorSetBatch[keccak256(abi.encode(operatorSets))];
     }
 
+    function isOperatorSetStrategy(OperatorSet calldata operatorSet, IStrategy strategy) external view returns (bool) {
+        return _isOperatorSetStrategy[keccak256(abi.encode(operatorSet, strategy))];
+    }
+
+    function isOperatorSetStrategyBatch(
+        OperatorSet calldata operatorSet,
+        IStrategy[] calldata strategies
+    ) external view returns (bool) {
+        return _isOperatorSetStrategyBatch[keccak256(abi.encode(operatorSet, strategies))];
+    }
+
+    // Setters
+
     function setIsOperatorSlashable(
         address operator, 
         OperatorSet memory operatorSet, 
         bool value
     ) public virtual {
-        _isOperatorSlashable[operator][bytes32(abi.encode(operatorSet))] = value;
+        _isOperatorSlashable[keccak256(abi.encode(operator, operatorSet))] = value;
     }
 
     function setIsOperatorSlashable(
@@ -43,4 +61,12 @@ contract AVSDirectoryMock is Test {
     function setIsOperatorSetBatch(OperatorSet[] memory operatorSets, bool value) public virtual {
         _isOperatorSetBatch[keccak256(abi.encode(operatorSets))] = value;
     }
+
+    function setIsOperatorSetStrategy(OperatorSet memory operatorSet, IStrategy strategy, bool value) public virtual {
+        _isOperatorSetStrategy[keccak256(abi.encode(operatorSet, strategy))] = value;
+    }
+
+    function setIsOperatorSetStrategyBatch(OperatorSet memory operatorSet, IStrategy[] memory strategies, bool value) public virtual {
+        _isOperatorSetStrategyBatch[keccak256(abi.encode(operatorSet, strategies))] = value;
+    }
 }

src/test/unit/AllocationManagerUnit.t.sol

@@ -111,6 +111,10 @@ contract AllocationManagerUnitTests is EigenLayerUnitTestSetup, IAllocationManag
         // Set operatorSet to being valid
         avsDirectoryMock.setIsOperatorSetBatch(operatorSets, true);
 
+        // Set strategy to being valid for operator set
+        avsDirectoryMock.setIsOperatorSetStrategy(operatorSets[0], strategy, true);
+        avsDirectoryMock.setIsOperatorSetStrategyBatch(operatorSets[0], _strategyMockArray(), true);
+
         uint64[] memory magnitudes = new uint64[](1);
         magnitudes[0] = magnitudeToSet;
 
@@ -285,6 +289,10 @@ contract AllocationManagerUnitTests is EigenLayerUnitTestSetup, IAllocationManag
         OperatorSet[] memory operatorSets = new OperatorSet[](1);
         operatorSets[0] = OperatorSet({avs: avs, operatorSetId: uint32(r)});
 
+        // set strategy to being valid for operator set
+        avsDirectoryMock.setIsOperatorSetStrategy(operatorSets[0], strategyMock, true);
+        avsDirectoryMock.setIsOperatorSetStrategyBatch(operatorSets[0], _strategyMockArray(), true);
+
         // Set operatorSet to being valid
         avsDirectoryMock.setIsOperatorSetBatch(operatorSets, true);
 
@@ -307,6 +315,7 @@ contract AllocationManagerUnitTests is EigenLayerUnitTestSetup, IAllocationManag
     /// Generate a random allocation for a single strategy and multiple operatorSets
     /// -----------------------------------------------------------------------
 
+    // TODO: optmize for compile/test time
     function _randomMagnitudeAllocation_singleStrat_multipleOpSets(
         uint256 r,
         uint256 salt,
@@ -343,12 +352,18 @@ contract AllocationManagerUnitTests is EigenLayerUnitTestSetup, IAllocationManag
         // Create magnitude allocation
         IAllocationManagerTypes.MagnitudeAllocation[] memory allocations =
             new IAllocationManagerTypes.MagnitudeAllocation[](1);
+
         allocations[0] = IAllocationManagerTypes.MagnitudeAllocation({
             strategy: strategyMock,
             expectedMaxMagnitude: 1e18, // magnitude starts at 100%
             operatorSets: operatorSets,
             magnitudes: magnitudes
         });
+
+        for (uint8 i = 0; i < numOpSets; i++) {
+            avsDirectoryMock.setIsOperatorSetStrategy(allocations[0].operatorSets[i], strategyMock, true);
+        }
+
         return allocations;
     }
 
@@ -375,6 +390,16 @@ contract AllocationManagerUnitTests is EigenLayerUnitTestSetup, IAllocationManag
         (MagnitudeAllocation[] memory allocations, MagnitudeAllocation[] memory deallocations) =
             _randomAllocationAndDeallocation_singleStrat_multipleOpSets(numOpSets, r, salt);
 
+        for (uint256 i; i < allocations.length; ++i) {
+            for (uint256 j; j < numOpSets; ++j) {
+                avsDirectoryMock.setIsOperatorSetStrategy(
+                    allocations[i].operatorSets[j], 
+                    allocations[i].strategy, 
+                    true
+                );
+            }
+        }
+
         // Allocate
         cheats.prank(operator);
         allocationManager.modifyAllocations(allocations);
@@ -418,13 +443,24 @@ contract AllocationManagerUnitTests is EigenLayerUnitTestSetup, IAllocationManag
         // Create deallocations
         IAllocationManagerTypes.MagnitudeAllocation[] memory deallocations =
             new IAllocationManagerTypes.MagnitudeAllocation[](1);
+
         deallocations[0] = IAllocationManagerTypes.MagnitudeAllocation({
             strategy: strategyMock,
             expectedMaxMagnitude: 1e18, // magnitude starts at 100%
             operatorSets: allocations[0].operatorSets,
             magnitudes: newMags
         });
 
+        for (uint256 i; i < allocations.length; ++i) {
+            for (uint256 j; j < numOpSets; ++j) {
+                avsDirectoryMock.setIsOperatorSetStrategy(
+                    allocations[i].operatorSets[j], 
+                    allocations[i].strategy, 
+                    true
+                );
+            }
+        }
+
         return (allocations, deallocations);
     }
 
@@ -532,6 +568,10 @@ contract AllocationManagerUnitTests_SlashOperator is AllocationManagerUnitTests
             slashingParams.operator, defaultAVS, slashingParams.operatorSetId, false
         );
 
+        avsDirectoryMock.setIsOperatorSetStrategyBatch(
+            _operatorSet(defaultAVS, slashingParams.operatorSetId), _strategyMockArray(), true
+        );
+
         cheats.expectRevert(IAllocationManagerErrors.InvalidOperator.selector);
         cheats.prank(defaultAVS);
         allocationManager.slashOperator(slashingParams);
@@ -541,6 +581,10 @@ contract AllocationManagerUnitTests_SlashOperator is AllocationManagerUnitTests
         SlashingParams memory slashingParams = _randomSlashingParams(defaultOperator, 0, 0);
         avsDirectoryMock.setIsOperatorSlashable(slashingParams.operator, defaultAVS, slashingParams.operatorSetId, true);
 
+        avsDirectoryMock.setIsOperatorSetStrategyBatch(
+            _operatorSet(defaultAVS, slashingParams.operatorSetId), _strategyMockArray(), true
+        );
+
         cheats.expectRevert(IAllocationManagerErrors.OperatorNotAllocated.selector);
         cheats.prank(defaultAVS);
         allocationManager.slashOperator(slashingParams);
@@ -560,6 +604,9 @@ contract AllocationManagerUnitTests_SlashOperator is AllocationManagerUnitTests
             description: "test"
         });
         avsDirectoryMock.setIsOperatorSlashable(slashingParams.operator, defaultAVS, slashingParams.operatorSetId, true);
+        avsDirectoryMock.setIsOperatorSetStrategyBatch(
+            _operatorSet(defaultAVS, slashingParams.operatorSetId), _strategyMockArray(), true
+        );
 
         // Expect revert
         cheats.expectRevert(IAllocationManagerErrors.OperatorNotAllocated.selector);
@@ -1200,6 +1247,7 @@ contract AllocationManagerUnitTests_SlashOperator is AllocationManagerUnitTests
             description: "test"
         });
         avsDirectoryMock.setIsOperatorSlashable(slashingParams.operator, defaultAVS, slashingParams.operatorSetId, true);
+        avsDirectoryMock.setIsOperatorSetStrategyBatch(OperatorSet(defaultAVS, slashingParams.operatorSetId), _strategyMockArray(), true);
 
         // Slash Operator
         cheats.prank(defaultAVS);
@@ -1498,6 +1546,10 @@ contract AllocationManagerUnitTests_ModifyAllocations is AllocationManagerUnitTe
 
         allocations[0].magnitudes[numOpSets - 1] = 1e18 + 1;
 
+        for (uint256 i; i < numOpSets; i++) {
+            avsDirectoryMock.setIsOperatorSetStrategy(allocations[0].operatorSets[i], allocations[0].strategy, true);
+        }
+
         // Overallocate
         cheats.expectRevert(IAllocationManagerErrors.InsufficientAllocatableMagnitude.selector);
         cheats.prank(defaultOperator);