Add Fee Collector

contracts/FeeProvider/FeeCollector.sol

@@ -0,0 +1,47 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.0;
+
+import "@openzeppelin/contracts/token/ERC721/IERC721Receiver.sol";
+import "@openzeppelin/contracts/token/ERC20/IERC20.sol";
+import "@openzeppelin/contracts/utils/introspection/ERC165Checker.sol";
+import "../interfaces/IFeeCollector.sol";
+import "../interfaces/ILockDealNFT.sol";
+import "../interfaces/IFeeProvider.sol";
+
+contract FeeCollector is IERC721Receiver, IFeeCollector {
+    ILockDealNFT public immutable lockDealNFT;
+    bool public feeCollected;
+
+    constructor(ILockDealNFT _lockDealNFT) {
+        lockDealNFT = _lockDealNFT;
+    }
+
+    function onERC721Received(
+        address,
+        address user,
+        uint256 poolId,
+        bytes calldata
+    ) external override returns (bytes4) {
+        require(msg.sender == address(lockDealNFT), "FeeCollector: invalid nft contract");
+        require(!feeCollected, "FeeCollector: fee already collected");
+        IProvider feeProvider = lockDealNFT.poolIdToProvider(poolId);
+        require(
+            ERC165Checker.supportsInterface(address(feeProvider), type(IFeeProvider).interfaceId),
+            "FeeCollector: wrong provider"
+        );
+        feeCollected = true;
+        uint256 amount = lockDealNFT.getWithdrawableAmount(poolId);
+        if (amount > 0) {
+            (address feeCollector, uint256 fee) = IERC2981(address(lockDealNFT)).royaltyInfo(poolId, amount);
+            lockDealNFT.safeTransferFrom(address(this), address(lockDealNFT), poolId);
+            IERC20 token = IERC20(lockDealNFT.tokenOf(poolId));
+            if (fee > 0) token.transfer(feeCollector, fee);
+            token.transfer(user, amount - fee);
+        }
+        if (lockDealNFT.ownerOf(poolId) == address(this)) {
+            lockDealNFT.transferFrom(address(this), user, poolId);
+        }
+        feeCollected = false;
+        return IERC721Receiver.onERC721Received.selector;
+    }
+}

contracts/FeeProvider/FeeDealProvider.sol

@@ -0,0 +1,29 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.0;
+
+import "../SimpleProviders/DealProvider/DealProvider.sol";
+import "./FeeProvider.sol";
+import "../interfaces/IFeeCollector.sol";
+
+contract FeeDealProvider is DealProvider, FeeProvider {
+    IFeeCollector public feeCollector;
+
+    constructor(IFeeCollector _feeCollector, ILockDealNFT _lockDealNFT) DealProvider(_lockDealNFT) {
+        feeCollector = _feeCollector;
+        name = "FeeDealProvider";
+    }
+
+    function _withdraw(
+        uint256 poolId,
+        uint256 amount
+    ) internal override firewallProtectedSig(0x9e2bf22c) returns (uint256 withdrawnAmount, bool isFinal) {
+        require(feeCollector.feeCollected(), "FeeDealProvider: fee not collected");
+        return super._withdraw(poolId, amount);
+    }
+
+    function supportsInterface(
+        bytes4 interfaceId
+    ) public view virtual override(FeeProvider, BasicProvider) returns (bool) {
+        return super.supportsInterface(interfaceId);
+    }
+}

contracts/FeeProvider/FeeLockProvider.sol

@@ -0,0 +1,18 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.0;
+
+import "../SimpleProviders/LockProvider/LockDealProvider.sol";
+import "./FeeProvider.sol";
+
+contract FeeLockProvider is LockDealProvider, FeeProvider {
+    constructor(ILockDealNFT _lockDealNFT, IProvider _provider) LockDealProvider(_lockDealNFT, address(_provider)) {
+        require(keccak256(bytes(_provider.name())) == keccak256(bytes("FeeDealProvider")), "invalid provider");
+        name = "FeeLockProvider";
+    }
+
+    function supportsInterface(
+        bytes4 interfaceId
+    ) public view virtual override(FeeProvider, BasicProvider) returns (bool) {
+        return super.supportsInterface(interfaceId);
+    }
+}

contracts/FeeProvider/FeeProvider.sol

@@ -0,0 +1,13 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.0;
+
+import "@openzeppelin/contracts/utils/introspection/ERC165.sol";
+import "../interfaces/IFeeProvider.sol";
+
+contract FeeProvider is IFeeProvider, ERC165 {
+    bool public constant feeProvider = true;
+
+    function supportsInterface(bytes4 interfaceId) public view virtual override returns (bool) {
+        return interfaceId == type(IFeeProvider).interfaceId || super.supportsInterface(interfaceId);
+    }
+}

