Add additional selector checks to modules

zkemail · Aug 1, 2024 · 5b26a9a · 5b26a9a
1 parent 88371b8
commit 5b26a9a
Show file tree

Hide file tree

Showing 5 changed files with 96 additions and 6 deletions.
diff --git a/src/interfaces/ISafe.sol b/src/interfaces/ISafe.sol
@@ -5,4 +5,6 @@ interface ISafe {
     function swapOwner(address prevOwner, address oldOwner, address newOwner) external;
     function isOwner(address owner) external view returns (bool);
     function getOwners() external view returns (address[] memory);
+    function setFallbackHandler(address handler) external;
+    function setGuard(address guard) external;
 }
diff --git a/src/modules/EmailRecoveryModule.sol b/src/modules/EmailRecoveryModule.sol
@@ -4,6 +4,7 @@ pragma solidity ^0.8.25;
 import { ERC7579ExecutorBase } from "@rhinestone/modulekit/src/Modules.sol";
 import { IERC7579Account } from "erc7579/interfaces/IERC7579Account.sol";
 import { IModule } from "erc7579/interfaces/IERC7579Module.sol";
+import { ISafe } from "../interfaces/ISafe.sol";
 import { IEmailRecoveryModule } from "../interfaces/IEmailRecoveryModule.sol";
 import { IEmailRecoveryManager } from "../interfaces/IEmailRecoveryManager.sol";
 
@@ -61,7 +62,9 @@ contract EmailRecoveryModule is ERC7579ExecutorBase, IEmailRecoveryModule {
         }
         if (
             _selector == IModule.onUninstall.selector || _selector == IModule.onInstall.selector
-                || _selector == bytes4(0)
+                || _selector == IERC7579Account.execute.selector
+                || _selector == ISafe.setFallbackHandler.selector
+                || _selector == ISafe.setGuard.selector || _selector == bytes4(0)
         ) {
             revert InvalidSelector(_selector);
         }

diff --git a/src/modules/UniversalEmailRecoveryModule.sol b/src/modules/UniversalEmailRecoveryModule.sol
@@ -4,6 +4,7 @@ pragma solidity ^0.8.25;
 import { ERC7579ExecutorBase } from "@rhinestone/modulekit/src/Modules.sol";
 import { IERC7579Account } from "erc7579/interfaces/IERC7579Account.sol";
 import { IModule } from "erc7579/interfaces/IERC7579Module.sol";
+import { ISafe } from "../interfaces/ISafe.sol";
 import { SentinelListLib, SENTINEL, ZERO_ADDRESS } from "sentinellist/SentinelList.sol";
 import { IUniversalEmailRecoveryModule } from "../interfaces/IUniversalEmailRecoveryModule.sol";
 import { IEmailRecoveryManager } from "../interfaces/IEmailRecoveryManager.sol";
@@ -71,14 +72,15 @@ contract UniversalEmailRecoveryModule is ERC7579ExecutorBase, IUniversalEmailRec
      * @notice Modifier to check whether the selector is safe. Reverts if the selector is for
      * "onInstall" or "onUninstall"
      */
-    modifier withoutUnsafeSelector(bytes4 recoverySelector) {
+    modifier withoutUnsafeSelector(bytes4 selector) {
         if (
-            recoverySelector == IModule.onUninstall.selector
-                || recoverySelector == IModule.onInstall.selector
+            selector == IModule.onUninstall.selector || selector == IModule.onInstall.selector
+                || selector == IERC7579Account.execute.selector
+                || selector == ISafe.setFallbackHandler.selector || selector == ISafe.setGuard.selector
+                || selector == bytes4(0)
         ) {
-            revert InvalidSelector(recoverySelector);
+            revert InvalidSelector(selector);
         }
-
         _;
     }
 

diff --git a/test/unit/modules/EmailRecoveryModule/constructor.t.sol b/test/unit/modules/EmailRecoveryModule/constructor.t.sol
@@ -3,6 +3,8 @@ pragma solidity ^0.8.25;
 
 import { console2 } from "forge-std/console2.sol";
 import { IModule } from "erc7579/interfaces/IERC7579Module.sol";
+import { IERC7579Account } from "erc7579/interfaces/IERC7579Account.sol";
+import { ISafe } from "src/interfaces/ISafe.sol";
 import { EmailRecoveryModuleBase } from "./EmailRecoveryModuleBase.t.sol";
 import { EmailRecoveryModule } from "src/modules/EmailRecoveryModule.sol";
 
@@ -47,6 +49,39 @@ contract EmailRecoveryManager_constructor_Test is EmailRecoveryModuleBase {
         );
     }
 
