feat: hook calls in function

rhinestonewtf · Feb 21, 2024 · 830dda3 · 830dda3
1 parent f1456fc
commit 830dda3
Show file tree

Hide file tree

Showing 4 changed files with 48 additions and 60 deletions.
diff --git a/accounts/safe7579/.solhintignore b/accounts/safe7579/.solhintignore
@@ -0,0 +1,2 @@
+node_modules/
+test/
diff --git a/accounts/safe7579/src/SafeERC7579.sol b/accounts/safe7579/src/SafeERC7579.sol
@@ -24,8 +24,6 @@ import {
 } from "@ERC4337/account-abstraction/contracts/core/UserOperationLib.sol";
 import { _packValidationData } from "@ERC4337/account-abstraction/contracts/core/Helpers.sol";
 
-import "forge-std/console2.sol";
-
 /**
  * @title ERC7579 Adapter for Safe accounts.
  * By using Safe's Fallback and Execution modules,
@@ -52,8 +50,8 @@ contract SafeERC7579 is ISafeOp, IERC7579Account, AccessControl, IMSA, HookManag
         external
         payable
         override
+        withHook // ! this modifier has side effects / external calls
         onlyEntryPointOrSelf
-        withHook
     {
         CallType callType = mode.getCallType();
 
@@ -80,7 +78,7 @@ contract SafeERC7579 is ISafeOp, IERC7579Account, AccessControl, IMSA, HookManag
         payable
         override
         onlyExecutorModule
-        withHook
+        withHook // ! this modifier has side effects / external calls
         returns (bytes[] memory returnData)
     {
         CallType callType = mode.getCallType();
@@ -107,7 +105,8 @@ contract SafeERC7579 is ISafeOp, IERC7579Account, AccessControl, IMSA, HookManag
         override
         onlyEntryPointOrSelf
     {
-        revert Unsupported();
+        (bool success, bytes memory ret) = address(this).delegatecall(userOp.callData[4:]);
+        if (!success) revert ExecutionFailed();
     }
 
     /**
@@ -125,11 +124,12 @@ contract SafeERC7579 is ISafeOp, IERC7579Account, AccessControl, IMSA, HookManag
     {
         address validator;
         uint256 nonce = userOp.nonce;
+        // solhint-disable-next-line no-inline-assembly
         assembly {
             validator := shr(96, nonce)
         }
 
-        // check if validator is enabled. If terminate the validation phase.
+        // check if validator is enabled. If not, use Safe's checkSignatures()
         if (!_isValidatorInstalled(validator)) return _validateSignatures(userOp);
 
         // bubble up the return value of the validator module
@@ -236,6 +236,7 @@ contract SafeERC7579 is ISafeOp, IERC7579Account, AccessControl, IMSA, HookManag
         if (moduleTypeId == MODULE_TYPE_VALIDATOR) return true;
         else if (moduleTypeId == MODULE_TYPE_EXECUTOR) return true;
         else if (moduleTypeId == MODULE_TYPE_FALLBACK) return true;
+        else if (moduleTypeId == MODULE_TYPE_HOOK) return true;
         else return false;
     }
 
@@ -355,7 +356,6 @@ contract SafeERC7579 is ISafeOp, IERC7579Account, AccessControl, IMSA, HookManag
         _initModuleManager();
 
         (address bootstrap, bytes memory bootstrapCall) = abi.decode(data, (address, bytes));
-        console2.log("bootstrap: ", bootstrap);
 
         (bool success,) = bootstrap.delegatecall(bootstrapCall);
         if (!success) revert AccountInitializationFailed();

diff --git a/accounts/safe7579/src/core/HookManager.sol b/accounts/safe7579/src/core/HookManager.sol
@@ -1,61 +1,59 @@
 // SPDX-License-Identifier: MIT
 pragma solidity ^0.8.23;
 
-import "./ModuleManager.sol";
-import "erc7579/interfaces/IERC7579Account.sol";
-import "erc7579/interfaces/IERC7579Module.sol";
+import { ModuleManager } from "./ModuleManager.sol";
+import { IHook, IModule } from "erc7579/interfaces/IERC7579Module.sol";
 
 /**
  * @title reference implementation of HookManager
  * @author zeroknots.eth | rhinestone.wtf
  */
 abstract contract HookManager is ModuleManager {
     /// @custom:storage-location erc7201:hookmanager.storage.msa
-    struct HookManagerStorage {
-        IHook _hook;
-    }
-
-    mapping(address smartAccount => HookManagerStorage) private _hookManagerStorage;
 
-    // keccak256("hookmanager.storage.msa");
-    bytes32 constant HOOKMANAGER_STORAGE_LOCATION =
-        0x36e05829dd1b9a4411d96a3549582172d7f071c1c0db5c573fcf94eb28431608;
+    mapping(address smartAccount => address hook) internal $hookManager;
 
     error HookPostCheckFailed();
     error HookAlreadyInstalled(address currentHook);
 
     modifier withHook() {
-        address hook = _getHook(msg.sender);
-        if (hook == address(0)) {
-            _;
-        } else {
-            bytes memory retData = _executeReturnData({
-                safe: msg.sender,
-                target: hook,
-                value: 0,
-                callData: abi.encodeCall(IHook.preCheck, (_msgSender(), msg.data))
-            });
-            bytes memory hookPreContext = abi.decode(retData, (bytes));
+        address hook = $hookManager[msg.sender];
+        bool isHookEnabled = hook != address(0);
+        bytes memory hookPreContext;
+        if (isHookEnabled) hookPreContext = _doPreHook(hook);
+
+        _; // <-- hooked Function Bytecode here
 
-            _;
-            retData = _executeReturnData({
+        if (isHookEnabled) _doPostHook(hook, hookPreContext);
+    }
+
+    function _doPreHook(address hook) internal returns (bytes memory hookPreContext) {
+        hookPreContext = abi.decode(
+            _executeReturnData({
                 safe: msg.sender,
                 target: hook,
                 value: 0,
-                callData: abi.encodeCall(IHook.postCheck, (hookPreContext))
-            });
-            bool success = abi.decode(retData, (bool));
+                callData: abi.encodeCall(IHook.preCheck, (_msgSender(), msg.data))
+            }),
+            (bytes)
+        );
+    }
 
-            if (!success) revert HookPostCheckFailed();
-        }
+    function _doPostHook(address hook, bytes memory hookPreContext) internal {
+        _execute({
+            safe: msg.sender,
+            target: hook,
+            value: 0,
+            callData: abi.encodeCall(IHook.postCheck, (hookPreContext))
+        });
     }
 
     function _setHook(address hook) internal virtual {
-        _hookManagerStorage[msg.sender]._hook = IHook(hook);
+        $hookManager[msg.sender] = hook;
     }
 
     function _installHook(address hook, bytes calldata data) internal virtual {
-        address currentHook = _getHook(msg.sender);
+        address currentHook = $hookManager[msg.sender];
         if (currentHook != address(0)) {
             revert HookAlreadyInstalled(currentHook);
         }
@@ -79,15 +77,11 @@ abstract contract HookManager is ModuleManager {
         });
     }
 
