feat: add execute user op

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 {EnumerableSet} from "@openzeppelin/contracts/utils/structs/EnumerableSet.sol";
@@ -29,7 +30,8 @@ contract UpgradeableModularAccount is
     IPluginExecutor,
     IStandardExecutor,
     PluginManagerInternals,
-    UUPSUpgradeable
+    UUPSUpgradeable,
+    IAccountExecute
 {
     using EnumerableSet for EnumerableSet.Bytes32Set;
     using FunctionReferenceLib for FunctionReference;
@@ -44,18 +46,22 @@ contract UpgradeableModularAccount is
     // As per the EIP-165 spec, no interface should ever match 0xffffffff
     bytes4 internal constant _INTERFACE_ID_INVALID = 0xffffffff;
     bytes4 internal constant _IERC165_INTERFACE_ID = 0x01ffc9a7;
+    uint8 internal constant _UO_CONTEXT_THRESHOLD = 128;
 
     event ModularAccountInitialized(IEntryPoint indexed entryPoint);
 
     error AlwaysDenyRule();
     error AuthorizeUpgradeReverted(bytes revertReason);
+    error CallerNotEntrypoint();
     error ExecFromPluginNotPermitted(address plugin, bytes4 selector);
     error ExecFromPluginExternalNotPermitted(address plugin, address target, uint256 value, bytes data);
     error InvalidConfiguration();
+    error InvalidContextProvided();
     error NativeTokenSpendingNotPermitted(address plugin);
     error PostExecHookReverted(address plugin, uint8 functionId, bytes revertReason);
     error PreExecHookReverted(address plugin, uint8 functionId, bytes revertReason);
     error PreRuntimeValidationHookFailed(address plugin, uint8 functionId, bytes revertReason);
+    error RequiresUserOperationContext(FunctionReference preExecHook);
     error RuntimeValidationFunctionMissing(bytes4 selector);
     error RuntimeValidationFunctionReverted(address plugin, uint8 functionId, bytes revertReason);
     error UnexpectedAggregator(address plugin, uint8 functionId, address aggregator);
@@ -67,7 +73,7 @@ contract UpgradeableModularAccount is
     modifier wrapNativeFunction() {
         _doRuntimeValidationIfNotFromEP();
 
-        PostExecToRun[] memory postExecHooks = _doPreExecHooks(msg.sig, msg.data);
+        PostExecToRun[] memory postExecHooks = _doPreExecHooks(msg.sig, abi.encode(msg.sender), msg.data);
 
         _;
 
@@ -121,7 +127,7 @@ contract UpgradeableModularAccount is
 
         PostExecToRun[] memory postExecHooks;
         // Cache post-exec hooks in memory
-        postExecHooks = _doPreExecHooks(msg.sig, msg.data);
+        postExecHooks = _doPreExecHooks(msg.sig, abi.encode(msg.sender), msg.data);
 
         // execute the function, bubbling up any reverts
         (bool execSuccess, bytes memory execReturnData) = execPlugin.call(msg.data);
@@ -165,6 +171,29 @@ contract UpgradeableModularAccount is
         }
     }
 
+    /// @inheritdoc IAccountExecute
+    /// @dev userOp.calldata = bytes4(executeUserOp.selector) || bytes callData || bytes
+    /// senderContext
+    function executeUserOp(PackedUserOperation calldata userOp, bytes32) external {
+        if (msg.sender != address(_ENTRY_POINT)) {
+            revert CallerNotEntrypoint();
+        }
+        (bytes memory callData, bytes memory senderContext) = abi.decode(userOp.callData, (bytes, bytes));
+
+        PostExecToRun[] memory postExecHooks = _doPreExecHooks(bytes4(callData), senderContext, callData);
+
+        (bool success, bytes memory result) = address(this).call(callData);
+
+        if (!success) {
+            // Directly bubble up revert messages
+            assembly ("memory-safe") {
+                revert(add(result, 32), mload(result))
+            }
+        }
+
+        _doCachedPostExecHooks(postExecHooks);
+    }
+
     /// @inheritdoc IPluginExecutor
     function executeFromPlugin(bytes calldata data) external payable override returns (bytes memory) {
         bytes4 selector = bytes4(data[:4]);
@@ -184,7 +213,7 @@ contract UpgradeableModularAccount is
             revert UnrecognizedFunction(selector);
         }
 
