fix: enable hooks for token receiver callbacks

alchemyplatform · Jan 27, 2024 · 7b02860 · 7b02860
1 parent 1d35562
commit 7b02860
Show file tree

Hide file tree

Showing 3 changed files with 65 additions and 68 deletions.
diff --git a/src/account/TokenReceiver.sol b/src/account/TokenReceiver.sol
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";
@@ -38,7 +41,6 @@ import {AccountExecutor} from "./AccountExecutor.sol";
 import {AccountLoupe} from "./AccountLoupe.sol";
 import {AccountStorageInitializable} from "./AccountStorageInitializable.sol";
 import {PluginManagerInternals} from "./PluginManagerInternals.sol";
-import {TokenReceiver} from "./TokenReceiver.sol";
 
 /// @title Upgradeable Modular Account
 /// @author Alchemy
@@ -48,10 +50,13 @@ contract UpgradeableModularAccount is
     AccountLoupe,
     AccountStorageInitializable,
     PluginManagerInternals,
-    TokenReceiver,
     IAccount,
     IAccountInitializable,
     IAccountView,
+    IERC165,
+    IERC721Receiver,
+    IERC777Recipient,
+    IERC1155Receiver,
     IPluginExecutor,
     IStandardExecutor,
     UUPSUpgradeable
@@ -371,6 +376,61 @@ contract UpgradeableModularAccount is
         _postNativeFunction(postExecHooks, postHookArgs);
     }
 
+    /// @inheritdoc IERC777Recipient
+    /// @dev Runtime validation is bypassed for this callback, 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
+    // solhint-disable-next-line no-empty-blocks
+    {
+        (FunctionReference[][] memory postExecHooks, bytes[] memory postHookArgs) =
+            _doPreExecHooks(_getAccountStorage().selectorData[msg.sig], "");
+        _postNativeFunction(postExecHooks, postHookArgs);
+    }
+
+    /// @inheritdoc IERC721Receiver
+    /// @dev Runtime validation is bypassed for this callback, 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 callback, 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 callback, 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
     function upgradeToAndCall(address newImplementation, bytes calldata data) public payable override onlyProxy {
         (FunctionReference[][] memory postExecHooks, bytes[] memory postHookArgs) = _preNativeFunction();
@@ -383,7 +443,8 @@ contract UpgradeableModularAccount is
         if (interfaceId == _INTERFACE_ID_INVALID) {
             return false;
         }
-        return interfaceId == _IERC165_INTERFACE_ID || super.supportsInterface(interfaceId)
+        return interfaceId == _IERC165_INTERFACE_ID || interfaceId == type(IERC721Receiver).interfaceId
+            || interfaceId == type(IERC777Recipient).interfaceId || interfaceId == type(IERC1155Receiver).interfaceId
             || _getAccountStorage().supportedInterfaces[interfaceId] > 0;
     }
 

diff --git a/src/helpers/KnownSelectors.sol b/src/helpers/KnownSelectors.sol
@@ -64,7 +64,7 @@ library KnownSelectors {
             || selector == IAccountLoupe.getExecutionHooks.selector
             || selector == IAccountLoupe.getPreValidationHooks.selector
             || selector == IAccountLoupe.getInstalledPlugins.selector
-        // check against TokenReceiver methods
+        // check against token receiver methods
         || selector == IERC777Recipient.tokensReceived.selector
             || selector == IERC721Receiver.onERC721Received.selector
             || selector == IERC1155Receiver.onERC1155Received.selector