contracts/FeeProvider/FeeTimedProvider.sol

@@ -0,0 +1,18 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.0;
+
+import "../SimpleProviders/TimedDealProvider/TimedDealProvider.sol";
+import "./FeeProvider.sol";
+
+contract FeeTimedProvider is TimedDealProvider, FeeProvider {
+    constructor(ILockDealNFT lockDealNFT, IProvider _provider) TimedDealProvider(lockDealNFT, address(_provider)) {
+        require(keccak256(bytes(_provider.name())) == keccak256(bytes("FeeLockProvider")), "invalid provider");
+        name = "FeeTimedProvider";
+    }
+
+    function supportsInterface(
+        bytes4 interfaceId
+    ) public view virtual override(FeeProvider, BasicProvider) returns (bool) {
+        return super.supportsInterface(interfaceId);
+    }
+}

contracts/SimpleProviders/DealProvider/DealProvider.sol

@@ -17,7 +17,7 @@ contract DealProvider is DealProviderState, BasicProvider {
     function _withdraw(
         uint256 poolId,
         uint256 amount
-    ) internal override firewallProtectedSig(0x9e2bf22c) returns (uint256 withdrawnAmount, bool isFinal) {
+    ) internal virtual override firewallProtectedSig(0x9e2bf22c) returns (uint256 withdrawnAmount, bool isFinal) {
         if (poolIdToAmount[poolId] >= amount) {
             poolIdToAmount[poolId] -= amount;
             withdrawnAmount = amount;

contracts/interfaces/IFeeCollector.sol

@@ -0,0 +1,7 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.0;
+
+interface IFeeCollector {
+    // Function to get the value of feeCollected
+    function feeCollected() external view returns (bool);
+}

contracts/interfaces/IFeeProvider.sol

@@ -0,0 +1,6 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.0;
+
+interface IFeeProvider {
+    function feeProvider() external returns (bool);
+}

contracts/mock/MockVaultManager.sol

@@ -1,47 +1,53 @@
 // SPDX-License-Identifier: MIT
 pragma solidity ^0.8.0;
 
+import "@openzeppelin/contracts/token/common/ERC2981.sol";
 import "@openzeppelin/contracts/token/ERC20/IERC20.sol";
 
-contract MockVaultManager {
+contract MockVaultManager is ERC2981 {
     mapping(address => uint) public tokenToVaultId;
     mapping(uint256 => address) vaultIdtoToken;
-    bool public transfers = true;
+    bool public transfers = false;
     uint256 public Id = 0;
 
     function setTransferStatus(bool status) external {
         transfers = status;
     }
 
-    function safeDeposit(address _tokenAddress, uint, address, bytes memory signature) external returns (uint vaultId) {
+    function safeDeposit(
+        address _tokenAddress,
+        uint amount,
+        address from,
+        bytes memory signature
+    ) external returns (uint vaultId) {
         require(keccak256(abi.encodePacked(signature)) == keccak256(abi.encodePacked("signature")), "wrong signature");
-        vaultId = _depositByToken(_tokenAddress);
+        vaultId = _depositByToken(_tokenAddress, from, amount);
     }
 
     function depositByToken(address _tokenAddress, uint256 amount) public returns (uint vaultId) {
-        IERC20(_tokenAddress).transferFrom(msg.sender, address(this), amount);
-        vaultId = _depositByToken(_tokenAddress);
+        vaultId = _depositByToken(_tokenAddress, msg.sender, amount);
     }
 
-    function _depositByToken(address _tokenAddress) internal returns (uint vaultId) {
+    function _depositByToken(address _tokenAddress, address from, uint256 amount) internal returns (uint vaultId) {
         vaultId = ++Id;
         vaultIdtoToken[vaultId] = _tokenAddress;
         tokenToVaultId[_tokenAddress] = vaultId;
+        if (transfers) IERC20(_tokenAddress).transferFrom(from, address(this), amount);
     }
 
     function withdrawByVaultId(uint _vaultId, address to, uint _amount) external {
-        // do nothing
+        if (_amount > 0 && transfers) IERC20(vaultIdtoToken[_vaultId]).transfer(to, _amount);
     }
 
     function vaultIdToTokenAddress(uint _vaultId) external view returns (address) {
         return vaultIdtoToken[_vaultId];
     }
 
-    function royaltyInfo(uint256, uint256) external pure returns (address receiver, uint256 royaltyAmount) {
-        return (address(0), 0);
-    }
-
     function vaultIdToTradeStartTime(uint256) external view returns (uint256) {
         return transfers ? block.timestamp - 1 : block.timestamp + 1;
     }
+
+    function setVaultRoyalty(uint _vaultId, address _royaltyReceiver, uint96 _feeNumerator) external {
+        _setTokenRoyalty(_vaultId, _royaltyReceiver, _feeNumerator);
+    }
 }

test/FeeCollector/FeeCollector.ts

@@ -0,0 +1,88 @@
+import { FeeTimedProvider, FeeDealProvider, FeeLockProvider, LockDealNFT, FeeCollector } from '../../typechain-types';
+import { MockVaultManager, DealProvider, ERC20Token } from '../../typechain-types';
+import { deployed } from '../helper';
+import { SignerWithAddress } from '@nomiclabs/hardhat-ethers/signers';
+import { time } from '@nomicfoundation/hardhat-network-helpers';
+import { expect } from 'chai';
+import { BigNumber, Bytes } from 'ethers';
+import { ethers } from 'hardhat';
+
+describe('Fee Collector', function () {
+  let lockDealNFT: LockDealNFT;
+  let feeDealProvider: FeeDealProvider;
+  let feeLockProvider: FeeLockProvider;
+  let feeTimeProvider: FeeTimedProvider;
+  let mockVaultManager: MockVaultManager;
+  let feeCollector: FeeCollector;
+  let poolId: number;
+  let collector: SignerWithAddress;
+  let owner: SignerWithAddress;
+  let addresses: string[];
+  let params: [BigNumber, number, number];
+  let token: ERC20Token;
+  const amount = ethers.utils.parseUnits('100', 18);
+  const fee = '1000'; // 10%
+  const ONE_DAY = 86400;
+  let startTime: number, finishTime: number;
+  const signature: Bytes = ethers.utils.toUtf8Bytes('signature');
+
+  before(async () => {
+    [owner, collector] = await ethers.getSigners();
+    mockVaultManager = await deployed('MockVaultManager');
+    lockDealNFT = await deployed('LockDealNFT', mockVaultManager.address, '');
+    feeCollector = await deployed('FeeCollector', lockDealNFT.address);
+    feeDealProvider = await deployed('FeeDealProvider', feeCollector.address, lockDealNFT.address);
+    feeLockProvider = await deployed('FeeLockProvider', lockDealNFT.address, feeDealProvider.address);
+    token = await deployed('ERC20Token', 'TestToken', 'TEST');
+    feeTimeProvider = await deployed('FeeTimedProvider', lockDealNFT.address, feeLockProvider.address);
+    await lockDealNFT.setApprovedContract(feeDealProvider.address, true);
+    await lockDealNFT.setApprovedContract(feeLockProvider.address, true);
+    await lockDealNFT.setApprovedContract(feeTimeProvider.address, true);
+    await lockDealNFT.setApprovedContract(feeCollector.address, true);
+    await mockVaultManager.setTransferStatus(true);
+    await token.approve(mockVaultManager.address, amount.mul(33));
+  });
+
+  beforeEach(async () => {
+    startTime = (await time.latest()) + ONE_DAY; // plus 1 day
+    finishTime = startTime + 7 * ONE_DAY; // plus 7 days from `startTime`
+    params = [amount, startTime, finishTime];
+    poolId = (await lockDealNFT.totalSupply()).toNumber();
+    addresses = [owner.address, token.address];
+    const vaultId = await mockVaultManager.Id();
+    await mockVaultManager.setVaultRoyalty(vaultId.add(1), collector.address, fee);
+  });
+
+  it('should revert withdraw wrong fee provider', async () => {
+    const nonFeeProvider: DealProvider = await deployed('DealProvider', lockDealNFT.address);
+    await lockDealNFT.setApprovedContract(nonFeeProvider.address, true);
+    await nonFeeProvider.createNewPool(addresses, params, signature);
+    await expect(
+      lockDealNFT['safeTransferFrom(address,address,uint256)'](owner.address, feeCollector.address, poolId),
+    ).to.be.revertedWith('FeeCollector: wrong provider');
+  });
+
+  it("should withdraw fee to fee collector's address from feeDealProvider", async () => {
+    await feeDealProvider.createNewPool(addresses, params, signature);
+    await lockDealNFT['safeTransferFrom(address,address,uint256)'](owner.address, feeCollector.address, poolId);
+    expect(await token.balanceOf(collector.address)).to.equal(amount.div(10));
+  });
+
+  it("should withdraw fee to fee collector's address from feeLockProvider", async () => {
+    await feeLockProvider.createNewPool(addresses, params, signature);
+    await time.setNextBlockTimestamp(startTime + 1);
+    const beforeBalance = await token.balanceOf(collector.address);
+    await lockDealNFT['safeTransferFrom(address,address,uint256)'](owner.address, feeCollector.address, poolId);
+    const afterBalance = await token.balanceOf(collector.address);
+    expect(afterBalance).to.equal(beforeBalance.add(amount.div(10)));
+  });
+
+  it("should withdraw fee to fee collector's address from feeTimeProvider", async () => {
+    await feeTimeProvider.createNewPool(addresses, params, signature);
+    await time.setNextBlockTimestamp(finishTime + 1);
+    const beforeBalance = await token.balanceOf(collector.address);
+    await lockDealNFT['safeTransferFrom(address,address,uint256)'](owner.address, feeCollector.address, poolId);
+    const afterBalance = await token.balanceOf(collector.address);
+    expect(afterBalance).to.equal(beforeBalance.add(amount.div(10)));
+  });
+});

test/FeeCollector/FeeDealProvider.ts

@@ -0,0 +1,73 @@
+import { FeeDealProvider, LockDealNFT, FeeCollector } from '../../typechain-types';
+import { MockVaultManager } from '../../typechain-types';
+import { deployed } from '../helper';
+import { ERC20Token } from '../../typechain-types';
+import { SignerWithAddress } from '@nomiclabs/hardhat-ethers/signers';
+import { expect } from 'chai';
+import { BigNumber, Bytes } from 'ethers';
+import { ethers } from 'hardhat';
+
+describe('Fee Deal Provider', function () {
+  let lockDealNFT: LockDealNFT;
+  let feeDealProvider: FeeDealProvider;
+  let mockVaultManager: MockVaultManager;
+  let feeCollector: FeeCollector;
+  let poolId: number;
+  let collector: SignerWithAddress;
+  let owner: SignerWithAddress;
+  let addresses: string[];
+  let params: [BigNumber];
+  let vaultId: BigNumber;
+  let token: ERC20Token;
+  const name: string = 'FeeDealProvider';
+  const amount = ethers.utils.parseUnits('100', 18);
+  const fee = '1000'; // 10%
+  const signature: Bytes = ethers.utils.toUtf8Bytes('signature');
+
+  before(async () => {
+    [owner, collector] = await ethers.getSigners();
+    mockVaultManager = await deployed('MockVaultManager');
+    lockDealNFT = await deployed('LockDealNFT', mockVaultManager.address, '');
+    feeCollector = await deployed('FeeCollector', lockDealNFT.address);
+    feeDealProvider = await deployed('FeeDealProvider', feeCollector.address, lockDealNFT.address);
+    token = await deployed('ERC20Token', 'TestToken', 'TEST');
+    await lockDealNFT.setApprovedContract(feeDealProvider.address, true);
+    await lockDealNFT.setApprovedContract(feeCollector.address, true);
+    await mockVaultManager.setTransferStatus(true);
+    await token.approve(mockVaultManager.address, amount.mul(33));
+  });
+
+  beforeEach(async () => {
+    poolId = (await lockDealNFT.totalSupply()).toNumber();
+    params = [amount];
+    addresses = [owner.address, token.address];
+    const vaultId = await mockVaultManager.Id();
+    await mockVaultManager.setVaultRoyalty(vaultId.add(1), collector.address, fee);
+  });
+
+  it("should return fee deal provider's name", async () => {
+    expect(await feeDealProvider.name()).to.equal(name);
+  });
+
+  it('should create a new fee deal', async () => {
+    await feeDealProvider.createNewPool(addresses, params, signature);
+    vaultId = await mockVaultManager.Id();
+    const data = await lockDealNFT.getData(poolId);
+    expect(data).to.deep.equal([feeDealProvider.address, name, poolId, vaultId, owner.address, token.address, params]);
+  });
+
+  it('should revert withdraw from LockDealNFT', async () => {
+    await feeDealProvider.createNewPool(addresses, params, signature);
+    await expect(
+      lockDealNFT['safeTransferFrom(address,address,uint256)'](owner.address, lockDealNFT.address, poolId),
+    ).to.be.revertedWith('FeeDealProvider: fee not collected');
+  });
+
+  it('should withdraw with fee calculation', async () => {
+    await feeDealProvider.createNewPool(addresses, params, signature);
+    const beforeBalance = await token.balanceOf(owner.address);
+    await lockDealNFT['safeTransferFrom(address,address,uint256)'](owner.address, feeCollector.address, poolId);
+    const afterBalance = await token.balanceOf(owner.address);
+    expect(afterBalance).to.equal(beforeBalance.add(amount.sub(amount.div(10))));
+  });
+});