Skip to content

Commit 80fa323

Browse files
nkrishangnkrishang
authored
Many accounts per EOA (#561)
* Enable multiple accounts per admin by not hardcoding _data * Update sender address in tests * initial unit test * Update unit test: create multiple accounts with same admin * Add test for Dynamic and Managed smart wallets * Use abi.encode instead of encodePacked * Move deposit fns to AccountExtension
1 parent a453dd8 commit 80fa323

File tree

13 files changed

+400
-65
lines changed

13 files changed

+400
-65
lines changed

‎contracts/prebuilts/account/dynamic/DynamicAccount.sol

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -31,9 +31,9 @@ contract DynamicAccount is AccountCore, BaseRouter {
3131
}
3232

3333
/// @notice Initializes the smart contract wallet.
34-
function initialize(address _defaultAdmin, bytes calldata) public override initializer {
34+
function initialize(address _defaultAdmin, bytes calldata _data) public override initializer {
3535
__BaseRouter_init();
36-
AccountCoreStorage.data().firstAdmin = _defaultAdmin;
36+
AccountCoreStorage.data().creationSalt = _generateSalt(_defaultAdmin, _data);
3737
_setAdmin(_defaultAdmin, true);
3838
}
3939

‎contracts/prebuilts/account/interface/IAccountFactory.sol

Lines changed: 2 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -9,16 +9,8 @@ interface IAccountFactory is IAccountFactoryCore {
99
//////////////////////////////////////////////////////////////*/
1010

1111
/// @notice Callback function for an Account to register its signers.
12-
function onSignerAdded(
13-
address signer,
14-
address creatorAdmin,
15-
bytes memory data
16-
) external;
12+
function onSignerAdded(address signer, bytes32 salt) external;
1713

1814
/// @notice Callback function for an Account to un-register its signers.
19-
function onSignerRemoved(
20-
address signer,
21-
address creatorAdmin,
22-
bytes memory data
23-
) external;
15+
function onSignerRemoved(address signer, bytes32 salt) external;
2416
}

‎contracts/prebuilts/account/non-upgradeable/Account.sol

Lines changed: 12 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -115,6 +115,17 @@ contract Account is AccountCore, ContractMetadata, ERC1271, ERC721Holder, ERC115
115115
}
116116
}
117117

118+
/// @notice Deposit funds for this account in Entrypoint.
119+
function addDeposit() public payable {
120+
entryPoint().depositTo{ value: msg.value }(address(this));
121+
}
122+
123+
/// @notice Withdraw funds for this account from Entrypoint.
124+
function withdrawDepositTo(address payable withdrawAddress, uint256 amount) public {
125+
_onlyAdmin();
126+
entryPoint().withdrawTo(withdrawAddress, amount);
127+
}
128+
118129
/*///////////////////////////////////////////////////////////////
119130
Internal functions
120131
//////////////////////////////////////////////////////////////*/
@@ -123,7 +134,7 @@ contract Account is AccountCore, ContractMetadata, ERC1271, ERC721Holder, ERC115
123134
function _registerOnFactory() internal virtual {
124135
BaseAccountFactory factoryContract = BaseAccountFactory(factory);
125136
if (!factoryContract.isRegistered(address(this))) {
126-
factoryContract.onRegister(AccountCoreStorage.data().firstAdmin, "");
137+
factoryContract.onRegister(AccountCoreStorage.data().creationSalt);
127138
}
128139
}
129140

‎contracts/prebuilts/account/utils/AccountCore.sol

Lines changed: 10 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -56,9 +56,9 @@ contract AccountCore is IAccountCore, Initializable, Multicall, BaseAccount, Acc
5656
}
5757

