[4/n permissions] feat: add native token spend limit plugin

.gitmodules

@@ -7,3 +7,6 @@
 [submodule "lib/forge-std"]
 	path = lib/forge-std
 	url = https://github.com/foundry-rs/forge-std
+[submodule "lib/modular-account-libs"]
+	path = lib/modular-account-libs
+	url = https://github.com/erc6900/modular-account-libs

remappings.txt

@@ -2,3 +2,4 @@ ds-test/=lib/forge-std/lib/ds-test/src/
 forge-std/=lib/forge-std/src/
 @eth-infinitism/account-abstraction/=lib/account-abstraction/contracts/
 @openzeppelin/=lib/openzeppelin-contracts/
+@modular-account-libs/=lib/modular-account-libs/src/

src/plugins/BasePlugin.sol

@@ -2,6 +2,8 @@
 pragma solidity ^0.8.25;
 
 import {ERC165, IERC165} from "@openzeppelin/contracts/utils/introspection/ERC165.sol";
+import {IAccountExecute} from "@eth-infinitism/account-abstraction/interfaces/IAccountExecute.sol";
+import {PackedUserOperation} from "@eth-infinitism/account-abstraction/interfaces/PackedUserOperation.sol";
 
 import {IPlugin} from "../interfaces/IPlugin.sol";
 
@@ -27,4 +29,26 @@ abstract contract BasePlugin is ERC165, IPlugin {
     function supportsInterface(bytes4 interfaceId) public view virtual override(ERC165, IERC165) returns (bool) {
         return interfaceId == type(IPlugin).interfaceId || super.supportsInterface(interfaceId);
     }
+
+    function _getSelectorAndCalldata(bytes calldata data) internal pure returns (bytes4, bytes memory) {
+        if (bytes4(data[:4]) == IAccountExecute.executeUserOp.selector) {
+            (PackedUserOperation memory uo,) = abi.decode(data[4:], (PackedUserOperation, bytes32));
+            bytes4 selector;
+            bytes memory callData = uo.callData;
+            // Bytes arr representation: [bytes32(len), bytes4(executeUserOp.selector), bytes4(actualSelector),
+            // bytes(actualCallData)]
+            // 1. Copy actualSelector into a new var
+            // 2. Shorten bytes arry by 8 by: store length - 8 into the new pointer location
+            // 3. Move the callData pointer by 8
+            assembly {
+                selector := mload(add(callData, 36))
+
+                let len := mload(callData)
+                mstore(add(callData, 8), sub(len, 8))
+                callData := add(callData, 8)
+            }
+            return (selector, callData);
+        }
+        return (bytes4(data[:4]), data[4:]);
+    }
 }

src/plugins/ERC20TokenLimitPlugin.sol

