Skip to content

Commit

Permalink
feat!: disallow using dependencies for hooks (#29)
Browse files Browse the repository at this point in the history
  • Loading branch information
@jaypaik
jaypaik authored Jan 23, 2024
1 parent 7846d38 commit 510b541
Show file tree
Hide file tree
Showing 2 changed files with 76 additions and 161 deletions.
27 changes: 17 additions & 10 deletions src/account/PluginManagerInternals.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -288,7 +288,6 @@ abstract contract PluginManagerInternals is IPluginManager {
_storage.pluginData[plugin].dependencies = dependencies;

// Update components according to the manifest.
// All conflicts should revert.

// Mark whether or not this plugin may spend native token amounts
if (manifest.canSpendNativeToken) {
Expand Down Expand Up @@ -380,6 +379,9 @@ abstract contract PluginManagerInternals is IPluginManager {
}
}

// Hooks are not allowed to be provided as dependencies, so we use an empty array for resolving them.
FunctionReference[] memory emptyDependencies;

length = manifest.preUserOpValidationHooks.length;
for (uint256 i = 0; i < length;) {
ManifestAssociatedFunction memory mh = manifest.preUserOpValidationHooks[i];
Expand All @@ -388,7 +390,7 @@ abstract contract PluginManagerInternals is IPluginManager {
_resolveManifestFunction(
mh.associatedFunction,
plugin,
dependencies,
emptyDependencies,
ManifestAssociatedFunctionType.PRE_HOOK_ALWAYS_DENY
)
);
Expand All @@ -406,7 +408,7 @@ abstract contract PluginManagerInternals is IPluginManager {
_resolveManifestFunction(
mh.associatedFunction,
plugin,
dependencies,
emptyDependencies,
ManifestAssociatedFunctionType.PRE_HOOK_ALWAYS_DENY
)
);
Expand All @@ -421,10 +423,10 @@ abstract contract PluginManagerInternals is IPluginManager {
_addExecHooks(
mh.executionSelector,
_resolveManifestFunction(
mh.preExecHook, plugin, dependencies, ManifestAssociatedFunctionType.PRE_HOOK_ALWAYS_DENY
mh.preExecHook, plugin, emptyDependencies, ManifestAssociatedFunctionType.PRE_HOOK_ALWAYS_DENY
),
_resolveManifestFunction(
mh.postExecHook, plugin, dependencies, ManifestAssociatedFunctionType.NONE
mh.postExecHook, plugin, emptyDependencies, ManifestAssociatedFunctionType.NONE
)
);

Expand Down Expand Up @@ -488,18 +490,20 @@ abstract contract PluginManagerInternals is IPluginManager {
}

// Remove components according to the manifest, in reverse order (by component type) of their installation.
// If any expected components are missing, revert.

// Hooks are not allowed to be provided as dependencies, so we use an empty array for resolving them.
FunctionReference[] memory emptyDependencies;

length = manifest.executionHooks.length;
for (uint256 i = 0; i < length;) {
ManifestExecutionHook memory mh = manifest.executionHooks[i];
_removeExecHooks(
mh.executionSelector,
_resolveManifestFunction(
mh.preExecHook, plugin, dependencies, ManifestAssociatedFunctionType.PRE_HOOK_ALWAYS_DENY
mh.preExecHook, plugin, emptyDependencies, ManifestAssociatedFunctionType.PRE_HOOK_ALWAYS_DENY
),
_resolveManifestFunction(
mh.postExecHook, plugin, dependencies, ManifestAssociatedFunctionType.NONE
mh.postExecHook, plugin, emptyDependencies, ManifestAssociatedFunctionType.NONE
)
);

Expand All @@ -516,7 +520,7 @@ abstract contract PluginManagerInternals is IPluginManager {
_resolveManifestFunction(
mh.associatedFunction,
plugin,
dependencies,
emptyDependencies,
ManifestAssociatedFunctionType.PRE_HOOK_ALWAYS_DENY
)
);
Expand All @@ -534,7 +538,7 @@ abstract contract PluginManagerInternals is IPluginManager {
_resolveManifestFunction(
mh.associatedFunction,
plugin,
dependencies,
emptyDependencies,
ManifestAssociatedFunctionType.PRE_HOOK_ALWAYS_DENY
)
);
Expand Down Expand Up @@ -698,6 +702,9 @@ abstract contract PluginManagerInternals is IPluginManager {
if (manifestFunction.functionType == ManifestAssociatedFunctionType.SELF) {
return FunctionReferenceLib.pack(plugin, manifestFunction.functionId);
} else if (manifestFunction.functionType == ManifestAssociatedFunctionType.DEPENDENCY) {
if (manifestFunction.dependencyIndex >= dependencies.length) {
revert InvalidPluginManifest();
}
return dependencies[manifestFunction.dependencyIndex];
} else if (manifestFunction.functionType == ManifestAssociatedFunctionType.RUNTIME_VALIDATION_ALWAYS_ALLOW)
{
Expand Down
210 changes: 59 additions & 151 deletions test/account/AccountExecHooks.t.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@ import {
PluginManifest
} from "../../src/interfaces/IPlugin.sol";
import {SingleOwnerPlugin} from "../../src/plugins/owner/SingleOwnerPlugin.sol";
import {PluginManagerInternals} from "../../src/account/PluginManagerInternals.sol";
import {UpgradeableModularAccount} from "../../src/account/UpgradeableModularAccount.sol";
import {FunctionReference, FunctionReferenceLib} from "../../src/helpers/FunctionReferenceLib.sol";

Expand Down Expand Up @@ -196,99 +197,47 @@ contract AccountExecHooksTest is OptimizedTest {
_uninstallPlugin(mockPlugin1);
}

function test_overlappingExecHookPairs_install() public {
function test_overlappingPreExecHooks_install() public {
// Install the first plugin.
_installPlugin1WithHooks(
_EXEC_SELECTOR,
ManifestFunction({
functionType: ManifestAssociatedFunctionType.SELF,
functionId: _PRE_HOOK_FUNCTION_ID_1,
functionType: ManifestAssociatedFunctionType.PRE_HOOK_ALWAYS_DENY,
functionId: 0,
dependencyIndex: 0
}),
ManifestFunction({
functionType: ManifestAssociatedFunctionType.SELF,
functionId: _POST_HOOK_FUNCTION_ID_2,
dependencyIndex: 0
})
ManifestFunction(ManifestAssociatedFunctionType.NONE, 0, 0)
);

// Install a second plugin that applies the first plugin's hook pair to the same selector.
FunctionReference[] memory dependencies = new FunctionReference[](2);
dependencies[0] = FunctionReferenceLib.pack(address(mockPlugin1), _PRE_HOOK_FUNCTION_ID_1);
dependencies[1] = FunctionReferenceLib.pack(address(mockPlugin1), _POST_HOOK_FUNCTION_ID_2);
_installPlugin2WithHooks(
// Install a second plugin that applies the same pre hook on the same selector.
_installPlugin2WithHooksExpectSuccess(
_EXEC_SELECTOR,
ManifestFunction({
functionType: ManifestAssociatedFunctionType.DEPENDENCY,
functionType: ManifestAssociatedFunctionType.PRE_HOOK_ALWAYS_DENY,
functionId: 0,
dependencyIndex: 0
}),
ManifestFunction({
functionType: ManifestAssociatedFunctionType.DEPENDENCY,
functionId: 0,
dependencyIndex: 1
}),
dependencies
ManifestFunction(ManifestAssociatedFunctionType.NONE, 0, 0),
new FunctionReference[](0)
);

vm.stopPrank();
}

/// @dev Plugin 1 hook pair: [1, 2]
/// Plugin 2 hook pair: [1, 2]
/// Expected execution: [1, 2]
function test_overlappingExecHookPairs_run() public {
test_overlappingExecHookPairs_install();

// Expect the pre hook to be called just once.
vm.expectCall(
address(mockPlugin1),
abi.encodeWithSelector(
IPlugin.preExecutionHook.selector,
_PRE_HOOK_FUNCTION_ID_1,
address(this), // caller
0, // msg.value in call to account
abi.encodeWithSelector(_EXEC_SELECTOR)
),
1
);

// Expect the post hook to be called just once, with the expected data.
vm.expectCall(
address(mockPlugin1),
abi.encodeWithSelector(IPlugin.postExecutionHook.selector, _POST_HOOK_FUNCTION_ID_2, ""),
1
);

function test_overlappingPreExecHooks_run() public {
(bool success,) = address(account).call(abi.encodeWithSelector(_EXEC_SELECTOR));
assertTrue(success);
assertFalse(success);
}

function test_overlappingExecHookPairs_uninstall() public {
test_overlappingExecHookPairs_install();
function test_overlappingPreExecHooks_uninstall() public {
test_overlappingPreExecHooks_install();

// Uninstall the second plugin.
_uninstallPlugin(mockPlugin2);

// Expect the pre/post hooks to still exist after uninstalling a plugin with a duplicate hook.
vm.expectCall(
address(mockPlugin1),
abi.encodeWithSelector(
IPlugin.preExecutionHook.selector,
_PRE_HOOK_FUNCTION_ID_1,
address(this), // caller
0, // msg.value in call to account
abi.encodeWithSelector(_EXEC_SELECTOR)
),
1
);
vm.expectCall(
address(mockPlugin1),
abi.encodeWithSelector(IPlugin.postExecutionHook.selector, _POST_HOOK_FUNCTION_ID_2, ""),
1
);
// Expect the pre hook to still exist.
(bool success,) = address(account).call(abi.encodeWithSelector(_EXEC_SELECTOR));
assertTrue(success);
assertFalse(success);

// Uninstall the first plugin.
_uninstallPlugin(mockPlugin1);
Expand All @@ -298,7 +247,7 @@ contract AccountExecHooksTest is OptimizedTest {
assertFalse(success);
}

function test_overlappingExecHookPairsOnPost_install() public {
function test_execHookDependencies_failInstall() public {
// Install the first plugin.
_installPlugin1WithHooks(
_EXEC_SELECTOR,
Expand All @@ -314,98 +263,29 @@ contract AccountExecHooksTest is OptimizedTest {
})
);

// Install the second plugin.
FunctionReference[] memory dependencies = new FunctionReference[](1);
dependencies[0] = FunctionReferenceLib.pack(address(mockPlugin1), _POST_HOOK_FUNCTION_ID_2);
_installPlugin2WithHooks(
// Attempt to install a second plugin that applies the first plugin's hook pair (as dependencies) to the
// same selector. This should revert.
FunctionReference[] memory dependencies = new FunctionReference[](2);
dependencies[0] = FunctionReferenceLib.pack(address(mockPlugin1), _PRE_HOOK_FUNCTION_ID_1);
dependencies[1] = FunctionReferenceLib.pack(address(mockPlugin1), _POST_HOOK_FUNCTION_ID_2);

_installPlugin2WithHooksExpectFail(
_EXEC_SELECTOR,
ManifestFunction({
functionType: ManifestAssociatedFunctionType.SELF,
functionId: _PRE_HOOK_FUNCTION_ID_3,
functionType: ManifestAssociatedFunctionType.DEPENDENCY,
functionId: 0,
dependencyIndex: 0
}),
ManifestFunction({
functionType: ManifestAssociatedFunctionType.DEPENDENCY,
functionId: 0,
dependencyIndex: 0
dependencyIndex: 1
}),
dependencies
);
}

/// @dev Plugin 1 hook pair: [1, 2]
/// Plugin 2 hook pair: [3, 2]
/// Expected execution: [1, 2], [3, 2]
function test_overlappingExecHookPairsOnPost_run() public {
test_overlappingExecHookPairsOnPost_install();

// Expect each pre hook to be called once.
vm.expectCall(
address(mockPlugin1),
abi.encodeWithSelector(
IPlugin.preExecutionHook.selector,
_PRE_HOOK_FUNCTION_ID_1,
address(this), // caller
0, // msg.value in call to account
abi.encodeWithSelector(_EXEC_SELECTOR)
),
1
);
vm.expectCall(
address(mockPlugin2),
abi.encodeWithSelector(
IPlugin.preExecutionHook.selector,
_PRE_HOOK_FUNCTION_ID_3,
address(this), // caller
0, // msg.value in call to account
abi.encodeWithSelector(_EXEC_SELECTOR)
),
1
);

// Expect the post hook to be called twice, with the expected data.
vm.expectCall(
address(mockPlugin1),
abi.encodeWithSelector(IPlugin.postExecutionHook.selector, _POST_HOOK_FUNCTION_ID_2, ""),
2
dependencies,
abi.encodePacked(PluginManagerInternals.InvalidPluginManifest.selector)
);

(bool success,) = address(account).call(abi.encodeWithSelector(_EXEC_SELECTOR));
assertTrue(success);
}

function test_overlappingExecHookPairsOnPost_uninstall() public {
test_overlappingExecHookPairsOnPost_install();

// Uninstall the second plugin.
_uninstallPlugin(mockPlugin2);

// Expect the pre/post hooks to still exist after uninstalling a plugin with a duplicate hook.
vm.expectCall(
address(mockPlugin1),
abi.encodeWithSelector(
IPlugin.preExecutionHook.selector,
_PRE_HOOK_FUNCTION_ID_1,
address(this), // caller
0, // msg.value in call to account
abi.encodeWithSelector(_EXEC_SELECTOR)
),
1
);
vm.expectCall(
address(mockPlugin1),
abi.encodeWithSelector(IPlugin.postExecutionHook.selector, _POST_HOOK_FUNCTION_ID_2, ""),
1
);
(bool success,) = address(account).call(abi.encodeWithSelector(_EXEC_SELECTOR));
assertTrue(success);

// Uninstall the first plugin.
_uninstallPlugin(mockPlugin1);

// Execution selector should no longer exist.
(success,) = address(account).call(abi.encodeWithSelector(_EXEC_SELECTOR));
assertFalse(success);
vm.stopPrank();
}

function _installPlugin1WithHooks(
Expand All @@ -430,7 +310,7 @@ contract AccountExecHooksTest is OptimizedTest {
});
}

function _installPlugin2WithHooks(
function _installPlugin2WithHooksExpectSuccess(
bytes4 selector,
ManifestFunction memory preHook,
ManifestFunction memory postHook,
Expand Down Expand Up @@ -461,6 +341,34 @@ contract AccountExecHooksTest is OptimizedTest {
});
}

function _installPlugin2WithHooksExpectFail(
bytes4 selector,
ManifestFunction memory preHook,
ManifestFunction memory postHook,
FunctionReference[] memory dependencies,
bytes memory revertData
) internal {
if (preHook.functionType == ManifestAssociatedFunctionType.DEPENDENCY) {
m2.dependencyInterfaceIds.push(type(IPlugin).interfaceId);
}
if (postHook.functionType == ManifestAssociatedFunctionType.DEPENDENCY) {
m2.dependencyInterfaceIds.push(type(IPlugin).interfaceId);
}

m2.executionHooks.push(ManifestExecutionHook(selector, preHook, postHook));

mockPlugin2 = new MockPlugin(m2);
manifestHash2 = keccak256(abi.encode(mockPlugin2.pluginManifest()));

vm.expectRevert(revertData);
account.installPlugin({
plugin: address(mockPlugin2),
manifestHash: manifestHash2,
pluginInitData: bytes(""),
dependencies: dependencies
});
}

function _uninstallPlugin(MockPlugin plugin) internal {
vm.expectEmit(true, true, true, true);
emit ReceivedCall(abi.encodeCall(IPlugin.onUninstall, (bytes(""))), 0);
Expand Down

0 comments on commit 510b541

Please sign in to comment.