Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[3/n permissions] feat: add execute user op #69

Merged
merged 4 commits into from
Jun 19, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
68 changes: 55 additions & 13 deletions src/account/PluginManager2.sol
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";

abstract contract PluginManager2 {
using EnumerableSet for EnumerableSet.Bytes32Set;
Expand All @@ -15,13 +16,15 @@ abstract contract PluginManager2 {
error ValidationAlreadySet(bytes4 selector, FunctionReference validationFunction);
error PreValidationAlreadySet(FunctionReference validationFunction, FunctionReference preValidationFunction);
error ValidationNotSet(bytes4 selector, FunctionReference validationFunction);
error PermissionAlreadySet(FunctionReference validationFunction, ExecutionHook hook);

function _installValidation(
FunctionReference validationFunction,
bool shared,
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 @@ -50,6 +53,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 (shared) {
if (_storage.validationData[validationFunction].isShared) {
revert DefaultValidationAlreadySet(validationFunction);
Expand All @@ -74,24 +97,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].isShared = 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
112 changes: 86 additions & 26 deletions src/account/UpgradeableModularAccount.sol
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
function execute(address target, uint256 value, bytes calldata data)
external
Expand Down Expand Up @@ -198,8 +231,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 @@ -208,6 +244,8 @@ contract UpgradeableModularAccount is
}
}

_doCachedPostExecHooks(postPermissionHooks);

return returnData;
}

Expand Down Expand Up @@ -247,7 +285,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 @@ -257,19 +295,27 @@ contract UpgradeableModularAccount is
bool shared,
bytes4[] memory selectors,
bytes calldata installData,
bytes calldata preValidationHooks
bytes calldata preValidationHooks,
bytes calldata permissionHooks
) external wrapNativeFunction {
_installValidation(validationFunction, shared, selectors, installData, preValidationHooks);
_installValidation(validationFunction, shared, selectors, installData, preValidationHooks, permissionHooks);
}

/// @inheritdoc IPluginManager
function uninstallValidation(
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 @@ -336,6 +382,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]);
}

FunctionReference userOpValidationFunction = FunctionReference.wrap(bytes21(userOp.signature[:21]));
bool isSharedValidation = uint8(userOp.signature[21]) == 1;
Expand All @@ -346,10 +395,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 @@ -444,26 +493,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 @@ -473,10 +515,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 @@ -486,15 +546,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
Original file line number Diff line number Diff line change
Expand Up @@ -32,12 +32,15 @@ interface IPluginManager {
/// @param shared Whether the validation function is shared across 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 shared,
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