Skip to content

Commit a2bd984

Browse files
jaypaikjaypaik
committed
refactor: fix ERC20TokenLimitModule and remove dependency on modular-account-libs
1 parent 9141114 commit a2bd984

File tree

5 files changed

+86
-71
lines changed

5 files changed

+86
-71
lines changed

‎.gitmodules‎

Lines changed: 0 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -7,9 +7,6 @@
77
[submodule "lib/forge-std"]
88
path = lib/forge-std
99
url = https://github.com/foundry-rs/forge-std
10-
[submodule "lib/modular-account-libs"]
11-
path = lib/modular-account-libs
12-
url = https://github.com/erc6900/modular-account-libs
1310
[submodule "lib/solady"]
1411
path = lib/solady
1512
url = https://github.com/vectorized/solady

‎lib/modular-account-libs‎

Lines changed: 0 additions & 1 deletion
This file was deleted.

‎remappings.txt‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,5 +2,4 @@ ds-test/=lib/forge-std/lib/ds-test/src/
22
forge-std/=lib/forge-std/src/
33
@eth-infinitism/account-abstraction/=lib/account-abstraction/contracts/
44
@openzeppelin/=lib/openzeppelin-contracts/
5-
@modular-account-libs/=lib/modular-account-libs/src/
65
solady=lib/solady/src/

‎src/modules/ERC20TokenLimitModule.sol‎

Lines changed: 53 additions & 53 deletions
Original file line numberDiff line numberDiff line change
@@ -3,14 +3,7 @@ pragma solidity ^0.8.20;
33

44
import {UserOperationLib} from "@eth-infinitism/account-abstraction/core/UserOperationLib.sol";
55
import {PackedUserOperation} from "@eth-infinitism/account-abstraction/interfaces/PackedUserOperation.sol";
6-
7-
import {
8-
AssociatedLinkedListSet,
9-
AssociatedLinkedListSetLib,
10-
SetValue
11-
} from "@modular-account-libs/libraries/AssociatedLinkedListSetLib.sol";
126
import {IERC20} from "@openzeppelin/contracts/token/ERC20/IERC20.sol";
13-
import {EnumerableSet} from "@openzeppelin/contracts/utils/structs/EnumerableSet.sol";
147

