Skip to content

Commit

Permalink
Allow direct plugin calls with validation & permission hooks (#90)
Browse files Browse the repository at this point in the history
  • Loading branch information
@Zer0dot @adamegyed
Zer0dot authored and adamegyed committed Jul 26, 2024
1 parent 24cc733 commit 39d327e
Show file tree
Hide file tree
Showing 9 changed files with 315 additions and 93 deletions.
36 changes: 18 additions & 18 deletions .solhint-test.json
View file
Original file line number Diff line number Diff line change
@@ -1,20 +1,20 @@
{
"extends": "solhint:recommended",
"rules": {
"func-name-mixedcase": "off",
"immutable-vars-naming": ["error"],
"no-unused-import": ["error"],
"compiler-version": ["error", ">=0.8.19"],
"custom-errors": "off",
"func-visibility": ["error", { "ignoreConstructors": true }],
"max-line-length": ["error", 120],
"max-states-count": ["warn", 30],
"modifier-name-mixedcase": ["error"],
"private-vars-leading-underscore": ["error"],
"no-inline-assembly": "off",
"avoid-low-level-calls": "off",
"one-contract-per-file": "off",
"no-empty-blocks": "off"
}
"extends": "solhint:recommended",
"rules": {
"func-name-mixedcase": "off",
"immutable-vars-naming": ["error"],
"no-unused-import": ["error"],
"compiler-version": ["error", ">=0.8.19"],
"custom-errors": "off",
"func-visibility": ["error", { "ignoreConstructors": true }],
"max-line-length": ["error", 120],
"max-states-count": ["warn", 30],
"modifier-name-mixedcase": ["error"],
"private-vars-leading-underscore": ["error"],
"no-inline-assembly": "off",
"avoid-low-level-calls": "off",
"one-contract-per-file": "off",
"no-empty-blocks": "off",
"reason-string": ["warn", { "maxLength": 64 }]
}
}
2 changes: 1 addition & 1 deletion src/account/AccountStorage.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -30,7 +30,7 @@ struct ValidationData {
bool isGlobal;
// Whether or not this validation is a signature validator.
bool isSignatureValidation;
// The pre validation hooks for this function selector.
// The pre validation hooks for this validation function.
PluginEntity[] preValidationHooks;
// Permission hooks for this validation function.
EnumerableSet.Bytes32Set permissionHooks;
Expand Down
35 changes: 19 additions & 16 deletions src/account/PluginManager2.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@ import {IPlugin} from "../interfaces/IPlugin.sol";
import {PluginEntity, ValidationConfig} from "../interfaces/IPluginManager.sol";
import {PluginEntityLib} from "../helpers/PluginEntityLib.sol";
import {ValidationConfigLib} from "../helpers/ValidationConfigLib.sol";
import {ValidationData, getAccountStorage, toSetValue, toPluginEntity} 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.
Expand All @@ -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;
uint32 internal constant _SELF_PERMIT_VALIDATION_FUNCTIONID = type(uint32).max;

error PreValidationAlreadySet(PluginEntity validationFunction, PluginEntity preValidationFunction);
error ValidationAlreadySet(bytes4 selector, PluginEntity validationFunction);
Expand All @@ -32,7 +33,7 @@ abstract contract PluginManager2 {
bytes memory permissionHooks
) internal {
ValidationData storage _validationData =
getAccountStorage().validationData[validationConfig.functionReference()];
getAccountStorage().validationData[validationConfig.pluginEntity()];

if (preValidationHooks.length > 0) {
(PluginEntity[] memory preValidationFunctions, bytes[] memory initDatas) =
Expand Down Expand Up @@ -63,7 +64,7 @@ abstract contract PluginManager2 {
ExecutionHook memory permissionFunction = permissionFunctions[i];

if (!_validationData.permissionHooks.add(toSetValue(permissionFunction))) {
revert PermissionAlreadySet(validationConfig.functionReference(), permissionFunction);
revert PermissionAlreadySet(validationConfig.pluginEntity(), permissionFunction);
}

if (initDatas[i].length > 0) {
Expand All @@ -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());
revert ValidationAlreadySet(selector, validationConfig.pluginEntity());
}
}

if (installData.length > 0) {
address plugin = validationConfig.plugin();
IPlugin(plugin).onInstall(installData);
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) {
IPlugin(validationConfig.plugin()).onInstall(installData);
}
}
}

Expand Down Expand Up @@ -120,12 +123,12 @@ abstract contract PluginManager2 {

// Clear permission hooks
EnumerableSet.Bytes32Set storage permissionHooks = _validationData.permissionHooks;
uint256 i = 0;
while (permissionHooks.length() > 0) {
PluginEntity permissionHook = toPluginEntity(permissionHooks.at(0));
permissionHooks.remove(toSetValue(permissionHook));
(address permissionHookPlugin,) = PluginEntityLib.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;
Expand Down
131 changes: 89 additions & 42 deletions src/account/UpgradeableModularAccount.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -78,22 +78,21 @@ contract UpgradeableModularAccount is
error SignatureValidationInvalid(address plugin, uint32 entityId);
error UnexpectedAggregator(address plugin, uint32 entityId, address aggregator);
error UnrecognizedFunction(bytes4 selector);
error UserOpValidationFunctionMissing(bytes4 selector);
error ValidationFunctionMissing(bytes4 selector);
error ValidationDoesNotApply(bytes4 selector, address plugin, uint32 entityId, bool isGlobal);
error ValidationSignatureSegmentMissing();
error SignatureSegmentOutOfOrder();

// 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) {
Expand Down Expand Up @@ -136,7 +135,7 @@ contract UpgradeableModularAccount is
revert UnrecognizedFunction(msg.sig);
}

_checkPermittedCallerIfNotFromEP();
_checkPermittedCallerAndAssociatedHooks();

PostExecToRun[] memory postExecHooks;
// Cache post-exec hooks in memory
Expand Down Expand Up @@ -500,17 +499,7 @@ contract UpgradeableModularAccount is
} else {
currentAuthData = "";
}

(address hookPlugin, uint32 hookEntityId) = preRuntimeValidationHooks[i].unpack();
try IValidationHook(hookPlugin).preRuntimeValidationHook(
hookEntityId, 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, hookEntityId, revertReason);
}
_doPreRuntimeValidationHook(preRuntimeValidationHooks[i], callData, currentAuthData);
}