-        PostExecToRun[] memory postExecHooks = _doPreExecHooks(selector, data);
+        PostExecToRun[] memory postExecHooks = _doPreExecHooks(selector, abi.encode(msg.sender), data);
 
         (bool success, bytes memory returnData) = execFunctionPlugin.call(data);
 
@@ -237,7 +266,7 @@ contract UpgradeableModularAccount is
 
         // Run any pre exec hooks for this selector
         PostExecToRun[] memory postExecHooks =
-            _doPreExecHooks(IPluginExecutor.executeFromPluginExternal.selector, msg.data);
+            _doPreExecHooks(IPluginExecutor.executeFromPluginExternal.selector, abi.encode(msg.sender), msg.data);
 
         // Perform the external call
         bytes memory returnData = _exec(target, value, data);
@@ -329,7 +358,22 @@ contract UpgradeableModularAccount is
 
         FunctionReference userOpValidationFunction = getAccountStorage().selectorData[selector].validation;
 
-        validationData = _doUserOpValidation(selector, userOpValidationFunction, userOp, userOpHash);
+        if (selector == this.executeUserOp.selector) {
+            // if executeUserOp is used, parse out the right calldata and verify senderContext
+            bytes memory validationContextData;
+            // userOp.calldata = bytes4(executeUserOp.selector) || bytes callData || bytes senderContext
+            // TODO: need to modify userOp.calldata here, it should point to `bytes callData` within
+            // userop.callData
+            (validationContextData, validationData) =
+                _doUserOpValidation(bytes4(userOp.callData[36:40]), userOpValidationFunction, userOp, userOpHash);
+
+            (, bytes memory sender) = abi.decode(userOp.callData[4:], (bytes, bytes));
+            if (keccak256(abi.encode(validationContextData)) != keccak256(abi.encode(sender))) {
+                revert InvalidContextProvided();
+            }
+        } else {
+            (, validationData) = _doUserOpValidation(selector, userOpValidationFunction, userOp, userOpHash);
+        }
     }
 
     // To support gas estimation, we don't fail early when the failure is caused by a signature failure
@@ -338,7 +382,7 @@ contract UpgradeableModularAccount is
         FunctionReference userOpValidationFunction,
         PackedUserOperation calldata userOp,
         bytes32 userOpHash
