Allow direct plugin calls with validation & permission hooks

src/account/AccountStorage.sol

@@ -32,7 +32,7 @@ struct ValidationData {
     bool isSignatureValidation;
     // How many execution hooks require the UO context.
     uint8 requireUOHookCount;
-    // The pre validation hooks for this function selector.
+    // The pre validation hooks for this validation function.
     FunctionReference[] preValidationHooks;
     // Permission hooks for this validation function.
     EnumerableSet.Bytes32Set permissionHooks;

src/account/PluginManager2.sol

@@ -7,7 +7,7 @@ import {IPlugin} from "../interfaces/IPlugin.sol";
 import {FunctionReference, ValidationConfig} from "../interfaces/IPluginManager.sol";
 import {FunctionReferenceLib} from "../helpers/FunctionReferenceLib.sol";
 import {ValidationConfigLib} from "../helpers/ValidationConfigLib.sol";
-import {ValidationData, getAccountStorage, toSetValue, toFunctionReference} from "./AccountStorage.sol";
+import {ValidationData, getAccountStorage, toSetValue} from "./AccountStorage.sol";
 import {ExecutionHook} from "../interfaces/IAccountLoupe.sol";
 
 // Temporary additional functions for a user-controlled install flow for validation functions.
@@ -17,6 +17,7 @@ abstract contract PluginManager2 {
 
     // Index marking the start of the data for the validation function.
     uint8 internal constant _RESERVED_VALIDATION_DATA_INDEX = 255;
+    uint8 internal constant _SELF_PERMIT_VALIDATION_FUNCTIONID = type(uint8).max;
 
     error PreValidationAlreadySet(FunctionReference validationFunction, FunctionReference preValidationFunction);
     error ValidationAlreadySet(bytes4 selector, FunctionReference validationFunction);
@@ -73,19 +74,21 @@ abstract contract PluginManager2 {
             }
         }
 
-        _validationData.isGlobal = validationConfig.isGlobal();
-        _validationData.isSignatureValidation = validationConfig.isSignatureValidation();
-
         for (uint256 i = 0; i < selectors.length; ++i) {
             bytes4 selector = selectors[i];
             if (!_validationData.selectors.add(toSetValue(selector))) {
                 revert ValidationAlreadySet(selector, validationConfig.functionReference());
             }
         }
 
-        if (installData.length > 0) {
-            address plugin = validationConfig.plugin();
-            IPlugin(plugin).onInstall(installData);
+        if (validationConfig.functionId() != _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) {
+                IPlugin(validationConfig.plugin()).onInstall(installData);
+            }
         }
     }
 
@@ -120,12 +123,13 @@ abstract contract PluginManager2 {
 
             // Clear permission hooks
             EnumerableSet.Bytes32Set storage permissionHooks = _validationData.permissionHooks;
-            uint256 i = 0;
-            while (permissionHooks.length() > 0) {
-                FunctionReference permissionHook = toFunctionReference(permissionHooks.at(0));
-                permissionHooks.remove(toSetValue(permissionHook));
-                (address permissionHookPlugin,) = FunctionReferenceLib.unpack(permissionHook);
-                IPlugin(permissionHookPlugin).onUninstall(permissionHookUninstallDatas[i++]);
+
+            uint256 len = permissionHooks.length();
+            for (uint256 i = 0; i < len; ++i) {
+                bytes32 permissionHook = permissionHooks.at(0);
+                permissionHooks.remove(permissionHook);
+                address permissionHookPlugin = address(uint160(bytes20(permissionHook)));
+                IPlugin(permissionHookPlugin).onUninstall(permissionHookUninstallDatas[i]);
             }
         }
         delete _validationData.preValidationHooks;

src/account/UpgradeableModularAccount.sol

@@ -86,14 +86,13 @@
     // Wraps execution of a native function with runtime validation and hooks
     // Used for upgradeTo, upgradeToAndCall, execute, executeBatch, installPlugin, uninstallPlugin
     modifier wrapNativeFunction() {
-        _checkPermittedCallerIfNotFromEP();
-
-        PostExecToRun[] memory postExecHooks =
-            _doPreHooks(getAccountStorage().selectorData[msg.sig].executionHooks, msg.data);
+        (PostExecToRun[] memory postPermissionHooks, PostExecToRun[] memory postExecHooks) =
+            _checkPermittedCallerAndAssociatedHooks();
 
         _;
 
         _doCachedPostExecHooks(postExecHooks);
+        _doCachedPostExecHooks(postPermissionHooks);
     }
 
     constructor(IEntryPoint anEntryPoint) {
@@ -136,7 +135,7 @@
             revert UnrecognizedFunction(msg.sig);
         }
 
-        _checkPermittedCallerIfNotFromEP();
+        _checkPermittedCallerAndAssociatedHooks();
 
         PostExecToRun[] memory postExecHooks;
         // Cache post-exec hooks in memory
@@ -500,17 +499,7 @@
             } else {
                 currentAuthData = "";
             }
