feat: move token receiver functions into account (#124)

alchemyplatform · Feb 2, 2024 · a8d57be · a8d57be
1 parent 3549c69
commit a8d57be
Show file tree

Hide file tree

Showing 14 changed files with 132 additions and 750 deletions.
diff --git a/README.md b/README.md
@@ -5,16 +5,17 @@ Alchemy's Modular Account is a maximally modular, upgradeable Smart Contract Acc
 ## Overview
 
 This repository contains:
-* [Modular Account implementation](src/account)
-* [Modular Account factories](src/factory)
-* 3 ERC-6900 compatible plugins:
-    * [MultiOwnerPlugin](src/plugins/owner) is a plugin supporting 1+ ECDSA owners.
-    * [TokenReceiverPlugin](src/plugins/TokenReceiverPlugin.sol) contains ERC721/ERC777/ERC1155 token receivers.
-    * [SessionKeyPlugin](src/plugins/session) enables session keys with optional permissions such as time ranges, token spend limits, and gas spend limits.
+
+- [Modular Account implementation](src/account)
+- [Modular Account factory](src/factory)
+- 2 ERC-6900 compatible plugins:
+  - [MultiOwnerPlugin](src/plugins/owner) is a plugin supporting 1+ ECDSA owners.
+  - [SessionKeyPlugin](src/plugins/session) enables session keys with optional permissions such as time ranges, token spend limits, and gas spend limits.
 
 The account and plugins conform to these ERC versions:
-* ERC-4337: 0.6.0
-* ERC-6900: 0.7.0
+
+- ERC-4337: 0.6.0
+- ERC-6900: 0.7.0
 
 ## Development
 
@@ -43,6 +44,7 @@ FOUNDRY_PROFILE=lite forge test -vvv
 ### Deployment
 
 A deployment script can be found in the `scripts/` folder
+
 ```bash
 forge script script/Deploy.s.sol --rpc-url $RPC_URL --broadcast
 ```

diff --git a/foundry.toml b/foundry.toml
@@ -7,7 +7,7 @@ out = 'out'
 test = 'test'
 libs = ['lib']
 optimizer = true
-optimizer_runs = 1600
+optimizer_runs = 900
 ignored_error_codes = []
 
 [fuzz]

diff --git a/script/Deploy.s.sol b/script/Deploy.s.sol
@@ -17,21 +17,16 @@
 
 pragma solidity ^0.8.22;
 
-import {console} from "forge-std/Test.sol";
 import {Script} from "forge-std/Script.sol";
+import {console} from "forge-std/Test.sol";
 
 import {IEntryPoint} from "@eth-infinitism/account-abstraction/interfaces/IEntryPoint.sol";
 
-import {IEntryPoint as IMSCAEntryPoint} from "../src/interfaces/erc4337/IEntryPoint.sol";
-
 import {UpgradeableModularAccount} from "../src/account/UpgradeableModularAccount.sol";
-
 import {MultiOwnerMSCAFactory} from "../src/factory/MultiOwnerMSCAFactory.sol";
-import {MultiOwnerTokenReceiverMSCAFactory} from "../src/factory/MultiOwnerTokenReceiverMSCAFactory.sol";
-
+import {IEntryPoint as IMSCAEntryPoint} from "../src/interfaces/erc4337/IEntryPoint.sol";
 import {BasePlugin} from "../src/plugins/BasePlugin.sol";
 import {MultiOwnerPlugin} from "../src/plugins/owner/MultiOwnerPlugin.sol";
-import {TokenReceiverPlugin} from "../src/plugins/TokenReceiverPlugin.sol";
 import {SessionKeyPlugin} from "../src/plugins/session/SessionKeyPlugin.sol";
 
 contract Deploy is Script {
@@ -47,13 +42,10 @@ contract Deploy is Script {
     address public ownerFactoryAddr = vm.envOr("OWNER_FACTORY", address(0));
     address public ownerAndTokenReceiverFactoryAddr = vm.envOr("OWNER_TOKEN_RECEIVER_FACTORY", address(0));
     MultiOwnerMSCAFactory ownerFactory;
-    MultiOwnerTokenReceiverMSCAFactory ownerAndTokenReceiverFactory;
 
     // Load plugins contract, if not in env, deploy new contract
     address public multiOwnerPlugin = vm.envOr("OWNER_PLUGIN", address(0));
     bytes32 public multiOwnerPluginManifestHash;
-    address public tokenReceiverPlugin = vm.envOr("TOKEN_RECEIVER_PLUGIN", address(0));
-    bytes32 public tokenReceiverPluginManifestHash;
     address public sessionKeyPlugin = vm.envOr("SESSION_KEY_PLUGIN", address(0));
 
     function run() public {
@@ -84,49 +76,19 @@ contract Deploy is Script {
         }
         multiOwnerPluginManifestHash = keccak256(abi.encode(BasePlugin(multiOwnerPlugin).pluginManifest()));
 
-        // Deploy multi owner plugin, and set plugin hash
-        if (tokenReceiverPlugin == address(0)) {
-            TokenReceiverPlugin trp = new TokenReceiverPlugin();
-            tokenReceiverPlugin = address(trp);
-            console.log("New TokenReceiverPlugin: ", tokenReceiverPlugin);
-        } else {
-            console.log("Exist TokenReceiverPlugin: ", tokenReceiverPlugin);
-        }
-        tokenReceiverPluginManifestHash = keccak256(abi.encode(BasePlugin(tokenReceiverPlugin).pluginManifest()));
-
         // Deploy MultiOwnerMSCAFactory, and add stake with EP
-        {
-            if (ownerFactoryAddr == address(0)) {
-                ownerFactory = new MultiOwnerMSCAFactory(
-                    owner, multiOwnerPlugin, mscaImpl, multiOwnerPluginManifestHash, entryPoint
-                );
-
-                ownerFactoryAddr = address(ownerFactory);
-                console.log("New MultiOwnerMSCAFactory: ", ownerFactoryAddr);
-            } else {
-                console.log("Exist MultiOwnerMSCAFactory: ", ownerFactoryAddr);
-            }
-            _addStakeForFactory(ownerFactoryAddr, entryPoint);
-        }
 
-        // Deploy MultiOwnerTokenReceiverMSCAFactory, and add stake with EP
-        if (ownerAndTokenReceiverFactoryAddr == address(0)) {
-            ownerAndTokenReceiverFactory = new MultiOwnerTokenReceiverMSCAFactory(
-                owner,
-                multiOwnerPlugin,
-                tokenReceiverPlugin,
-                mscaImpl,
-                multiOwnerPluginManifestHash,
-                tokenReceiverPluginManifestHash,
-                entryPoint
+        if (ownerFactoryAddr == address(0)) {
+            ownerFactory = new MultiOwnerMSCAFactory(
+                owner, multiOwnerPlugin, mscaImpl, multiOwnerPluginManifestHash, entryPoint
             );
 
-            ownerAndTokenReceiverFactoryAddr = address(ownerAndTokenReceiverFactory);
-            console.log("New MultiOwnerTokenReceiverMSCAFactory: ", ownerAndTokenReceiverFactoryAddr);
+            ownerFactoryAddr = address(ownerFactory);
+            console.log("New MultiOwnerMSCAFactory: ", ownerFactoryAddr);
         } else {
-            console.log("Exist MultiOwnerTokenReceiverMSCAFactory: ", ownerAndTokenReceiverFactoryAddr);
+            console.log("Exist MultiOwnerMSCAFactory: ", ownerFactoryAddr);
         }
-        _addStakeForFactory(ownerAndTokenReceiverFactoryAddr, entryPoint);
+        _addStakeForFactory(ownerFactoryAddr, entryPoint);
 
         // Deploy SessionKeyPlugin impl
         if (sessionKeyPlugin == address(0)) {

diff --git a/src/account/PluginManagerInternals.sol b/src/account/PluginManagerInternals.sol
@@ -53,6 +53,9 @@ abstract contract PluginManagerInternals is IPluginManager, AccountStorageV1 {
         uint256 callbackGasLimit;
     }
 
+    // As per the EIP-165 spec, no interface should ever match 0xffffffff
+    bytes4 internal constant _INVALID_INTERFACE_ID = 0xffffffff;
+
     // These flags are used in LinkedListSet values to optimize lookups.
     // It's important that they don't overlap with bit 1 and bit 2, which are reserved bits used to indicate
     // the sentinel value and the existence of a next value, respectively.
@@ -64,10 +67,10 @@ abstract contract PluginManagerInternals is IPluginManager, AccountStorageV1 {
     error DuplicatePreUserOpValidationHookLimitExceeded(bytes4 selector, FunctionReference hook);
     error Erc4337FunctionNotAllowed(bytes4 selector);
     error ExecutionFunctionAlreadySet(bytes4 selector);
-    error IPluginFunctionNotAllowed(bytes4 selector);
-    error IPluginInterfaceNotAllowed();
+    error InterfaceNotAllowed();
     error InvalidDependenciesProvided();
     error InvalidPluginManifest();
+    error IPluginFunctionNotAllowed(bytes4 selector);
     error MissingPluginDependency(address dependency);
     error NativeFunctionNotAllowed(bytes4 selector);
     error NullFunctionReference();
@@ -452,8 +455,8 @@ abstract contract PluginManagerInternals is IPluginManager, AccountStorageV1 {
         length = manifest.interfaceIds.length;
         for (uint256 i = 0; i < length; ++i) {
             bytes4 interfaceId = manifest.interfaceIds[i];
-            if (interfaceId == type(IPlugin).interfaceId) {
-                revert IPluginInterfaceNotAllowed();
+            if (interfaceId == type(IPlugin).interfaceId || interfaceId == _INVALID_INTERFACE_ID) {
+                revert InterfaceNotAllowed();
             }
             storage_.supportedInterfaces[interfaceId] += 1;
         }

diff --git a/src/account/UpgradeableModularAccount.sol b/src/account/UpgradeableModularAccount.sol
@@ -17,6 +17,9 @@
 
 pragma solidity ^0.8.22;
 
+import {IERC1155Receiver} from "@openzeppelin/contracts/interfaces/IERC1155Receiver.sol";
+import {IERC777Recipient} from "@openzeppelin/contracts/interfaces/IERC777Recipient.sol";
+import {IERC721Receiver} from "@openzeppelin/contracts/token/ERC721/IERC721Receiver.sol";
 import {IERC165} from "@openzeppelin/contracts/utils/introspection/IERC165.sol";
 
 import {UUPSUpgradeable} from "../../ext/UUPSUpgradeable.sol";
@@ -51,6 +54,9 @@ contract UpgradeableModularAccount is
     IAccountInitializable,
     IAccountView,
     IERC165,
+    IERC721Receiver,
+    IERC777Recipient,
+    IERC1155Receiver,
     IPluginExecutor,
     IStandardExecutor,
     UUPSUpgradeable
@@ -80,8 +86,6 @@ contract UpgradeableModularAccount is
     // ERC-4337 v0.6.0 entrypoint address only
     IEntryPoint private immutable _ENTRY_POINT;
 
-    // As per the EIP-165 spec, no interface should ever match 0xffffffff
-    bytes4 internal constant _INTERFACE_ID_INVALID = 0xffffffff;
     bytes4 internal constant _IERC165_INTERFACE_ID = 0x01ffc9a7;
 
     event ModularAccountInitialized(IEntryPoint indexed entryPoint);
@@ -370,16 +374,58 @@ contract UpgradeableModularAccount is
         _postNativeFunction(postExecHooks, postHookArgs);
     }
 
-    /// @inheritdoc IERC165
-    function supportsInterface(bytes4 interfaceId) external view override returns (bool) {
-        if (interfaceId == _INTERFACE_ID_INVALID) {
-            return false;
-        }
-        if (interfaceId == _IERC165_INTERFACE_ID) {
-            return true;
-        }
+    /// @inheritdoc IERC777Recipient
+    /// @dev Runtime validation is bypassed for this function, but we still allow pre and post exec hooks to be
+    /// assigned and run.
+    function tokensReceived(address, address, address, uint256, bytes calldata, bytes calldata)
+        external
+        override
+    {
+        (FunctionReference[][] memory postExecHooks, bytes[] memory postHookArgs) =
+            _doPreExecHooks(_getAccountStorage().selectorData[msg.sig], "");
+        _postNativeFunction(postExecHooks, postHookArgs);
+    }
 
-        return _getAccountStorage().supportedInterfaces[interfaceId] > 0;
+    /// @inheritdoc IERC721Receiver
+    /// @dev Runtime validation is bypassed for this function, but we still allow pre and post exec hooks to be
+    /// assigned and run.
+    function onERC721Received(address, address, uint256, bytes calldata)
+        external
+        override
+        returns (bytes4 selector)
+    {
+        (FunctionReference[][] memory postExecHooks, bytes[] memory postHookArgs) =
+            _doPreExecHooks(_getAccountStorage().selectorData[msg.sig], "");
+        selector = IERC721Receiver.onERC721Received.selector;
+        _postNativeFunction(postExecHooks, postHookArgs);
+    }
+
+    /// @inheritdoc IERC1155Receiver
+    /// @dev Runtime validation is bypassed for this function, but we still allow pre and post exec hooks to be
+    /// assigned and run.
+    function onERC1155Received(address, address, uint256, uint256, bytes calldata)
+        external
+        override
+        returns (bytes4 selector)
+    {
+        (FunctionReference[][] memory postExecHooks, bytes[] memory postHookArgs) =
+            _doPreExecHooks(_getAccountStorage().selectorData[msg.sig], "");
+        selector = IERC1155Receiver.onERC1155Received.selector;
+        _postNativeFunction(postExecHooks, postHookArgs);
+    }
+
+    /// @inheritdoc IERC1155Receiver
+    /// @dev Runtime validation is bypassed for this function, but we still allow pre and post exec hooks to be
+    /// assigned and run.
+    function onERC1155BatchReceived(address, address, uint256[] calldata, uint256[] calldata, bytes calldata)
+        external
+        override
+        returns (bytes4 selector)
+    {
+        (FunctionReference[][] memory postExecHooks, bytes[] memory postHookArgs) =
+            _doPreExecHooks(_getAccountStorage().selectorData[msg.sig], "");
+        selector = IERC1155Receiver.onERC1155BatchReceived.selector;
+        _postNativeFunction(postExecHooks, postHookArgs);
     }
 
     /// @inheritdoc UUPSUpgradeable
@@ -389,6 +435,16 @@ contract UpgradeableModularAccount is
         _postNativeFunction(postExecHooks, postHookArgs);
     }
 
+    /// @inheritdoc IERC165
+    function supportsInterface(bytes4 interfaceId) public view override returns (bool) {
+        if (interfaceId == _INVALID_INTERFACE_ID) {
+            return false;
+        }
+        return interfaceId == _IERC165_INTERFACE_ID || interfaceId == type(IERC721Receiver).interfaceId
+            || interfaceId == type(IERC777Recipient).interfaceId || interfaceId == type(IERC1155Receiver).interfaceId
+            || _getAccountStorage().supportedInterfaces[interfaceId] > 0;
+    }
+
     /// @inheritdoc IAccountView
     function entryPoint() public view override returns (IEntryPoint) {
         return _ENTRY_POINT;
@@ -417,7 +473,7 @@ contract UpgradeableModularAccount is
     }
 
     /// @dev Wraps execution of a native function with runtime validation and hooks. Used for upgradeToAndCall,
-    /// execute, executeBatch, installPlugin, uninstallPlugin.
+    /// execute, executeBatch, installPlugin, uninstallPlugin, and the token receiver functions.
     function _postNativeFunction(FunctionReference[][] memory postExecHooks, bytes[] memory postHookArgs)
         internal
     {