feat: add execute user op

src/account/AccountStorage.sol

@@ -38,8 +38,6 @@ struct ValidationData {
     bool isDefault;
     // Whether or not this validation is a signature validator.
     bool isSignatureValidation;
-    // How many execution hooks require the UO context.
-    uint8 requireUOHookCount;
     // The pre validation hooks for this function selector.
     EnumerableSet.Bytes32Set preValidationHooks;
     // Permission hooks for this validation function.

src/account/PluginManager2.sol

@@ -7,6 +7,7 @@ import {IPlugin} from "../interfaces/IPlugin.sol";
 import {FunctionReference} from "../interfaces/IPluginManager.sol";
 import {FunctionReferenceLib} from "../helpers/FunctionReferenceLib.sol";
 import {AccountStorage, getAccountStorage, toSetValue, toFunctionReference} from "./AccountStorage.sol";
+import {ExecutionHook} from "../interfaces/IAccountLoupe.sol";
 
 // Temporary additional functions for a user-controlled install flow for validation functions.
 abstract contract PluginManager2 {
@@ -16,13 +17,15 @@ abstract contract PluginManager2 {
     error PreValidationAlreadySet(FunctionReference validationFunction, FunctionReference preValidationFunction);
     error ValidationAlreadySet(bytes4 selector, FunctionReference validationFunction);
     error ValidationNotSet(bytes4 selector, FunctionReference validationFunction);
+    error PermissionAlreadySet(FunctionReference validationFunction, ExecutionHook hook);
 
     function _installValidation(
         FunctionReference validationFunction,
         bool isDefault,
         bytes4[] memory selectors,
         bytes calldata installData,
-        bytes memory preValidationHooks
+        bytes memory preValidationHooks,
+        bytes memory permissionHooks
     )
         // TODO: flag for signature validation
         internal
@@ -51,6 +54,26 @@ abstract contract PluginManager2 {
             }
         }
 
+        if (permissionHooks.length > 0) {
+            (ExecutionHook[] memory permissionFunctions, bytes[] memory initDatas) =
+                abi.decode(permissionHooks, (ExecutionHook[], bytes[]));
+
+            for (uint256 i = 0; i < permissionFunctions.length; ++i) {
+                ExecutionHook memory permissionFunction = permissionFunctions[i];
+
+                if (
+                    !_storage.validationData[validationFunction].permissionHooks.add(toSetValue(permissionFunction))
+                ) {
+                    revert PermissionAlreadySet(validationFunction, permissionFunction);
+                }
+
+                if (initDatas[i].length > 0) {
+                    (address executionPlugin,) = FunctionReferenceLib.unpack(permissionFunction.hookFunction);
+                    IPlugin(executionPlugin).onInstall(initDatas[i]);
+                }
+            }
+        }
+
         if (isDefault) {
             if (_storage.validationData[validationFunction].isDefault) {
                 revert DefaultValidationAlreadySet(validationFunction);
@@ -75,27 +98,44 @@ abstract contract PluginManager2 {
         FunctionReference validationFunction,
         bytes4[] calldata selectors,
         bytes calldata uninstallData,
-        bytes calldata preValidationHookUninstallData
+        bytes calldata preValidationHookUninstallData,
+        bytes calldata permissionHookUninstallData
     ) internal {
         AccountStorage storage _storage = getAccountStorage();
 
         _storage.validationData[validationFunction].isDefault = false;
         _storage.validationData[validationFunction].isSignatureValidation = false;
 
-        bytes[] memory preValidationHookUninstallDatas = abi.decode(preValidationHookUninstallData, (bytes[]));
+        {
+            bytes[] memory preValidationHookUninstallDatas = abi.decode(preValidationHookUninstallData, (bytes[]));
 
-        // Clear pre validation hooks
-        EnumerableSet.Bytes32Set storage preValidationHooks =
-            _storage.validationData[validationFunction].preValidationHooks;
-        while (preValidationHooks.length() > 0) {
-            FunctionReference preValidationFunction = toFunctionReference(preValidationHooks.at(0));
-            preValidationHooks.remove(toSetValue(preValidationFunction));
-            (address preValidationPlugin,) = FunctionReferenceLib.unpack(preValidationFunction);
-            if (preValidationHookUninstallDatas[0].length > 0) {
+            // Clear pre validation hooks
+            EnumerableSet.Bytes32Set storage preValidationHooks =
+                _storage.validationData[validationFunction].preValidationHooks;
+            while (preValidationHooks.length() > 0) {
+                FunctionReference preValidationFunction = toFunctionReference(preValidationHooks.at(0));
+                preValidationHooks.remove(toSetValue(preValidationFunction));
+                (address preValidationPlugin,) = FunctionReferenceLib.unpack(preValidationFunction);
                 IPlugin(preValidationPlugin).onUninstall(preValidationHookUninstallDatas[0]);
             }
         }
 
+        {
+            bytes[] memory permissionHookUninstallDatas = abi.decode(permissionHookUninstallData, (bytes[]));
+
+            // Clear permission hooks
+            EnumerableSet.Bytes32Set storage permissionHooks =
+                _storage.validationData[validationFunction].permissionHooks;
+            while (permissionHooks.length() > 0) {
+                FunctionReference permissionHook = toFunctionReference(permissionHooks.at(0));
+                permissionHooks.remove(toSetValue(permissionHook));
+                (address permissionHookPlugin,) = FunctionReferenceLib.unpack(permissionHook);
+                if (permissionHookUninstallDatas[0].length > 0) {
+                    IPlugin(permissionHookPlugin).onUninstall(permissionHookUninstallDatas[0]);
+                }
+            }
+        }
+
         // Because this function also calls `onUninstall`, and removes the default flag from validation, we must
         // assume these selectors passed in to be exhaustive.
         // TODO: consider enforcing this from user-supplied install config.

src/account/UpgradeableModularAccount.sol

@@ -4,6 +4,7 @@ pragma solidity ^0.8.25;
 import {BaseAccount} from "@eth-infinitism/account-abstraction/core/BaseAccount.sol";
 import {IEntryPoint} from "@eth-infinitism/account-abstraction/interfaces/IEntryPoint.sol";
 import {PackedUserOperation} from "@eth-infinitism/account-abstraction/interfaces/PackedUserOperation.sol";
+import {IAccountExecute} from "@eth-infinitism/account-abstraction/interfaces/IAccountExecute.sol";
 import {UUPSUpgradeable} from "@openzeppelin/contracts/proxy/utils/UUPSUpgradeable.sol";
 import {IERC165} from "@openzeppelin/contracts/utils/introspection/IERC165.sol";
 import {IERC1271} from "@openzeppelin/contracts/interfaces/IERC1271.sol";
@@ -38,6 +39,7 @@ contract UpgradeableModularAccount is
     IERC165,
     IERC1271,
     IStandardExecutor,
+    IAccountExecute,
     PluginManagerInternals,
     PluginManager2,
     UUPSUpgradeable
@@ -66,6 +68,7 @@ contract UpgradeableModularAccount is
     error ExecFromPluginNotPermitted(address plugin, bytes4 selector);
     error ExecFromPluginExternalNotPermitted(address plugin, address target, uint256 value, bytes data);
     error NativeTokenSpendingNotPermitted(address plugin);
+    error NotEntryPoint();
     error PostExecHookReverted(address plugin, uint8 functionId, bytes revertReason);
     error PreExecHookReverted(address plugin, uint8 functionId, bytes revertReason);
     error PreRuntimeValidationHookFailed(address plugin, uint8 functionId, bytes revertReason);
@@ -84,7 +87,7 @@ contract UpgradeableModularAccount is
         _checkPermittedCallerIfNotFromEP();
 
         PostExecToRun[] memory postExecHooks =
-            _doPreExecHooks(getAccountStorage().selectorData[msg.sig].executionHooks, msg.data);
+            _doPreHooks(getAccountStorage().selectorData[msg.sig].executionHooks, msg.data, false);
 
         _;
 
@@ -138,7 +141,7 @@ contract UpgradeableModularAccount is
 
         PostExecToRun[] memory postExecHooks;
         // Cache post-exec hooks in memory
-        postExecHooks = _doPreExecHooks(getAccountStorage().selectorData[msg.sig].executionHooks, msg.data);
+        postExecHooks = _doPreHooks(getAccountStorage().selectorData[msg.sig].executionHooks, msg.data, false);
 
         // execute the function, bubbling up any reverts
         (bool execSuccess, bytes memory execReturnData) = execPlugin.call(msg.data);
@@ -155,6 +158,41 @@ contract UpgradeableModularAccount is
         return execReturnData;
     }
 
+    /// @notice Execution function that allows UO context to be passed to execution hooks
+    /// @dev This function is only callable by the EntryPoint
+    function executeUserOp(PackedUserOperation calldata userOp, bytes32) external {
+        if (msg.sender != address(_ENTRY_POINT)) {
+            revert NotEntryPoint();
+        }
+
+        FunctionReference userOpValidationFunction = FunctionReference.wrap(bytes21(userOp.signature[:21]));
+
+        // remove first 4 bytes which is executeUserOp.selector
+        PackedUserOperation memory uo;
+        bytes memory callData = uo.callData;
+        assembly {
+            let len := mload(callData)
+            mstore(add(callData, 4), sub(len, 4))
+            callData := add(callData, 4)
+        }
+        uo.callData = callData;
+
+        PostExecToRun[] memory postPermissionHooks = _doPreHooks(
+            getAccountStorage().validationData[userOpValidationFunction].permissionHooks, abi.encode(uo), true
+        );
+
+        (bool success, bytes memory result) = address(this).call(userOp.callData[4:]);
+
+        if (!success) {
+            // Directly bubble up revert messages
+            assembly ("memory-safe") {
+                revert(add(result, 32), mload(result))
+            }
+        }
+
+        _doCachedPostExecHooks(postPermissionHooks);
+    }
+
     /// @inheritdoc IStandardExecutor
     /// @notice May be validated by a default validation.
     function execute(address target, uint256 value, bytes calldata data)
@@ -201,8 +239,11 @@ contract UpgradeableModularAccount is
 
         _doRuntimeValidation(runtimeValidationFunction, data, authorization[22:]);
 
-        // If runtime validation passes, execute the call
+        // If runtime validation passes, do runtime permission checks
+        PostExecToRun[] memory postPermissionHooks =
+            _doPreHooks(getAccountStorage().validationData[runtimeValidationFunction].permissionHooks, data, false);
 
+        // Execute the call
         (bool success, bytes memory returnData) = address(this).call(data);
 
         if (!success) {
@@ -211,6 +252,8 @@ contract UpgradeableModularAccount is
             }
         }
 
+        _doCachedPostExecHooks(postPermissionHooks);
+
         return returnData;
     }
 
@@ -252,7 +295,7 @@ contract UpgradeableModularAccount is
         external
         initializer
     {
-        _installValidation(validationFunction, true, new bytes4[](0), installData, bytes(""));
+        _installValidation(validationFunction, true, new bytes4[](0), installData, bytes(""), bytes(""));
         emit ModularAccountInitialized(_ENTRY_POINT);
     }
 
@@ -263,9 +306,12 @@ contract UpgradeableModularAccount is
         bool isDefault,
         bytes4[] memory selectors,
         bytes calldata installData,
-        bytes calldata preValidationHooks
+        bytes calldata preValidationHooks,
+        bytes calldata permissionHooks
     ) external wrapNativeFunction {
-        _installValidation(validationFunction, isDefault, selectors, installData, preValidationHooks);
+        _installValidation(
+            validationFunction, isDefault, selectors, installData, preValidationHooks, permissionHooks
+        );
     }
 
     /// @inheritdoc IPluginManager
@@ -274,9 +320,16 @@ contract UpgradeableModularAccount is
         FunctionReference validationFunction,
         bytes4[] calldata selectors,
         bytes calldata uninstallData,
-        bytes calldata preValidationHookUninstallData
+        bytes calldata preValidationHookUninstallData,
+        bytes calldata permissionHookUninstallData
     ) external wrapNativeFunction {
-        _uninstallValidation(validationFunction, selectors, uninstallData, preValidationHookUninstallData);
+        _uninstallValidation(
+            validationFunction,
+            selectors,
+            uninstallData,
+            preValidationHookUninstallData,
+            permissionHookUninstallData
+        );
     }
 
     /// @notice ERC165 introspection
@@ -344,21 +397,24 @@ contract UpgradeableModularAccount is
             revert UnrecognizedFunction(bytes4(userOp.callData));
         }
         bytes4 selector = bytes4(userOp.callData);
+        if (selector == this.executeUserOp.selector) {
+            selector = bytes4(userOp.callData[4:8]);
+        }
 
         // Revert if the provided `authorization` less than 21 bytes long, rather than right-padding.
         FunctionReference userOpValidationFunction = FunctionReference.wrap(bytes21(userOp.signature[:21]));
         bool isDefaultValidation = uint8(userOp.signature[21]) == 1;
 
         _checkIfValidationApplies(selector, userOpValidationFunction, isDefaultValidation);
 
-        // Check if there are exec hooks associated with the validator that require UO context, and revert if the
-        // call isn't to `executeUserOp`
-        // This check must be here because if context isn't passed, we wouldn't be able to get the exec hooks
-        // associated with the validator
-        if (getAccountStorage().validationData[userOpValidationFunction].requireUOHookCount > 0) {
-            /**
-             * && msg.sig != this.executeUserOp.selector
-             */
+        // Check if there are permission hooks associated with the validator, and revert if the call isn't to
+        // `executeUserOp`
+        // This check must be here because if context isn't passed, we can't tell in execution which hooks should
+        // have ran
+        if (
+            getAccountStorage().validationData[userOpValidationFunction].permissionHooks.length() > 0
+                && bytes4(userOp.callData[:4]) != this.executeUserOp.selector
+        ) {
             revert RequireUserOperationContext();
         }
 
@@ -453,12 +509,11 @@ contract UpgradeableModularAccount is
         }
     }
 
