Skip to content

Commit

Permalink
Proxy pattern smart wallet factory contracts (#562)
Browse files Browse the repository at this point in the history 
* 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

* Create BaseAccountFactoryStorage

* Make account factory contracts initializable

* Fix build errors in tests

* Store factory in initialize fn

* Update tests
  • Loading branch information
@nkrishang
nkrishang authored Oct 30, 2023
1 parent 80fa323 commit c778b3c
Show file tree
Hide file tree
Showing 17 changed files with 236 additions and 74 deletions.
11 changes: 9 additions & 2 deletions contracts/prebuilts/account/dynamic/DynamicAccount.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -24,16 +24,23 @@ contract DynamicAccount is AccountCore, BaseRouter {
//////////////////////////////////////////////////////////////*/

constructor(IEntryPoint _entrypoint, Extension[] memory _defaultExtensions)
AccountCore(_entrypoint, msg.sender)
AccountCore(_entrypoint)
BaseRouter(_defaultExtensions)
{
_disableInitializers();
}

/// @notice Initializes the smart contract wallet.
function initialize(address _defaultAdmin, bytes calldata _data) public override initializer {
function initialize(
address _defaultAdmin,
address _factory,
bytes calldata _data
) public override initializer {
__BaseRouter_init();

// This is passed as data in the `_registerOnFactory()` call in `AccountExtension` / `Account`.
AccountCoreStorage.data().creationSalt = _generateSalt(_defaultAdmin, _data);
AccountCoreStorage.data().factory = _factory;
_setAdmin(_defaultAdmin, true);
}

Expand Down
19 changes: 12 additions & 7 deletions contracts/prebuilts/account/dynamic/DynamicAccountFactory.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,10 +4,11 @@ pragma solidity ^0.8.12;
// Utils
import "../utils/BaseAccountFactory.sol";
import "@thirdweb-dev/dynamic-contracts/src/interface/IExtension.sol";
import "../../../extension/upgradeable/Initializable.sol";

// Extensions
import "../../../extension/upgradeable//PermissionsEnumerable.sol";
import "../../../extension/upgradeable//ContractMetadata.sol";
import "../../../extension/upgradeable/PermissionsEnumerable.sol";
import "../../../extension/upgradeable/ContractMetadata.sol";

// Smart wallet implementation
import { DynamicAccount, IEntryPoint } from "./DynamicAccount.sol";
Expand All @@ -21,20 +22,24 @@ import { DynamicAccount, IEntryPoint } from "./DynamicAccount.sol";
// \$$$$ |$$ | $$ |$$ |$$ | \$$$$$$$ |\$$$$$\$$$$ |\$$$$$$$\ $$$$$$$ |
// \____/ \__| \__|\__|\__| \_______| \_____\____/ \_______|\_______/

contract DynamicAccountFactory is BaseAccountFactory, ContractMetadata, PermissionsEnumerable {
contract DynamicAccountFactory is Initializable, BaseAccountFactory, ContractMetadata, PermissionsEnumerable {
address public constant ENTRYPOINT_ADDRESS = 0x5FF137D4b0FDCD49DcA30c7CF57E578a026d2789;

/*///////////////////////////////////////////////////////////////
Constructor
//////////////////////////////////////////////////////////////*/

constructor(address _defaultAdmin, IExtension.Extension[] memory _defaultExtensions)
constructor(IExtension.Extension[] memory _defaultExtensions)
BaseAccountFactory(
payable(address(new DynamicAccount(IEntryPoint(ENTRYPOINT_ADDRESS), _defaultExtensions))),
address(new DynamicAccount(IEntryPoint(ENTRYPOINT_ADDRESS), _defaultExtensions)),
ENTRYPOINT_ADDRESS
)
{
{}

/// @notice Initializes the factory contract.
function initialize(address _defaultAdmin, string memory _contractURI) external initializer {
_setupRole(DEFAULT_ADMIN_ROLE, _defaultAdmin);
_setupContractURI(_contractURI);
}

/*///////////////////////////////////////////////////////////////
Expand All @@ -47,7 +52,7 @@ contract DynamicAccountFactory is BaseAccountFactory, ContractMetadata, Permissi
address _admin,
bytes calldata _data
) internal override {
DynamicAccount(payable(_account)).initialize(_admin, _data);
DynamicAccount(payable(_account)).initialize(_admin, address(this), _data);
}

/// @dev Returns whether contract metadata can be set in the given execution context.
Expand Down
6 changes: 3 additions & 3 deletions contracts/prebuilts/account/managed/ManagedAccount.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -19,15 +19,15 @@ import "@thirdweb-dev/dynamic-contracts/src/core/Router.sol";
import "@thirdweb-dev/dynamic-contracts/src/interface/IRouterState.sol";

contract ManagedAccount is AccountCore, Router, IRouterState {
constructor(IEntryPoint _entrypoint, address _factory) AccountCore(_entrypoint, _factory) {}
constructor(IEntryPoint _entrypoint) AccountCore(_entrypoint) {}

/// @notice Returns the implementation contract address for a given function signature.
function getImplementationForFunction(bytes4 _functionSelector) public view virtual override returns (address) {
return Router(payable(factory)).getImplementationForFunction(_functionSelector);
return Router(payable(AccountCoreStorage.data().factory)).getImplementationForFunction(_functionSelector);
}

/// @notice Returns all extensions of the Router.
function getAllExtensions() external view returns (Extension[] memory) {
return IRouterState(payable(factory)).getAllExtensions();
return IRouterState(payable(AccountCoreStorage.data().factory)).getAllExtensions();
}
}
27 changes: 18 additions & 9 deletions contracts/prebuilts/account/managed/ManagedAccountFactory.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@ pragma solidity ^0.8.12;
// Utils
import "@thirdweb-dev/dynamic-contracts/src/presets/BaseRouter.sol";
import "../utils/BaseAccountFactory.sol";
import "../../../extension/upgradeable/Initializable.sol";

// Extensions
import "../../../extension/upgradeable//PermissionsEnumerable.sol";
Expand All @@ -21,25 +22,33 @@ import { ManagedAccount, IEntryPoint } from "./ManagedAccount.sol";
// \$$$$ |$$ | $$ |$$ |$$ | \$$$$$$$ |\$$$$$\$$$$ |\$$$$$$$\ $$$$$$$ |
// \____/ \__| \__|\__|\__| \_______| \_____\____/ \_______|\_______/

contract ManagedAccountFactory is BaseAccountFactory, ContractMetadata, PermissionsEnumerable, BaseRouter {
contract ManagedAccountFactory is
Initializable,
BaseAccountFactory,
ContractMetadata,
PermissionsEnumerable,
BaseRouter
{
/*///////////////////////////////////////////////////////////////
Constructor
//////////////////////////////////////////////////////////////*/

constructor(
address _defaultAdmin,
IEntryPoint _entrypoint,
Extension[] memory _defaultExtensions
)
constructor(IEntryPoint _entrypoint, Extension[] memory _defaultExtensions)
BaseRouter(_defaultExtensions)
BaseAccountFactory(payable(address(new ManagedAccount(_entrypoint, address(this)))), address(_entrypoint))
{
BaseAccountFactory(address(new ManagedAccount(_entrypoint)), address(_entrypoint))
{}

/// @notice Initializes the factory contract.
function initialize(address _defaultAdmin, string memory _contractURI) external initializer {
__BaseRouter_init();

_setupRole(DEFAULT_ADMIN_ROLE, _defaultAdmin);

bytes32 _extensionRole = keccak256("EXTENSION_ROLE");
_setupRole(_extensionRole, _defaultAdmin);
_setRoleAdmin(_extensionRole, _extensionRole);

_setupContractURI(_contractURI);
}

/*///////////////////////////////////////////////////////////////
Expand All @@ -52,7 +61,7 @@ contract ManagedAccountFactory is BaseAccountFactory, ContractMetadata, Permissi
address _admin,
bytes calldata _data
) internal override {
ManagedAccount(payable(_account)).initialize(_admin, _data);
ManagedAccount(payable(_account)).initialize(_admin, address(this), _data);
}

/// @dev Returns whether all relevant permission and other checks are met before any upgrade.
Expand Down
4 changes: 2 additions & 2 deletions contracts/prebuilts/account/non-upgradeable/Account.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -37,7 +37,7 @@ contract Account is AccountCore, ContractMetadata, ERC1271, ERC721Holder, ERC115
Constructor, Initializer, Modifiers
//////////////////////////////////////////////////////////////*/

constructor(IEntryPoint _entrypoint, address _factory) AccountCore(_entrypoint, _factory) {}
constructor(IEntryPoint _entrypoint) AccountCore(_entrypoint) {}

/// @notice Checks whether the caller is the EntryPoint contract or the admin.
modifier onlyAdminOrEntrypoint() virtual {
Expand Down Expand Up @@ -132,7 +132,7 @@ contract Account is AccountCore, ContractMetadata, ERC1271, ERC721Holder, ERC115

/// @dev Registers the account on the factory if it hasn't been registered yet.
function _registerOnFactory() internal virtual {
BaseAccountFactory factoryContract = BaseAccountFactory(factory);
BaseAccountFactory factoryContract = BaseAccountFactory(AccountCoreStorage.data().factory);
if (!factoryContract.isRegistered(address(this))) {
factoryContract.onRegister(AccountCoreStorage.data().creationSalt);
}
Expand Down
13 changes: 8 additions & 5 deletions contracts/prebuilts/account/non-upgradeable/AccountFactory.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@ pragma solidity ^0.8.12;
import "../utils/BaseAccountFactory.sol";
import "../utils/BaseAccount.sol";
import "../../../external-deps/openzeppelin/proxy/Clones.sol";
import "../../../extension/upgradeable/Initializable.sol";

// Extensions
import "../../../extension/upgradeable//PermissionsEnumerable.sol";
Expand All @@ -25,15 +26,17 @@ import { Account } from "./Account.sol";
// \$$$$ |$$ | $$ |$$ |$$ | \$$$$$$$ |\$$$$$\$$$$ |\$$$$$$$\ $$$$$$$ |
// \____/ \__| \__|\__|\__| \_______| \_____\____/ \_______|\_______/

contract AccountFactory is BaseAccountFactory, ContractMetadata, PermissionsEnumerable {
contract AccountFactory is Initializable, BaseAccountFactory, ContractMetadata, PermissionsEnumerable {
/*///////////////////////////////////////////////////////////////
Constructor
//////////////////////////////////////////////////////////////*/

constructor(address _defaultAdmin, IEntryPoint _entrypoint)
BaseAccountFactory(address(new Account(_entrypoint, address(this))), address(_entrypoint))
{
constructor(IEntryPoint _entrypoint) BaseAccountFactory(address(new Account(_entrypoint)), address(_entrypoint)) {}

/// @notice Initializes the factory contract.
function initialize(address _defaultAdmin, string memory _contractURI) external initializer {
_setupRole(DEFAULT_ADMIN_ROLE, _defaultAdmin);
_setupContractURI(_contractURI);
}

/*///////////////////////////////////////////////////////////////
Expand All @@ -46,7 +49,7 @@ contract AccountFactory is BaseAccountFactory, ContractMetadata, PermissionsEnum
address _admin,
bytes calldata _data
) internal override {
Account(payable(_account)).initialize(_admin, _data);
Account(payable(_account)).initialize(_admin, address(this), _data);
}

/// @dev Returns whether contract metadata can be set in the given execution context.
Expand Down
31 changes: 20 additions & 11 deletions contracts/prebuilts/account/utils/AccountCore.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -39,33 +39,39 @@ contract AccountCore is IAccountCore, Initializable, Multicall, BaseAccount, Acc
State
//////////////////////////////////////////////////////////////*/

/// @notice EIP 4337 factory for this contract.
address public immutable factory;

/// @notice EIP 4337 Entrypoint contract.
IEntryPoint private immutable entrypointContract;

/*///////////////////////////////////////////////////////////////
Constructor, Initializer, Modifiers
//////////////////////////////////////////////////////////////*/

constructor(IEntryPoint _entrypoint, address _factory) EIP712("Account", "1") {
constructor(IEntryPoint _entrypoint) EIP712("Account", "1") {
_disableInitializers();
factory = _factory;
entrypointContract = _entrypoint;
}

/// @notice Initializes the smart contract wallet.
function initialize(address _defaultAdmin, bytes calldata _data) public virtual initializer {
function initialize(
address _defaultAdmin,
address _factory,
bytes calldata _data
) public virtual initializer {
// This is passed as data in the `_registerOnFactory()` call in `AccountExtension` / `Account`.
AccountCoreStorage.data().creationSalt = _generateSalt(_defaultAdmin, _data);
AccountCoreStorage.data().factory = _factory;
_setAdmin(_defaultAdmin, true);
}

/*///////////////////////////////////////////////////////////////
View functions
//////////////////////////////////////////////////////////////*/

/// @notice Returns the address of the account factory.
function factory() public view virtual override returns (address) {
return AccountCoreStorage.data().factory;
}

/// @notice Returns the EIP 4337 entrypoint contract.
function entryPoint() public view virtual override returns (IEntryPoint) {
address entrypointOverride = AccountCoreStorage.data().entrypointOverride;
Expand Down Expand Up @@ -235,19 +241,22 @@ contract AccountCore is IAccountCore, Initializable, Multicall, BaseAccount, Acc
/// @notice Makes the given account an admin.
function _setAdmin(address _account, bool _isAdmin) internal virtual override {
super._setAdmin(_account, _isAdmin);
if (factory.code.length > 0) {

address factoryAddr = factory();
if (factoryAddr.code.length > 0) {
if (_isAdmin) {
BaseAccountFactory(factory).onSignerAdded(_account, AccountCoreStorage.data().creationSalt);
BaseAccountFactory(factoryAddr).onSignerAdded(_account, AccountCoreStorage.data().creationSalt);
} else {
BaseAccountFactory(factory).onSignerRemoved(_account, AccountCoreStorage.data().creationSalt);
BaseAccountFactory(factoryAddr).onSignerRemoved(_account, AccountCoreStorage.data().creationSalt);
}
}
}

/// @notice Runs after every `changeRole` run.
function _afterSignerPermissionsUpdate(SignerPermissionRequest calldata _req) internal virtual override {
if (factory.code.length > 0) {
BaseAccountFactory(factory).onSignerAdded(_req.signer, AccountCoreStorage.data().creationSalt);
address factoryAddr = factory();
if (factoryAddr.code.length > 0) {
BaseAccountFactory(factoryAddr).onSignerAdded(_req.signer, AccountCoreStorage.data().creationSalt);
}
}
}
1 change: 1 addition & 0 deletions contracts/prebuilts/account/utils/AccountCoreStorage.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@ library AccountCoreStorage {

struct Data {
address entrypointOverride;
address factory;
bytes32 creationSalt;
}

Expand Down
26 changes: 16 additions & 10 deletions contracts/prebuilts/account/utils/BaseAccountFactory.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@
pragma solidity ^0.8.12;

// Utils
import "./BaseAccountFactoryStorage.sol";
import "../../../extension/Multicall.sol";
import "../../../external-deps/openzeppelin/proxy/Clones.sol";
import "../../../external-deps/openzeppelin/utils/structs/EnumerableSet.sol";
Expand Down Expand Up @@ -32,9 +33,6 @@ abstract contract BaseAccountFactory is IAccountFactory, Multicall {
address public immutable accountImplementation;
address public immutable entrypoint;

EnumerableSet.AddressSet private allAccounts;
mapping(address => EnumerableSet.AddressSet) internal accountsOfSigner;

/*///////////////////////////////////////////////////////////////
Constructor
//////////////////////////////////////////////////////////////*/
Expand All @@ -61,7 +59,10 @@ abstract contract BaseAccountFactory is IAccountFactory, Multicall {
account = Clones.cloneDeterministic(impl, salt);

if (msg.sender != entrypoint) {
require(allAccounts.add(account), "AccountFactory: account already registered");
require(
_baseAccountFactoryStorage().allAccounts.add(account),
"AccountFactory: account already registered"
);
}

_initializeAccount(account, _admin, _data);
Expand All @@ -76,14 +77,14 @@ abstract contract BaseAccountFactory is IAccountFactory, Multicall {
address account = msg.sender;
require(_isAccountOfFactory(account, _salt), "AccountFactory: not an account.");

require(allAccounts.add(account), "AccountFactory: account already registered");
require(_baseAccountFactoryStorage().allAccounts.add(account), "AccountFactory: account already registered");
}

function onSignerAdded(address _signer, bytes32 _salt) external {
address account = msg.sender;
require(_isAccountOfFactory(account, _salt), "AccountFactory: not an account.");

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

if (isNewSigner) {
emit SignerAdded(account, _signer);
Expand All @@ -95,7 +96,7 @@ abstract contract BaseAccountFactory is IAccountFactory, Multicall {
address account = msg.sender;
require(_isAccountOfFactory(account, _salt), "AccountFactory: not an account.");

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

if (isAccount) {
emit SignerRemoved(account, _signer);
Expand All @@ -108,12 +109,12 @@ abstract contract BaseAccountFactory is IAccountFactory, Multicall {

/// @notice Returns whether an account is registered on this factory.
function isRegistered(address _account) external view returns (bool) {
return allAccounts.contains(_account);
return _baseAccountFactoryStorage().allAccounts.contains(_account);
}

/// @notice Returns all accounts created on the factory.
function getAllAccounts() external view returns (address[] memory) {
return allAccounts.values();
return _baseAccountFactoryStorage().allAccounts.values();
}

/// @notice Returns the address of an Account that would be deployed with the given admin signer.
Expand All @@ -124,7 +125,7 @@ abstract contract BaseAccountFactory is IAccountFactory, Multicall {

/// @notice Returns all accounts that the given address is a signer of.
function getAccountsOfSigner(address signer) external view returns (address[] memory accounts) {
return accountsOfSigner[signer].values();
return _baseAccountFactoryStorage().accountsOfSigner[signer].values();
}

/*///////////////////////////////////////////////////////////////
Expand All @@ -147,6 +148,11 @@ abstract contract BaseAccountFactory is IAccountFactory, Multicall {
return keccak256(abi.encode(_admin, _data));
}

/// @dev Returns the BaseAccountFactory contract's storage.
function _baseAccountFactoryStorage() internal pure returns (BaseAccountFactoryStorage.Data storage) {
return BaseAccountFactoryStorage.data();
}

/// @dev Called in `createAccount`. Initializes the account contract created in `createAccount`.
function _initializeAccount(
address _account,
Expand Down
Loading

0 comments on commit c778b3c

Please sign in to comment.