-    ) internal returns (uint256 validationData) {
+    ) internal returns (bytes memory contextData, uint256 validationData) {
         if (userOpValidationFunction.isEmptyOrMagicValue()) {
             // If the validation function is empty, then the call cannot proceed.
             // Alternatively, the validation function may be set to the RUNTIME_VALIDATION_ALWAYS_ALLOW magic
@@ -370,7 +414,8 @@ contract UpgradeableModularAccount is
         // Run the user op validationFunction
         {
             (address plugin, uint8 functionId) = userOpValidationFunction.unpack();
-            currentValidationData = IPlugin(plugin).userOpValidationFunction(functionId, userOp, userOpHash);
+            (contextData, currentValidationData) =
+                IPlugin(plugin).userOpValidationFunction(functionId, userOp, userOpHash);
 
             if (preUserOpValidationHooksLength != 0) {
                 // If we have other validation data we need to coalesce with
@@ -426,7 +471,7 @@ contract UpgradeableModularAccount is
         }
     }
 
-    function _doPreExecHooks(bytes4 selector, bytes calldata data)
+    function _doPreExecHooks(bytes4 selector, bytes memory senderContext, bytes memory data)
         internal
         returns (PostExecToRun[] memory postHooksToRun)
     {
@@ -436,7 +481,7 @@ contract UpgradeableModularAccount is
         uint256 postOnlyHooksLength = selectorData.postOnlyHooks.length();
 
         // Overallocate on length - not all of this may get filled up. We set the correct length later.
-        postHooksToRun = new PostExecToRun[](preExecHooksLength + postOnlyHooksLength);
+        postHooksToRun = new PostExecToRun[](selectorData.preHooks.length() + selectorData.postOnlyHooks.length());
 
         // Copy all post hooks to the array. This happens before any pre hooks are run, so we can
         // be sure that the set of hooks to run will not be affected by state changes mid-execution.
@@ -461,10 +506,29 @@ contract UpgradeableModularAccount is
         // Run the pre hooks and copy their return data to the post hooks array, if an associated post-exec hook
         // exists.
         for (uint256 i = 0; i < preExecHooksLength; ++i) {
-            bytes32 key = selectorData.preHooks.at(i);
-            FunctionReference preExecHook = _toFunctionReference(key);
+            address plugin;
+            uint8 functionId;
+            {
+                bytes32 key = selectorData.preHooks.at(i);
+                FunctionReference preExecHook = _toFunctionReference(key);
+                (plugin, functionId) = preExecHook.unpack();
+
+                // functionId > _UO_CONTEXT_THRESHOLD means that the plugin wants UO context
+                if (msg.sender == address(_ENTRY_POINT) && functionId > _UO_CONTEXT_THRESHOLD) {
+                    if (abi.decode(senderContext, (address)) == msg.sender) {
+                        revert RequiresUserOperationContext(preExecHook);
+                    }
+                }
+            }
 
-            bytes memory preExecHookReturnData = _runPreExecHook(preExecHook, data);
+            bytes memory preExecHookReturnData;
+            try IPlugin(plugin).preExecutionHook(functionId, senderContext, msg.value, data) returns (
+                bytes memory returnData
+            ) {
+                preExecHookReturnData = returnData;
+            } catch (bytes memory revertReason) {
+                revert PreExecHookReverted(plugin, functionId, revertReason);
+            }
 
             // If there is an associated post-exec hook, save the return data.
             PostExecToRun memory postExecToRun = postHooksToRun[i + postOnlyHooksLength];
@@ -474,20 +538,6 @@ contract UpgradeableModularAccount is
         }
     }
 
-    function _runPreExecHook(FunctionReference preExecHook, bytes calldata data)
-        internal
-        returns (bytes memory preExecHookReturnData)
-    {
-        (address plugin, uint8 functionId) = preExecHook.unpack();
-        try IPlugin(plugin).preExecutionHook(functionId, msg.sender, msg.value, data) returns (
-            bytes memory returnData
-        ) {
-            preExecHookReturnData = returnData;
-        } catch (bytes memory revertReason) {
-            revert PreExecHookReverted(plugin, functionId, revertReason);
-        }
-    }
-
     /// @dev Associated post hooks are run in reverse order of their pre hooks.
     function _doCachedPostExecHooks(PostExecToRun[] memory postHooksToRun) internal {
         uint256 postHooksToRunLength = postHooksToRun.length;

src/interfaces/IPlugin.sol

@@ -126,7 +126,7 @@ interface IPlugin {
     /// @return Packed validation data for validAfter (6 bytes), validUntil (6 bytes), and authorizer (20 bytes).
     function userOpValidationFunction(uint8 functionId, PackedUserOperation calldata userOp, bytes32 userOpHash)
         external
-        returns (uint256);
+        returns (bytes memory, uint256);
 
     /// @notice Run the pre runtime validation hook specified by the `functionId`.
     /// @dev To indicate the entire call should revert, the function MUST revert.
@@ -152,11 +152,11 @@ interface IPlugin {
     /// @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 sender The caller address.
+    /// @param senderContext The caller address.
     /// @param value The call value.
     /// @param data The calldata sent.
     /// @return Context to pass to a post execution hook, if present. An empty bytes array MAY be returned.
-    function preExecutionHook(uint8 functionId, address sender, uint256 value, bytes calldata data)
+    function preExecutionHook(uint8 functionId, bytes calldata senderContext, uint256 value, bytes calldata data)
         external
         returns (bytes memory);
 

src/plugins/BasePlugin.sol

@@ -56,7 +56,7 @@ abstract contract BasePlugin is ERC165, IPlugin {
     function userOpValidationFunction(uint8 functionId, PackedUserOperation calldata userOp, bytes32 userOpHash)
         external
         virtual
-        returns (uint256)
+        returns (bytes memory, uint256)
     {
         (functionId, userOp, userOpHash);
         revert NotImplemented();
@@ -96,16 +96,16 @@ abstract contract BasePlugin is ERC165, IPlugin {
     /// @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 sender The caller address.
+    /// @param senderContext Senders context. It'll be an address except for some UOs
     /// @param value The call value.
     /// @param data The calldata sent.
     /// @return Context to pass to a post execution hook, if present. An empty bytes array MAY be returned.
-    function preExecutionHook(uint8 functionId, address sender, uint256 value, bytes calldata data)
+    function preExecutionHook(uint8 functionId, bytes calldata senderContext, uint256 value, bytes calldata data)
         external
         virtual
         returns (bytes memory)
     {
-        (functionId, sender, value, data);
+        (functionId, senderContext, value, data);
         revert NotImplemented();
     }
 

src/plugins/owner/SingleOwnerPlugin.sol

@@ -97,15 +97,15 @@ contract SingleOwnerPlugin is BasePlugin, ISingleOwnerPlugin, IERC1271 {
         external
         view
         override
-        returns (uint256)
+        returns (bytes memory, uint256)
     {
         if (functionId == uint8(FunctionId.VALIDATION_OWNER_OR_SELF)) {
             // Validate the user op signature against the owner.
             (address signer,,) = (userOpHash.toEthSignedMessageHash()).tryRecover(userOp.signature);
             if (signer == address(0) || signer != _owners[msg.sender]) {
-                return _SIG_VALIDATION_FAILED;
+                return ("", _SIG_VALIDATION_FAILED);
             }
-            return _SIG_VALIDATION_PASSED;
+            return (abi.encode(signer), _SIG_VALIDATION_PASSED);
         }
         revert NotImplemented();
     }

test/account/AccountExecHooks.t.sol

@@ -75,7 +75,7 @@ contract AccountExecHooksTest is AccountTestBase {
             abi.encodeWithSelector(
                 IPlugin.preExecutionHook.selector,
                 _PRE_HOOK_FUNCTION_ID_1,
-                address(this), // caller
+                abi.encode(address(this)), // caller context
                 0, // msg.value in call to account
                 abi.encodeWithSelector(_EXEC_SELECTOR)
             ),
@@ -119,7 +119,7 @@ contract AccountExecHooksTest is AccountTestBase {
             abi.encodeWithSelector(
                 IPlugin.preExecutionHook.selector,
                 _PRE_HOOK_FUNCTION_ID_1,
-                address(this), // caller
+                abi.encode(address(this)), // caller context
                 0, // msg.value in call to account
                 abi.encodeWithSelector(_EXEC_SELECTOR)
             ),

test/mocks/plugins/ComprehensivePlugin.sol

@@ -61,10 +61,10 @@ contract ComprehensivePlugin is BasePlugin {
         external
         pure
         override
-        returns (uint256)
+        returns (bytes memory, uint256)
     {
         if (functionId == uint8(FunctionId.VALIDATION)) {
-            return 0;
+            return ("", 0);
         }
         revert NotImplemented();
     }
@@ -89,7 +89,7 @@ contract ComprehensivePlugin is BasePlugin {
         revert NotImplemented();
     }
 
-    function preExecutionHook(uint8 functionId, address, uint256, bytes calldata)
+    function preExecutionHook(uint8 functionId, bytes calldata, uint256, bytes calldata)
         external
         pure
         override

test/mocks/plugins/ValidationPluginMocks.sol

@@ -48,10 +48,10 @@ abstract contract MockBaseUserOpValidationPlugin is BaseTestPlugin {
         external
         view
         override
-        returns (uint256)
+        returns (bytes memory, uint256)
     {
         if (functionId == uint8(FunctionId.USER_OP_VALIDATION)) {
-            return _userOpValidationFunctionData;
+            return ("", _userOpValidationFunctionData);
         }
         revert NotImplemented();
     }

test/plugin/SingleOwnerPlugin.t.sol

@@ -137,7 +137,7 @@ contract SingleOwnerPluginTest is OptimizedTest {
         userOp.signature = abi.encodePacked(r, s, v);
 
         // sig check should fail
-        uint256 success = plugin.userOpValidationFunction(
+        (, uint256 success) = plugin.userOpValidationFunction(
             uint8(ISingleOwnerPlugin.FunctionId.VALIDATION_OWNER_OR_SELF), userOp, userOpHash
         );
         assertEq(success, 1);
@@ -147,7 +147,7 @@ contract SingleOwnerPluginTest is OptimizedTest {
         assertEq(signer, plugin.owner());
 
         // sig check should pass
-        success = plugin.userOpValidationFunction(
+        (, success) = plugin.userOpValidationFunction(
             uint8(ISingleOwnerPlugin.FunctionId.VALIDATION_OWNER_OR_SELF), userOp, userOpHash
         );
         assertEq(success, 0);