feat: add native token spend limit plugin

src/plugins/NativeTokenLimitPlugin.sol

@@ -0,0 +1,152 @@
+// 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 {IAccountExecute} from "@eth-infinitism/account-abstraction/interfaces/IAccountExecute.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 paymaster is used, UO gas would not cause the limit to decrease.
+contract NativeTokenLimitPlugin is BasePlugin, IExecutionHook, IValidationHook {
+    using UserOperationLib for PackedUserOperation;
+    using EnumerableSet for EnumerableSet.Bytes32Set;
+
+    string public constant NAME = "Native Token Limit";
+    string public constant VERSION = "1.0.0";
+    string public constant AUTHOR = "ERC-6900 Authors";
+
+    mapping(address account => uint256[] limits) public limits;
+
+    error ExceededNativeTokenLimit();
+    error ExceededNumberOfEntities();
+
+    function updateLimits(uint8 functionId, uint256 newLimit) external {
+        limits[msg.sender][functionId] = newLimit;
+    }
+
+    /// @inheritdoc IExecutionHook
+    function preExecutionHook(uint8 functionId, bytes calldata data) external override returns (bytes memory) {
+        bytes calldata callData;
+        bytes4 execSelector;
+
+        // TODO: plugins should never have to do these gymnastics
+        execSelector = bytes4(data[52:56]);
+        if (execSelector == IAccountExecute.executeUserOp.selector) {
+            callData = data[56:];
+            execSelector = bytes4(callData);
+        } else {
+            callData = data[52:];
+        }
+
+        uint256 value;
+        // Get value being sent
+        if (execSelector == IStandardExecutor.execute.selector) {
+            value = uint256(bytes32(callData[36:68]));
+        } else if (execSelector == IStandardExecutor.executeBatch.selector) {
+            Call[] memory calls = abi.decode(callData[4:], (Call[]));
+            for (uint256 i = 0; i < calls.length; i++) {
+                value += calls[i].value;
+            }
+        }
+
+        uint256 limit = limits[msg.sender][functionId];
+        if (value > limit) {
+            revert ExceededNativeTokenLimit();
+        }
+        limits[msg.sender][functionId] = limit - value;
+
+        return "";
+    }
+
+    /// @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 functionId, address, uint256, bytes calldata) external pure override {
+        // silence warnings
+        (functionId);
+    }
+
+    /// @inheritdoc IValidationHook
+    function preUserOpValidationHook(uint8 functionId, PackedUserOperation calldata userOp, bytes32)
+        external
+        returns (uint256)
+    {
+        // Decrease limit only if no paymaster is used
+        if (userOp.paymasterAndData.length == 0) {
+            uint256 vgl = UserOperationLib.unpackVerificationGasLimit(userOp);
+            uint256 cgl = UserOperationLib.unpackCallGasLimit(userOp);
+            uint256 totalGas = userOp.preVerificationGas + vgl + cgl;
+            uint256 usage = totalGas * UserOperationLib.unpackMaxFeePerGas(userOp);
+
+            uint256 limit = limits[msg.sender][functionId];
+            if (usage > limit) {
+                revert ExceededNativeTokenLimit();
+            }
+            limits[msg.sender][functionId] = limit - usage;
+        }
+        return 0;
+    }
+
+    /// @inheritdoc IPlugin
+    function onInstall(bytes calldata data) external override {
+        uint256[] memory spendLimits = abi.decode(data, (uint256[]));
+
+        for (uint256 i = 0; i < spendLimits.length; i++) {
+            limits[msg.sender].push(spendLimits[i]);
+        }
+
+        if (limits[msg.sender].length > type(uint8).max) {
+            revert ExceededNumberOfEntities();
+        }
+    }
+
+    /// @inheritdoc IPlugin
+    function onUninstall(bytes calldata data) external override {
+        uint8 functionId = abi.decode(data, (uint8));
+        delete limits[msg.sender][functionId];
+    }
+
+    /// @inheritdoc IPlugin
+    function pluginManifest() external pure override returns (PluginManifest memory) {
+        // silence warnings
+        PluginManifest memory manifest;
+        return manifest;
+    }
+
+    /// @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 super.supportsInterface(interfaceId);
+    }
+}

test/plugin/NativeTokenLimitPlugin.t.sol

