Proxy pattern smart wallet factory contracts

contracts/prebuilts/account/dynamic/DynamicAccount.sol

@@ -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);
     }
 

contracts/prebuilts/account/dynamic/DynamicAccountFactory.sol

@@ -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";
@@ -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);
     }
 
     /*///////////////////////////////////////////////////////////////
@@ -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.

contracts/prebuilts/account/managed/ManagedAccount.sol

@@ -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();
     }
 }

contracts/prebuilts/account/managed/ManagedAccountFactory.sol

@@ -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";
@@ -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);
     }
 
     /*///////////////////////////////////////////////////////////////
@@ -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.

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

@@ -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 {
@@ -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);
         }

contracts/prebuilts/account/non-upgradeable/AccountFactory.sol

@@ -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";
@@ -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);
     }
 
     /*///////////////////////////////////////////////////////////////
@@ -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.

contracts/prebuilts/account/utils/AccountCore.sol

@@ -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;
@@ -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);
         }
     }
 }

contracts/prebuilts/account/utils/AccountCoreStorage.sol

@@ -9,6 +9,7 @@ library AccountCoreStorage {
 
     struct Data {
         address entrypointOverride;
+        address factory;
         bytes32 creationSalt;
     }
 

contracts/prebuilts/account/utils/BaseAccountFactory.sol

@@ -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";
@@ -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
     //////////////////////////////////////////////////////////////*/
@@ -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);
@@ -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);
@@ -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);
@@ -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.
@@ -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();
     }
 
     /*///////////////////////////////////////////////////////////////
@@ -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,