-    function _doPreExecHooks(EnumerableSet.Bytes32Set storage executionHooks, bytes calldata data)
+    function _doPreHooks(EnumerableSet.Bytes32Set storage executionHooks, bytes memory data, bool isPackedUO)
         internal
         returns (PostExecToRun[] memory postHooksToRun)
     {
         uint256 hooksLength = executionHooks.length();
-
         // Overallocate on length - not all of this may get filled up. We set the correct length later.
         postHooksToRun = new PostExecToRun[](hooksLength);
 
@@ -479,7 +534,14 @@ contract UpgradeableModularAccount is
             (FunctionReference hookFunction, bool isPreHook, bool isPostHook) = toExecutionHook(key);
 
             if (isPreHook) {
-                bytes memory preExecHookReturnData = _runPreExecHook(hookFunction, data);
+                bytes memory preExecHookReturnData;
+
+                // isPackedUO implies it is a permission hook
+                if (isPackedUO) {
+                    preExecHookReturnData = _runPreUserOpExecHook(hookFunction, data);
+                } else {
+                    preExecHookReturnData = _runPreExecHook(hookFunction, data);
+                }
 
                 // If there is an associated post-exec hook, save the return data.
                 if (isPostHook) {
@@ -489,7 +551,22 @@ contract UpgradeableModularAccount is
         }
     }
 
-    function _runPreExecHook(FunctionReference preExecHook, bytes calldata data)
+    function _runPreUserOpExecHook(FunctionReference preExecHook, bytes memory data)
+        internal
+        returns (bytes memory preExecHookReturnData)
+    {
+        (address plugin, uint8 functionId) = preExecHook.unpack();
+        try IExecutionHook(plugin).preExecutionHook(functionId, msg.sender, msg.value, data) returns (
+            bytes memory returnData
+        ) {
+            preExecHookReturnData = returnData;
+        } catch (bytes memory revertReason) {
+            // TODO: same issue with EP0.6 - we can't do bytes4 error codes in plugins
+            revert PreExecHookReverted(plugin, functionId, revertReason);
+        }
+    }
+
+    function _runPreExecHook(FunctionReference preExecHook, bytes memory data)
         internal
         returns (bytes memory preExecHookReturnData)
     {
@@ -499,6 +576,7 @@ contract UpgradeableModularAccount is
         ) {
             preExecHookReturnData = returnData;
         } catch (bytes memory revertReason) {
+            // TODO: same issue with EP0.6 - we can't do bytes4 error codes in plugins
             revert PreExecHookReverted(plugin, functionId, revertReason);
         }
     }

src/interfaces/IPermissionHook.sol

@@ -0,0 +1,18 @@
+// SPDX-License-Identifier: CC0-1.0
+pragma solidity ^0.8.25;
+
+import {PackedUserOperation} from "@eth-infinitism/account-abstraction/interfaces/PackedUserOperation.sol";
+import {IExecutionHook} from "./IExecutionHook.sol";
+
+interface IPermissionHook is IExecutionHook {
+    /// @notice Run the pre execution permission hook specified by the `functionId`, passing in the whole user
+    /// operation.
+    /// @dev To indicate the entire call should revert, the function MUST revert.
+    /// @param functionId An identifier that routes the call to different internal implementations, should there be
+    /// more than one.
+    /// @param uo The packed user operation
+    /// @return Context to pass to a post execution hook, if present. An empty bytes array MAY be returned.
+    function preUserOpExecutionHook(uint8 functionId, PackedUserOperation calldata uo)
+        external
+        returns (bytes memory);
+}