if (authSegment.getIndex() != _RESERVED_VALIDATION_DATA_INDEX) {
Expand Down Expand Up @@ -605,9 +594,78 @@ contract UpgradeableModularAccount is
}
}

function _doPreRuntimeValidationHook(
PluginEntity validationHook,
bytes memory callData,
bytes memory currentAuthData
) internal {
(address hookPlugin, uint32 hookEntityId) = validationHook.unpack();
try IValidationHook(hookPlugin).preRuntimeValidationHook(
hookEntityId, 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, hookEntityId, revertReason);
}
}

// solhint-disable-next-line no-empty-blocks
function _authorizeUpgrade(address newImplementation) internal override {}

/**
* 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
) {
return (new PostExecToRun[](0), new PostExecToRun[](0));
}

PluginEntity directCallValidationKey = PluginEntityLib.pack(msg.sender, _SELF_PERMIT_VALIDATION_FUNCTIONID);

_checkIfValidationAppliesCallData(msg.data, directCallValidationKey, false);

// Direct call is allowed, run associated permission & validation hooks

// Validation hooks
PluginEntity[] 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);
}

function _checkIfValidationAppliesCallData(
bytes calldata callData,
PluginEntity validationFunction,
Expand Down Expand Up @@ -661,25 +719,6 @@ contract UpgradeableModularAccount is
}
}

function _checkIfValidationAppliesSelector(bytes4 selector, PluginEntity validationFunction, bool isGlobal)
internal
view
{
AccountStorage storage _storage = getAccountStorage();

// Check that the provided validation function is applicable to the selector
if (isGlobal) {
if (!_globalValidationAllowed(selector) || !_storage.validationData[validationFunction].isGlobal) {
revert UserOpValidationFunctionMissing(selector);
}
} else {
// Not global validation, but per-selector
if (!getAccountStorage().validationData[validationFunction].selectors.contains(toSetValue(selector))) {
revert UserOpValidationFunctionMissing(selector);
}
}
}

function _globalValidationAllowed(bytes4 selector) internal view returns (bool) {
if (
selector == this.execute.selector || selector == this.executeBatch.selector
Expand All @@ -693,14 +732,22 @@ contract UpgradeableModularAccount is
return getAccountStorage().selectorData[selector].allowGlobalValidation;
}

function _checkPermittedCallerIfNotFromEP() internal view {
function _checkIfValidationAppliesSelector(bytes4 selector, PluginEntity validationFunction, bool isGlobal)
internal
view
{
AccountStorage storage _storage = getAccountStorage();

if (
msg.sender != address(_ENTRY_POINT) && msg.sender != address(this)
&& !_storage.selectorData[msg.sig].isPublic
) {
revert ExecFromPluginNotPermitted(msg.sender, msg.sig);
// Check that the provided validation function is applicable to the selector
if (isGlobal) {
if (!_globalValidationAllowed(selector) || !_storage.validationData[validationFunction].isGlobal) {
revert ValidationFunctionMissing(selector);
}
} else {
// Not global validation, but per-selector
if (!getAccountStorage().validationData[validationFunction].selectors.contains(toSetValue(selector))) {
revert ValidationFunctionMissing(selector);
}
}
}
}
2 changes: 1 addition & 1 deletion src/helpers/ValidationConfigLib.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -78,7 +78,7 @@ library ValidationConfigLib {
return uint32(bytes4(ValidationConfig.unwrap(config) << 160));
}

function functionReference(ValidationConfig config) internal pure returns (PluginEntity) {
function pluginEntity(ValidationConfig config) internal pure returns (PluginEntity) {
return PluginEntity.wrap(bytes24(ValidationConfig.unwrap(config)));
}

Expand Down
Loading

0 comments on commit 39d327e

Please sign in to comment.