From eb191850f70300ef52a72cf80e7e3523b660756d Mon Sep 17 00:00:00 2001 From: howydev <132113803+howydev@users.noreply.github.com> Date: Mon, 17 Jun 2024 19:31:35 -0400 Subject: [PATCH] feat: permission hooks installation via installValidation --- src/account/PluginManager2.sol | 68 ++++++++++++++++++----- src/account/UpgradeableModularAccount.sol | 29 ++++++++-- src/interfaces/IPluginManager.sol | 11 +++- test/account/AccountLoupe.t.sol | 7 ++- test/account/ValidationIntersection.t.sol | 14 ++++- 5 files changed, 105 insertions(+), 24 deletions(-) diff --git a/src/account/PluginManager2.sol b/src/account/PluginManager2.sol index 1b1745b0..afa2963e 100644 --- a/src/account/PluginManager2.sol +++ b/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"; abstract contract PluginManager2 { using EnumerableSet for EnumerableSet.Bytes32Set; @@ -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 @@ -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); @@ -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]); + } } } diff --git a/src/account/UpgradeableModularAccount.sol b/src/account/UpgradeableModularAccount.sol index 8c8aaf90..aa74ecd2 100644 --- a/src/account/UpgradeableModularAccount.sol +++ b/src/account/UpgradeableModularAccount.sol @@ -231,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) { @@ -241,6 +244,8 @@ contract UpgradeableModularAccount is } } + _doCachedPostExecHooks(postPermissionHooks); + return returnData; } @@ -280,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); } @@ -290,9 +295,10 @@ 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 @@ -300,9 +306,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 @@ -369,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; @@ -538,6 +554,7 @@ contract UpgradeableModularAccount is 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); } } diff --git a/src/interfaces/IPluginManager.sol b/src/interfaces/IPluginManager.sol index f3809211..f7e7220b 100644 --- a/src/interfaces/IPluginManager.sol +++ b/src/interfaces/IPluginManager.sol @@ -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. @@ -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. diff --git a/test/account/AccountLoupe.t.sol b/test/account/AccountLoupe.t.sol index e21f232f..aebc65af 100644 --- a/test/account/AccountLoupe.t.sol +++ b/test/account/AccountLoupe.t.sol @@ -43,7 +43,12 @@ contract AccountLoupeTest is AccountTestBase { bytes[] memory installDatas = new bytes[](2); vm.prank(address(entryPoint)); account1.installValidation( - ownerValidation, true, new bytes4[](0), bytes(""), abi.encode(preValidationHooks, installDatas) + ownerValidation, + true, + new bytes4[](0), + bytes(""), + abi.encode(preValidationHooks, installDatas), + bytes("") ); } diff --git a/test/account/ValidationIntersection.t.sol b/test/account/ValidationIntersection.t.sol index 877f8ff9..6f451a16 100644 --- a/test/account/ValidationIntersection.t.sol +++ b/test/account/ValidationIntersection.t.sol @@ -67,7 +67,12 @@ contract ValidationIntersectionTest is AccountTestBase { }); bytes[] memory installDatas = new bytes[](1); account1.installValidation( - oneHookValidation, true, new bytes4[](0), bytes(""), abi.encode(preValidationHooks, installDatas) + oneHookValidation, + true, + new bytes4[](0), + bytes(""), + abi.encode(preValidationHooks, installDatas), + bytes("") ); account1.installPlugin({ plugin: address(twoHookPlugin), @@ -87,7 +92,12 @@ contract ValidationIntersectionTest is AccountTestBase { }); installDatas = new bytes[](2); account1.installValidation( - twoHookValidation, true, new bytes4[](0), bytes(""), abi.encode(preValidationHooks, installDatas) + twoHookValidation, + true, + new bytes4[](0), + bytes(""), + abi.encode(preValidationHooks, installDatas), + bytes("") ); vm.stopPrank(); }