158
import {IExecutionHookModule} from "../interfaces/IExecutionHookModule.sol";
169
import {Call, IModularAccount} from "../interfaces/IModularAccount.sol";
@@ -27,29 +20,24 @@ import {BaseModule, IERC165} from "./BaseModule.sol";
2720
/// token contract
2821
contract ERC20TokenLimitModule is BaseModule, IExecutionHookModule {
2922
using UserOperationLib for PackedUserOperation;
30-
using EnumerableSet for EnumerableSet.AddressSet;
31-
using AssociatedLinkedListSetLib for AssociatedLinkedListSet;
3223

3324
struct ERC20SpendLimit {
3425
address token;
35-
uint256[] limits;
26+
uint256 limit;
3627
}
3728

38-
string internal constant _NAME = "ERC20 Token Limit Module";
39-
string internal constant _VERSION = "1.0.0";
40-
string internal constant _AUTHOR = "ERC-6900 Authors";
29+
struct SpendLimit {
30+
bool hasLimit;
31+
uint256 limit;
32+
}
4133

42-
mapping(uint32 entityId => mapping(address token => mapping(address account => uint256 limit))) public limits;
43-
AssociatedLinkedListSet internal _tokenList;
34+
mapping(uint32 entityId => mapping(address token => mapping(address account => SpendLimit))) public limits;
4435

36+
error ERC20NotAllowed(address);
4537
error ExceededTokenLimit();
46-
error ExceededNumberOfEntities();
38+
error InvalidCalldataLength();
4739
error SelectorNotAllowed();
48-
49-
function updateLimits(uint32 entityId, address token, uint256 newLimit) external {
50-
_tokenList.tryAdd(msg.sender, SetValue.wrap(bytes30(bytes20(token))));
51-
limits[entityId][token][msg.sender] = newLimit;
52-
}
40+
error SpendingRequestNotAllowed(bytes4);
5341

5442
/// @inheritdoc IExecutionHookModule
5543
function preExecutionHook(uint32 entityId, address, uint256, bytes calldata data)
@@ -60,54 +48,40 @@ contract ERC20TokenLimitModule is BaseModule, IExecutionHookModule {
6048
(bytes4 selector, bytes memory callData) = _getSelectorAndCalldata(data);
6149

6250
if (selector == IModularAccount.execute.selector) {
51+
// when calling execute or ERC20 functions directly
6352
(address token,, bytes memory innerCalldata) = abi.decode(callData, (address, uint256, bytes));
64-
if (_tokenList.contains(msg.sender, SetValue.wrap(bytes30(bytes20(token))))) {
65-
_decrementLimit(entityId, token, innerCalldata);
66-
}
53+
_decrementLimitIfApplies(entityId, token, innerCalldata);
6754
} else if (selector == IModularAccount.executeBatch.selector) {
6855
Call[] memory calls = abi.decode(callData, (Call[]));
6956
for (uint256 i = 0; i < calls.length; i++) {
70-
if (_tokenList.contains(msg.sender, SetValue.wrap(bytes30(bytes20(calls[i].target))))) {
71-
_decrementLimit(entityId, calls[i].target, calls[i].data);
72-
}
57+
_decrementLimitIfApplies(entityId, calls[i].target, calls[i].data);
7358
}
59+
} else {
60+
revert SpendingRequestNotAllowed(selector);
7461
}
75-
7662
return "";
7763
}
7864

7965
/// @inheritdoc IModule
66+
/// @param data should be encoded with the entityId of the validation and a list of ERC20 spend limits
8067
function onInstall(bytes calldata data) external override {
81-
(uint32 startEntityId, ERC20SpendLimit[] memory spendLimits) =
82-
abi.decode(data, (uint32, ERC20SpendLimit[]));
83-
84-
if (startEntityId + spendLimits.length > type(uint32).max) {
85-
revert ExceededNumberOfEntities();
86-
}
68+
(uint32 entityId, ERC20SpendLimit[] memory spendLimits) = abi.decode(data, (uint32, ERC20SpendLimit[]));
8769

8870
for (uint8 i = 0; i < spendLimits.length; i++) {
89-
_tokenList.tryAdd(msg.sender, SetValue.wrap(bytes30(bytes20(spendLimits[i].token))));
90-
for (uint256 j = 0; j < spendLimits[i].limits.length; j++) {
91-
limits[i + startEntityId][spendLimits[i].token][msg.sender] = spendLimits[i].limits[j];
92-
}
71+
address token = spendLimits[i].token;
72+
updateLimits(entityId, token, true, spendLimits[i].limit);
9373
}
9474
}
9575

9676
/// @inheritdoc IModule
77+
/// @notice uninstall this module can only clear limit for one token of one entity. To clear all limits, users
78+
/// are recommended to use updateLimit for each token and entityId.
79+
/// @param data should be encoded with the entityId of the validation and the token address to be uninstalled
9780
function onUninstall(bytes calldata data) external override {
9881
(address token, uint32 entityId) = abi.decode(data, (address, uint32));
9982
delete limits[entityId][token][msg.sender];
10083
}
10184

102-
function getTokensForAccount(address account) external view returns (address[] memory tokens) {
103-
SetValue[] memory set = _tokenList.getAll(account);
104-
tokens = new address[](set.length);
105-
for (uint256 i = 0; i < tokens.length; i++) {
106-
tokens[i] = address(bytes20(bytes32(SetValue.unwrap(set[i]))));
107-
}
108-
return tokens;
109-
}
110-
11185
/// @inheritdoc IExecutionHookModule
11286
function postExecutionHook(uint32, bytes calldata) external pure override {
11387
revert NotImplemented();
@@ -118,27 +92,53 @@ contract ERC20TokenLimitModule is BaseModule, IExecutionHookModule {
11892
return "erc6900.erc20-token-limit-module.1.0.0";
11993
}
12094

95+
/// @notice Update the token limit of a validation
96+
/// @param entityId The validation entityId to update
97+
/// @param token The token address whose limit will be updated
98+
/// @param newLimit The new limit of the token for the validation
99+
function updateLimits(uint32 entityId, address token, bool hasLimit, uint256 newLimit) public {
100+
if (token == address(0)) {
101+
revert ERC20NotAllowed(address(0));
102+
}
103+
limits[entityId][token][msg.sender] = SpendLimit({hasLimit: hasLimit, limit: newLimit});
104+
}
105+
121106
/// @inheritdoc BaseModule
122107
function supportsInterface(bytes4 interfaceId) public view override(BaseModule, IERC165) returns (bool) {
123108
return interfaceId == type(IExecutionHookModule).interfaceId || super.supportsInterface(interfaceId);
124109
}
125110

126-
function _decrementLimit(uint32 entityId, address token, bytes memory innerCalldata) internal {
111+
function _decrementLimitIfApplies(uint32 entityId, address token, bytes memory innerCalldata) internal {
112+
SpendLimit storage spendLimit = limits[entityId][token][msg.sender];
113+
114+
if (!spendLimit.hasLimit) {
115+
return;
116+
}
117+
118+
if (innerCalldata.length < 68) {
119+
revert InvalidCalldataLength();
120+
}
121+
127122
bytes4 selector;
128123
uint256 spend;
129-
assembly {
124+
assembly ("memory-safe") {
130125
selector := mload(add(innerCalldata, 32)) // 0:32 is arr len, 32:36 is selector
131126
spend := mload(add(innerCalldata, 68)) // 36:68 is recipient, 68:100 is spend
132127
}
133-
if (selector == IERC20.transfer.selector || selector == IERC20.approve.selector) {
134-
uint256 limit = limits[entityId][token][msg.sender];
128+
if (_isAllowedERC20Function(selector)) {
129+
uint256 limit = spendLimit.limit;
135130
if (spend > limit) {
136131
revert ExceededTokenLimit();
137132
}
138-
// solhint-disable-next-line reentrancy
139-
limits[entityId][token][msg.sender] = limit - spend;
133+
unchecked {
134+
spendLimit.limit = limit - spend;
135+
}
140136
} else {
141137
revert SelectorNotAllowed();
142138
}
143139
}
140+
141+
function _isAllowedERC20Function(bytes4 selector) internal pure returns (bool) {
142+
return selector == IERC20.transfer.selector || selector == IERC20.approve.selector;
143+
}
144144
}

‎test/module/ERC20TokenLimitModule.t.sol‎

Lines changed: 33 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -43,7 +43,7 @@ contract ERC20TokenLimitModuleTest is AccountTestBase {
4343
limits[0] = spendLimit;
4444

4545
ERC20TokenLimitModule.ERC20SpendLimit[] memory limit = new ERC20TokenLimitModule.ERC20SpendLimit[](1);
46-
limit[0] = ERC20TokenLimitModule.ERC20SpendLimit({token: address(erc20), limits: limits});
46+
limit[0] = ERC20TokenLimitModule.ERC20SpendLimit({token: address(erc20), limit: spendLimit});
4747

4848
bytes[] memory hooks = new bytes[](1);
4949
hooks[0] = abi.encodePacked(
@@ -82,9 +82,14 @@ contract ERC20TokenLimitModuleTest is AccountTestBase {
8282

8383
function test_userOp_executeLimit() public {
8484
vm.startPrank(address(entryPoint));
85-
assertEq(module.limits(0, address(erc20), address(acct)), 10 ether);
85+
86+
(, uint256 limit) = module.limits(0, address(erc20), address(acct));
87+
88+
assertEq(limit, 10 ether);
8689
acct.executeUserOp(_getPackedUO(_getExecuteWithSpend(5 ether)), bytes32(0));
87-
assertEq(module.limits(0, address(erc20), address(acct)), 5 ether);
90+
91+
(, limit) = module.limits(0, address(erc20), address(acct));
92+
assertEq(limit, 5 ether);
8893
}
8994

9095
function test_userOp_executeBatchLimit() public {
@@ -100,9 +105,12 @@ contract ERC20TokenLimitModuleTest is AccountTestBase {
100105
});
101106

102107
vm.startPrank(address(entryPoint));
103-
assertEq(module.limits(0, address(erc20), address(acct)), 10 ether);
108+
(, uint256 limit) = module.limits(0, address(erc20), address(acct));
109+
assertEq(limit, 10 ether);
104110
acct.executeUserOp(_getPackedUO(abi.encodeCall(IModularAccount.executeBatch, (calls))), bytes32(0));
105-
assertEq(module.limits(0, address(erc20), address(acct)), 10 ether - 6 ether - 100_001);
111+
112+
(, limit) = module.limits(0, address(erc20), address(acct));
113+
assertEq(limit, 10 ether - 6 ether - 100_001);
106114
}
107115

108116
function test_userOp_executeBatch_approveAndTransferLimit() public {
@@ -118,9 +126,12 @@ contract ERC20TokenLimitModuleTest is AccountTestBase {
118126
});
119127

120128
vm.startPrank(address(entryPoint));
121-
assertEq(module.limits(0, address(erc20), address(acct)), 10 ether);
129+
(, uint256 limit) = module.limits(0, address(erc20), address(acct));
130+
assertEq(limit, 10 ether);
122131
acct.executeUserOp(_getPackedUO(abi.encodeCall(IModularAccount.executeBatch, (calls))), bytes32(0));
123-
assertEq(module.limits(0, address(erc20), address(acct)), 10 ether - 6 ether - 100_001);
132+
133+
(, limit) = module.limits(0, address(erc20), address(acct));
134+
assertEq(limit, 10 ether - 6 ether - 100_001);
124135
}
125136

126137
function test_userOp_executeBatch_approveAndTransferLimit_fail() public {
@@ -136,21 +147,27 @@ contract ERC20TokenLimitModuleTest is AccountTestBase {
136147
});
137148

138149
vm.startPrank(address(entryPoint));
139-
assertEq(module.limits(0, address(erc20), address(acct)), 10 ether);
150+
(, uint256 limit) = module.limits(0, address(erc20), address(acct));
151+
assertEq(limit, 10 ether);
140152
PackedUserOperation[] memory uos = new PackedUserOperation[](1);
141153
uos[0] = _getPackedUO(abi.encodeCall(IModularAccount.executeBatch, (calls)));
142154
entryPoint.handleOps(uos, bundler);
143155
// no spend consumed
144-
assertEq(module.limits(0, address(erc20), address(acct)), 10 ether);
156+
157+
(, limit) = module.limits(0, address(erc20), address(acct));
158+
assertEq(limit, 10 ether);
145159
}
146160

147161
function test_runtime_executeLimit() public {
148-
assertEq(module.limits(0, address(erc20), address(acct)), 10 ether);
162+
(, uint256 limit) = module.limits(0, address(erc20), address(acct));
163+
assertEq(limit, 10 ether);
149164
acct.executeWithRuntimeValidation(
150165
_getExecuteWithSpend(5 ether),
151166
_encodeSignature(ModuleEntityLib.pack(address(validationModule), 0), 1, "")
152167
);
153-
assertEq(module.limits(0, address(erc20), address(acct)), 5 ether);
168+
169+
(, limit) = module.limits(0, address(erc20), address(acct));
170+
assertEq(limit, 5 ether);
154171
}
155172

156173
function test_runtime_executeBatchLimit() public {
@@ -165,11 +182,14 @@ contract ERC20TokenLimitModuleTest is AccountTestBase {
165182
data: abi.encodeCall(IERC20.approve, (recipient, 5 ether + 100_000))
166183
});
167184

168-
assertEq(module.limits(0, address(erc20), address(acct)), 10 ether);
185+
(, uint256 limit) = module.limits(0, address(erc20), address(acct));
186+
assertEq(limit, 10 ether);
169187
acct.executeWithRuntimeValidation(
170188
abi.encodeCall(IModularAccount.executeBatch, (calls)),
171189
_encodeSignature(ModuleEntityLib.pack(address(validationModule), 0), 1, "")
172190
);
173-
assertEq(module.limits(0, address(erc20), address(acct)), 10 ether - 6 ether - 100_001);
191+
192+
(, limit) = module.limits(0, address(erc20), address(acct));
193+
assertEq(limit, 10 ether - 6 ether - 100_001);
174194
}
175195
}

0 commit comments

Comments
 (0)