fix: check if rollup chain ID is valid

alt-research · May 16, 2024 · 5c8e488 · 5c8e488
1 parent d64fa64
commit 5c8e488
Show file tree

Hide file tree

Showing 2 changed files with 50 additions and 4 deletions.
diff --git a/contracts/src/core/MachServiceManager.sol b/contracts/src/core/MachServiceManager.sol
@@ -74,6 +74,14 @@ contract MachServiceManager is
         _;
     }
 
+    /// @notice when applied to a function, ensures that the `rollupChainID` is valid.
+    modifier onlyValidRollupChainID(uint256 rollupChainID) {
+        if (!rollupChainIDs[rollupChainID]) {
+            revert InvalidRollupChainID();
+        }
+        _;
+    }
+
     constructor(
         IAVSDirectory __avsDirectory,
         IRegistryCoordinator __registryCoordinator,
@@ -179,7 +187,11 @@ contract MachServiceManager is
      * @notice Remove an Alert.
      * @param messageHash The message hash of the alert
      */
-    function removeAlert(uint256 rollupChainId, bytes32 messageHash) external onlyOwner {
+    function removeAlert(uint256 rollupChainId, bytes32 messageHash)
+        external
+        onlyValidRollupChainID(rollupChainId)
+        onlyOwner
+    {
         bool ret = _messageHashes[rollupChainId].remove(messageHash);
         if (ret) {
             _resolvedMessageHashes[rollupChainId].add(messageHash);
@@ -257,7 +269,7 @@ contract MachServiceManager is
         uint256 rollupChainId,
         AlertHeader calldata alertHeader,
         NonSignerStakesAndSignature memory nonSignerStakesAndSignature
-    ) external whenNotPaused onlyAlertConfirmer {
+    ) external whenNotPaused onlyAlertConfirmer onlyValidRollupChainID(rollupChainId) {
         // make sure the information needed to derive the non-signers and batch is in calldata to avoid emitting events
         if (tx.origin != msg.sender) {
             revert InvalidSender();

diff --git a/contracts/test/MachServiceManager.t.sol b/contracts/test/MachServiceManager.t.sol
@@ -520,7 +520,33 @@ contract MachServiceManagerTest is BLSAVSDeployer {
         vm.stopPrank();
     }
 
-    function test_removeAlert() public {
+    function test_confirmAlert_RevertIfInvalidRollupChainID() public {
+        vm.startPrank(proxyAdminOwner);
+        serviceManager.disableAllowlist();
+        vm.stopPrank();
+
+        (
+            uint32 referenceBlockNumber,
+            BLSSignatureChecker.NonSignerStakesAndSignature memory nonSignerStakesAndSignature
+        ) = _registerSignatoriesAndGetNonSignerStakeAndSignatureRandom(nonRandomNumber, 1, quorumBitmap);
+
+        bytes memory quorumThresholdPercentages = new bytes(1);
+        quorumThresholdPercentages[0] = bytes1(uint8(67));
+
+        IMachServiceManager.AlertHeader memory alertHeader = IMachServiceManager.AlertHeader({
+            messageHash: "foo",
+            quorumNumbers: quorumNumbers,
+            quorumThresholdPercentages: quorumThresholdPercentages,
+            referenceBlockNumber: referenceBlockNumber
+        });
+
+        vm.startPrank(proxyAdminOwner);
+        vm.expectRevert(InvalidRollupChainID.selector);
+        serviceManager.confirmAlert(99, alertHeader, nonSignerStakesAndSignature);
+        vm.stopPrank();
+    }
+
+    function test_RemoveAlert() public {
         test_confirmAlert();
         vm.startPrank(proxyAdminOwner);
         assertEq(serviceManager.totalAlerts(1), 1);
@@ -535,7 +561,15 @@ contract MachServiceManagerTest is BLSAVSDeployer {
         vm.stopPrank();
     }
 
-    function test_removeAlert_RevertIfNotOwner() public {
+    function test_RemoveAlert_RevertIfInvalidRollupChainID() public {
+        test_confirmAlert();
+        vm.startPrank(proxyAdminOwner);
+        vm.expectRevert(InvalidRollupChainID.selector);
+        serviceManager.removeAlert(42, "foo");
+        vm.stopPrank();
+    }
+
+    function test_RemoveAlert_RevertIfNotOwner() public {
         test_confirmAlert();
         vm.expectRevert("Ownable: caller is not the owner");
         serviceManager.removeAlert(1, "foo");