Skip to content

Commit

Permalink
feat: add execute user op
Browse files Browse the repository at this point in the history
  • Loading branch information
@howydev
howydev committed Jun 24, 2024
1 parent c4fa5d4 commit 7346d52
Show file tree
Hide file tree
Showing 5 changed files with 170 additions and 44 deletions.
68 changes: 55 additions & 13 deletions src/account/PluginManager2.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand All @@ -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
Expand Down Expand Up @@ -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);
Expand All @@ -75,24 +98,43 @@ 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[]));

// 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) {
IPlugin(preValidationPlugin).onUninstall(preValidationHookUninstallDatas[0]);
{
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) {
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]);
}
}
}

Expand Down
114 changes: 88 additions & 26 deletions src/account/UpgradeableModularAccount.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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";
Expand Down Expand Up @@ -38,6 +39,7 @@ contract UpgradeableModularAccount is
IERC165,
IERC1271,
IStandardExecutor,
IAccountExecute,
PluginManagerInternals,
PluginManager2,
UUPSUpgradeable
Expand Down Expand Up @@ -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);
Expand All @@ -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);

_;

Expand Down Expand Up @@ -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);
Expand All @@ -155,6 +158,36 @@ 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]));

PostExecToRun[] memory postPermissionHooks = _doPreHooks(
getAccountStorage().validationData[userOpValidationFunction].permissionHooks, abi.encode(userOp), true
);

PostExecToRun[] memory postExecHooks = _doPreHooks(
getAccountStorage().selectorData[bytes4(userOp.callData[4:8])].executionHooks, abi.encode(userOp), 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(postExecHooks);
_doCachedPostExecHooks(postPermissionHooks);
}

/// @inheritdoc IStandardExecutor
/// @notice May be validated by a default validation.
function execute(address target, uint256 value, bytes calldata data)
Expand Down Expand Up @@ -201,8 +234,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) {
Expand All @@ -211,6 +247,8 @@ contract UpgradeableModularAccount is
}
}

_doCachedPostExecHooks(postPermissionHooks);

return returnData;
}

Expand Down Expand Up @@ -252,7 +290,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);
}

Expand All @@ -263,9 +301,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
Expand All @@ -274,9 +315,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
Expand Down Expand Up @@ -344,6 +392,9 @@ 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]));
Expand All @@ -355,10 +406,10 @@ contract UpgradeableModularAccount is
// 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
*/
if (
getAccountStorage().validationData[userOpValidationFunction].requireUOHookCount > 0
&& bytes4(userOp.callData[:4]) != this.executeUserOp.selector
) {
revert RequireUserOperationContext();
}

Expand Down Expand Up @@ -453,26 +504,19 @@ 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);

// 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.
for (uint256 i = 0; i < hooksLength; ++i) {
bytes32 key = executionHooks.at(i);
(FunctionReference hookFunction,, bool isPostHook, bool requireUOContext) = toExecutionHook(key);
if (requireUOContext) {
/**
* && msg.sig != this.executeUserOp.selector
*/
revert RequireUserOperationContext();
}
(FunctionReference hookFunction,, bool isPostHook,) = toExecutionHook(key);
if (isPostHook) {
postHooksToRun[i].postExecHook = hookFunction;
}
Expand All @@ -482,10 +526,28 @@ contract UpgradeableModularAccount is
// exists.
for (uint256 i = 0; i < hooksLength; ++i) {
bytes32 key = executionHooks.at(i);
(FunctionReference hookFunction, bool isPreHook, bool isPostHook,) = toExecutionHook(key);
(FunctionReference hookFunction, bool isPreHook, bool isPostHook, bool requireUOContext) =
toExecutionHook(key);

if (!isPackedUO && requireUOContext) {
revert RequireUserOperationContext();
}

if (isPreHook) {
bytes memory preExecHookReturnData = _runPreExecHook(hookFunction, data);
bytes memory preExecHookReturnData;

if (isPackedUO) {
if (requireUOContext) {
preExecHookReturnData = _runPreExecHook(hookFunction, data);
} else {
PackedUserOperation memory uo = abi.decode(data, (PackedUserOperation));
preExecHookReturnData =
_runPreExecHook(hookFunction, abi.encodePacked(msg.sender, msg.value, uo.callData));
}
} else {
preExecHookReturnData =
_runPreExecHook(hookFunction, abi.encodePacked(msg.sender, msg.value, data));
}

// If there is an associated post-exec hook, save the return data.
if (isPostHook) {
Expand All @@ -495,15 +557,15 @@ contract UpgradeableModularAccount is
}
}

function _runPreExecHook(FunctionReference preExecHook, bytes calldata data)
function _runPreExecHook(FunctionReference preExecHook, bytes memory data)
internal
returns (bytes memory preExecHookReturnData)
{
(address plugin, uint8 functionId) = preExecHook.unpack();
try IExecutionHook(plugin).preExecutionHook(functionId, abi.encodePacked(msg.sender, msg.value, data))
returns (bytes memory returnData) {
try IExecutionHook(plugin).preExecutionHook(functionId, 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);
}
}
Expand Down
11 changes: 9 additions & 2 deletions src/interfaces/IPluginManager.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -32,12 +32,15 @@ interface IPluginManager {
/// @param isDefault Whether the validation function applies for all selectors in the default pool.
/// @param selectors The selectors to install the validation function for.
/// @param installData Optional data to be decoded and used by the plugin to setup initial plugin state.
/// @param preValidationHooks Optional pre-validation hooks to install for the validation function.
/// @param permissionHooks Optional permission hooks to install for the validation function.
function installValidation(
FunctionReference validationFunction,
bool isDefault,
bytes4[] memory selectors,
bytes calldata installData,
bytes calldata preValidationHooks
bytes calldata preValidationHooks,
bytes calldata permissionHooks
) external;

/// @notice Uninstall a validation function from a set of execution selectors.
Expand All @@ -46,11 +49,15 @@ interface IPluginManager {
/// @param selectors The selectors to uninstall the validation function for.
/// @param uninstallData Optional data to be decoded and used by the plugin to clear plugin data for the
/// account.
/// @param preValidationHookUninstallData Optional data to be decoded and used by the plugin to clear account
/// data
/// @param permissionHookUninstallData Optional data to be decoded and used by the plugin to clear account data
function uninstallValidation(
FunctionReference validationFunction,
bytes4[] calldata selectors,
bytes calldata uninstallData,
bytes calldata preValidationHookUninstallData
bytes calldata preValidationHookUninstallData,
bytes calldata permissionHookUninstallData
) external;

/// @notice Uninstall a plugin from the modular account.
Expand Down
Loading

0 comments on commit 7346d52

Please sign in to comment.