-    function _getHook(address smartAccount) internal view returns (address _hook) {
-        return address(_hookManagerStorage[smartAccount]._hook);
-    }
-
     function _isHookInstalled(address module) internal view returns (bool) {
-        return _getHook(msg.sender) == module;
+        return $hookManager[msg.sender] == module;
     }
 
     function getActiveHook() external view returns (address hook) {
-        return _getHook(msg.sender);
+        return $hookManager[msg.sender];
     }
 }
diff --git a/accounts/safe7579/src/core/ModuleManager.sol b/accounts/safe7579/src/core/ModuleManager.sol
@@ -1,9 +1,9 @@
 // SPDX-License-Identifier: MIT
 pragma solidity ^0.8.23;
 
-import { SentinelListLib, SENTINEL } from "sentinellist/SentinelList.sol";
+import { SentinelListLib } from "sentinellist/SentinelList.sol";
 import { SentinelList4337Lib } from "sentinellist/SentinelList4337.sol";
-import { IModule, IExecutor, IValidator, IFallback } from "erc7579/interfaces/IERC7579Module.sol";
+import { IModule } from "erc7579/interfaces/IERC7579Module.sol";
 import { ExecutionHelper } from "./ExecutionHelper.sol";
 import { Receiver } from "erc7579/core/Receiver.sol";
 import { AccessControl } from "./AccessControl.sol";
@@ -29,11 +29,12 @@ abstract contract ModuleManager is AccessControl, Receiver, ExecutionHelper {
     error CannotRemoveLastValidator();
     error InitializerError();
     error ValidatorStorageHelperError();
+    error NoFallbackHandler();
 
+    mapping(address smartAccount => ModuleManagerStorage moduleManagerStorage) internal
+        $moduleManager;
 
-    mapping(address smartAccount => ModuleManagerStorage) private $moduleManager;
-
-    SentinelList4337Lib.SentinelList $validators;
+    SentinelList4337Lib.SentinelList internal $validators;
 
     modifier onlyExecutorModule() {
         if (!_isExecutorInstalled(_msgSender())) revert InvalidModule(_msgSender());
@@ -42,14 +43,9 @@ abstract contract ModuleManager is AccessControl, Receiver, ExecutionHelper {
 
     /**
      * Initializes linked list that handles installed Validator and Executor
-     * For Validators:
-     *      The Safe Account will call VALIDATOR_STORAGE via DELEGTATECALL.
-     *      Due to the storage restrictions of ERC-4337 of the validation phase,
-     *      Validators are stored within the Safe's account storage.
      */
     function _initModuleManager() internal {
         ModuleManagerStorage storage $mms = $moduleManager[msg.sender];
-
         // this will revert if list is already initialized
         $validators.init({ account: msg.sender });
         $mms._executors.init();
@@ -75,7 +71,6 @@ abstract contract ModuleManager is AccessControl, Receiver, ExecutionHelper {
 
     /**
      * Uninstall and de-initialize validator module
-     * @dev this function Write into the Safe account storage (validator linked) list via
      */
     function _uninstallValidator(address validator, bytes memory data) internal {
         (address prev, bytes memory disableModuleData) = abi.decode(data, (address, bytes));
@@ -104,10 +99,6 @@ abstract contract ModuleManager is AccessControl, Receiver, ExecutionHelper {
         isInstalled = $validators.contains({ account: msg.sender, entry: validator });
     }
 
-    /**
-     * THIS IS NOT PART OF THE STANDARD
-     * Helper Function to access linked list
-     */
     function getValidatorPaginated(
         address start,
         uint256 pageSize
@@ -217,9 +208,10 @@ abstract contract ModuleManager is AccessControl, Receiver, ExecutionHelper {
     }
 
     // FALLBACK
+    // solhint-disable-next-line no-complex-fallback
     fallback() external payable override(Receiver) receiverFallback {
         address handler = _getFallbackHandler();
-        if (handler == address(0)) revert();
+        if (handler == address(0)) revert NoFallbackHandler();
         /* solhint-disable no-inline-assembly */
         /// @solidity memory-safe-assembly
         // solhint-disable-next-line no-inline-assembly