@@ -0,0 +1,154 @@
+// SPDX-License-Identifier: GPL-3.0
+pragma solidity ^0.8.25;
+
+import {PackedUserOperation} from "@eth-infinitism/account-abstraction/interfaces/PackedUserOperation.sol";
+import {UserOperationLib} from "@eth-infinitism/account-abstraction/core/UserOperationLib.sol";
+import {EnumerableSet} from "@openzeppelin/contracts/utils/structs/EnumerableSet.sol";
+import {IERC20} from "@openzeppelin/contracts/token/ERC20/IERC20.sol";
+import {
+    SetValue,
+    AssociatedLinkedListSet,
+    AssociatedLinkedListSetLib
+} from "@modular-account-libs/libraries/AssociatedLinkedListSetLib.sol";
+
+import {PluginManifest, PluginMetadata} from "../interfaces/IPlugin.sol";
+import {IStandardExecutor, Call} from "../interfaces/IStandardExecutor.sol";
+import {IPlugin} from "../interfaces/IPlugin.sol";
+import {IExecutionHook} from "../interfaces/IExecutionHook.sol";
+import {BasePlugin, IERC165} from "./BasePlugin.sol";
+
+/// @title ERC20 Token Limit Plugin
+/// @author ERC-6900 Authors
+/// @notice This plugin supports an ERC20 token spend limit. This should be combined with a contract whitelist
+/// plugin to make sure that token transfers not tracked by the plugin don't happen.
+/// Note: this plugin is opinionated on what selectors can be called for token contracts to guard against weird
+/// edge cases like DAI. You wouldn't be able to use uni v2 pairs directly as the pair contract is also the LP
+/// token contract
+contract ERC20TokenLimitPlugin is BasePlugin, IExecutionHook {
+    using UserOperationLib for PackedUserOperation;
+    using EnumerableSet for EnumerableSet.AddressSet;
+    using AssociatedLinkedListSetLib for AssociatedLinkedListSet;
+
+    struct ERC20SpendLimit {
+        address token;
+        uint256[] limits;
+    }
+
+    string internal constant _NAME = "ERC20 Token Limit Plugin";
+    string internal constant _VERSION = "1.0.0";
+    string internal constant _AUTHOR = "ERC-6900 Authors";
+
+    mapping(uint8 functionId => mapping(address token => mapping(address account => uint256 limit))) public limits;
+    AssociatedLinkedListSet internal _tokenList;
+
+    error ExceededTokenLimit();
+    error ExceededNumberOfEntities();
+    error SelectorNotAllowed();
+
+    function updateLimits(uint8 functionId, address token, uint256 newLimit) external {
+        _tokenList.tryAdd(msg.sender, SetValue.wrap(bytes30(bytes20(token))));
+        limits[functionId][token][msg.sender] = newLimit;
+    }
+
+    /// @inheritdoc IExecutionHook
+    function preExecutionHook(uint8 functionId, address, uint256, bytes calldata data)
+        external
+        override
+        returns (bytes memory)
+    {
+        (bytes4 selector, bytes memory callData) = _getSelectorAndCalldata(data);
+
+        if (selector == IStandardExecutor.execute.selector) {
+            (address token,, bytes memory innerCalldata) = abi.decode(callData, (address, uint256, bytes));
+            if (_tokenList.contains(msg.sender, SetValue.wrap(bytes30(bytes20(token))))) {
+                _decrementLimit(functionId, token, innerCalldata);
+            }
+        } else if (selector == IStandardExecutor.executeBatch.selector) {
+            Call[] memory calls = abi.decode(callData, (Call[]));
+            for (uint256 i = 0; i < calls.length; i++) {
+                if (_tokenList.contains(msg.sender, SetValue.wrap(bytes30(bytes20(calls[i].target))))) {
+                    _decrementLimit(functionId, calls[i].target, calls[i].data);
+                }
+            }
+        }
+
+        return "";
+    }
+
+    /// @inheritdoc IPlugin
+    function onInstall(bytes calldata data) external override {
+        (uint8 startFunctionId, ERC20SpendLimit[] memory spendLimits) =
+            abi.decode(data, (uint8, ERC20SpendLimit[]));
+
+        if (startFunctionId + spendLimits.length > type(uint8).max) {
+            revert ExceededNumberOfEntities();
+        }
+
+        for (uint8 i = 0; i < spendLimits.length; i++) {
+            _tokenList.tryAdd(msg.sender, SetValue.wrap(bytes30(bytes20(spendLimits[i].token))));
+            for (uint256 j = 0; j < spendLimits[i].limits.length; j++) {
+                limits[i + startFunctionId][spendLimits[i].token][msg.sender] = spendLimits[i].limits[j];
+            }
+        }
+    }
+
+    /// @inheritdoc IPlugin
+    function onUninstall(bytes calldata data) external override {
+        (address token, uint8 functionId) = abi.decode(data, (address, uint8));
+        delete limits[functionId][token][msg.sender];
+    }
+
+    function getTokensForAccount(address account) external view returns (address[] memory tokens) {
+        SetValue[] memory set = _tokenList.getAll(account);
+        tokens = new address[](set.length);
+        for (uint256 i = 0; i < tokens.length; i++) {
+            tokens[i] = address(bytes20(bytes32(SetValue.unwrap(set[i]))));
+        }
+        return tokens;
+    }
+
+    /// @inheritdoc IExecutionHook
+    function postExecutionHook(uint8, bytes calldata) external pure override {
+        revert NotImplemented();
+    }
+
+    /// @inheritdoc IPlugin
+    // solhint-disable-next-line no-empty-blocks
+    function pluginManifest() external pure override returns (PluginManifest memory) {}
+
+    /// @inheritdoc IPlugin
+    function pluginMetadata() external pure virtual override returns (PluginMetadata memory) {
+        PluginMetadata memory metadata;
+        metadata.name = _NAME;
+        metadata.version = _VERSION;
+        metadata.author = _AUTHOR;
+
+        metadata.permissionRequest = new string[](1);
+        metadata.permissionRequest[0] = "erc20-token-limit";
+        return metadata;
+    }
+
+    /// @inheritdoc BasePlugin
+    function supportsInterface(bytes4 interfaceId) public view override(BasePlugin, IERC165) returns (bool) {
+        return super.supportsInterface(interfaceId);
+    }
+
+    function _decrementLimit(uint8 functionId, address token, bytes memory innerCalldata) internal {
+        bytes4 selector;
+        uint256 spend;
+        assembly {
+            selector := mload(add(innerCalldata, 32)) // 0:32 is arr len, 32:36 is selector
+            spend := mload(add(innerCalldata, 68)) // 36:68 is recipient, 68:100 is spend
+        }
+        if (selector == IERC20.transfer.selector || selector == IERC20.approve.selector) {
+            uint256 limit = limits[functionId][token][msg.sender];
+            if (spend > limit) {
+                revert ExceededTokenLimit();
+            }
+            // solhint-disable-next-line reentrancy
+            limits[functionId][token][msg.sender] = limit - spend;
+        } else {
+            revert SelectorNotAllowed();
+        }
+    }
+}