@@ -0,0 +1,206 @@
+// SPDX-License-Identifier: UNLICENSED
+pragma solidity ^0.8.19;
+
+import {EntryPoint} from "@eth-infinitism/account-abstraction/core/EntryPoint.sol";
+import {PackedUserOperation} from "@eth-infinitism/account-abstraction/interfaces/PackedUserOperation.sol";
+
+import {UpgradeableModularAccount} from "../../src/account/UpgradeableModularAccount.sol";
+import {FunctionReference} from "../../src/helpers/FunctionReferenceLib.sol";
+import {NativeTokenLimitPlugin} from "../../src/plugins/NativeTokenLimitPlugin.sol";
+import {MockPlugin} from "../mocks/MockPlugin.sol";
+import {ExecutionHook} from "../../src/interfaces/IAccountLoupe.sol";
+import {FunctionReferenceLib} from "../../src/helpers/FunctionReferenceLib.sol";
+import {IStandardExecutor, Call} from "../../src/interfaces/IStandardExecutor.sol";
+import {PluginManifest} from "../../src/interfaces/IPlugin.sol";
+
+import {MSCAFactoryFixture} from "../mocks/MSCAFactoryFixture.sol";
+import {OptimizedTest} from "../utils/OptimizedTest.sol";
+
+contract NativeTokenLimitPluginTest is OptimizedTest {
+    EntryPoint public entryPoint = new EntryPoint();
+    address public recipient = address(1);
+    address payable public bundler = payable(address(2));
+    PluginManifest internal _m;
+    MockPlugin public validationPlugin = new MockPlugin(_m);
+    FunctionReference public validationFunction;
+
+    UpgradeableModularAccount public acct;
+    NativeTokenLimitPlugin public plugin = new NativeTokenLimitPlugin();
+    uint256 public spendLimit = 10 ether;
+
+    function setUp() public {
+        // Set up a validator with hooks from the gas spend limit plugin attached
+
+        MSCAFactoryFixture factory = new MSCAFactoryFixture(entryPoint, _deploySingleOwnerPlugin());
+
+        acct = factory.createAccount(address(this), 0);
+
+        vm.deal(address(acct), 10 ether);
+
+        FunctionReference[] memory preValidationHooks = new FunctionReference[](1);
+        preValidationHooks[0] = FunctionReferenceLib.pack(address(plugin), 0);
+
+        ExecutionHook[] memory permissionHooks = new ExecutionHook[](1);
+        permissionHooks[0] = ExecutionHook({
+            hookFunction: FunctionReferenceLib.pack(address(plugin), 0),
+            isPreHook: true,
+            isPostHook: false,
+            requireUOContext: false
+        });
+
+        uint256[] memory spendLimits = new uint256[](1);
+        spendLimits[0] = spendLimit;
+
+        bytes[] memory preValHooksInitDatas = new bytes[](1);
+        preValHooksInitDatas[0] = "";
+
+        bytes[] memory permissionInitDatas = new bytes[](1);
+        permissionInitDatas[0] = abi.encode(spendLimits);
+
+        vm.prank(address(acct));
+        acct.installValidation(
+            FunctionReferenceLib.pack(address(validationPlugin), 0),
+            true,
+            new bytes4[](0),
+            new bytes(0),
+            abi.encode(preValidationHooks, preValHooksInitDatas),
+            abi.encode(permissionHooks, permissionInitDatas)
+        );
+
+        validationFunction = FunctionReferenceLib.pack(address(validationPlugin), 0);
+    }
+
+    function _getExecuteWithValue(uint256 value) internal view returns (bytes memory) {
+        return abi.encodeCall(UpgradeableModularAccount.execute, (recipient, value, ""));
+    }
+
+    function _getPackedUO(uint256 gas1, uint256 gas2, uint256 gas3, uint256 gasPrice, bytes memory callData)
+        internal
+        view
+        returns (PackedUserOperation memory uo)
+    {
+        uo = PackedUserOperation({
+            sender: address(acct),
+            nonce: 0,
+            initCode: "",
+            callData: abi.encodePacked(UpgradeableModularAccount.executeUserOp.selector, callData),
+            accountGasLimits: bytes32(bytes16(uint128(gas1))) | bytes32(uint256(gas2)),
+            preVerificationGas: gas3,
+            gasFees: bytes32(uint256(uint128(gasPrice))),
+            paymasterAndData: "",
+            signature: abi.encodePacked(FunctionReferenceLib.pack(address(validationPlugin), 0), uint8(1))
+        });
+    }
+
+    function test_userOp_gasLimit() public {
+        vm.startPrank(address(entryPoint));
+
+        // uses 10e - 200000 of gas
+        assertEq(plugin.limits(address(acct), 0), 10 ether);
+        uint256 result = acct.validateUserOp(
+            _getPackedUO(100000, 100000, 10 ether - 400000, 1, _getExecuteWithValue(0)), bytes32(0), 0
+        );
+        assertEq(plugin.limits(address(acct), 0), 200000);
+
+        uint256 expected = uint256(type(uint48).max) << 160;
+        assertEq(result, expected);
+
+        // uses 200k + 1 wei of gas
+        vm.expectRevert(NativeTokenLimitPlugin.ExceededNativeTokenLimit.selector);
+        result = acct.validateUserOp(_getPackedUO(100000, 100000, 1, 1, _getExecuteWithValue(0)), bytes32(0), 0);
+    }
+
+    function test_userOp_executeLimit() public {
+        vm.startPrank(address(entryPoint));
+
+        // uses 5e of native tokens
+        assertEq(plugin.limits(address(acct), 0), 10 ether);
+        acct.executeUserOp(_getPackedUO(0, 0, 0, 0, _getExecuteWithValue(5 ether)), bytes32(0));
+        assertEq(plugin.limits(address(acct), 0), 5 ether);
+
+        // uses 5e + 1wei of native tokens
+        vm.expectRevert(
+            abi.encodePacked(
+                UpgradeableModularAccount.PreExecHookReverted.selector,
+                abi.encode(
+                    address(plugin),
+                    uint8(0),
+                    abi.encodePacked(NativeTokenLimitPlugin.ExceededNativeTokenLimit.selector)
+                )
+            )
+        );
+        acct.executeUserOp(_getPackedUO(0, 0, 0, 0, _getExecuteWithValue(5 ether + 1)), bytes32(0));
+    }
+
+    function test_userOp_executeBatchLimit() public {
+        Call[] memory calls = new Call[](3);
+        calls[0] = Call({target: recipient, value: 1, data: ""});
+        calls[1] = Call({target: recipient, value: 1 ether, data: ""});
+        calls[2] = Call({target: recipient, value: 5 ether + 100000, data: ""});
+
+        vm.startPrank(address(entryPoint));
+        assertEq(plugin.limits(address(acct), 0), 10 ether);
+        acct.executeUserOp(
+            _getPackedUO(0, 0, 0, 0, abi.encodeCall(IStandardExecutor.executeBatch, (calls))), bytes32(0)
+        );
+        assertEq(plugin.limits(address(acct), 0), 10 ether - 6 ether - 100001);
+        assertEq(recipient.balance, 6 ether + 100001);
+    }
+
+    function test_userOp_combinedExecLimit_success() public {
+        assertEq(plugin.limits(address(acct), 0), 10 ether);
+        PackedUserOperation[] memory uos = new PackedUserOperation[](1);
+        uos[0] = _getPackedUO(100000, 100000, 100000, 1, _getExecuteWithValue(5 ether));
+        entryPoint.handleOps(uos, bundler);
+
+        assertEq(plugin.limits(address(acct), 0), 5 ether - 300000);
+        assertEq(recipient.balance, 5 ether);
+    }
+
+    function test_userOp_combinedExecBatchLimit_success() public {
+        Call[] memory calls = new Call[](3);
+        calls[0] = Call({target: recipient, value: 1, data: ""});
+        calls[1] = Call({target: recipient, value: 1 ether, data: ""});
+        calls[2] = Call({target: recipient, value: 5 ether + 100000, data: ""});
+
+        vm.startPrank(address(entryPoint));
+        assertEq(plugin.limits(address(acct), 0), 10 ether);
+        PackedUserOperation[] memory uos = new PackedUserOperation[](1);
+        uos[0] = _getPackedUO(200000, 200000, 200000, 1, abi.encodeCall(IStandardExecutor.executeBatch, (calls)));
+        entryPoint.handleOps(uos, bundler);
+
+        assertEq(plugin.limits(address(acct), 0), 10 ether - 6 ether - 700001);
+        assertEq(recipient.balance, 6 ether + 100001);
+    }
+
+    function test_userOp_combinedExecLimit_failExec() public {
+        assertEq(plugin.limits(address(acct), 0), 10 ether);
+        PackedUserOperation[] memory uos = new PackedUserOperation[](1);
+        uos[0] = _getPackedUO(100000, 100000, 100000, 1, _getExecuteWithValue(10 ether));
+        entryPoint.handleOps(uos, bundler);
+
+        assertEq(plugin.limits(address(acct), 0), 10 ether - 300000);
+        assertEq(recipient.balance, 0);
+    }
+
+    function test_runtime_executeLimit() public {
+        assertEq(plugin.limits(address(acct), 0), 10 ether);
+        acct.executeWithAuthorization(
+            _getExecuteWithValue(5 ether), abi.encodePacked(validationFunction, uint8(1))
+        );
+        assertEq(plugin.limits(address(acct), 0), 5 ether);
+    }
+
+    function test_runtime_executeBatchLimit() public {
+        Call[] memory calls = new Call[](3);
+        calls[0] = Call({target: recipient, value: 1, data: ""});
+        calls[1] = Call({target: recipient, value: 1 ether, data: ""});
+        calls[2] = Call({target: recipient, value: 5 ether + 100000, data: ""});
+
+        assertEq(plugin.limits(address(acct), 0), 10 ether);
+        acct.executeWithAuthorization(
+            abi.encodeCall(IStandardExecutor.executeBatch, (calls)), abi.encodePacked(validationFunction, uint8(1))
+        );
+        assertEq(plugin.limits(address(acct), 0), 4 ether - 100001);
+    }
+}