-
-            (address hookPlugin, uint8 hookFunctionId) = preRuntimeValidationHooks[i].unpack();
-            try IValidationHook(hookPlugin).preRuntimeValidationHook(
-                hookFunctionId, msg.sender, msg.value, callData, currentAuthData
-            )
-            // forgefmt: disable-start
-            // solhint-disable-next-line no-empty-blocks
-            {} catch (bytes memory revertReason) {
-            // forgefmt: disable-end
-                revert PreRuntimeValidationHookFailed(hookPlugin, hookFunctionId, revertReason);
-            }
+            _doPreRuntimeValidationHook(preRuntimeValidationHooks[i], callData, currentAuthData);
         }
 
         if (authSegment.getIndex() != _RESERVED_VALIDATION_DATA_INDEX) {
@@ -603,6 +592,23 @@
         }
     }
 
+    function _doPreRuntimeValidationHook(
+        FunctionReference validationHook,
+        bytes memory callData,
+        bytes memory currentAuthData
+    ) internal {
+        (address hookPlugin, uint8 hookFunctionId) = validationHook.unpack();
+        try IValidationHook(hookPlugin).preRuntimeValidationHook(
+            hookFunctionId, msg.sender, msg.value, callData, currentAuthData
+        )
+        // forgefmt: disable-start
+        // solhint-disable-next-line no-empty-blocks
+        {} catch (bytes memory revertReason) {
+        // forgefmt: disable-end
+            revert PreRuntimeValidationHookFailed(hookPlugin, hookFunctionId, revertReason);
+        }
+    }
+
     // solhint-disable-next-line no-empty-blocks
     function _authorizeUpgrade(address newImplementation) internal override {}
 
@@ -669,7 +675,8 @@
         // Check that the provided validation function is applicable to the selector
         if (isGlobal) {
             if (!_globalValidationAllowed(selector) || !_storage.validationData[validationFunction].isGlobal) {
-                revert UserOpValidationFunctionMissing(selector);
+                revert UserOpValidationFunctionMissing(selector); //TODO: Update this error as it can be runtime
+                    // validation too
             }
         } else {
             // Not global validation, but per-selector
@@ -692,14 +699,56 @@
         return getAccountStorage().selectorData[selector].allowGlobalValidation;
     }
 
