feat: add erc20 token limit module

src/modules/ERC20TokenLimitModule.sol

@@ -0,0 +1,142 @@
+// SPDX-License-Identifier: GPL-3.0
+pragma solidity ^0.8.26;
+
+import {UserOperationLib} from "@eth-infinitism/account-abstraction/core/UserOperationLib.sol";
+import {PackedUserOperation} from "@eth-infinitism/account-abstraction/interfaces/PackedUserOperation.sol";
+
+import {IERC20} from "@openzeppelin/contracts/token/ERC20/IERC20.sol";
+
+import {IExecutionHookModule} from "@erc6900/reference-implementation/interfaces/IExecutionHookModule.sol";
+import {Call, IModularAccount} from "@erc6900/reference-implementation/interfaces/IModularAccount.sol";
+import {IModule} from "@erc6900/reference-implementation/interfaces/IModule.sol";
+
+import {BaseModule, IERC165} from "./BaseModule.sol";
+
+/// @title ERC20 Token Limit Module
+/// @author Alchemy & ERC-6900 Authors
+/// @notice This module supports ERC20 token spend limits. A few key features/restrictions features:
+///     - This module only provide hooks associated with validations.
+///     - For any validation function with this hook installed, all ERC20s without limits specified here will be
+/// reverted.
+///     - Only spending request through the following native execution functions are supported:
+/// IModularAccount.execute, IModularAccount.executeWithAuthorization, IAccountExecute.executeUserOp,
+/// IModularAccount.executeBatch. All other spending request will be reverted.
+///     - this module is opinionated on what selectors (transfer and approve only) 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 ERC20TokenLimitModule is BaseModule, IExecutionHookModule {
+    using UserOperationLib for PackedUserOperation;
+
+    struct ERC20SpendLimit {
+        address token;
+        uint256 limit;
+    }
+
+    mapping(uint32 entityId => mapping(address token => mapping(address account => uint256 limit))) public limits;
+
+    error ExceededTokenLimit();
+    error SelectorNotAllowed();
+    error SpendingRequestNotAllowed(bytes4);
+    error ERC20NotAllowed(address);
+    error InvalidCalldataLength();
+
+    /// @notice Update the token limit of a validation
+    /// @param entityId The validation entityId to update
+    /// @param token The token address whose limit will be updated
+    /// @param newLimit The new limit of the token for the validation
+    function updateLimits(uint32 entityId, address token, uint256 newLimit) external {
+        if (token == address(0)) {
+            revert ERC20NotAllowed(address(0));
+        }
+        limits[entityId][token][msg.sender] = newLimit;
+    }
+
+    /// @inheritdoc IExecutionHookModule
+    function preExecutionHook(uint32 entityId, address, uint256, bytes calldata data)
+        external
+        override
+        returns (bytes memory)
+    {
+        (bytes4 selector, bytes memory callData) = _getSelectorAndCalldata(data);
+
+        if (selector == IModularAccount.execute.selector) {
+            // when calling execute or ERC20 functions directly
+            (address token,, bytes memory innerCalldata) = abi.decode(callData, (address, uint256, bytes));
+            _decrementLimit(entityId, token, innerCalldata);
+        } else if (selector == IModularAccount.executeBatch.selector) {
+            Call[] memory calls = abi.decode(callData, (Call[]));
+            for (uint256 i = 0; i < calls.length; i++) {
+                _decrementLimit(entityId, calls[i].target, calls[i].data);
+            }
+        } else {
+            revert SpendingRequestNotAllowed(selector);
+        }
+        return "";
+    }
+
+    /// @inheritdoc IModule
+    /// @param data should be encoded with the entityId of the validation and a list of ERC20 spend limits
+    function onInstall(bytes calldata data) external override {
+        (uint32 entityId, ERC20SpendLimit[] memory spendLimits) = abi.decode(data, (uint32, ERC20SpendLimit[]));
+
+        for (uint8 i = 0; i < spendLimits.length; i++) {
+            address token = spendLimits[i].token;
+            if (token == address(0)) {
+                revert ERC20NotAllowed(address(0));
+            }
+            limits[entityId][token][msg.sender] = spendLimits[i].limit;
+        }
+    }
+
+    /// @inheritdoc IModule
+    /// @notice uninstall this module can only clear limit for one token of one entity. To clear all limits, users
+    /// are recommended to use updateLimit for each token and entityId.
+    /// @param data should be encoded with the entityId of the validation and the token address to be uninstalled
+    function onUninstall(bytes calldata data) external override {
+        (address token, uint32 entityId) = abi.decode(data, (address, uint32));
+        delete limits[entityId][token][msg.sender];
+    }
+
+    /// @inheritdoc IExecutionHookModule
+    function postExecutionHook(uint32, bytes calldata) external pure override {
+        revert NotImplemented();
+    }
+
+    /// @inheritdoc IModule
+    function moduleId() external pure returns (string memory) {
+        return "alchemy.erc20-token-limit-module.0.0.1";
+    }
+
+    /// @inheritdoc BaseModule
+    function supportsInterface(bytes4 interfaceId) public view override(BaseModule, IERC165) returns (bool) {
+        return interfaceId == type(IExecutionHookModule).interfaceId || super.supportsInterface(interfaceId);
+    }
+
+    function _decrementLimit(uint32 entityId, address token, bytes memory innerCalldata) internal {
+        if (innerCalldata.length < 68) {
+            revert InvalidCalldataLength();
+        }
+
+        bytes4 selector;
+        uint256 spend;
+        assembly ("memory-safe") {
+            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 (_isAllowedERC20Function(selector)) {
+            uint256 limit = limits[entityId][token][msg.sender];
+            if (spend > limit) {
+                revert ExceededTokenLimit();
+            }
+            unchecked {
+                limits[entityId][token][msg.sender] = limit - spend;
+            }
+        } else {
+            revert SelectorNotAllowed();
+        }
+    }
+
+    function _isAllowedERC20Function(bytes4 selector) internal pure returns (bool) {
+        return selector == IERC20.transfer.selector || selector == IERC20.approve.selector;
+    }
+}

