diff --git a/src/plugins/ColdStoragePlugin.sol b/src/plugins/ColdStoragePlugin.sol new file mode 100644 index 00000000..1090658d --- /dev/null +++ b/src/plugins/ColdStoragePlugin.sol @@ -0,0 +1,214 @@ +// SPDX-License-Identifier: GPL-3.0 +pragma solidity ^0.8.25; + +import {SignatureChecker} from "@openzeppelin/contracts/utils/cryptography/SignatureChecker.sol"; +import {IAccountExecute} from "@eth-infinitism/account-abstraction/interfaces/IAccountExecute.sol"; +import {IERC721} from "@openzeppelin/contracts/token/ERC721/IERC721.sol"; +import {EnumerableSet} from "@openzeppelin/contracts/utils/structs/EnumerableSet.sol"; + +import {IPlugin, PluginManifest, ManifestExecutionHook, PluginMetadata} from "../interfaces/IPlugin.sol"; +import {BasePlugin} from "./BasePlugin.sol"; +import {IStandardExecutor, Call} from "../interfaces/IStandardExecutor.sol"; + +/// @title Cold Storage Plugin +/// @author ERC-6900 Authors +/// @notice This plugin allows modular accounts to add additional restrictions on transferring certain NFTs +contract ColdStoragePlugin is BasePlugin { + using EnumerableSet for EnumerableSet.UintSet; + + struct ColdStorageStruct { + address guardian; + address isApprovedFor; + } + + struct ToStoreStruct { + address guardian; + address nft; + uint256 tokenId; + } + + string public constant NAME = "Cold Storage Plugin"; + string public constant VERSION = "1.0.0"; + string public constant AUTHOR = "ERC-6900 Authors"; + + bytes4 internal constant _SAFE_TRANSFER_FROM_SELECTOR = + bytes4(keccak256("safeTransferFrom(address,address,uint256)")); + bytes4 internal constant _SAFE_TRANSFER_FROM_WITH_DATA_SELECTOR = + bytes4(keccak256("safeTransferFrom(address,address,uint256,bytes)")); + + mapping(address account => mapping(address nft => EnumerableSet.UintSet tokenIds)) internal _coldStorage; + mapping(bytes32 id => ColdStorageStruct) public coldStorage; + mapping(address guardian => uint256 nonce) public nonces; + + error ActionNotAllowed(); + error AlreadyInColdStorage(); + error NotGuardian(); + error SelectorNotAllowed(); + + function addToColdStorage(address nft, uint256 tokenId, address guardian) external { + bytes32 id = encodeId(msg.sender, nft, tokenId); + ColdStorageStruct storage cs = coldStorage[id]; + if (cs.guardian != address(0)) { + revert AlreadyInColdStorage(); + } + _coldStorage[msg.sender][nft].add(tokenId); + cs.guardian = guardian; + } + + function unlock(address account, address nft, uint256 tokenId, address isApprovedFor, bytes calldata signature) + external + { + bytes32 id = encodeId(account, nft, tokenId); + ColdStorageStruct storage cs = coldStorage[id]; + address guardian = cs.guardian; + if ( + guardian != msg.sender + || SignatureChecker.isValidSignatureNow( + guardian, + encodeSignatureDigest(account, nft, tokenId, isApprovedFor, nonces[guardian]++), + signature + ) + ) { + revert NotGuardian(); + } + cs.isApprovedFor = isApprovedFor; + } + + function _checkColdStorage(address nft, bytes memory callData) internal view { + bytes4 selector = bytes4(callData); + if (selector == IERC721.setApprovalForAll.selector && _coldStorage[msg.sender][nft].length() > 0) { + revert ActionNotAllowed(); + } + if (selector == IERC721.approve.selector) { + address to; + uint256 tokenId; + assembly { + // callData = 32b length | 4b selector | 32b to | 32b tokenId + to := mload(add(callData, 36)) + tokenId := mload(add(callData, 68)) + } + bytes32 id = encodeId(msg.sender, nft, tokenId); + if (coldStorage[id].isApprovedFor != to) { + revert ActionNotAllowed(); + } + } else if ( + selector == IERC721.transferFrom.selector || selector == _SAFE_TRANSFER_FROM_SELECTOR + || selector == _SAFE_TRANSFER_FROM_WITH_DATA_SELECTOR + ) { + address to; + uint256 tokenId; + assembly { + // callData = 32b length | 4b selector | 32b from | 32b to | 32b tokenId + to := mload(add(callData, 68)) + tokenId := mload(add(callData, 100)) + } + bytes32 id = encodeId(msg.sender, nft, tokenId); + if (coldStorage[id].isApprovedFor != to) { + revert ActionNotAllowed(); + } + } + } + + function preExecutionHook(uint8, bytes calldata data) external view returns (bytes memory) { + bytes calldata topLevelCallData; + bytes4 topLevelSelector; + + topLevelSelector = bytes4(data[52:56]); + if (topLevelSelector == IAccountExecute.executeUserOp.selector) { + topLevelCallData = data[56:]; + topLevelSelector = bytes4(topLevelCallData); + } else { + topLevelCallData = data[52:]; + } + + if (topLevelSelector == IStandardExecutor.execute.selector) { + address token = address(uint160(uint256(bytes32(topLevelCallData[4:36])))); + + bytes calldata executeCalldata; + uint256 offset = uint256(bytes32(topLevelCallData[68:100])); + + assembly { + let relativeOffset := add(add(topLevelCallData.offset, offset), 4) + executeCalldata.offset := add(relativeOffset, 32) + executeCalldata.length := calldataload(relativeOffset) + } + + _checkColdStorage(token, executeCalldata); + } else if (topLevelSelector == IStandardExecutor.executeBatch.selector) { + Call[] memory calls = abi.decode(topLevelCallData[4:], (Call[])); + for (uint256 i = 0; i < calls.length; i++) { + _checkColdStorage(calls[i].target, calls[i].data); + } + } + return ""; + } + + function encodeId(address account, address nft, uint256 tokenId) public pure returns (bytes32) { + return keccak256(abi.encode(account, nft, tokenId)); + } + + function encodeSignatureDigest( + address account, + address nft, + uint256 tokenId, + address isApprovedFor, + uint256 nonce + ) public pure returns (bytes32) { + return keccak256(abi.encode(account, nft, tokenId, isApprovedFor, nonce)); + } + + /// @inheritdoc IPlugin + function onInstall(bytes calldata data) external override { + ToStoreStruct[] memory toStore = abi.decode(data, (ToStoreStruct[])); + for (uint256 i = 0; i < toStore.length; i++) { + _coldStorage[msg.sender][toStore[i].nft].add(toStore[i].tokenId); + coldStorage[encodeId(msg.sender, toStore[i].nft, toStore[i].tokenId)] = + ColdStorageStruct({guardian: toStore[i].guardian, isApprovedFor: address(0)}); + } + } + + /// @inheritdoc IPlugin + /// @dev Accounts are expected to provide all the necessary data to remove the restrictions + function onUninstall(bytes calldata data) external override { + ToStoreStruct[] memory toStore = abi.decode(data, (ToStoreStruct[])); + for (uint256 i = 0; i < toStore.length; i++) { + _coldStorage[msg.sender][toStore[i].nft].remove(toStore[i].tokenId); + delete coldStorage[encodeId(msg.sender, toStore[i].nft, toStore[i].tokenId)]; + } + } + + /// @inheritdoc IPlugin + function pluginManifest() external pure override returns (PluginManifest memory) { + PluginManifest memory manifest; + + // TODO: think it will be more efficient to attach hooks to `transfer.selector` etc + // in execute and executeBatch, instead of calling hooks associated with those selectors, call them based + // on selectors of msg.data[:4] + manifest.executionHooks = new ManifestExecutionHook[](2); + manifest.executionHooks[0] = ManifestExecutionHook({ + executionSelector: IStandardExecutor.execute.selector, + functionId: 0, + isPreHook: true, + isPostHook: false, + requireUOContext: false + }); + manifest.executionHooks[1] = ManifestExecutionHook({ + executionSelector: IStandardExecutor.executeBatch.selector, + functionId: 0, + isPreHook: true, + isPostHook: false, + requireUOContext: false + }); + + return manifest; + } + + /// @inheritdoc IPlugin + function pluginMetadata() external pure virtual override returns (PluginMetadata memory) { + PluginMetadata memory metadata; + metadata.name = NAME; + metadata.version = VERSION; + metadata.author = AUTHOR; + return metadata; + } +}