-    function _checkPermittedCallerIfNotFromEP() internal view {
+    /**
+     * Order of operations:
+     *      1. Check if the sender is the entry point, the account itself, or the selector called is public.
+     *          - Yes: Return an empty array, there are no post-permissionHooks.
+     *          - No: Continue
+     *      2. Check if the called selector (msg.sig) is included in the set of selectors the msg.sender can
+     *         directly call.
+     *          - Yes: Continue
+     *          - No: Revert, the caller is not allowed to call this selector
+     *      3. If there are runtime validation hooks associated with this caller-sig combination, run them.
+     *      4. Run the pre-permissionHooks associated with this caller-sig combination, and return the
+     *         post-permissionHooks to run later.
+     */
+    function _checkPermittedCallerAndAssociatedHooks()
+        internal
+        returns (PostExecToRun[] memory, PostExecToRun[] memory)
+    {
         AccountStorage storage _storage = getAccountStorage();
 
         if (
-            msg.sender != address(_ENTRY_POINT) && msg.sender != address(this)
-                && !_storage.selectorData[msg.sig].isPublic
+            msg.sender == address(_ENTRY_POINT) || msg.sender == address(this)
+                || _storage.selectorData[msg.sig].isPublic
         ) {
-            revert ExecFromPluginNotPermitted(msg.sender, msg.sig);
+            return (new PostExecToRun[](0), new PostExecToRun[](0));
         }
+
+        FunctionReference directCallValidationKey =
+            FunctionReferenceLib.pack(msg.sender, _SELF_PERMIT_VALIDATION_FUNCTIONID);
+
+        _checkIfValidationAppliesCallData(msg.data, directCallValidationKey, false);
+
+        // Direct call is allowed, run associated permission & validation hooks
+
+        // Validation hooks
+        FunctionReference[] memory preRuntimeValidationHooks =
+            _storage.validationData[directCallValidationKey].preValidationHooks;
+
+        uint256 hookLen = preRuntimeValidationHooks.length;
+        for (uint256 i = 0; i < hookLen; ++i) {
+            _doPreRuntimeValidationHook(preRuntimeValidationHooks[i], msg.data, "");
+        }
+
+        // Permission hooks
+        PostExecToRun[] memory postPermissionHooks =
+            _doPreHooks(_storage.validationData[directCallValidationKey].permissionHooks, msg.data);
+
+        // Exec hooks
+        PostExecToRun[] memory postExecutionHooks =
+            _doPreHooks(_storage.selectorData[msg.sig].executionHooks, msg.data);
+
+        return (postPermissionHooks, postExecutionHooks);
     }
 }

test/account/DirectCallsFromPlugin.t.sol