5858
/// @notice Initializes the smart contract wallet.
59-
function initialize(address _defaultAdmin, bytes calldata) public virtual initializer {
59+
function initialize(address _defaultAdmin, bytes calldata _data) public virtual initializer {
6060
// This is passed as data in the `_registerOnFactory()` call in `AccountExtension` / `Account`.
61-
AccountCoreStorage.data().firstAdmin = _defaultAdmin;
61+
AccountCoreStorage.data().creationSalt = _generateSalt(_defaultAdmin, _data);
6262
_setAdmin(_defaultAdmin, true);
6363
}
6464

@@ -168,17 +168,6 @@ contract AccountCore is IAccountCore, Initializable, Multicall, BaseAccount, Acc
168168
External functions
169169
//////////////////////////////////////////////////////////////*/
170170

171-
/// @notice Deposit funds for this account in Entrypoint.
172-
function addDeposit() public payable {
173-
entryPoint().depositTo{ value: msg.value }(address(this));
174-
}
175-
176-
/// @notice Withdraw funds for this account from Entrypoint.
177-
function withdrawDepositTo(address payable withdrawAddress, uint256 amount) public {
178-
_onlyAdmin();
179-
entryPoint().withdrawTo(withdrawAddress, amount);
180-
}
181-
182171
/// @notice Overrides the Entrypoint contract being used.
183172
function setEntrypointOverride(IEntryPoint _entrypointOverride) public virtual {
184173
_onlyAdmin();
@@ -189,6 +178,11 @@ contract AccountCore is IAccountCore, Initializable, Multicall, BaseAccount, Acc
189178
Internal functions
190179
//////////////////////////////////////////////////////////////*/
191180

181+
/// @dev Returns the salt used when deploying an Account.
182+
function _generateSalt(address _admin, bytes memory _data) internal view virtual returns (bytes32) {
183+
return keccak256(abi.encode(_admin, _data));
184+
}
185+
192186
function getFunctionSignature(bytes calldata data) internal pure returns (bytes4 functionSelector) {
193187
require(data.length >= 4, "!Data");
194188
return bytes4(data[:4]);
@@ -243,17 +237,17 @@ contract AccountCore is IAccountCore, Initializable, Multicall, BaseAccount, Acc
243237
super._setAdmin(_account, _isAdmin);
244238
if (factory.code.length > 0) {
245239
if (_isAdmin) {
246-
BaseAccountFactory(factory).onSignerAdded(_account, AccountCoreStorage.data().firstAdmin, "");
240+
BaseAccountFactory(factory).onSignerAdded(_account, AccountCoreStorage.data().creationSalt);
247241
} else {
248-
BaseAccountFactory(factory).onSignerRemoved(_account, AccountCoreStorage.data().firstAdmin, "");
242+
BaseAccountFactory(factory).onSignerRemoved(_account, AccountCoreStorage.data().creationSalt);
249243
}
250244
}
251245
}
252246

253247
/// @notice Runs after every `changeRole` run.
254248
function _afterSignerPermissionsUpdate(SignerPermissionRequest calldata _req) internal virtual override {
255249
if (factory.code.length > 0) {
256-
BaseAccountFactory(factory).onSignerAdded(_req.signer, AccountCoreStorage.data().firstAdmin, "");
250+
BaseAccountFactory(factory).onSignerAdded(_req.signer, AccountCoreStorage.data().creationSalt);
257251
}
258252
}
259253
}

‎contracts/prebuilts/account/utils/AccountCoreStorage.sol

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,7 @@ library AccountCoreStorage {
99

1010
struct Data {
1111
address entrypointOverride;
12-
address firstAdmin;
12+
bytes32 creationSalt;
1313
}
1414

1515
function data() internal pure returns (Data storage acountCoreData) {

‎contracts/prebuilts/account/utils/AccountExtension.sol

Lines changed: 12 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -116,6 +116,17 @@ contract AccountExtension is ContractMetadata, ERC1271, AccountPermissions, ERC7
116116
}
117117
}
118118

119+
/// @notice Deposit funds for this account in Entrypoint.
120+
function addDeposit() public payable {
121+
AccountCore(payable(address(this))).entryPoint().depositTo{ value: msg.value }(address(this));
122+
}
123+
124+
/// @notice Withdraw funds for this account from Entrypoint.
125+
function withdrawDepositTo(address payable withdrawAddress, uint256 amount) public {
126+
_onlyAdmin();
127+
AccountCore(payable(address(this))).entryPoint().withdrawTo(withdrawAddress, amount);
128+
}
129+
119130
/*///////////////////////////////////////////////////////////////
120131
Internal functions
121132
//////////////////////////////////////////////////////////////*/
@@ -125,7 +136,7 @@ contract AccountExtension is ContractMetadata, ERC1271, AccountPermissions, ERC7
125136
address factory = AccountCore(payable(address(this))).factory();
126137
BaseAccountFactory factoryContract = BaseAccountFactory(factory);
127138
if (!factoryContract.isRegistered(address(this))) {
128-
factoryContract.onRegister(AccountCoreStorage.data().firstAdmin, "");
139+
factoryContract.onRegister(AccountCoreStorage.data().creationSalt);
129140
}
130141
}
131142

