Add getAllWallets to BaseAccountFactory (#412)

* Add getAllWallets to BaseAccountFactory

* Remove redundant local isRegistered flag

* Move execute and executeBatch from AccountCore -> AccountExtension

* rename getAllWallets -> getAllAccounts + tests

* Remove unused functions from BytesLib

* docs update

* dev build

* Add getAllAccounts to IAccountFactory

* dev build

* pkg update

Author: nkrishang

contracts/lib/BytesLib.sol

@@ -0,0 +1,25 @@
+// SPDX-License-Identifier: Apache 2.0
+/*
+ * @title Solidity Bytes Arrays Utils
+ * @author Gonçalo Sá <[email protected]>
+ *
+ * @dev Bytes tightly packed arrays utility library for ethereum contracts written in Solidity.
+ *      The library lets you concatenate, slice and type cast bytes arrays both in memory and storage.
+ *
+ *  Credits: https://github.com/GNSPS/solidity-bytes-utils/blob/master/contracts/BytesLib.sol
+ */
+
+pragma solidity >=0.8.0 <0.9.0;
+
+library BytesLib {
+    function toAddress(bytes memory _bytes, uint256 _start) internal pure returns (address) {
+        require(_bytes.length >= _start + 20, "toAddress_outOfBounds");
+        address tempAddress;
+
+        assembly {
+            tempAddress := div(mload(add(add(_bytes, 0x20), _start)), 0x1000000000000000000000000)
+        }
+
+        return tempAddress;
+    }
+}

contracts/package.json

@@ -1,7 +1,7 @@
 {
   "name": "@thirdweb-dev/contracts",
   "description": "Collection of smart contracts deployable via the thirdweb SDK, dashboard and CLI",
-  "version": "3.5.8",
+  "version": "3.5.9",
   "license": "Apache-2.0",
   "repository": {
     "type": "git",

contracts/smart-wallet/dynamic/DynamicAccountFactory.sol

@@ -27,7 +27,7 @@ contract DynamicAccountFactory is BaseAccountFactory, ContractMetadata, Permissi
     //////////////////////////////////////////////////////////////*/
 
     constructor(IEntryPoint _entrypoint, IExtension.Extension[] memory _defaultExtensions)
-        BaseAccountFactory(payable(address(new DynamicAccount(_entrypoint, _defaultExtensions))))
+        BaseAccountFactory(payable(address(new DynamicAccount(_entrypoint, _defaultExtensions))), address(_entrypoint))
     {
         _setupRole(DEFAULT_ADMIN_ROLE, msg.sender);
     }

contracts/smart-wallet/interfaces/IAccountFactory.sol

@@ -35,6 +35,9 @@ interface IAccountFactory {
     /// @notice Returns the address of the Account implementation.
     function accountImplementation() external view returns (address);
 
+    /// @notice Returns all accounts created on the factory.
+    function getAllAccounts() external view returns (address[] memory);
+
     /// @notice Returns the address of an Account that would be deployed with the given admin signer.
     function getAddress(address adminSigner, bytes calldata data) external view returns (address);
 

contracts/smart-wallet/managed/ManagedAccountFactory.sol

@@ -28,7 +28,7 @@ contract ManagedAccountFactory is BaseAccountFactory, ContractMetadata, Permissi
 
     constructor(IEntryPoint _entrypoint, Extension[] memory _defaultExtensions)
         BaseRouter(_defaultExtensions)
-        BaseAccountFactory(payable(address(new ManagedAccount(_entrypoint, address(this)))))
+        BaseAccountFactory(payable(address(new ManagedAccount(_entrypoint, address(this)))), address(_entrypoint))
     {
         _setupRole(DEFAULT_ADMIN_ROLE, msg.sender);
     }

contracts/smart-wallet/non-upgradeable/Account.sol

@@ -6,7 +6,7 @@ pragma solidity ^0.8.11;
 /* solhint-disable reason-string */
 
 // Base
-import "./../utils/BaseAccount.sol";
+import "../utils/BaseAccount.sol";
 
 // Extensions
 import "../../extension/Multicall.sol";
@@ -19,7 +19,7 @@ import "../../eip/ERC1271.sol";
 
 // Utils
 import "../../openzeppelin-presets/utils/cryptography/ECDSA.sol";
-import "./AccountFactory.sol";
+import "../utils/BaseAccountFactory.sol";
 
 //   $$\     $$\       $$\                 $$\                         $$\
 //   $$ |    $$ |      \__|                $$ |                        $$ |
@@ -176,6 +176,7 @@ contract Account is
         uint256 _value,
         bytes calldata _calldata
     ) external virtual onlyAdminOrEntrypoint {
+        _registerOnFactory();
         _call(_target, _value, _calldata);
     }
 
@@ -185,6 +186,8 @@ contract Account is
         uint256[] calldata _value,
         bytes[] calldata _calldata
     ) external virtual onlyAdminOrEntrypoint {
+        _registerOnFactory();
+
         require(_target.length == _calldata.length && _target.length == _value.length, "Account: wrong array lengths.");
         for (uint256 i = 0; i < _target.length; i++) {
             _call(_target[i], _value[i], _calldata[i]);
@@ -205,6 +208,14 @@ contract Account is
                         Internal functions
     //////////////////////////////////////////////////////////////*/
 
+    /// @dev Registers the account on the factory if it hasn't been registered yet.
+    function _registerOnFactory() internal virtual {
+        BaseAccountFactory factoryContract = BaseAccountFactory(factory);
+        if (!factoryContract.isRegistered(address(this))) {
+            factoryContract.onRegister();
+        }
+    }
+
     /// @dev Calls a target contract and reverts if it fails.
     function _call(
         address _target,
@@ -268,9 +279,9 @@ contract Account is
         super._setAdmin(_account, _isAdmin);
         if (factory.code.length > 0) {
             if (_isAdmin) {
-                AccountFactory(factory).onSignerAdded(_account);
+                BaseAccountFactory(factory).onSignerAdded(_account);
             } else {
-                AccountFactory(factory).onSignerRemoved(_account);
+                BaseAccountFactory(factory).onSignerRemoved(_account);
             }
         }
     }
@@ -279,9 +290,9 @@ contract Account is
     function _afterChangeRole(RoleRequest calldata _req) internal virtual override {
         if (factory.code.length > 0) {
             if (_req.action == RoleAction.GRANT) {
-                AccountFactory(factory).onSignerAdded(_req.target);
+                BaseAccountFactory(factory).onSignerAdded(_req.target);
             } else if (_req.action == RoleAction.REVOKE) {
-                AccountFactory(factory).onSignerRemoved(_req.target);
+                BaseAccountFactory(factory).onSignerRemoved(_req.target);
             }
         }
     }

contracts/smart-wallet/non-upgradeable/AccountFactory.sol

@@ -30,7 +30,9 @@ contract AccountFactory is BaseAccountFactory, ContractMetadata, PermissionsEnum
                             Constructor
     //////////////////////////////////////////////////////////////*/
 
-    constructor(IEntryPoint _entrypoint) BaseAccountFactory(address(new Account(_entrypoint, address(this)))) {
+    constructor(IEntryPoint _entrypoint)
+        BaseAccountFactory(address(new Account(_entrypoint, address(this))), address(_entrypoint))
+    {
         _setupRole(DEFAULT_ADMIN_ROLE, msg.sender);
     }
 

contracts/smart-wallet/utils/AccountExtension.sol

@@ -13,6 +13,8 @@ import "../../openzeppelin-presets/token/ERC1155/utils/ERC1155Holder.sol";
 
 // Utils
 import "../../openzeppelin-presets/utils/cryptography/ECDSA.sol";
+import "./BaseAccountFactory.sol";
+import "./AccountCore.sol";
 
 //   $$\     $$\       $$\                 $$\                         $$\
 //   $$ |    $$ |      \__|                $$ |                        $$ |
@@ -37,19 +39,19 @@ contract AccountExtension is ContractMetadata, AccountPermissions, ERC721Holder,
                     Constructor, Initializer, Modifiers
     //////////////////////////////////////////////////////////////*/
 
+    /// @notice Checks whether the caller is the EntryPoint contract or the admin.
+    modifier onlyAdminOrEntrypoint() virtual {
+        require(msg.sender == address(entrypointContract) || isAdmin(msg.sender), "Account: not admin or EntryPoint.");
+        _;
+    }
+
     // solhint-disable-next-line no-empty-blocks
     receive() external payable virtual {}
 
     constructor(address _entrypoint) EIP712("Account", "1") {
         entrypointContract = _entrypoint;
     }
 
-    /// @notice Checks whether the caller is the EntryPoint contract or the admin.
-    modifier onlyAdminOrEntrypoint() virtual {
-        require(msg.sender == entrypointContract || isAdmin(msg.sender), "Account: not admin or EntryPoint.");
-        _;
-    }
-
     /*///////////////////////////////////////////////////////////////
                             View functions
     //////////////////////////////////////////////////////////////*/
@@ -72,6 +74,7 @@ contract AccountExtension is ContractMetadata, AccountPermissions, ERC721Holder,
         uint256 _value,
         bytes calldata _calldata
     ) external virtual onlyAdminOrEntrypoint {
+        _registerOnFactory();
         _call(_target, _value, _calldata);
     }
 
@@ -81,6 +84,7 @@ contract AccountExtension is ContractMetadata, AccountPermissions, ERC721Holder,
         uint256[] calldata _value,
         bytes[] calldata _calldata
     ) external virtual onlyAdminOrEntrypoint {
+        _registerOnFactory();
         require(_target.length == _calldata.length && _target.length == _value.length, "Account: wrong array lengths.");
         for (uint256 i = 0; i < _target.length; i++) {
             _call(_target[i], _value[i], _calldata[i]);
@@ -91,6 +95,15 @@ contract AccountExtension is ContractMetadata, AccountPermissions, ERC721Holder,
                         Internal functions
     //////////////////////////////////////////////////////////////*/
 
+    /// @dev Registers the account on the factory if it hasn't been registered yet.
+    function _registerOnFactory() internal virtual {
+        address factory = AccountCore(payable(address(this))).factory();
+        BaseAccountFactory factoryContract = BaseAccountFactory(factory);
+        if (!factoryContract.isRegistered(address(this))) {
+            factoryContract.onRegister();
+        }
+    }
+
     /// @dev Calls a target contract and reverts if it fails.
     function _call(
         address _target,

contracts/smart-wallet/utils/BaseAccountFactory.sol

@@ -7,6 +7,7 @@ import "../../openzeppelin-presets/proxy/Clones.sol";
 import "../../openzeppelin-presets/utils/structs/EnumerableSet.sol";
 import "../utils/BaseAccount.sol";
 import "../../extension/interface/IAccountPermissions.sol";
+import "../../lib/BytesLib.sol";
 
 // Interface
 import "../interfaces/IEntrypoint.sol";
@@ -29,16 +30,19 @@ abstract contract BaseAccountFactory is IAccountFactory, Multicall {
     //////////////////////////////////////////////////////////////*/
 
     address public immutable accountImplementation;
+    address public immutable entrypoint;
 
+    EnumerableSet.AddressSet private allAccounts;
     mapping(address => EnumerableSet.AddressSet) internal accountsOfSigner;
     mapping(address => EnumerableSet.AddressSet) internal signersOfAccount;
 
     /*///////////////////////////////////////////////////////////////
                             Constructor
     //////////////////////////////////////////////////////////////*/
 
-    constructor(address _accountImpl) {
+    constructor(address _accountImpl, address _entrypoint) {
         accountImplementation = _accountImpl;
+        entrypoint = _entrypoint;
     }
 
     /*///////////////////////////////////////////////////////////////
@@ -57,18 +61,28 @@ abstract contract BaseAccountFactory is IAccountFactory, Multicall {
 
         account = Clones.cloneDeterministic(impl, salt);
 
+        if (msg.sender != entrypoint) {
+            require(allAccounts.add(account), "AccountFactory: account already registered");
+        }
+
         _initializeAccount(account, _admin, _data);
 
         emit AccountCreated(account, _admin);
 
         return account;
     }
 
-    /// @notice Callback function for an Account to register its signers.
-    function onSignerAdded(address _signer) external {
+    /// @notice Callback function for an Account to register itself on the factory.
+    function onRegister() external {
         address account = msg.sender;
+        require(_isAccountOfFactory(account), "AccountFactory: not an account.");
+
+        require(allAccounts.add(account), "AccountFactory: account already registered");
+    }
 
-        require(account.code.length > 0, "AccountFactory: not an account.");
+    function onSignerAdded(address _signer) external {
+        address account = msg.sender;
+        require(_isAccountOfFactory(account), "AccountFactory: not an account.");
 
         bool isAdmin = IAccountPermissions(account).isAdmin(_signer);
 
@@ -94,8 +108,7 @@ abstract contract BaseAccountFactory is IAccountFactory, Multicall {
     /// @notice Callback function for an Account to un-register its signers.
     function onSignerRemoved(address _signer) external {
         address account = msg.sender;
-
-        require(account.code.length > 0, "AccountFactory: not an account.");
+        require(_isAccountOfFactory(account), "AccountFactory: not an account.");
 
         bool isAccount = accountsOfSigner[_signer].remove(account);
         bool isSigner = signersOfAccount[account].remove(_signer);
@@ -111,6 +124,16 @@ abstract contract BaseAccountFactory is IAccountFactory, Multicall {
                             View functions
     //////////////////////////////////////////////////////////////*/
 
+    /// @notice Returns whether an account is registered on this factory.
+    function isRegistered(address _account) external view returns (bool) {
+        return allAccounts.contains(_account);
+    }
+
+    /// @notice Returns all accounts created on the factory.
+    function getAllAccounts() external view returns (address[] memory) {
+        return allAccounts.values();
+    }
+
     /// @notice Returns the address of an Account that would be deployed with the given admin signer.
     function getAddress(address _adminSigner, bytes calldata _data) public view returns (address) {
         bytes32 salt = _generateSalt(_adminSigner, _data);
@@ -131,6 +154,17 @@ abstract contract BaseAccountFactory is IAccountFactory, Multicall {
                             Internal functions
     //////////////////////////////////////////////////////////////*/
 
+    /// @dev Returns whether the caller is an account deployed by this factory.
+    function _isAccountOfFactory(address _account) internal view virtual returns (bool) {
+        address impl = _getImplementation(_account);
+        return _account.code.length > 0 && impl == accountImplementation;
+    }
+
+    function _getImplementation(address cloneAddress) internal view returns (address) {
+        bytes memory code = cloneAddress.code;
+        return BytesLib.toAddress(code, 10);
+    }
+
     /// @dev Returns the salt used when deploying an Account.
     function _generateSalt(address _admin, bytes calldata) internal view virtual returns (bytes32) {
         return keccak256(abi.encode(_admin));