add fees on escrow

contracts/escrow/Escrow.sol

@@ -8,6 +8,7 @@
 import "@openzeppelin/contracts/token/ERC20/utils/SafeERC20.sol";
 import "@openzeppelin/contracts/utils/math/SafeMath.sol";
 import "@openzeppelin/contracts/security/ReentrancyGuard.sol";
+import "../interfaces/IFactoryRouter.sol";
 
 /**
  * @title Escrow contract
@@ -32,6 +33,9 @@
     using SafeMath for uint256;
     using SafeERC20 for IERC20;
 
+    // OPC fee router
+    address immutable public factoryRouter; 
+    address immutable public opcCollector;
 
     /*  User funds are stored per user and per token */
     struct userFunds{
@@ -40,7 +44,12 @@
     }
 
     mapping(address => mapping(address => userFunds)) private funds; // user -> token -> userFunds
-
+    // Mapping from user to an array of token addresses they have funds in
+    mapping(address => address[]) private userTokens;
+
+    // A helper mapping to avoid duplicates: user => token => hasTokenFunds
+    mapping(address => mapping(address => bool)) private hasFundsInToken;
+
     /*  Payee authorizations are stored per user and per token */
     struct auth{
         address payee;
@@ -74,6 +83,12 @@
     event Claimed(address indexed payee,uint256 jobId,address token,address indexed payer,uint256 amount,bytes proof);
     event Canceled(address indexed payee,uint256 jobId,address token,address indexed payer,uint256 amount);
 
+    // Add constructor to set router
+    constructor(address _factoryRouter,address _opcCollector) {
+        require(_factoryRouter != address(0), "Invalid router");
+        factoryRouter = _factoryRouter;
+        opcCollector = _opcCollector;
+    }
     /* Payer actions  */
 
     /**
@@ -102,6 +117,10 @@
     function _deposit(address token,uint256 amount) internal{
         require(token!=address(0),"Invalid token address");
         funds[msg.sender][token].available+=amount;
+        if (!hasFundsInToken[msg.sender][token]) {
+            userTokens[msg.sender].push(token);
+            hasFundsInToken[msg.sender][token] = true;
+        }
         emit Deposit(msg.sender,token,amount);
         uint256 balanceBefore = IERC20(token).balanceOf(address(this));
         IERC20(token).safeTransferFrom(msg.sender, address(this), amount);
@@ -112,6 +131,7 @@
 
     }
 
+
     /**
      * @dev withdraw
      *      Called by payer to withdraw available (not locked) funds from the contract
@@ -128,13 +148,30 @@
     function _withdraw(address token,uint256 amount) internal{
         if(funds[msg.sender][token].available>=amount){
             funds[msg.sender][token].available-=amount;
+            if(funds[msg.sender][token].available==0){
+                address[] storage tokens = userTokens[msg.sender];
+                for (uint256 i = 0; i < tokens.length; i++) {
+                    if (tokens[i] == token) {
+                        tokens[i] = tokens[tokens.length - 1]; // overwrite with last element
+                        tokens.pop(); // remove last element
+                        break;
+                    }
+                }
+                // Update interaction status
+                hasFundsInToken[msg.sender][token] = false;
+            }
             emit Withdraw(msg.sender,token,amount);
             IERC20(token).safeTransfer(
                 msg.sender,
                 amount
             );
+
         }
     }
+    function getUserTokens(address user) external view returns (address[] memory) {
+        return userTokens[user];
+    }
+
 
     /**
      * @dev authorize
@@ -475,11 +512,24 @@
                 userAuths[payer][token][i].currentLocks-=1;
             }
         }
+        // OPC fee logic
+        uint256 opcFee = IFactoryRouter(factoryRouter).getOPCFee(token);
+        uint256 feeAmount = amount.mul(opcFee).div(1e18);
+        uint256 payout = amount.sub(feeAmount);
+        // Transfer OPC fee to collector if any
+        if(feeAmount > 0){
+            if(opcCollector==address(0)){
+                IERC20(token).safeTransfer(IFactoryRouter(factoryRouter).getOPCCollector(), feeAmount);
+            }
+            else{
+                IERC20(token).safeTransfer(opcCollector, feeAmount);
+            }
+        }
         //update user funds
         funds[payer][token].available+=tempLock.amount-amount;
         funds[payer][token].locked-=tempLock.amount;
         //update payee balance
-        funds[msg.sender][token].available+=amount;
+        funds[msg.sender][token].available+=payout;
         //delete the lock
         if(index<locks.length-1){
             locks[index]=locks[locks.length-1];

scripts/deploy-contracts.js

@@ -11,6 +11,7 @@ const { UV_FS_O_FILEMAP } = require("constants");
 const ethers = hre.ethers;
 require("dotenv").config();
 const DEAD_ADDRESS = "0x000000000000000000000000000000000000dEaD"
+const ZERO_ADDRESS = '0x0000000000000000000000000000000000000000';
 let shouldDeployV4 = true;
 let shouldDeployDF = true;
 let shouldDeployVE = true;
@@ -67,6 +68,7 @@ async function main() {
   let sleepAmount = 10;
   let additionalApprovedTokens = []
   let pdrTrueValSubmiter = null
+  let router
   console.log("Using chain " + networkDetails.chainId);
   switch (networkDetails.chainId) {
     case 1:
@@ -442,7 +444,6 @@ async function main() {
 
     if (logging) console.log("Deploying Router");
     const Router = await ethers.getContractFactory("FactoryRouter", owner);
-    let router
     if (options) router = await Router.connect(owner).deploy(
       owner.address,
       addresses.Ocean,
@@ -933,7 +934,7 @@ async function main() {
       owner
     );
 
-    const deployEscrow = await Escrow.connect(owner).deploy(options)
+    const deployEscrow = await Escrow.connect(owner).deploy(router.address,ZERO_ADDRESS,options)
     await deployEscrow.deployTransaction.wait();
     if (show_verify) {
       console.log("\tRun the following to verify on etherscan");

test/unit/escrow/Escrow.test.js

@@ -2,7 +2,8 @@ const { assert,expect } = require('chai');
 const { ethers } = require("hardhat");
 const { json } = require('hardhat/internal/core/params/argumentTypes');
 const { web3 } = require("@openzeppelin/test-helpers/src/setup");
-const { getEventFromTx } = require("../../helpers/utils")
+const { getEventFromTx } = require("../../helpers/utils");
+
 
 const addressZero = '0x0000000000000000000000000000000000000000';
 
@@ -19,8 +20,9 @@ describe('Escrow tests', function () {
   let Mock20Contract;
   let Mock20DecimalsContract;
   let EscrowContract;
+  let FactoryRouter
   let signers;
-  let payee1,payee2,payee3,payer1,payer2,payer3;
+  let payee1,payee2,payee3,payer1,payer2,payer3,opcCollector
   before(async function () {
     // Get the contractOwner and collector address
     signers = await ethers.getSigners();
@@ -30,14 +32,25 @@ describe('Escrow tests', function () {
     payer1=signers[4]
     payer2=signers[5]
     payer3=signers[6]
+    opcCollector=signers[7]
+    const Router = await ethers.getContractFactory("FactoryRouter");
     const MockErc20 = await ethers.getContractFactory('MockERC20');
     const MockErc20Decimals = await ethers.getContractFactory('MockERC20Decimals');
     const Escrow = await ethers.getContractFactory('Escrow');
     Mock20Contract = await MockErc20.deploy(signers[0].address,"MockERC20", 'MockERC20');
     Mock20DecimalsContract = await MockErc20Decimals.deploy("Mock6Digits", 'Mock6Digits', 6);
-    EscrowContract = await Escrow.deploy();
     await Mock20Contract.deployed();
     await Mock20DecimalsContract.deployed();
+    // DEPLOY ROUTER, SETTING OWNER
+    FactoryRouter = await Router.deploy(
+      signers[0].address,
+      Mock20Contract.address,
+      '0x000000000000000000000000000000000000dead',
+      opcCollector.address,
+      []
+    );
+    await FactoryRouter.deployed();
+    EscrowContract = await Escrow.deploy(FactoryRouter.address,addressZero);
     await EscrowContract.deployed();
     // top up accounts
     await Mock20Contract.transfer(payer1.address,web3.utils.toWei("10000"))
@@ -57,11 +70,14 @@ describe('Escrow tests', function () {
 });
 
 it('Escrow - deposit', async function () {
+  let fundTokens=await EscrowContract.connect(payer1).getUserTokens(payer1.address)
+  expect(fundTokens).to.be.empty;
   expect(await Mock20Contract.balanceOf(EscrowContract.address)).to.equal(0);
   expect(await Mock20DecimalsContract.balanceOf(EscrowContract.address)).to.equal(0);
   await Mock20Contract.connect(payer1).approve(EscrowContract.address, web3.utils.toWei("10000"));
   await EscrowContract.connect(payer1).deposit(Mock20Contract.address,web3.utils.toWei("100"));
-
+  fundTokens=await EscrowContract.connect(payer1).getUserTokens(payer1.address)
+  expect(fundTokens).to.include(Mock20Contract.address);
   expect(await Mock20Contract.balanceOf(EscrowContract.address)).to.equal(web3.utils.toWei("100"));
   expect(await Mock20DecimalsContract.balanceOf(EscrowContract.address)).to.equal(0);
   const funds=await EscrowContract.connect(payer1).getFunds(Mock20Contract.address)
@@ -81,8 +97,10 @@ it('Escrow - withdraw', async function () {
     expect(await Mock20Contract.balanceOf(EscrowContract.address)).to.equal(balanceMock20);
     await EscrowContract.connect(payer1).withdraw([Mock20Contract.address],[web3.utils.toWei("10")]);
     expect(await Mock20Contract.balanceOf(EscrowContract.address)).to.equal(web3.utils.toWei("90"));
+    expect(await EscrowContract.connect(payer1).getUserTokens(payer1.address)).to.include(Mock20Contract.address);
 });
 
+
 it('Escrow - auth', async function () {
     await EscrowContract.connect(payer1).authorizeMultiple([Mock20Contract.address],[payee1.address],[web3.utils.toWei("50")],[100],[2]);
     const auths=await EscrowContract.connect(payer1).getAuthorizations(Mock20Contract.address,payer1.address,payee1.address)
@@ -138,6 +156,8 @@ it('Escrow - lock', async function () {
     const payer1Available=payer1Funds.available
     const payer1Locked=payer1Funds.locked
     const payee1Balance=await Mock20Contract.balanceOf(payee1.address)
+    const opcBalance=await Mock20Contract.balanceOf(opcCollector.address)
+
     // claim jobId
     let jobId=1 // full claim
     let lock
@@ -152,13 +172,19 @@ it('Escrow - lock', async function () {
     const txReceipt = await tx.wait();
     const event = getEventFromTx(txReceipt, 'Claimed')
     assert(event, "Cannot find Claimed event")
+    const opcBalanceAfter=await Mock20Contract.balanceOf(opcCollector.address)
+
+    const opcCollectorFee=await FactoryRouter.getOPCFee(Mock20Contract.address)
+    const expectedPayee=lock.amount.sub(lock.amount.mul(opcCollectorFee).div(web3.utils.toWei("1")));
+    expect(event.args.amount).to.equal(lock.amount)
+    expect(opcBalanceAfter).to.equal(opcBalance.add(lock.amount.sub(expectedPayee)))
     const afterpayer1Funds=await EscrowContract.connect(payer1).getFunds(Mock20Contract.address)
     const afterpayer1Available=afterpayer1Funds.available
     const afterpayer1Locked=afterpayer1Funds.locked
     const afterpayee1Balance=await Mock20Contract.balanceOf(payee1.address)
     expect(afterpayer1Available).to.equal(payer1Available)
     expect(afterpayer1Locked).to.equal(payer1Locked.sub(lock.amount))
-    expect(afterpayee1Balance).to.equal(payee1Balance.add(lock.amount))
+    expect(afterpayee1Balance).to.equal(payee1Balance.add(expectedPayee))
     // make sure lock is gone
     for( oneLock of await EscrowContract.connect(payee1).getLocks(Mock20Contract.address,payer1.address,payee1.address)){
       expect(oneLock.jobId).to.not.equal(jobId)
@@ -171,6 +197,8 @@ it('Escrow - lock', async function () {
     const payer1Available=payer1Funds.available
     const payer1Locked=payer1Funds.locked
     const payee1Balance=await Mock20Contract.balanceOf(payee1.address)
+    const opcBalance=await Mock20Contract.balanceOf(opcCollector.address)
+
     // claim jobId
     let jobId=2 // partial claim
     let lock
@@ -181,20 +209,26 @@ it('Escrow - lock', async function () {
       }
     }
     expect(lock.jobId).to.equal(jobId)
+
     const claimedAmount=web3.utils.toWei("1")
+    const bnClaimedAmount=ethers.BigNumber.from(claimedAmount)
     const returnAmount=lock.amount.sub(claimedAmount)
     const tx=await EscrowContract.connect(payee1).claimLocksAndWithdraw([lock.jobId],[lock.token],[lock.payer],[claimedAmount],[0]);
     const txReceipt = await tx.wait();
     const event = getEventFromTx(txReceipt, 'Claimed')
     assert(event, "Cannot find Claimed event")
-
+    const opcBalanceAfter=await Mock20Contract.balanceOf(opcCollector.address)
+
+    const opcCollectorFee=await FactoryRouter.getOPCFee(Mock20Contract.address)
+    const expectedPayee=bnClaimedAmount.sub(bnClaimedAmount.mul(opcCollectorFee).div(web3.utils.toWei("1")));
+    expect(opcBalanceAfter).to.equal(opcBalance.add(bnClaimedAmount.sub(expectedPayee)))
     const afterpayer1Funds=await EscrowContract.connect(payer1).getFunds(Mock20Contract.address)
     const afterpayer1Available=afterpayer1Funds.available
     const afterpayer1Locked=afterpayer1Funds.locked
     const afterpayee1Balance=await Mock20Contract.balanceOf(payee1.address)
     expect(afterpayer1Available).to.equal(payer1Available.add(returnAmount))
     expect(afterpayer1Locked).to.equal(payer1Locked.sub(lock.amount))
-    expect(afterpayee1Balance).to.equal(payee1Balance.add(claimedAmount))
+    expect(afterpayee1Balance).to.equal(payee1Balance.add(expectedPayee))
     // make sure lock is gone
     for( oneLock of await EscrowContract.connect(payee1).getLocks(Mock20Contract.address,payer1.address,payee1.address)){
       expect(oneLock.jobId).to.not.equal(jobId)
@@ -271,4 +305,33 @@ it('Escrow - lock', async function () {
       expect(oneLock.jobId).to.not.equal(jobId)
     }
   });
+  it('Escrow - deposit with decimals', async function () {
+    expect(await Mock20DecimalsContract.balanceOf(EscrowContract.address)).to.equal(0);
+    await Mock20DecimalsContract.connect(payer1).approve(EscrowContract.address, ethers.utils.parseUnits("10000", 6));
+    await EscrowContract.connect(payer1).deposit(Mock20DecimalsContract.address,ethers.utils.parseUnits("100", 6));
+
+    expect(await Mock20DecimalsContract.balanceOf(EscrowContract.address)).to.equal(ethers.utils.parseUnits("100", 6));
+    const funds=await EscrowContract.connect(payer1).getFunds(Mock20DecimalsContract.address)
+    expect(funds.available).to.equal(ethers.utils.parseUnits("100", 6))
+    expect(funds.locked).to.equal(0)
+    const locks=await EscrowContract.connect(payer1).getLocks(addressZero,addressZero,addressZero)
+    expect(locks.length).to.equal(0)
+    const auths=await EscrowContract.connect(payer1).getAuthorizations(Mock20DecimalsContract.address,payer1.address,addressZero)
+    expect(auths.length).to.equal(0)
+
+  });
+
+  it('Escrow - withdraw with decimals', async function () {
+      const balanceMock20=await Mock20DecimalsContract.balanceOf(EscrowContract.address);
+      await EscrowContract.connect(payer1).withdraw([Mock20DecimalsContract.address],[ethers.utils.parseUnits("10000", 6)]);
+      expect(await Mock20DecimalsContract.balanceOf(EscrowContract.address)).to.equal(balanceMock20);
+      await EscrowContract.connect(payer1).withdraw([Mock20DecimalsContract.address],[ethers.utils.parseUnits("10", 6)]);
+      expect(await Mock20DecimalsContract.balanceOf(EscrowContract.address)).to.equal(ethers.utils.parseUnits("90", 6));
+  });
+  it('Escrow - withdraw all funds', async function () {
+    expect(await EscrowContract.connect(payer1).getUserTokens(payer1.address)).to.include(Mock20Contract.address);
+    const payer1Funds=await EscrowContract.connect(payer1).getFunds(Mock20Contract.address)
+    await EscrowContract.connect(payer1).withdraw([Mock20Contract.address],[payer1Funds.available]);
+    expect(await EscrowContract.connect(payer1).getUserTokens(payer1.address)).does.not.include(Mock20Contract.address);
+  });
 });