src/plugins/NativeTokenLimitPlugin.sol

@@ -0,0 +1,167 @@
+// SPDX-License-Identifier: GPL-3.0
+pragma solidity ^0.8.25;
+
+import {PackedUserOperation} from "@eth-infinitism/account-abstraction/interfaces/PackedUserOperation.sol";
+import {UserOperationLib} from "@eth-infinitism/account-abstraction/core/UserOperationLib.sol";
+import {EnumerableSet} from "@openzeppelin/contracts/utils/structs/EnumerableSet.sol";
+
+import {PluginManifest, PluginMetadata} from "../interfaces/IPlugin.sol";
+import {IStandardExecutor, Call} from "../interfaces/IStandardExecutor.sol";
+import {IPlugin} from "../interfaces/IPlugin.sol";
+import {IExecutionHook} from "../interfaces/IExecutionHook.sol";
+import {IValidationHook} from "../interfaces/IValidationHook.sol";
+import {BasePlugin, IERC165} from "./BasePlugin.sol";
+
+/// @title Native Token Limit Plugin
+/// @author ERC-6900 Authors
+/// @notice This plugin supports a single total native token spend limit.
+/// It tracks a total spend limit across UserOperation gas limits and native token transfers.
+/// If a non whitelisted paymaster is used, UO gas would not cause the limit to decrease.
+/// If a whitelisted paymaster is used, gas is still counted towards the limit
+contract NativeTokenLimitPlugin is BasePlugin, IExecutionHook, IValidationHook {
+    using UserOperationLib for PackedUserOperation;
+    using EnumerableSet for EnumerableSet.Bytes32Set;
+
+    string internal constant _NAME = "Native Token Limit";
+    string internal constant _VERSION = "1.0.0";
+    string internal constant _AUTHOR = "ERC-6900 Authors";
+
+    mapping(uint256 funcIds => mapping(address account => uint256 limit)) public limits;
+    // Accounts should add paymasters that still use the accounts tokens here
+    // E.g. ERC20 paymasters that pull funds from the account
+    mapping(address paymaster => mapping(address account => bool allowed)) public specialPaymasters;
+
+    error ExceededNativeTokenLimit();
+    error ExceededNumberOfEntities();
+
+    function updateLimits(uint8 functionId, uint256 newLimit) external {
+        limits[functionId][msg.sender] = newLimit;
+    }
+
+    function updateSpecialPaymaster(address paymaster, bool allowed) external {
+        specialPaymasters[paymaster][msg.sender] = allowed;
+    }
+
+    /// @inheritdoc IValidationHook
+    function preUserOpValidationHook(uint8 functionId, PackedUserOperation calldata userOp, bytes32)
+        external
+        returns (uint256)
+    {
+        // Decrease limit only if no paymaster is used, or if its a special paymaster
+        if (
+            userOp.paymasterAndData.length == 0
+                || specialPaymasters[address(bytes20(userOp.paymasterAndData[:20]))][msg.sender]
+        ) {
+            uint256 vgl = UserOperationLib.unpackVerificationGasLimit(userOp);
+            uint256 cgl = UserOperationLib.unpackCallGasLimit(userOp);
+            uint256 pvgl;
+            uint256 ppogl;
+            if (userOp.paymasterAndData.length > 0) {
+                // Can skip the EP length check here since it would have reverted there if it was invalid
+                (, pvgl, ppogl) = UserOperationLib.unpackPaymasterStaticFields(userOp.paymasterAndData);
+            }
+            uint256 totalGas = userOp.preVerificationGas + vgl + cgl + pvgl + ppogl;
+            uint256 usage = totalGas * UserOperationLib.unpackMaxFeePerGas(userOp);
+
+            uint256 limit = limits[functionId][msg.sender];
+            if (usage > limit) {
+                revert ExceededNativeTokenLimit();
+            }
+            limits[functionId][msg.sender] = limit - usage;
+        }
+        return 0;
+    }
+
+    /// @inheritdoc IExecutionHook
+    function preExecutionHook(uint8 functionId, address, uint256, bytes calldata data)
+        external
+        override
+        returns (bytes memory)
+    {
+        return _checkAndDecrementLimit(functionId, data);
+    }
+
+    /// @inheritdoc IPlugin
+    function onInstall(bytes calldata data) external override {
+        (uint8 startFunctionId, uint256[] memory spendLimits) = abi.decode(data, (uint8, uint256[]));
+
+        if (startFunctionId + spendLimits.length > type(uint8).max) {
+            revert ExceededNumberOfEntities();
+        }
+
+        for (uint256 i = 0; i < spendLimits.length; i++) {
+            limits[i + startFunctionId][msg.sender] = spendLimits[i];
+        }
+    }
+
+    /// @inheritdoc IPlugin
+    function onUninstall(bytes calldata data) external override {
+        // This is the highest functionId that's being used by the account
+        uint8 functionId = abi.decode(data, (uint8));
+        for (uint256 i = 0; i < functionId; i++) {
+            delete limits[i][msg.sender];
+        }
+    }
+
+    /// @inheritdoc IExecutionHook
+    function postExecutionHook(uint8, bytes calldata) external pure override {
+        revert NotImplemented();
+    }
+
+    // No implementation, no revert
+    // Runtime spends no account gas, and we check native token spend limits in exec hooks
+    function preRuntimeValidationHook(uint8, address, uint256, bytes calldata, bytes calldata)
+        external
+        pure
+        override
+    {} // solhint-disable-line no-empty-blocks
+
+    /// @inheritdoc IPlugin
+    // solhint-disable-next-line no-empty-blocks
+    function pluginManifest() external pure override returns (PluginManifest memory) {}
+
+    /// @inheritdoc IPlugin
+    function pluginMetadata() external pure virtual override returns (PluginMetadata memory) {
+        PluginMetadata memory metadata;
+        metadata.name = _NAME;
+        metadata.version = _VERSION;
+        metadata.author = _AUTHOR;
+
+        metadata.permissionRequest = new string[](2);
+        metadata.permissionRequest[0] = "native-token-limit";
+        metadata.permissionRequest[1] = "gas-limit";
+        return metadata;
+    }
+
+    // ┏━━━━━━━━━━━━━━━┓
+    // ┃    EIP-165    ┃
+    // ┗━━━━━━━━━━━━━━━┛
+
+    /// @inheritdoc BasePlugin
+    function supportsInterface(bytes4 interfaceId) public view override(BasePlugin, IERC165) returns (bool) {
+        return interfaceId == type(IExecutionHook).interfaceId || super.supportsInterface(interfaceId);
+    }
+
+    function _checkAndDecrementLimit(uint8 functionId, bytes calldata data) internal returns (bytes memory) {
+        (bytes4 selector, bytes memory callData) = _getSelectorAndCalldata(data);
+
+        uint256 value;
+        // Get value being sent
+        if (selector == IStandardExecutor.execute.selector) {
+            (, value) = abi.decode(callData, (address, uint256));
+        } else if (selector == IStandardExecutor.executeBatch.selector) {
+            Call[] memory calls = abi.decode(callData, (Call[]));
+            for (uint256 i = 0; i < calls.length; i++) {
+                value += calls[i].value;
+            }
+        }
+
+        uint256 limit = limits[functionId][msg.sender];
+        if (value > limit) {
+            revert ExceededNativeTokenLimit();
+        }
+        limits[functionId][msg.sender] = limit - value;
+
+        return "";
+    }
+}

test/mocks/MockERC20.sol

@@ -0,0 +1,12 @@
+// SPDX-License-Identifier: UNLICENSED
+pragma solidity ^0.8.19;
+
+import {ERC20} from "@openzeppelin/contracts/token/ERC20/ERC20.sol";
+
+contract MockERC20 is ERC20 {
+    constructor() ERC20("MockERC20", "MERC") {}
+
+    function mint(address to, uint256 amount) external {
+        _mint(to, amount);
+    }
+}