From 1ba0233d2e0591c313c0ce08a74f67a01cbe8162 Mon Sep 17 00:00:00 2001 From: adam Date: Wed, 17 Jul 2024 21:29:10 -0400 Subject: [PATCH] merge plugin manager contracts --- src/account/ModuleManager2.sol | 148 --------------- src/account/ModuleManagerInternals.sol | 181 ++++++++++++++++--- src/account/UpgradeableModularAccount.sol | 2 - src/interfaces/IModule.sol | 2 +- test/mocks/modules/ComprehensiveModule.sol | 2 +- test/mocks/modules/ReturnDataModuleMocks.sol | 2 +- test/mocks/modules/ValidationModuleMocks.sol | 6 +- 7 files changed, 162 insertions(+), 181 deletions(-) delete mode 100644 src/account/ModuleManager2.sol diff --git a/src/account/ModuleManager2.sol b/src/account/ModuleManager2.sol deleted file mode 100644 index e1868403..00000000 --- a/src/account/ModuleManager2.sol +++ /dev/null @@ -1,148 +0,0 @@ -// SPDX-License-Identifier: UNLICENSED -pragma solidity ^0.8.25; - -import {EnumerableSet} from "@openzeppelin/contracts/utils/structs/EnumerableSet.sol"; - -import {ModuleEntityLib} from "../helpers/ModuleEntityLib.sol"; -import {ValidationConfigLib} from "../helpers/ValidationConfigLib.sol"; - -import {ExecutionHook} from "../interfaces/IAccountLoupe.sol"; -import {IModule} from "../interfaces/IModule.sol"; -import {ModuleEntity, ValidationConfig} from "../interfaces/IModuleManager.sol"; -import {ValidationData, getAccountStorage, toSetValue} from "./AccountStorage.sol"; - -// Temporary additional functions for a user-controlled install flow for validation functions. -abstract contract ModuleManager2 { - using EnumerableSet for EnumerableSet.Bytes32Set; - using ValidationConfigLib for ValidationConfig; - - // Index marking the start of the data for the validation function. - uint8 internal constant _RESERVED_VALIDATION_DATA_INDEX = 255; - uint32 internal constant _SELF_PERMIT_VALIDATION_FUNCTIONID = type(uint32).max; - - error PreValidationAlreadySet(ModuleEntity validationFunction, ModuleEntity preValidationFunction); - error ValidationAlreadySet(bytes4 selector, ModuleEntity validationFunction); - error ValidationNotSet(bytes4 selector, ModuleEntity validationFunction); - error PermissionAlreadySet(ModuleEntity validationFunction, ExecutionHook hook); - error PreValidationHookLimitExceeded(); - - function _installValidation( - ValidationConfig validationConfig, - bytes4[] memory selectors, - bytes calldata installData, - bytes memory preValidationHooks, - bytes memory permissionHooks - ) internal { - ValidationData storage _validationData = - getAccountStorage().validationData[validationConfig.moduleEntity()]; - - if (preValidationHooks.length > 0) { - (ModuleEntity[] memory preValidationFunctions, bytes[] memory initDatas) = - abi.decode(preValidationHooks, (ModuleEntity[], bytes[])); - - for (uint256 i = 0; i < preValidationFunctions.length; ++i) { - ModuleEntity preValidationFunction = preValidationFunctions[i]; - - _validationData.preValidationHooks.push(preValidationFunction); - - if (initDatas[i].length > 0) { - (address preValidationModule,) = ModuleEntityLib.unpack(preValidationFunction); - IModule(preValidationModule).onInstall(initDatas[i]); - } - } - - // Avoid collision between reserved index and actual indices - if (_validationData.preValidationHooks.length > _RESERVED_VALIDATION_DATA_INDEX) { - revert PreValidationHookLimitExceeded(); - } - } - - 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 (!_validationData.permissionHooks.add(toSetValue(permissionFunction))) { - revert PermissionAlreadySet(validationConfig.moduleEntity(), permissionFunction); - } - - if (initDatas[i].length > 0) { - (address executionModule,) = ModuleEntityLib.unpack(permissionFunction.hookFunction); - IModule(executionModule).onInstall(initDatas[i]); - } - } - } - - for (uint256 i = 0; i < selectors.length; ++i) { - bytes4 selector = selectors[i]; - if (!_validationData.selectors.add(toSetValue(selector))) { - revert ValidationAlreadySet(selector, validationConfig.moduleEntity()); - } - } - - if (validationConfig.entityId() != _SELF_PERMIT_VALIDATION_FUNCTIONID) { - // Only allow global validations and signature validations if they're not direct-call validations. - - _validationData.isGlobal = validationConfig.isGlobal(); - _validationData.isSignatureValidation = validationConfig.isSignatureValidation(); - if (installData.length > 0) { - IModule(validationConfig.module()).onInstall(installData); - } - } - } - - function _uninstallValidation( - ModuleEntity validationFunction, - bytes calldata uninstallData, - bytes calldata preValidationHookUninstallData, - bytes calldata permissionHookUninstallData - ) internal { - ValidationData storage _validationData = getAccountStorage().validationData[validationFunction]; - - _validationData.isGlobal = false; - _validationData.isSignatureValidation = false; - - { - bytes[] memory preValidationHookUninstallDatas = abi.decode(preValidationHookUninstallData, (bytes[])); - - // Clear pre validation hooks - ModuleEntity[] storage preValidationHooks = _validationData.preValidationHooks; - for (uint256 i = 0; i < preValidationHooks.length; ++i) { - ModuleEntity preValidationFunction = preValidationHooks[i]; - if (preValidationHookUninstallDatas[0].length > 0) { - (address preValidationModule,) = ModuleEntityLib.unpack(preValidationFunction); - IModule(preValidationModule).onUninstall(preValidationHookUninstallDatas[0]); - } - } - delete _validationData.preValidationHooks; - } - - { - bytes[] memory permissionHookUninstallDatas = abi.decode(permissionHookUninstallData, (bytes[])); - - // Clear permission hooks - EnumerableSet.Bytes32Set storage permissionHooks = _validationData.permissionHooks; - uint256 permissionHookLen = permissionHooks.length(); - for (uint256 i = 0; i < permissionHookLen; ++i) { - bytes32 permissionHook = permissionHooks.at(0); - permissionHooks.remove(permissionHook); - address permissionHookModule = address(uint160(bytes20(permissionHook))); - IModule(permissionHookModule).onUninstall(permissionHookUninstallDatas[i]); - } - } - - // Clear selectors - uint256 selectorLen = _validationData.selectors.length(); - for (uint256 i = 0; i < selectorLen; ++i) { - bytes32 selectorSetValue = _validationData.selectors.at(0); - _validationData.selectors.remove(selectorSetValue); - } - - if (uninstallData.length > 0) { - (address module,) = ModuleEntityLib.unpack(validationFunction); - IModule(module).onUninstall(uninstallData); - } - } -} diff --git a/src/account/ModuleManagerInternals.sol b/src/account/ModuleManagerInternals.sol index a6a275b6..9fd60c45 100644 --- a/src/account/ModuleManagerInternals.sol +++ b/src/account/ModuleManagerInternals.sol @@ -7,14 +7,22 @@ import {EnumerableSet} from "@openzeppelin/contracts/utils/structs/EnumerableSet import {KnownSelectors} from "../helpers/KnownSelectors.sol"; import {ModuleEntityLib} from "../helpers/ModuleEntityLib.sol"; +import {ValidationConfigLib} from "../helpers/ValidationConfigLib.sol"; import {ExecutionHook} from "../interfaces/IAccountLoupe.sol"; import {IModule, ManifestExecutionHook, ManifestValidation, ModuleManifest} from "../interfaces/IModule.sol"; -import {IModuleManager, ModuleEntity} from "../interfaces/IModuleManager.sol"; -import {AccountStorage, SelectorData, getAccountStorage, toSetValue} from "./AccountStorage.sol"; +import {IModuleManager, ModuleEntity, ValidationConfig} from "../interfaces/IModuleManager.sol"; +import {AccountStorage, SelectorData, ValidationData, getAccountStorage, toSetValue} from "./AccountStorage.sol"; abstract contract ModuleManagerInternals is IModuleManager { using EnumerableSet for EnumerableSet.Bytes32Set; using ModuleEntityLib for ModuleEntity; + using ValidationConfigLib for ValidationConfig; + + // Index marking the start of the data for the validation function. + uint8 internal constant _RESERVED_VALIDATION_DATA_INDEX = 255; + + // Magic value for the Entity ID of direct call validation. + uint32 internal constant _SELF_PERMIT_VALIDATION_FUNCTIONID = type(uint32).max; error ArrayLengthMismatch(); error Erc4337FunctionNotAllowed(bytes4 selector); @@ -22,10 +30,12 @@ abstract contract ModuleManagerInternals is IModuleManager { error IModuleFunctionNotAllowed(bytes4 selector); error NativeFunctionNotAllowed(bytes4 selector); error NullModule(); + error PermissionAlreadySet(ModuleEntity validationFunction, ExecutionHook hook); error ModuleInstallCallbackFailed(address module, bytes revertReason); error ModuleInterfaceNotSupported(address module); error ModuleNotInstalled(address module); - error ValidationFunctionAlreadySet(bytes4 selector, ModuleEntity validationFunction); + error PreValidationHookLimitExceeded(); + error ValidationAlreadySet(bytes4 selector, ModuleEntity validationFunction); // Storage update operations @@ -71,39 +81,35 @@ abstract contract ModuleManagerInternals is IModuleManager { _selectorData.allowGlobalValidation = false; } - function _addValidationFunction(address module, ManifestValidation memory mv) internal { - AccountStorage storage _storage = getAccountStorage(); + function _addValidationFunction(ValidationConfig validationConfig, bytes4[] memory selectors) internal { + ValidationData storage _validationData = + getAccountStorage().validationData[validationConfig.moduleEntity()]; - ModuleEntity validationFunction = ModuleEntityLib.pack(module, mv.entityId); - - if (mv.isDefault) { - _storage.validationData[validationFunction].isGlobal = true; + if (validationConfig.isGlobal()) { + _validationData.isGlobal = true; } - if (mv.isSignatureValidation) { - _storage.validationData[validationFunction].isSignatureValidation = true; + if (validationConfig.isSignatureValidation()) { + _validationData.isSignatureValidation = true; } // Add the validation function to the selectors. - uint256 length = mv.selectors.length; + uint256 length = selectors.length; for (uint256 i = 0; i < length; ++i) { - bytes4 selector = mv.selectors[i]; - _storage.validationData[validationFunction].selectors.add(toSetValue(selector)); + _validationData.selectors.add(toSetValue(selectors[i])); } } - function _removeValidationFunction(address module, ManifestValidation memory mv) internal { - AccountStorage storage _storage = getAccountStorage(); - - ModuleEntity validationFunction = ModuleEntityLib.pack(module, mv.entityId); + function _removeValidationFunction(ModuleEntity validationFunction) internal { + ValidationData storage _validationData = getAccountStorage().validationData[validationFunction]; - _storage.validationData[validationFunction].isGlobal = false; - _storage.validationData[validationFunction].isSignatureValidation = false; + _validationData.isGlobal = false; + _validationData.isSignatureValidation = false; // Clear the selectors - while (_storage.validationData[validationFunction].selectors.length() > 0) { - bytes32 selector = _storage.validationData[validationFunction].selectors.at(0); - _storage.validationData[validationFunction].selectors.remove(selector); + uint256 length = _validationData.selectors.length(); + for (uint256 i = 0; i < length; ++i) { + _validationData.selectors.remove(_validationData.selectors.at(0)); } } @@ -161,7 +167,11 @@ abstract contract ModuleManagerInternals is IModuleManager { for (uint256 i = 0; i < length; ++i) { // Todo: limit this to only "direct runtime call" validation path (old EFP), // and add a way for the user to specify permission/pre-val hooks here. - _addValidationFunction(module, manifest.validationFunctions[i]); + ManifestValidation memory mv = manifest.validationFunctions[i]; + + ValidationConfig validationConfig = + ValidationConfigLib.pack(module, mv.entityId, mv.isGlobal, mv.isSignatureValidation); + _addValidationFunction(validationConfig, mv.selectors); } length = manifest.executionHooks.length; @@ -204,7 +214,9 @@ abstract contract ModuleManagerInternals is IModuleManager { length = manifest.validationFunctions.length; for (uint256 i = 0; i < length; ++i) { - _removeValidationFunction(module, manifest.validationFunctions[i]); + ModuleEntity validationFunction = + ModuleEntityLib.pack(module, manifest.validationFunctions[i].entityId); + _removeValidationFunction(validationFunction); } length = manifest.executionFunctions.length; @@ -228,4 +240,123 @@ abstract contract ModuleManagerInternals is IModuleManager { emit ModuleUninstalled(module, onUninstallSuccess); } + + function _installValidation( + ValidationConfig validationConfig, + bytes4[] memory selectors, + bytes calldata installData, + bytes memory preValidationHooks, + bytes memory permissionHooks + ) internal { + ValidationData storage _validationData = + getAccountStorage().validationData[validationConfig.moduleEntity()]; + + if (preValidationHooks.length > 0) { + (ModuleEntity[] memory preValidationFunctions, bytes[] memory initDatas) = + abi.decode(preValidationHooks, (ModuleEntity[], bytes[])); + + for (uint256 i = 0; i < preValidationFunctions.length; ++i) { + ModuleEntity preValidationFunction = preValidationFunctions[i]; + + _validationData.preValidationHooks.push(preValidationFunction); + + if (initDatas[i].length > 0) { + (address preValidationPlugin,) = ModuleEntityLib.unpack(preValidationFunction); + IModule(preValidationPlugin).onInstall(initDatas[i]); + } + } + + // Avoid collision between reserved index and actual indices + if (_validationData.preValidationHooks.length > _RESERVED_VALIDATION_DATA_INDEX) { + revert PreValidationHookLimitExceeded(); + } + } + + 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 (!_validationData.permissionHooks.add(toSetValue(permissionFunction))) { + revert PermissionAlreadySet(validationConfig.moduleEntity(), permissionFunction); + } + + if (initDatas[i].length > 0) { + (address executionPlugin,) = ModuleEntityLib.unpack(permissionFunction.hookFunction); + IModule(executionPlugin).onInstall(initDatas[i]); + } + } + } + + for (uint256 i = 0; i < selectors.length; ++i) { + bytes4 selector = selectors[i]; + if (!_validationData.selectors.add(toSetValue(selector))) { + revert ValidationAlreadySet(selector, validationConfig.moduleEntity()); + } + } + + if (validationConfig.entityId() != _SELF_PERMIT_VALIDATION_FUNCTIONID) { + // Only allow global validations and signature validations if they're not direct-call validations. + + _validationData.isGlobal = validationConfig.isGlobal(); + _validationData.isSignatureValidation = validationConfig.isSignatureValidation(); + if (installData.length > 0) { + IModule(validationConfig.module()).onInstall(installData); + } + } + } + + function _uninstallValidation( + ModuleEntity validationFunction, + bytes calldata uninstallData, + bytes calldata preValidationHookUninstallData, + bytes calldata permissionHookUninstallData + ) internal { + ValidationData storage _validationData = getAccountStorage().validationData[validationFunction]; + + _removeValidationFunction(validationFunction); + + { + bytes[] memory preValidationHookUninstallDatas = abi.decode(preValidationHookUninstallData, (bytes[])); + + // Clear pre validation hooks + ModuleEntity[] storage preValidationHooks = _validationData.preValidationHooks; + for (uint256 i = 0; i < preValidationHooks.length; ++i) { + ModuleEntity preValidationFunction = preValidationHooks[i]; + if (preValidationHookUninstallDatas[0].length > 0) { + (address preValidationPlugin,) = ModuleEntityLib.unpack(preValidationFunction); + IModule(preValidationPlugin).onUninstall(preValidationHookUninstallDatas[0]); + } + } + delete _validationData.preValidationHooks; + } + + { + bytes[] memory permissionHookUninstallDatas = abi.decode(permissionHookUninstallData, (bytes[])); + + // Clear permission hooks + EnumerableSet.Bytes32Set storage permissionHooks = _validationData.permissionHooks; + uint256 permissionHookLen = permissionHooks.length(); + for (uint256 i = 0; i < permissionHookLen; ++i) { + bytes32 permissionHook = permissionHooks.at(0); + permissionHooks.remove(permissionHook); + address permissionHookPlugin = address(uint160(bytes20(permissionHook))); + IModule(permissionHookPlugin).onUninstall(permissionHookUninstallDatas[i]); + } + } + + // Clear selectors + uint256 selectorLen = _validationData.selectors.length(); + for (uint256 i = 0; i < selectorLen; ++i) { + bytes32 selectorSetValue = _validationData.selectors.at(0); + _validationData.selectors.remove(selectorSetValue); + } + + if (uninstallData.length > 0) { + (address plugin,) = ModuleEntityLib.unpack(validationFunction); + IModule(plugin).onUninstall(uninstallData); + } + } } diff --git a/src/account/UpgradeableModularAccount.sol b/src/account/UpgradeableModularAccount.sol index 2f1da16d..737beebb 100644 --- a/src/account/UpgradeableModularAccount.sol +++ b/src/account/UpgradeableModularAccount.sol @@ -29,7 +29,6 @@ import {AccountLoupe} from "./AccountLoupe.sol"; import {AccountStorage, getAccountStorage, toExecutionHook, toSetValue} from "./AccountStorage.sol"; import {AccountStorageInitializable} from "./AccountStorageInitializable.sol"; -import {ModuleManager2} from "./ModuleManager2.sol"; import {ModuleManagerInternals} from "./ModuleManagerInternals.sol"; contract UpgradeableModularAccount is @@ -42,7 +41,6 @@ contract UpgradeableModularAccount is IStandardExecutor, IAccountExecute, ModuleManagerInternals, - ModuleManager2, UUPSUpgradeable { using EnumerableSet for EnumerableSet.Bytes32Set; diff --git a/src/interfaces/IModule.sol b/src/interfaces/IModule.sol index 8851acaf..ab78b222 100644 --- a/src/interfaces/IModule.sol +++ b/src/interfaces/IModule.sol @@ -16,7 +16,7 @@ struct ManifestExecutionFunction { // todo: do we need these at all? Or do we fully switch to `installValidation`? struct ManifestValidation { uint32 entityId; - bool isDefault; + bool isGlobal; bool isSignatureValidation; bytes4[] selectors; } diff --git a/test/mocks/modules/ComprehensiveModule.sol b/test/mocks/modules/ComprehensiveModule.sol index bb73b5db..aca09f92 100644 --- a/test/mocks/modules/ComprehensiveModule.sol +++ b/test/mocks/modules/ComprehensiveModule.sol @@ -146,7 +146,7 @@ contract ComprehensiveModule is IValidation, IValidationHook, IExecutionHook, Ba manifest.validationFunctions = new ManifestValidation[](1); manifest.validationFunctions[0] = ManifestValidation({ entityId: uint32(EntityId.VALIDATION), - isDefault: true, + isGlobal: true, isSignatureValidation: false, selectors: validationSelectors }); diff --git a/test/mocks/modules/ReturnDataModuleMocks.sol b/test/mocks/modules/ReturnDataModuleMocks.sol index 8fe3241d..1f219e8b 100644 --- a/test/mocks/modules/ReturnDataModuleMocks.sol +++ b/test/mocks/modules/ReturnDataModuleMocks.sol @@ -126,7 +126,7 @@ contract ResultConsumerModule is BaseModule, IValidation { manifest.validationFunctions = new ManifestValidation[](1); manifest.validationFunctions[0] = ManifestValidation({ entityId: 0, - isDefault: true, + isGlobal: true, isSignatureValidation: false, selectors: validationSelectors }); diff --git a/test/mocks/modules/ValidationModuleMocks.sol b/test/mocks/modules/ValidationModuleMocks.sol index ed152af2..4c37190f 100644 --- a/test/mocks/modules/ValidationModuleMocks.sol +++ b/test/mocks/modules/ValidationModuleMocks.sol @@ -118,7 +118,7 @@ contract MockUserOpValidationModule is MockBaseUserOpValidationModule { manifest.validationFunctions = new ManifestValidation[](1); manifest.validationFunctions[0] = ManifestValidation({ entityId: uint32(EntityId.USER_OP_VALIDATION), - isDefault: false, + isGlobal: false, isSignatureValidation: false, selectors: validationSelectors }); @@ -161,7 +161,7 @@ contract MockUserOpValidation1HookModule is MockBaseUserOpValidationModule { manifest.validationFunctions = new ManifestValidation[](2); manifest.validationFunctions[0] = ManifestValidation({ entityId: uint32(EntityId.USER_OP_VALIDATION), - isDefault: false, + isGlobal: false, isSignatureValidation: false, selectors: validationSelectors }); @@ -207,7 +207,7 @@ contract MockUserOpValidation2HookModule is MockBaseUserOpValidationModule { manifest.validationFunctions = new ManifestValidation[](1); manifest.validationFunctions[0] = ManifestValidation({ entityId: uint32(EntityId.USER_OP_VALIDATION), - isDefault: false, + isGlobal: false, isSignatureValidation: false, selectors: validationSelectors });