diff --git a/foundry.toml b/foundry.toml index ff7f6c4..6c48258 100644 --- a/foundry.toml +++ b/foundry.toml @@ -1,5 +1,5 @@ [profile.default] -evm_version = "cancun" +evm_version = "prague" src = "src" out = "out" libs = ["node_modules"] diff --git a/src/MSAAdvanced.sol b/src/MSAAdvanced.sol index c9db2cb..714322c 100644 --- a/src/MSAAdvanced.sol +++ b/src/MSAAdvanced.sol @@ -419,6 +419,7 @@ contract MSAAdvanced is if (isERC7702) { _addStorageBase(MODULEMANAGER_STORAGE_LOCATION); _addStorageBase(HOOKMANAGER_STORAGE_LOCATION); + _addStorageBase(PREVALIDATION_HOOKMANAGER_STORAGE_LOCATION); } // bootstrap the account @@ -438,10 +439,30 @@ contract MSAAdvanced is if (!success) revert(); } - function _onRedelegation() internal override { + function _onRedelegation(bytes calldata context) internal override { + // Decode the context data which should contain arrays of uninstall data for validators and + // executors + // (bytes[] memory validatorUninstallData, bytes[] memory executorUninstallData) = + // abi.decode(context, (bytes[], bytes[])); + + // Call the uninstall functions with the decoded data + // _tryUninstallValidators(validatorUninstallData); + // _tryUninstallExecutors(executorUninstallData); + _tryUninstallValidators(); _tryUninstallExecutors(); + + // Continue with other uninstallations _tryUninstallHook(_getHook()); + _tryUninstallPreValidationHook( + _getPreValidationHook(MODULE_TYPE_PREVALIDATION_HOOK_ERC1271), + MODULE_TYPE_PREVALIDATION_HOOK_ERC1271 + ); + _tryUninstallPreValidationHook( + _getPreValidationHook(MODULE_TYPE_PREVALIDATION_HOOK_ERC4337), + MODULE_TYPE_PREVALIDATION_HOOK_ERC4337 + ); + _tryUninstallFallbacks(); _initModuleManager(); } } diff --git a/src/core/ERC7779Adapter.sol b/src/core/ERC7779Adapter.sol index e35fdec..ab971f9 100644 --- a/src/core/ERC7779Adapter.sol +++ b/src/core/ERC7779Adapter.sol @@ -42,6 +42,12 @@ abstract contract ERC7779Adapter is IERC7779 { assembly { $.slot := ERC7779_STORAGE_BASE } + // Check if the storageBase already exists to avoid duplication + for (uint256 i = 0; i < $.storageBases.length; i++) { + if ($.storageBases[i] == storageBase) { + return; // Exit if the storageBase is already present + } + } $.storageBases.push(storageBase); } @@ -54,14 +60,17 @@ abstract contract ERC7779Adapter is IERC7779 { * It should uninitialize storages if needed and execute wallet-specific logic to prepare for redelegation. * msg.sender should be the owner of the account. + * @param context - The context of the redelegation. Additional data required to uninitialize + storages. */ - function onRedelegation() external returns (bool) { + function onRedelegation(bytes calldata context) external returns (bool) { require(msg.sender == address(this), NonAuthorizedOnRedelegationCaller()); - _onRedelegation(); + _onRedelegation(context); return true; } /// @dev This function is called before redelegation. /// @dev Account should override this function to implement the specific logic. - function _onRedelegation() internal virtual; + /// @param context The context is the additional data required to uninitialize storages. + function _onRedelegation(bytes calldata context) internal virtual; } diff --git a/src/core/ModuleManager.sol b/src/core/ModuleManager.sol index 24ffd64..029ba6c 100644 --- a/src/core/ModuleManager.sol +++ b/src/core/ModuleManager.sol @@ -46,6 +46,8 @@ abstract contract ModuleManager is AccountBase, Receiver { // single fallback handler for all fallbacks // account vendors may implement this differently. This is just a reference implementation mapping(bytes4 selector => FallbackHandler fallbackHandler) $fallbacks; + ///< List of fallback selectors, initialized upon contract deployment. + bytes4[] fallbackSelectors; } function $moduleManager() internal pure virtual returns (ModuleManagerStorage storage $ims) { @@ -95,36 +97,35 @@ abstract contract ModuleManager is AccountBase, Receiver { IValidator(validator).onUninstall(disableModuleData); } - /* - function _tryUninstallValidators(bytes[] calldata data) internal { + function _tryUninstallValidators() internal virtual { SentinelListLib.SentinelList storage $valdiators = $moduleManager().$valdiators; - uint256 length = data.length; - uint256 index; address validator = $valdiators.getNext(SENTINEL); while (validator != SENTINEL) { - bytes memory uninstallData; - if (index < length) { - uninstallData = data[index]; - } - try IValidator(validator).onUninstall(uninstallData) {} catch { - emit ValidatorUninstallFailed(validator, uninstallData); + try IValidator(validator).onUninstall("") { } + catch { + emit ValidatorUninstallFailed(validator, ""); } validator = $valdiators.getNext(validator); - index++; } $valdiators.popAll(); } - */ - function _tryUninstallValidators() internal { + function _tryUninstallValidators(bytes[] memory data) internal virtual { SentinelListLib.SentinelList storage $valdiators = $moduleManager().$valdiators; + uint256 length = data.length; + uint256 index; address validator = $valdiators.getNext(SENTINEL); while (validator != SENTINEL) { - try IValidator(validator).onUninstall("") { } + bytes memory uninstallData; + if (index < length) { + uninstallData = data[index]; + } + try IValidator(validator).onUninstall(uninstallData) { } catch { - emit ValidatorUninstallFailed(validator, ""); + emit ValidatorUninstallFailed(validator, uninstallData); } validator = $valdiators.getNext(validator); + index++; } $valdiators.popAll(); } @@ -168,36 +169,35 @@ abstract contract ModuleManager is AccountBase, Receiver { IExecutor(executor).onUninstall(disableModuleData); } - /* - function _tryUninstallExecutors(bytes[] calldata data) internal { + function _tryUninstallExecutors() internal virtual { SentinelListLib.SentinelList storage $executors = $moduleManager().$executors; - uint256 length = data.length; - uint256 index; address executor = $executors.getNext(SENTINEL); while (executor != SENTINEL) { - bytes memory uninstallData; - if (index < length) { - uninstallData = data[index]; - } - try IExecutor(executor).onUninstall(uninstallData) {} catch { - emit ExecutorUninstallFailed(executor, uninstallData); + try IExecutor(executor).onUninstall("") { } + catch { + emit ExecutorUninstallFailed(executor, ""); } executor = $executors.getNext(executor); - index++; } $executors.popAll(); } - */ - function _tryUninstallExecutors() internal { + function _tryUninstallExecutors(bytes[] memory data) internal virtual { SentinelListLib.SentinelList storage $executors = $moduleManager().$executors; + uint256 length = data.length; + uint256 index; address executor = $executors.getNext(SENTINEL); while (executor != SENTINEL) { - try IExecutor(executor).onUninstall("") { } + bytes memory uninstallData; + if (index < length) { + uninstallData = data[index]; + } + try IExecutor(executor).onUninstall(uninstallData) { } catch { - emit ExecutorUninstallFailed(executor, ""); + emit ExecutorUninstallFailed(executor, uninstallData); } executor = $executors.getNext(executor); + index++; } $executors.popAll(); } @@ -237,6 +237,8 @@ abstract contract ModuleManager is AccountBase, Receiver { revert("Function selector already used"); } $moduleManager().$fallbacks[selector] = FallbackHandler(handler, calltype); + // Add the selector to the maintained list of fallback selectors + $moduleManager().fallbackSelectors.push(selector); IFallback(handler).onInstall(initData); } @@ -262,9 +264,37 @@ abstract contract ModuleManager is AccountBase, Receiver { $moduleManager().$fallbacks[selector] = FallbackHandler(address(0), CallType.wrap(0x00)); + // Remove selector from fallbackSelectors via swap-and-pop + uint256 len = $moduleManager().fallbackSelectors.length; + for (uint256 i = 0; i < len; i++) { + if ($moduleManager().fallbackSelectors[i] == selector) { + $moduleManager().fallbackSelectors[i] = $moduleManager().fallbackSelectors[len - 1]; + $moduleManager().fallbackSelectors.pop(); + break; + } + } + IFallback(handler).onUninstall(_deInitData); } + // Review: if uninstalling selectors also need some data. + function _tryUninstallFallbacks() internal { + uint256 len = $moduleManager().fallbackSelectors.length; + + for (uint256 i = 0; i < len; i++) { + bytes4 selector = $moduleManager().fallbackSelectors[i]; + FallbackHandler memory fallbackHandler = $moduleManager().$fallbacks[selector]; + + if (address(fallbackHandler.handler) == address(0)) continue; + + IFallback(fallbackHandler.handler).onUninstall(abi.encodePacked(selector)); + + $moduleManager().$fallbacks[selector] = FallbackHandler(address(0), CallType.wrap(0x00)); + } + + delete $moduleManager().fallbackSelectors; + } + function _isFallbackHandlerInstalled(bytes4 functionSig) internal view virtual returns (bool) { FallbackHandler storage $fallback = $moduleManager().$fallbacks[functionSig]; return $fallback.handler != address(0); diff --git a/src/core/PreValidationHookManager.sol b/src/core/PreValidationHookManager.sol index f672d82..991d813 100644 --- a/src/core/PreValidationHookManager.sol +++ b/src/core/PreValidationHookManager.sol @@ -92,6 +92,7 @@ abstract contract PreValidationHookManager { } function _tryUninstallPreValidationHook(address hook, uint256 hookType) internal virtual { + if (hook == address(0)) return; PreValidationHookManagerStorage storage $ = _getStorage(); if (hookType == MODULE_TYPE_PREVALIDATION_HOOK_ERC1271) { try $.hook1271.onUninstall("") { } diff --git a/src/interfaces/IERC7779.sol b/src/interfaces/IERC7779.sol index 74dc068..0e57f6f 100644 --- a/src/interfaces/IERC7779.sol +++ b/src/interfaces/IERC7779.sol @@ -27,6 +27,8 @@ interface IERC7779 { * It should uninitialize storages if needed and execute wallet-specific logic to prepare for redelegation. * msg.sender should be the owner of the account. + * @param context - The context of the redelegation. Additional data required to uninitialize + storages. */ - function onRedelegation() external returns (bool); + function onRedelegation(bytes calldata context) external returns (bool); } diff --git a/test/advanced/EIP7702.t.sol b/test/advanced/EIP7702.t.sol index e3b5433..ca262a7 100644 --- a/test/advanced/EIP7702.t.sol +++ b/test/advanced/EIP7702.t.sol @@ -308,7 +308,7 @@ contract EIP7702 is TestBaseUtilAdvanced { assertTrue(IMSA(account).isModuleInstalled(MODULE_TYPE_HOOK, address(hook), "")); // storage is cleared vm.prank(address(account)); - IMSA(account).onRedelegation(); + IMSA(account).onRedelegation(bytes("")); assertFalse( IMSA(account).isModuleInstalled(MODULE_TYPE_VALIDATOR, address(defaultValidator), "") ); diff --git a/test/mocks/MockERC7779.sol b/test/mocks/MockERC7779.sol index fc89ca6..c4c7b3a 100644 --- a/test/mocks/MockERC7779.sol +++ b/test/mocks/MockERC7779.sol @@ -8,7 +8,7 @@ contract MockERC7779 is ERC7779Adapter { _addStorageBase(storageBase); } - function _onRedelegation() internal override { + function _onRedelegation(bytes calldata context) internal override { // do nothing } }