+    function test_Constructor_RevertWhen_UnsafeExecuteSelector() public {
+        vm.expectRevert(
+            abi.encodeWithSelector(
+                EmailRecoveryModule.InvalidSelector.selector, IERC7579Account.execute.selector
+            )
+        );
+        new EmailRecoveryModule(
+            emailRecoveryManagerAddress, validatorAddress, IERC7579Account.execute.selector
+        );
+    }
+
+    function test_Constructor_RevertWhen_UnsafeSetFallbackHandlerSelector() public {
+        vm.expectRevert(
+            abi.encodeWithSelector(
+                EmailRecoveryModule.InvalidSelector.selector, ISafe.setFallbackHandler.selector
+            )
+        );
+        new EmailRecoveryModule(
+            emailRecoveryManagerAddress, validatorAddress, ISafe.setFallbackHandler.selector
+        );
+    }
+
+    function test_Constructor_RevertWhen_UnsafeSetGuardSelector() public {
+        vm.expectRevert(
+            abi.encodeWithSelector(
+                EmailRecoveryModule.InvalidSelector.selector, ISafe.setGuard.selector
+            )
+        );
+        new EmailRecoveryModule(
+            emailRecoveryManagerAddress, validatorAddress, ISafe.setGuard.selector
+        );
+    }
+
     function test_Constructor_RevertWhen_InvalidSelector() public {
         vm.expectRevert(
             abi.encodeWithSelector(EmailRecoveryModule.InvalidSelector.selector, bytes4(0))

diff --git a/test/unit/modules/UniversalEmailRecoveryModule/allowValidatorRecovery.t.sol b/test/unit/modules/UniversalEmailRecoveryModule/allowValidatorRecovery.t.sol
@@ -5,6 +5,8 @@ import { console2 } from "forge-std/console2.sol";
 import { ModuleKitHelpers } from "modulekit/ModuleKit.sol";
 import { MODULE_TYPE_EXECUTOR, MODULE_TYPE_VALIDATOR } from "modulekit/external/ERC7579.sol";
 import { IModule } from "erc7579/interfaces/IERC7579Module.sol";
+import { IERC7579Account } from "erc7579/interfaces/IERC7579Account.sol";
+import { ISafe } from "src/interfaces/ISafe.sol";
 import { SentinelListLib } from "sentinellist/SentinelList.sol";
 import { OwnableValidator } from "src/test/OwnableValidator.sol";
 import { UniversalEmailRecoveryModule } from "src/modules/UniversalEmailRecoveryModule.sol";
@@ -50,6 +52,52 @@ contract UniversalEmailRecoveryModule_allowValidatorRecovery_Test is UnitBase {
         );
     }
 
+    function test_AllowValidatorRecovery_RevertWhen_UnsafeExecuteSelector() public {
+        vm.expectRevert(
+            abi.encodeWithSelector(
+                UniversalEmailRecoveryModule.InvalidSelector.selector,
+                IERC7579Account.execute.selector
+            )
+        );
+        vm.startPrank(accountAddress);
+        emailRecoveryModule.allowValidatorRecovery(
+            validatorAddress, bytes("0"), IERC7579Account.execute.selector
+        );
+    }
+
+    function test_AllowValidatorRecovery_RevertWhen_UnsafeSetFallbackHandlerSelector() public {
+        vm.expectRevert(
+            abi.encodeWithSelector(
+                UniversalEmailRecoveryModule.InvalidSelector.selector,
+                ISafe.setFallbackHandler.selector
+            )
+        );
+        vm.startPrank(accountAddress);
+        emailRecoveryModule.allowValidatorRecovery(
+            validatorAddress, bytes("0"), ISafe.setFallbackHandler.selector
+        );
+    }
+
+    function test_AllowValidatorRecovery_RevertWhen_UnsafeSetGuardSelector() public {
+        vm.expectRevert(
+            abi.encodeWithSelector(
+                UniversalEmailRecoveryModule.InvalidSelector.selector, ISafe.setGuard.selector
+            )
+        );
+        vm.startPrank(accountAddress);
+        emailRecoveryModule.allowValidatorRecovery(
+            validatorAddress, bytes("0"), ISafe.setGuard.selector
+        );
+    }
+
+    function test_AllowValidatorRecovery_RevertWhen_InvalidSelector() public {
+        vm.expectRevert(
+            abi.encodeWithSelector(UniversalEmailRecoveryModule.InvalidSelector.selector, bytes4(0))
+        );
+        vm.startPrank(accountAddress);
+        emailRecoveryModule.allowValidatorRecovery(validatorAddress, bytes("0"), bytes4(0));
+    }
+
     function test_AllowValidatorRecovery_RevertWhen_InvalidValidator() public {
         OwnableValidator newValidator = new OwnableValidator();
         address newValidatorAddress = address(newValidator);