@@ -0,0 +1,139 @@
+pragma solidity ^0.8.19;
+
+import {DirectCallPlugin} from "../mocks/plugins/DirectCallPlugin.sol";
+import {ExecutionHook} from "../../src/interfaces/IAccountLoupe.sol";
+import {IStandardExecutor, Call} from "../../src/interfaces/IStandardExecutor.sol";
+import {FunctionReferenceLib, FunctionReference} from "../../src/helpers/FunctionReferenceLib.sol";
+import {ValidationConfig, ValidationConfigLib} from "../../src/helpers/ValidationConfigLib.sol";
+import {UpgradeableModularAccount} from "../../src/account/UpgradeableModularAccount.sol";
+
+import {AccountTestBase} from "../utils/AccountTestBase.sol";
+
+contract DirectCallsFromPluginTest is AccountTestBase {
+    using ValidationConfigLib for ValidationConfig;
+
+    DirectCallPlugin internal _plugin;
+    FunctionReference internal _pluginFunctionReference;
+
+    function setUp() public {
+        _plugin = new DirectCallPlugin();
+        assertFalse(_plugin.preHookRan());
+        assertFalse(_plugin.postHookRan());
+        _pluginFunctionReference = FunctionReferenceLib.pack(address(_plugin), type(uint8).max);
+    }
+
+    /* -------------------------------------------------------------------------- */
+    /*                                  Negatives                                 */
+    /* -------------------------------------------------------------------------- */
+
+    function test_Fail_DirectCallPluginNotInstalled() external {
+        vm.prank(address(_plugin));
+        vm.expectRevert(_buildDirectCallDisallowedError(IStandardExecutor.execute.selector));
+        account1.execute(address(0), 0, "");
+    }
+
+    function test_Fail_DirectCallPluginUninstalled() external {
+        _installPlugin();
+
+        _uninstallPlugin();
+
+        vm.prank(address(_plugin));
+        vm.expectRevert(_buildDirectCallDisallowedError(IStandardExecutor.execute.selector));
+        account1.execute(address(0), 0, "");
+    }
+
+    function test_Fail_DirectCallPluginCallOtherSelector() external {
+        _installPlugin();
+
+        Call[] memory calls = new Call[](0);
+
+        vm.prank(address(_plugin));
+        vm.expectRevert(_buildDirectCallDisallowedError(IStandardExecutor.executeBatch.selector));
+        account1.executeBatch(calls);
+    }
+
+    /* -------------------------------------------------------------------------- */
+    /*                                  Positives                                 */
+    /* -------------------------------------------------------------------------- */
+
+    function test_Pass_DirectCallFromPluginPrank() external {
+        _installPlugin();
+
+        vm.prank(address(_plugin));
+        account1.execute(address(0), 0, "");
+
+        assertTrue(_plugin.preHookRan());
+        assertTrue(_plugin.postHookRan());
+    }
+
+    function test_Pass_DirectCallFromPluginCallback() external {
+        _installPlugin();
+
+        bytes memory encodedCall = abi.encodeCall(DirectCallPlugin.directCall, ());
+
+        vm.prank(address(entryPoint));
+        bytes memory result = account1.execute(address(_plugin), 0, encodedCall);
+
+        assertTrue(_plugin.preHookRan());
+        assertTrue(_plugin.postHookRan());
+
+        // the directCall() function in the _plugin calls back into `execute()` with an encoded call back into the
+        // _plugin's getData() function.
+        assertEq(abi.decode(result, (bytes)), abi.encode(_plugin.getData()));
+    }
+
+    function test_Flow_DirectCallFromPluginSequence() external {
+        // Install => Succeesfully call => uninstall => fail to call
+
+        _installPlugin();
+
+        vm.prank(address(_plugin));
+        account1.execute(address(0), 0, "");
+
+        assertTrue(_plugin.preHookRan());
+        assertTrue(_plugin.postHookRan());
+
+        _uninstallPlugin();
+
+        vm.prank(address(_plugin));
+        vm.expectRevert(_buildDirectCallDisallowedError(IStandardExecutor.execute.selector));
+        account1.execute(address(0), 0, "");
+    }
+
+    /* -------------------------------------------------------------------------- */
+    /*                                  Internals                                 */
+    /* -------------------------------------------------------------------------- */
+
+    function _installPlugin() internal {
+        bytes4[] memory selectors = new bytes4[](1);
+        selectors[0] = IStandardExecutor.execute.selector;
+
+        ExecutionHook[] memory permissionHooks = new ExecutionHook[](1);
+        bytes[] memory permissionHookInitDatas = new bytes[](1);
+
+        permissionHooks[0] = ExecutionHook({
+            hookFunction: FunctionReferenceLib.pack(address(_plugin), 0xff),
+            isPreHook: true,
+            isPostHook: true
+        });
+
+        bytes memory encodedPermissionHooks = abi.encode(permissionHooks, permissionHookInitDatas);
+
+        vm.prank(address(entryPoint));
+
+        ValidationConfig validationConfig = ValidationConfigLib.pack(_pluginFunctionReference, false, false);
+
+        account1.installValidation(validationConfig, selectors, "", "", encodedPermissionHooks);
+    }
+
+    function _uninstallPlugin() internal {
+        vm.prank(address(entryPoint));
+        account1.uninstallValidation(
+            _pluginFunctionReference, "", abi.encode(new bytes[](0)), abi.encode(new bytes[](1))
+        );
+    }
+
+    function _buildDirectCallDisallowedError(bytes4 selector) internal pure returns (bytes memory) {
+        return abi.encodeWithSelector(UpgradeableModularAccount.UserOpValidationFunctionMissing.selector, selector);
+    }
+}

test/account/PermittedCallPermissions.t.sol

@@ -48,8 +48,7 @@ contract PermittedCallPermissionsTest is AccountTestBase {
     function test_permittedCall_NotAllowed() public {
         vm.expectRevert(
             abi.encodeWithSelector(
-                UpgradeableModularAccount.ExecFromPluginNotPermitted.selector,
-                address(permittedCallerPlugin),
+                UpgradeableModularAccount.UserOpValidationFunctionMissing.selector,
                 ResultCreatorPlugin.bar.selector
             )
         );