‎contracts/prebuilts/account/utils/BaseAccountFactory.sol

Lines changed: 10 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -72,20 +72,16 @@ abstract contract BaseAccountFactory is IAccountFactory, Multicall {
7272
}
7373

7474
/// @notice Callback function for an Account to register itself on the factory.
75-
function onRegister(address _defaultAdmin, bytes memory _data) external {
75+
function onRegister(bytes32 _salt) external {
7676
address account = msg.sender;
77-
require(_isAccountOfFactory(account, _defaultAdmin, _data), "AccountFactory: not an account.");
77+
require(_isAccountOfFactory(account, _salt), "AccountFactory: not an account.");
7878

7979
require(allAccounts.add(account), "AccountFactory: account already registered");
8080
}
8181

82-
function onSignerAdded(
83-
address _signer,
84-
address _defaultAdmin,
85-
bytes memory _data
86-
) external {
82+
function onSignerAdded(address _signer, bytes32 _salt) external {
8783
address account = msg.sender;
88-
require(_isAccountOfFactory(account, _defaultAdmin, _data), "AccountFactory: not an account.");
84+
require(_isAccountOfFactory(account, _salt), "AccountFactory: not an account.");
8985

9086
bool isNewSigner = accountsOfSigner[_signer].add(account);
9187

@@ -95,13 +91,9 @@ abstract contract BaseAccountFactory is IAccountFactory, Multicall {
9591
}
9692

9793
/// @notice Callback function for an Account to un-register its signers.
98-
function onSignerRemoved(
99-
address _signer,
100-
address _defaultAdmin,
101-
bytes memory _data
102-
) external {
94+
function onSignerRemoved(address _signer, bytes32 _salt) external {
10395
address account = msg.sender;
104-
require(_isAccountOfFactory(account, _defaultAdmin, _data), "AccountFactory: not an account.");
96+
require(_isAccountOfFactory(account, _salt), "AccountFactory: not an account.");
10597

10698
bool isAccount = accountsOfSigner[_signer].remove(account);
10799

@@ -140,13 +132,8 @@ abstract contract BaseAccountFactory is IAccountFactory, Multicall {
140132
//////////////////////////////////////////////////////////////*/
141133

142134
/// @dev Returns whether the caller is an account deployed by this factory.
143-
function _isAccountOfFactory(
144-
address _account,
145-
address _admin,
146-
bytes memory _data
147-
) internal view virtual returns (bool) {
148-
bytes32 salt = _generateSalt(_admin, _data);
149-
address predicted = Clones.predictDeterministicAddress(accountImplementation, salt);
135+
function _isAccountOfFactory(address _account, bytes32 _salt) internal view virtual returns (bool) {
136+
address predicted = Clones.predictDeterministicAddress(accountImplementation, _salt);
150137
return _account == predicted;
151138
}
152139

@@ -156,8 +143,8 @@ abstract contract BaseAccountFactory is IAccountFactory, Multicall {
156143
}
157144

158145
/// @dev Returns the salt used when deploying an Account.
159-
function _generateSalt(address _admin, bytes memory) internal view virtual returns (bytes32) {
160-
return keccak256(abi.encode(_admin));
146+
function _generateSalt(address _admin, bytes memory _data) internal view virtual returns (bytes32) {
147+
return keccak256(abi.encode(_admin, _data));
161148
}
162149

163150
/// @dev Called in `createAccount`. Initializes the account contract created in `createAccount`.

‎src/test/benchmark/AccountBenchmark.t.sol

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -49,7 +49,7 @@ contract AccountBenchmarkTest is BaseTest {
4949
address private nonSigner;
5050

5151
// UserOp terminology: `sender` is the smart wallet.
52-
address private sender = 0xBB956D56140CA3f3060986586A2631922a4B347E;
52+
address private sender = 0x0df2C3523703d165Aa7fA1a552f3F0B56275DfC6;
5353
address payable private beneficiary = payable(address(0x45654));
5454

5555
bytes32 private uidCache = bytes32("random uid");

‎src/test/smart-wallet/Account.t.sol

Lines changed: 110 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -49,7 +49,7 @@ contract SimpleAccountTest is BaseTest {
4949
address private nonSigner;
5050

5151
// UserOp terminology: `sender` is the smart wallet.
52-
address private sender = 0xBB956D56140CA3f3060986586A2631922a4B347E;
52+
address private sender = 0x0df2C3523703d165Aa7fA1a552f3F0B56275DfC6;
5353
address payable private beneficiary = payable(address(0x45654));
5454

5555
bytes32 private uidCache = bytes32("random uid");
@@ -141,6 +141,63 @@ contract SimpleAccountTest is BaseTest {
141141
ops[0] = op;
142142
}
143143

144+
function _setupUserOpWithSender(
145+
bytes memory _initCode,
146+
bytes memory _callDataForEntrypoint,
147+
address _sender
148+
) internal returns (UserOperation[] memory ops) {
149+
uint256 nonce = entrypoint.getNonce(_sender, 0);
150+
151+
// Get user op fields
152+
UserOperation memory op = UserOperation({
153+
sender: _sender,
154+
nonce: nonce,
155+
initCode: _initCode,
156+
callData: _callDataForEntrypoint,
157+
callGasLimit: 500_000,
158+
verificationGasLimit: 500_000,
159+
preVerificationGas: 500_000,
160+
maxFeePerGas: 0,
161+
maxPriorityFeePerGas: 0,
162+
paymasterAndData: bytes(""),
163+
signature: bytes("")
164+
});
165+
166+
// Sign UserOp
167+
bytes32 opHash = EntryPoint(entrypoint).getUserOpHash(op);
168+
bytes32 msgHash = ECDSA.toEthSignedMessageHash(opHash);
169+
170+
(uint8 v, bytes32 r, bytes32 s) = vm.sign(accountAdminPKey, msgHash);
171+
bytes memory userOpSignature = abi.encodePacked(r, s, v);
172+
173+
address recoveredSigner = ECDSA.recover(msgHash, v, r, s);
174+
address expectedSigner = vm.addr(accountAdminPKey);
175+
assertEq(recoveredSigner, expectedSigner);
176+
177+
op.signature = userOpSignature;
178+
179+
// Store UserOp
180+
ops = new UserOperation[](1);
181+
ops[0] = op;
182+
}
183+
184+
function _setupUserOpExecuteWithSender(
185+
bytes memory _initCode,
186+
address _target,
187+
uint256 _value,
188+
bytes memory _callData,
189+
address _sender
190+
) internal returns (UserOperation[] memory) {
191+
bytes memory callDataForEntrypoint = abi.encodeWithSignature(
192+
"execute(address,uint256,bytes)",
193+
_target,
194+
_value,
195+
_callData
196+
);
197+
198+
return _setupUserOpWithSender(_initCode, callDataForEntrypoint, _sender);
199+
}
200+
144201
function _setupUserOpExecute(
145202
uint256 _signerPKey,
146203
bytes memory _initCode,
@@ -175,6 +232,11 @@ contract SimpleAccountTest is BaseTest {
175232
return _setupUserOp(_signerPKey, _initCode, callDataForEntrypoint);
176233
}
177234

235+
/// @dev Returns the salt used when deploying an Account.
236+
function _generateSalt(address _admin, bytes memory _data) internal view virtual returns (bytes32) {
237+
return keccak256(abi.encode(_admin, _data));
238+
}
239+
178240
function setUp() public override {
179241
super.setUp();
180242

@@ -234,7 +296,53 @@ contract SimpleAccountTest is BaseTest {
234296
function test_revert_onRegister_nonFactoryChildContract() public {
235297
vm.prank(address(0x12345));
236298
vm.expectRevert("AccountFactory: not an account.");
237-
accountFactory.onRegister(accountAdmin, "");
299+
accountFactory.onRegister(_generateSalt(accountAdmin, ""));
300+
}
301+
302+
/// @dev Create more than one accounts with the same admin.
303+
function test_state_createAccount_viaEntrypoint_multipleAccountSameAdmin() public {
304+
uint256 amount = 100;
305+
306+
for (uint256 i = 0; i < amount; i += 1) {
307+
bytes memory initCallData = abi.encodeWithSignature(
308+
"createAccount(address,bytes)",
309+
accountAdmin,
310+
bytes(abi.encode(i))
311+
);
312+
bytes memory initCode = abi.encodePacked(abi.encodePacked(address(accountFactory)), initCallData);
313+
314+
address expectedSenderAddress = Clones.predictDeterministicAddress(
315+
accountFactory.accountImplementation(),
316+
_generateSalt(accountAdmin, bytes(abi.encode(i))),
317+
address(accountFactory)
318+
);
319+
320+
UserOperation[] memory userOpCreateAccount = _setupUserOpExecuteWithSender(
321+
initCode,
322+
address(0),
323+
0,
324+
bytes(abi.encode(i)),
325+
expectedSenderAddress
326+
);
327+
328+
vm.expectEmit(true, true, false, true);
329+
emit AccountCreated(expectedSenderAddress, accountAdmin);
330+
EntryPoint(entrypoint).handleOps(userOpCreateAccount, beneficiary);
331+
}
332+
333+
address[] memory allAccounts = accountFactory.getAllAccounts();
334+
assertEq(allAccounts.length, amount);
335+
336+
for (uint256 i = 0; i < amount; i += 1) {
337+
assertEq(
338+
allAccounts[i],
339+
Clones.predictDeterministicAddress(
340+
accountFactory.accountImplementation(),
341+
_generateSalt(accountAdmin, bytes(abi.encode(i))),
342+
address(accountFactory)
343+
)
344+
);
345+
}
238346
}
239347

240348
/*///////////////////////////////////////////////////////////////

‎src/test/smart-wallet/AccountVulnPOC.t.sol

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -73,7 +73,7 @@ contract SimpleAccountVulnPOCTest is BaseTest {
7373
address private nonSigner;
7474

7575
// UserOp terminology: `sender` is the smart wallet.
76-
address private sender = 0xBB956D56140CA3f3060986586A2631922a4B347E;
76+
address private sender = 0x0df2C3523703d165Aa7fA1a552f3F0B56275DfC6;
7777
address payable private beneficiary = payable(address(0x45654));
7878

7979
bytes32 private uidCache = bytes32("random uid");

0 commit comments

Comments
 (0)