refactor: reuse internal install func

erc6900 · Aug 22, 2024 · 9d8e1ff · 9d8e1ff
1 parent f95611a
commit 9d8e1ff
Show file tree

Hide file tree

Showing 4 changed files with 41 additions and 35 deletions.
diff --git a/src/account/ModuleManagerInternals.sol b/src/account/ModuleManagerInternals.sol
@@ -129,19 +129,17 @@ abstract contract ModuleManagerInternals is IModularAccount {
         hooks.remove(toSetValue(hookConfig));
     }
 
-    function _installExecution(address module, ExecutionManifest calldata manifest, bytes memory moduleInstallData)
-        internal
-    {
+    function _installExecution(
+        address module,
+        ExecutionManifest calldata manifest,
+        bytes calldata moduleInstallData
+    ) internal {
         AccountStorage storage _storage = getAccountStorage();
 
         if (module == address(0)) {
             revert NullModule();
         }
 
-        if (!ERC165Checker.supportsInterface(module, type(IModule).interfaceId)) {
-            revert InterfaceNotSupported(module);
-        }
-
         // Update components according to the manifest.
         uint256 length = manifest.executionFunctions.length;
         for (uint256 i = 0; i < length; ++i) {
@@ -170,18 +168,12 @@ abstract contract ModuleManagerInternals is IModularAccount {
             _storage.supportedIfaces[manifest.interfaceIds[i]] += 1;
         }
 
-        // Initialize the module storage for the account.
-        // solhint-disable-next-line no-empty-blocks
-        try IModule(module).onInstall(moduleInstallData) {}
-        catch {
-            bytes memory revertReason = collectReturnData();
-            revert ModuleInstallCallbackFailed(module, revertReason);
-        }
+        _onInstall(module, moduleInstallData, type(IModule).interfaceId);
 
         emit ExecutionInstalled(module, manifest);
     }
 
-    function _uninstallExecution(address module, ExecutionManifest calldata manifest, bytes memory uninstallData)
+    function _uninstallExecution(address module, ExecutionManifest calldata manifest, bytes calldata uninstallData)
         internal
     {
         AccountStorage storage _storage = getAccountStorage();
@@ -214,29 +206,22 @@ abstract contract ModuleManagerInternals is IModularAccount {
         }
 
         // Clear the module storage for the account.
-        bool onUninstallSuccess = true;
-        // solhint-disable-next-line no-empty-blocks
-        try IModule(module).onUninstall(uninstallData) {}
-        catch {
-            onUninstallSuccess = false;
-        }
+        bool onUninstallSuccess = _onUninstall(module, uninstallData);
 
         emit ExecutionUninstalled(module, onUninstallSuccess, manifest);
     }
 
-    function _onInstall(address module, bytes calldata data) internal {
-        if (data.length > 0) {
-            IModule(module).onInstall(data);
-        }
-    }
-
     function _onInstall(address module, bytes calldata data, bytes4 interfaceId) internal {
         if (data.length > 0) {
             if (!ERC165Checker.supportsInterface(module, interfaceId)) {
                 revert InterfaceNotSupported(module);
             }
-
-            IModule(module).onInstall(data);
+            // solhint-disable-next-line no-empty-blocks
+            try IModule(module).onInstall(data) {}
+            catch {
+                bytes memory revertReason = collectReturnData();
+                revert ModuleInstallCallbackFailed(module, revertReason);
+            }
         }
     }
 

diff --git a/test/account/AccountExecHooks.t.sol b/test/account/AccountExecHooks.t.sol
@@ -162,27 +162,27 @@ contract AccountExecHooksTest is AccountTestBase {
         mockModule1 = new MockModule(_m1);
 
         vm.expectEmit(true, true, true, true);
-        emit ReceivedCall(abi.encodeCall(IModule.onInstall, (bytes(""))), 0);
+        emit ReceivedCall(abi.encodeCall(IModule.onInstall, (bytes("a"))), 0);
         vm.expectEmit(true, true, true, true);
         emit ExecutionInstalled(address(mockModule1), _m1);
 
         vm.startPrank(address(entryPoint));
         account1.installExecution({
             module: address(mockModule1),
             manifest: mockModule1.executionManifest(),
-            moduleInstallData: bytes("")
+            moduleInstallData: bytes("a")
         });
         vm.stopPrank();
     }
 
     function _uninstallExecution(MockModule module) internal {
         vm.expectEmit(true, true, true, true);
-        emit ReceivedCall(abi.encodeCall(IModule.onUninstall, (bytes(""))), 0);
+        emit ReceivedCall(abi.encodeCall(IModule.onUninstall, (bytes("b"))), 0);
         vm.expectEmit(true, true, true, true);
         emit ExecutionUninstalled(address(module), true, module.executionManifest());
 
         vm.startPrank(address(entryPoint));
-        account1.uninstallExecution(address(module), module.executionManifest(), bytes(""));
+        account1.uninstallExecution(address(module), module.executionManifest(), bytes("b"));
         vm.stopPrank();
     }
 }
diff --git a/test/account/DirectCallsFromModule.t.sol b/test/account/DirectCallsFromModule.t.sol
@@ -11,6 +11,8 @@ import {DirectCallModule} from "../mocks/modules/DirectCallModule.sol";
 
 import {AccountTestBase} from "../utils/AccountTestBase.sol";
 
+import {DIRECT_CALL_VALIDATION_ENTITYID} from "../../src/helpers/Constants.sol";
+
 contract DirectCallsFromModuleTest is AccountTestBase {
     using ValidationConfigLib for ValidationConfig;
 
@@ -23,7 +25,7 @@ contract DirectCallsFromModuleTest is AccountTestBase {
         _module = new DirectCallModule();
         assertFalse(_module.preHookRan());
         assertFalse(_module.postHookRan());
-        _moduleEntity = ModuleEntityLib.pack(address(_module), type(uint32).max);
+        _moduleEntity = ModuleEntityLib.pack(address(_module), DIRECT_CALL_VALIDATION_ENTITYID);
     }
 
     /* -------------------------------------------------------------------------- */
@@ -104,6 +106,25 @@ contract DirectCallsFromModuleTest is AccountTestBase {
         account1.execute(address(0), 0, "");
     }
 
+    function test_directCallsFromEOA() external {
+        address extraOwner = makeAddr("extraOwner");
+
+        bytes4[] memory selectors = new bytes4[](1);
+        selectors[0] = IModularAccount.execute.selector;
+
+        vm.prank(address(entryPoint));
+
+        account1.installValidation(
+            ValidationConfigLib.pack(extraOwner, DIRECT_CALL_VALIDATION_ENTITYID, false, false),
+            selectors,
+            "",
+            new bytes[](0)
+        );
+
+        vm.prank(extraOwner);
+        account1.execute(makeAddr("dead"), 0, "");
+    }
+
     /* -------------------------------------------------------------------------- */
     /*                                  Internals                                 */
     /* -------------------------------------------------------------------------- */

diff --git a/test/account/UpgradeableModularAccount.t.sol b/test/account/UpgradeableModularAccount.t.sol
@@ -293,7 +293,7 @@ contract UpgradeableModularAccountTest is AccountTestBase {
 
         ExecutionManifest memory m;
 
-        account1.installExecution({module: address(badModule), manifest: m, moduleInstallData: ""});
+        account1.installExecution({module: address(badModule), manifest: m, moduleInstallData: "a"});
     }
 
     function test_installExecution_alreadyInstalled() public {