contracts-imtbl/contracts/game/NFTLock.sol

// SPDX-License-Identifier: MIT
pragma solidity 0.8.19;
import {IERC721} from "@openzeppelin/contracts/token/ERC721/IERC721.sol";
import {ERC721Holder} from "@openzeppelin/contracts/token/ERC721/utils/ERC721Holder.sol";
import {EnumerableSet} from "@openzeppelin/contracts/utils/structs/EnumerableSet.sol";
import {AccessControl} from "@openzeppelin/contracts/access/AccessControl.sol";
contract NFTLock is AccessControl, ERC721Holder {
    using EnumerableSet for EnumerableSet.UintSet;

    mapping(address => mapping(uint256 => address)) lockedOriginal;
    mapping(address => mapping(address => EnumerableSet.UintSet)) lockedRecords;
    mapping(address => bool) supportNftList;
    bytes32 public constant OPERATOR_ROLE = keccak256("OPERATOR_ROLE");
    bytes32 public constant UNLOCK_ROLE = keccak256("UNLOCK_ROLE");
    bytes32 public constant RELEASE_ROLE = keccak256("RELEASE_ROLE");

    event Lock(address indexed nft, address indexed user, uint256 indexed tokenId);
    event UnLock(address indexed nft, address indexed user, uint256 indexed tokenId);
    event BatchLock(address indexed nft, address indexed user, uint256[] tokenId);
    event Release(address indexed nft, uint256[] tokenIds, string serverId);

    constructor() {
        _setupRole(DEFAULT_ADMIN_ROLE, msg.sender);
        _setupRole(OPERATOR_ROLE, msg.sender);
        _setupRole(UNLOCK_ROLE, msg.sender);
        _setupRole(RELEASE_ROLE, msg.sender);
    }

    function lock(address nft, uint256 tokenId) external {
        require(supportNftList[nft], "can't lock this nft");
        IERC721(nft).transferFrom(msg.sender, address(this), tokenId);
        lockedOriginal[nft][tokenId] = msg.sender;
        lockedRecords[nft][msg.sender].add(tokenId);
        emit Lock(nft, msg.sender, tokenId);
    }

    function unlock(address nft, address to, uint256 tokenId) external onlyUnlocker {
        IERC721(nft).transferFrom(address(this), to, tokenId);
        lockedRecords[nft][lockedOriginal[nft][tokenId]].remove(tokenId);
        delete lockedOriginal[nft][tokenId];
        emit UnLock(nft, to, tokenId);
    }

    function batchLock(address nft, uint256[] calldata tokenIds) external {
        require(tokenIds.length <= 100, "tokenIds too many");
        for (uint256 i = 0; i < tokenIds.length; i++) {
            IERC721(nft).transferFrom(msg.sender, address(this), tokenIds[i]);
            lockedOriginal[nft][tokenIds[i]] = msg.sender;
            lockedRecords[nft][msg.sender].add(tokenIds[i]);
        }
        emit BatchLock(nft, msg.sender, tokenIds);
    }

    function release(address nft, uint256[] calldata tokenIds, string calldata serverId) external onlyReleaser {
        require(tokenIds.length <= 100, "tokenIds too many");
        for (uint256 i = 0; i < tokenIds.length; i++) {
            IERC721(nft).transferFrom(msg.sender, address(this), tokenIds[i]);
            lockedRecords[nft][msg.sender].add(tokenIds[i]);
        }
        emit Release(nft, tokenIds, serverId);
    }

    function originalOwner(address token, uint256 tokenId) public view returns (address) {
        return lockedOriginal[token][tokenId];
    }

    function lockedNum(address token, address user) public view returns (uint256) {
        return lockedRecords[token][user].length();
    }

    function lockedNft(address token, address user) public view returns (uint256[] memory) {
        return lockedRecords[token][user].values();
    }

    /** ------set------- **/
    function setOperatorRole(address to) external {
        grantRole(OPERATOR_ROLE, to);
    }

    function removeOperatorRole(address to) external {
        revokeRole(OPERATOR_ROLE, to);
    }

    function setReleaseRole(address to) external {
        grantRole(RELEASE_ROLE, to);
    }

    function removeReleaseRole(address to) external {
        revokeRole(RELEASE_ROLE, to);
    }

    function setUnlockRole(address to) external {
        grantRole(UNLOCK_ROLE, to);
    }

    function removeUnlockRole(address to) external {
        revokeRole(UNLOCK_ROLE, to);
    }

    function addSupportNftList(address[] calldata nftList) external onlyOperator {
        for (uint256 i = 0; i < nftList.length; i++) {
            supportNftList[nftList[i]] = true;
        }
    }
    function removeSupportNft(address nftAddress) external onlyOperator {
        require(supportNftList[nftAddress], "can't remove");
        delete supportNftList[nftAddress];
    }

    /** ------modifier------- **/
    modifier onlyOperator() {
        require(hasRole(OPERATOR_ROLE, msg.sender), "not operator role");
        _;
    }

    modifier onlyUnlocker() {
        require(hasRole(UNLOCK_ROLE, msg.sender), "not unlocker role");
        _;
    }

    modifier onlyReleaser() {
        require(hasRole(RELEASE_ROLE, msg.sender), "not releaser role");
        _;
    }
}