fix: [spearbit-15] IPlugin function prevention in installing executio…

…n functions
alchemyplatform · Dec 12, 2023 · 9082c75 · 9082c75
1 parent 84f4b3d
commit 9082c75
Show file tree

Hide file tree

Showing 4 changed files with 62 additions and 2 deletions.
diff --git a/src/account/PluginManagerInternals.sol b/src/account/PluginManagerInternals.sol
@@ -50,6 +50,7 @@ abstract contract PluginManagerInternals is IPluginManager, AccountStorageV1 {
     error Erc4337FunctionNotAllowed(bytes4 selector);
     error ExecutionFunctionAlreadySet(bytes4 selector);
     error ExecutionFunctionNotSet(bytes4 selector);
+    error IPluginFunctionNotAllowed(bytes4 selector);
     error InvalidDependenciesProvided();
     error InvalidPluginManifest();
     error MissingPluginDependency(address dependency);
@@ -75,11 +76,17 @@ abstract contract PluginManagerInternals is IPluginManager, AccountStorageV1 {
             revert ExecutionFunctionAlreadySet(selector);
         }
 
-        // make sure incoming execution function does not collide with any native functions (data are stored on the
+        // Make sure incoming execution function does not collide with any native functions (data are stored on the
         // account implementation contract)
         if (KnownSelectors.isNativeFunction(selector)) {
             revert NativeFunctionNotAllowed(selector);
         }
+
+        // Make sure incoming execution function is not any IPlugin functions
+        if (KnownSelectors.isIPluginFunction(selector)) {
+            revert IPluginFunctionNotAllowed(selector);
+        }
+
         // Also make sure it doesn't collide with functions defined by ERC-4337
         // and called by the entry point. This prevents a malicious plugin from
         // sneaking in a function with the same selector as e.g.

diff --git a/src/helpers/KnownSelectors.sol b/src/helpers/KnownSelectors.sol
@@ -10,6 +10,7 @@ import {IAccountLoupe} from "../../src/interfaces/IAccountLoupe.sol";
 import {IAccountView} from "../../src/interfaces/IAccountView.sol";
 import {IAggregator} from "../../src/interfaces/erc4337/IAggregator.sol";
 import {IPaymaster} from "../../src/interfaces/erc4337/IPaymaster.sol";
+import {IPlugin} from "../interfaces/IPlugin.sol";
 import {IPluginExecutor} from "../interfaces/IPluginExecutor.sol";
 import {IPluginManager} from "../../src/interfaces/IPluginManager.sol";
 import {IStandardExecutor} from "../interfaces/IStandardExecutor.sol";
@@ -53,4 +54,15 @@ library KnownSelectors {
             || selector == IAggregator.aggregateSignatures.selector
             || selector == IPaymaster.validatePaymasterUserOp.selector || selector == IPaymaster.postOp.selector;
     }
+
+    function isIPluginFunction(bytes4 selector) internal pure returns (bool) {
+        return selector == IPlugin.onInstall.selector || selector == IPlugin.onUninstall.selector
+            || selector == IPlugin.preUserOpValidationHook.selector
+            || selector == IPlugin.userOpValidationFunction.selector
+            || selector == IPlugin.preRuntimeValidationHook.selector
+            || selector == IPlugin.runtimeValidationFunction.selector || selector == IPlugin.preExecutionHook.selector
+            || selector == IPlugin.postExecutionHook.selector || selector == IPlugin.onHookApply.selector
+            || selector == IPlugin.onHookUnapply.selector || selector == IPlugin.pluginManifest.selector
+            || selector == IPlugin.pluginMetadata.selector;
+    }
 }
diff --git a/test/account/UpgradeableModularAccountPluginManager.t.sol b/test/account/UpgradeableModularAccountPluginManager.t.sol
@@ -263,6 +263,29 @@ contract UpgradeableModularAccountPluginManagerTest is Test {
         });
     }
 
+    function test_installPlugin_failWithIPluginFunctionSelector() public {
+        vm.startPrank(owner2);
+
+        PluginManifest memory manifestBad;
+        manifestBad.executionFunctions = new bytes4[](1);
+        manifestBad.executionFunctions[0] = IPlugin.onInstall.selector;
+        MockPlugin mockPluginBad = new MockPlugin(manifestBad);
+        bytes32 manifestHashBad = keccak256(abi.encode(mockPluginBad.pluginManifest()));
+
+        vm.expectRevert(
+            abi.encodeWithSelector(
+                PluginManagerInternals.IPluginFunctionNotAllowed.selector, IPlugin.onInstall.selector
+            )
+        );
+        IPluginManager(account2).installPlugin({
+            plugin: address(mockPluginBad),
+            manifestHash: manifestHashBad,
+            pluginInitData: bytes(""),
+            dependencies: new FunctionReference[](0),
+            injectedHooks: new IPluginManager.InjectedHook[](0)
+        });
+    }
+
     function test_installPlugin_failWtihErc4337FunctionSelector() public {
         vm.startPrank(owner2);
 

diff --git a/test/helpers/KnownSelectors.t.sol b/test/helpers/KnownSelectors.t.sol
@@ -12,9 +12,10 @@ import {IPaymaster} from "@eth-infinitism/account-abstraction/interfaces/IPaymas
 import {KnownSelectors} from "../../src/helpers/KnownSelectors.sol";
 import {IAccountLoupe} from "../../src/interfaces/IAccountLoupe.sol";
 import {IAccountInitializable} from "../../src/interfaces/IAccountInitializable.sol";
-import {IStandardExecutor} from "../../src/interfaces/IStandardExecutor.sol";
+import {IPlugin} from "../../src/interfaces/IPlugin.sol";
 import {IPluginExecutor} from "../../src/interfaces/IPluginExecutor.sol";
 import {IPluginManager} from "../../src/interfaces/IPluginManager.sol";
+import {IStandardExecutor} from "../../src/interfaces/IStandardExecutor.sol";
 
 contract KnownSelectorsTest is Test {
     function test_isNativeFunction() public {
@@ -64,4 +65,21 @@ contract KnownSelectorsTest is Test {
 
         assertFalse(KnownSelectors.isErc4337Function(BaseAccount.validateUserOp.selector));
     }
+
+    function test_isIPluginFunction() public {
+        assertTrue(KnownSelectors.isIPluginFunction(IPlugin.onInstall.selector));
+        assertTrue(KnownSelectors.isIPluginFunction(IPlugin.onUninstall.selector));
+        assertTrue(KnownSelectors.isIPluginFunction(IPlugin.preUserOpValidationHook.selector));
+        assertTrue(KnownSelectors.isIPluginFunction(IPlugin.userOpValidationFunction.selector));
+        assertTrue(KnownSelectors.isIPluginFunction(IPlugin.preRuntimeValidationHook.selector));
+        assertTrue(KnownSelectors.isIPluginFunction(IPlugin.runtimeValidationFunction.selector));
+        assertTrue(KnownSelectors.isIPluginFunction(IPlugin.preExecutionHook.selector));
+        assertTrue(KnownSelectors.isIPluginFunction(IPlugin.postExecutionHook.selector));
+        assertTrue(KnownSelectors.isIPluginFunction(IPlugin.onHookApply.selector));
+        assertTrue(KnownSelectors.isIPluginFunction(IPlugin.onHookUnapply.selector));
+        assertTrue(KnownSelectors.isIPluginFunction(IPlugin.pluginManifest.selector));
+        assertTrue(KnownSelectors.isIPluginFunction(IPlugin.pluginMetadata.selector));
+
+        assertFalse(KnownSelectors.isIPluginFunction(IPaymaster.postOp.selector));
+    }
 }