test/module/ERC20TokenLimitModule.t.sol

@@ -0,0 +1,169 @@
+// SPDX-License-Identifier: UNLICENSED
+pragma solidity ^0.8.26;
+
+import {MockERC20} from "../mocks/MockERC20.sol";
+import {PackedUserOperation} from "@eth-infinitism/account-abstraction/interfaces/PackedUserOperation.sol";
+import {IERC20} from "@openzeppelin/contracts/token/ERC20/IERC20.sol";
+
+import {ExecutionManifest} from "@erc6900/reference-implementation/interfaces/IExecutionModule.sol";
+import {Call, IModularAccount} from "@erc6900/reference-implementation/interfaces/IModularAccount.sol";
+
+import {ModularAccount} from "../../src/account/ModularAccount.sol";
+import {HookConfigLib} from "../../src/libraries/HookConfigLib.sol";
+import {ModuleEntity} from "../../src/libraries/ModuleEntityLib.sol";
+import {ModuleEntityLib} from "../../src/libraries/ModuleEntityLib.sol";
+import {ValidationConfigLib} from "../../src/libraries/ValidationConfigLib.sol";
+import {ERC20TokenLimitModule} from "../../src/modules/ERC20TokenLimitModule.sol";
+
+import {MockModule} from "../mocks/modules/MockModule.sol";
+import {AccountTestBase} from "../utils/AccountTestBase.sol";
+
+contract ERC20TokenLimitModuleTest is AccountTestBase {
+    address public recipient = address(1);
+    MockERC20 public erc20;
+    address payable public bundler = payable(address(2));
+    ExecutionManifest internal _m;
+    MockModule public validationModule = new MockModule(_m);
+    ModuleEntity public validationFunction;
+
+    ModularAccount public acct;
+    ERC20TokenLimitModule public module = new ERC20TokenLimitModule();
+    uint256 public spendLimit = 10 ether;
+
+    function setUp() public {
+        // Set up a validator with hooks from the erc20 spend limit module attached
+        acct = factory.createAccount(address(this), 0, 0);
+
+        erc20 = new MockERC20();
+        erc20.mint(address(acct), 10 ether);
+
+        ERC20TokenLimitModule.ERC20SpendLimit[] memory limit = new ERC20TokenLimitModule.ERC20SpendLimit[](1);
+        limit[0] = ERC20TokenLimitModule.ERC20SpendLimit({token: address(erc20), limit: spendLimit});
+
+        bytes[] memory hooks = new bytes[](1);
+        hooks[0] = abi.encodePacked(
+            HookConfigLib.packExecHook({_module: address(module), _entityId: 0, _hasPre: true, _hasPost: false}),
+            abi.encode(uint32(0), limit)
+        );
+
+        vm.prank(address(acct));
+        acct.installValidation(
+            ValidationConfigLib.pack(address(validationModule), 0, true, true, true), new bytes4[](0), "", hooks
+        );
+
+        validationFunction = ModuleEntityLib.pack(address(validationModule), 0);
+    }
+
+    function _getPackedUO(bytes memory callData) internal view returns (PackedUserOperation memory uo) {
+        uo = PackedUserOperation({
+            sender: address(acct),
+            nonce: 0,
+            initCode: "",
+            callData: abi.encodePacked(ModularAccount.executeUserOp.selector, callData),
+            accountGasLimits: bytes32(bytes16(uint128(200_000))) | bytes32(uint256(200_000)),
+            preVerificationGas: 200_000,
+            gasFees: bytes32(uint256(uint128(0))),
+            paymasterAndData: "",
+            signature: _encodeSignature(ModuleEntityLib.pack(address(validationModule), 0), 1, "")
+        });
+    }
+
+    function _getExecuteWithSpend(uint256 value) internal view returns (bytes memory) {
+        return abi.encodeCall(
+            ModularAccount.execute, (address(erc20), 0, abi.encodeCall(IERC20.transfer, (recipient, value)))
+        );
+    }
+
+    function test_userOp_executeLimit() public {
+        vm.startPrank(address(entryPoint));
+        assertEq(module.limits(0, address(erc20), address(acct)), 10 ether);
+        acct.executeUserOp(_getPackedUO(_getExecuteWithSpend(5 ether)), bytes32(0));
+        assertEq(module.limits(0, address(erc20), address(acct)), 5 ether);
+    }
+
+    function test_userOp_executeBatchLimit() public {
+        Call[] memory calls = new Call[](3);
+        calls[0] =
+            Call({target: address(erc20), value: 0, data: abi.encodeCall(IERC20.transfer, (recipient, 1 wei))});
+        calls[1] =
+            Call({target: address(erc20), value: 0, data: abi.encodeCall(IERC20.transfer, (recipient, 1 ether))});
+        calls[2] = Call({
+            target: address(erc20),
+            value: 0,
+            data: abi.encodeCall(IERC20.transfer, (recipient, 5 ether + 100_000))
+        });
+
+        vm.startPrank(address(entryPoint));
+        assertEq(module.limits(0, address(erc20), address(acct)), 10 ether);
+        acct.executeUserOp(_getPackedUO(abi.encodeCall(IModularAccount.executeBatch, (calls))), bytes32(0));
+        assertEq(module.limits(0, address(erc20), address(acct)), 10 ether - 6 ether - 100_001);
+    }
+
+    function test_userOp_executeBatch_approveAndTransferLimit() public {
+        Call[] memory calls = new Call[](3);
+        calls[0] =
+            Call({target: address(erc20), value: 0, data: abi.encodeCall(IERC20.approve, (recipient, 1 wei))});
+        calls[1] =
+            Call({target: address(erc20), value: 0, data: abi.encodeCall(IERC20.transfer, (recipient, 1 ether))});
+        calls[2] = Call({
+            target: address(erc20),
+            value: 0,
+            data: abi.encodeCall(IERC20.approve, (recipient, 5 ether + 100_000))
+        });
+
+        vm.startPrank(address(entryPoint));
+        assertEq(module.limits(0, address(erc20), address(acct)), 10 ether);
+        acct.executeUserOp(_getPackedUO(abi.encodeCall(IModularAccount.executeBatch, (calls))), bytes32(0));
+        assertEq(module.limits(0, address(erc20), address(acct)), 10 ether - 6 ether - 100_001);
+    }
+
+    function test_userOp_executeBatch_approveAndTransferLimit_fail() public {
+        Call[] memory calls = new Call[](3);
+        calls[0] =
+            Call({target: address(erc20), value: 0, data: abi.encodeCall(IERC20.approve, (recipient, 1 wei))});
+        calls[1] =
+            Call({target: address(erc20), value: 0, data: abi.encodeCall(IERC20.transfer, (recipient, 1 ether))});
+        calls[2] = Call({
+            target: address(erc20),
+            value: 0,
+            data: abi.encodeCall(IERC20.approve, (recipient, 9 ether + 100_000))
+        });
+
+        vm.startPrank(address(entryPoint));
+        assertEq(module.limits(0, address(erc20), address(acct)), 10 ether);
+        PackedUserOperation[] memory uos = new PackedUserOperation[](1);
+        uos[0] = _getPackedUO(abi.encodeCall(IModularAccount.executeBatch, (calls)));
+        entryPoint.handleOps(uos, bundler);
+        // no spend consumed
+        assertEq(module.limits(0, address(erc20), address(acct)), 10 ether);
+    }
+
+    function test_runtime_executeLimit() public {
+        assertEq(module.limits(0, address(erc20), address(acct)), 10 ether);
+        acct.executeWithAuthorization(
+            _getExecuteWithSpend(5 ether),
+            _encodeSignature(ModuleEntityLib.pack(address(validationModule), 0), 1, "")
+        );
+        assertEq(module.limits(0, address(erc20), address(acct)), 5 ether);
+    }
+
+    function test_runtime_executeBatchLimit() public {
+        Call[] memory calls = new Call[](3);
+        calls[0] =
+            Call({target: address(erc20), value: 0, data: abi.encodeCall(IERC20.approve, (recipient, 1 wei))});
+        calls[1] =
+            Call({target: address(erc20), value: 0, data: abi.encodeCall(IERC20.transfer, (recipient, 1 ether))});
+        calls[2] = Call({
+            target: address(erc20),
+            value: 0,
+            data: abi.encodeCall(IERC20.approve, (recipient, 5 ether + 100_000))
+        });
+
+        assertEq(module.limits(0, address(erc20), address(acct)), 10 ether);
+        acct.executeWithAuthorization(
+            abi.encodeCall(IModularAccount.executeBatch, (calls)),
+            _encodeSignature(ModuleEntityLib.pack(address(validationModule), 0), 1, "")
+        );
+        assertEq(module.limits(0, address(erc20), address(acct)), 10 ether - 6 ether - 100_001);
+    }
+}