add L1 fee from user

.gitignore

@@ -4,3 +4,4 @@ build
 cache
 artifacts
 src/types
+.vscode

.solhint.json

@@ -7,7 +7,8 @@
         "printWidth": 110
       }
     ],
-    "quotes": ["error", "double"]
+    "quotes": ["error", "double"],
+    "compiler-version": ["error", "^0.7.0"]
   },
   "plugins": ["prettier"]
 }

contracts/Mocks/WETH.sol

@@ -0,0 +1,23 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.7.0;
+
+import "@openzeppelin/contracts/token/ERC20/ERC20.sol";
+import "@openzeppelin/contracts/access/Ownable.sol";
+
+contract WETH is ERC20 {
+  
+  constructor(
+    string memory name,
+    string memory ticker
+  ) ERC20(name, ticker) {}
+
+  function deposit() external payable {
+    _mint(msg.sender, msg.value);
+  }
+
+  function withdraw(uint256 value) external {
+    _burn(msg.sender, value);     
+    (bool success, ) = msg.sender.call{value: value}("");
+    require(success, "WETH: ETH transfer failed");
+  }
+}

contracts/TornadoPool.sol

@@ -46,6 +46,7 @@ contract TornadoPool is MerkleTreeWithHistory, IERC20Receiver, ReentrancyGuard,
     bytes encryptedOutput1;
     bytes encryptedOutput2;
     bool isL1Withdrawal;
+    uint256 l1Fee;
   }
 
   struct Proof {
@@ -275,7 +276,7 @@ contract TornadoPool is MerkleTreeWithHistory, IERC20Receiver, ReentrancyGuard,
     if (_extData.extAmount < 0) {
       require(_extData.recipient != address(0), "Can't withdraw to zero address");
       if (_extData.isL1Withdrawal) {
-        token.transferAndCall(omniBridge, uint256(-_extData.extAmount), abi.encodePacked(l1Unwrapper, _extData.recipient));
+        token.transferAndCall(omniBridge, uint256(-_extData.extAmount), abi.encodePacked(l1Unwrapper, _extData.recipient, _extData.l1Fee));
       } else {
         token.transfer(_extData.recipient, uint256(-_extData.extAmount));
       }

contracts/bridge/L1Unwrapper.sol

@@ -14,9 +14,13 @@ pragma solidity ^0.7.0;
 pragma abicoder v2;
 
 import "omnibridge/contracts/helpers/WETHOmnibridgeRouter.sol";
+import "@openzeppelin/contracts/math/SafeMath.sol";
+import { BytesHelper } from "../libraries/Bytes.sol";
 
 /// @dev Extension for original WETHOmnibridgeRouter that stores TornadoPool account registrations.
-contract L1Helper is WETHOmnibridgeRouter {
+contract L1Unwrapper is WETHOmnibridgeRouter {
+  using SafeMath for uint256;
+
   event PublicKey(address indexed owner, bytes key);
 
   struct Account {
@@ -61,4 +65,29 @@ contract L1Helper is WETHOmnibridgeRouter {
   function _register(Account memory _account) internal {
     emit PublicKey(_account.owner, _account.publicKey);
   }
+
+  /**
+  * @dev Bridged callback function used for unwrapping received tokens.
+  * Can only be called by the associated Omnibridge contract.
+  * @param _token bridged token contract address, should be WETH.
+  * @param _value amount of bridged/received tokens.
+  * @param _data extra data passed alongside with relayTokensAndCall on the other side of the bridge.
+  * Should contain coins receiver address and L1 executer fee amount.
+  */
+  function onTokenBridged(
+    address _token,
+    uint256 _value,
+    bytes memory _data
+  ) override external {
+    require(_token == address(WETH));
+    require(msg.sender == address(bridge));
+    require(_data.length == 52);
+    
+    WETH.withdraw(_value);
+
+    uint256 l1Fee = BytesHelper.sliceToUint(_data, 20);
+
+    AddressHelper.safeSendValue(payable(BytesHelper.bytesToAddress(_data)), _value.sub(l1Fee));
+    AddressHelper.safeSendValue(payable(tx.origin), l1Fee);
+  }
 }

contracts/libraries/Bytes.sol

@@ -0,0 +1,36 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.7.0;
+
+/**
+ * @title Bytes
+ * @dev Helper methods to transform bytes to other solidity types.
+ */
+library BytesHelper {
+    /**
+     * @dev Truncate bytes array if its size is more than 20 bytes.
+     * NOTE: This function does not perform any checks on the received parameter.
+     * Make sure that the _bytes argument has a correct length, not less than 20 bytes.
+     * A case when _bytes has length less than 20 will lead to the undefined behaviour,
+     * since assembly will read data from memory that is not related to the _bytes argument.
+     * @param _bytes to be converted to address type
+     * @return addr address included in the firsts 20 bytes of the bytes array in parameter.
+     */
+    function bytesToAddress(bytes memory _bytes) internal pure returns (address addr) {
+        assembly {
+            addr := mload(add(_bytes, 20))
+        }
+    }
+
+    /**
+     * @param _bytes it's 32 length slice to be converted to uint type
+     * @param _start start index of slice
+     * @return x uint included in the 32 length slice of the bytes array in parameter.
+     */
+    function sliceToUint(bytes memory _bytes, uint _start) internal pure returns (uint x)
+    {
+        require(_bytes.length >= _start + 32, "slicing out of range");
+        assembly {
+            x := mload(add(_bytes, add(0x20, _start)))
+        }
+    }
+}

package.json

@@ -44,7 +44,7 @@
     "fixed-merkle-tree": "^0.5.1",
     "hardhat": "^2.3.0",
     "mocha": "^9.1.0",
-    "omnibridge": "git+https://github.com/peppersec/omnibridge.git#aa3a970c29752a4da5f3fc7ccf0733783c1acf0b",
+    "omnibridge": "git+https://github.com/peppersec/omnibridge.git#1f0baaa34bbfdc8f2ddb37c0554ad7d964a96803",
     "prompt-sync": "^4.2.0",
     "snarkjs": "git+https://github.com/tornadocash/snarkjs.git#f37f146948f3b28086493e71512006b030588fc2",
     "tmp-promise": "^3.0.2",

src/index.js

@@ -16,7 +16,7 @@ async function buildMerkleTree({ tornadoPool }) {
   return new MerkleTree(MERKLE_TREE_HEIGHT, leaves, { hashFunction: poseidonHash2 })
 }
 
-async function getProof({ inputs, outputs, tree, extAmount, fee, recipient, relayer, isL1Withdrawal }) {
+async function getProof({ inputs, outputs, tree, extAmount, fee, recipient, relayer, isL1Withdrawal, l1Fee }) {
   inputs = shuffle(inputs)
   outputs = shuffle(outputs)
 
@@ -45,6 +45,7 @@ async function getProof({ inputs, outputs, tree, extAmount, fee, recipient, rela
     encryptedOutput1: outputs[0].encrypt(),
     encryptedOutput2: outputs[1].encrypt(),
     isL1Withdrawal,
+    l1Fee,
   }
 
   const extDataHash = getExtDataHash(extData)
@@ -94,6 +95,7 @@ async function prepareTransaction({
   recipient = 0,
   relayer = 0,
   isL1Withdrawal = false,
+  l1Fee = 0,
 }) {
   if (inputs.length > 16 || outputs.length > 2) {
     throw new Error('Incorrect inputs/outputs count')
@@ -118,6 +120,7 @@ async function prepareTransaction({
     recipient,
     relayer,
     isL1Withdrawal,
+    l1Fee,
   })
 
   return {

src/utils.js

@@ -22,12 +22,13 @@ function getExtDataHash({
   encryptedOutput1,
   encryptedOutput2,
   isL1Withdrawal,
+  l1Fee,
 }) {
   const abi = new ethers.utils.AbiCoder()
 
   const encodedData = abi.encode(
     [
-      'tuple(address recipient,int256 extAmount,address relayer,uint256 fee,bytes encryptedOutput1,bytes encryptedOutput2,bool isL1Withdrawal)',
+      'tuple(address recipient,int256 extAmount,address relayer,uint256 fee,bytes encryptedOutput1,bytes encryptedOutput2,bool isL1Withdrawal,uint256 l1Fee)',
     ],
     [
       {
@@ -38,6 +39,7 @@ function getExtDataHash({
         encryptedOutput1: encryptedOutput1,
         encryptedOutput2: encryptedOutput2,
         isL1Withdrawal: isL1Withdrawal,
+        l1Fee: l1Fee,
       },
     ],
   )

test/full.test.js

@@ -26,7 +26,7 @@ describe('TornadoPool', function () {
 
   async function fixture() {
     require('../scripts/compileHasher')
-    const [sender, gov, l1Unwrapper, multisig] = await ethers.getSigners()
+    const [sender, gov, multisig] = await ethers.getSigners()
     const verifier2 = await deploy('Verifier2')
     const verifier16 = await deploy('Verifier16')
     const hasher = await deploy('Hasher')
@@ -34,8 +34,12 @@ describe('TornadoPool', function () {
     const token = await deploy('PermittableToken', 'Wrapped ETH', 'WETH', 18, l1ChainId)
     await token.mint(sender.address, utils.parseEther('10000'))
 
+    const l1Token = await deploy('WETH', 'Wrapped ETH', 'WETH')
+    await l1Token.deposit({value: utils.parseEther('3')})
+
     const amb = await deploy('MockAMB', gov.address, l1ChainId)
     const omniBridge = await deploy('MockOmniBridge', amb.address)
+    const l1Unwrapper = await deploy('L1Unwrapper', amb.address, l1Token.address, gov.address)
 
     /** @type {TornadoPool} */
     const tornadoPoolImpl = await deploy(
@@ -69,7 +73,7 @@ describe('TornadoPool', function () {
 
     await token.approve(tornadoPool.address, utils.parseEther('10000'))
 
-    return { tornadoPool, token, proxy, omniBridge, amb, gov, multisig }
+    return { tornadoPool, token, proxy, omniBridge, amb, gov, multisig, l1Unwrapper, sender, l1Token }
   }
 
   describe('Upgradeability tests', () => {
@@ -271,6 +275,93 @@ describe('TornadoPool', function () {
     expect(omniBridgeBalance).to.be.equal(aliceWithdrawAmount)
   })
 
+  it('should withdraw with L1 fee', async function () {
+    const { tornadoPool, token, omniBridge, amb, l1Unwrapper, sender, l1Token } = await loadFixture(fixture)
+    const aliceKeypair = new Keypair() // contains private and public keys
+
+    // regular L1 deposit -------------------------------------------
+    const aliceDepositAmount = utils.parseEther('0.07')
+    const aliceDepositUtxo = new Utxo({ amount: aliceDepositAmount, keypair: aliceKeypair })
+    const { args, extData } = await prepareTransaction({
+      tornadoPool,
+      outputs: [aliceDepositUtxo],
+    })
+
+    let onTokenBridgedData = encodeDataForBridge({
+      proof: args,
+      extData,
+    })
+
+    let onTokenBridgedTx = await tornadoPool.populateTransaction.onTokenBridged(
+      token.address,
+      aliceDepositUtxo.amount,
+      onTokenBridgedData,
+    )
+    // emulating bridge. first it sends tokens to omnibridge mock then it sends to the pool
+    await token.transfer(omniBridge.address, aliceDepositAmount)
+    let transferTx = await token.populateTransaction.transfer(tornadoPool.address, aliceDepositAmount)
+
+    await omniBridge.execute([
+      { who: token.address, callData: transferTx.data }, // send tokens to pool
+      { who: tornadoPool.address, callData: onTokenBridgedTx.data }, // call onTokenBridgedTx
+    ])
+
+    // withdrawal with L1 fee ---------------------------------------
+    // withdraws a part of his funds from the shielded pool
+    const aliceWithdrawAmount = utils.parseEther('0.06')
+    const l1Fee = utils.parseEther('0.01')
+    // sum of desired withdraw amount and L1 fee are stored in extAmount
+    const extAmount = aliceWithdrawAmount.add(l1Fee)
+    const recipient = '0xDeaD00000000000000000000000000000000BEEf'
+    const aliceChangeUtxo = new Utxo({
+      amount: aliceDepositAmount.sub(extAmount),
+      keypair: aliceKeypair,
+    })
+    await transaction({
+      tornadoPool,
+      inputs: [aliceDepositUtxo],
+      outputs: [aliceChangeUtxo],
+      recipient: recipient,
+      isL1Withdrawal: true,
+      l1Fee: l1Fee,
+    })
+
+    const filter = omniBridge.filters.OnTokenTransfer()
+    const fromBlock = await ethers.provider.getBlock()
+    const events = await omniBridge.queryFilter(filter, fromBlock.number)
+    onTokenBridgedData = events[0].args.data
+    const hexL1Fee = '0x' + events[0].args.data.toString().slice(42)
+    expect(ethers.BigNumber.from(hexL1Fee)).to.be.equal(l1Fee)
+
+    const recipientBalance = await token.balanceOf(recipient)
+    expect(recipientBalance).to.be.equal(0)
+    const omniBridgeBalance = await token.balanceOf(omniBridge.address)
+    expect(omniBridgeBalance).to.be.equal(extAmount)
+
+    // L1 transactions:
+    onTokenBridgedTx = await l1Unwrapper.populateTransaction.onTokenBridged(
+      l1Token.address,
+      extAmount,
+      onTokenBridgedData,
+    )
+    // emulating bridge. first it sends tokens to amb mock then it sends to the recipient
+    await l1Token.transfer(amb.address, extAmount)
+    transferTx = await l1Token.populateTransaction.transfer(l1Unwrapper.address, extAmount)
+
+    const senderBalanceBefore = await ethers.provider.getBalance(sender.address)
+
+    let tx = await amb.execute([
+      { who: l1Token.address, callData: transferTx.data }, // send tokens to L1Unwrapper
+      { who: l1Unwrapper.address, callData: onTokenBridgedTx.data }, // call onTokenBridged on L1Unwrapper
+    ])
+
+    let receipt = await tx.wait()
+    let txFee = receipt.cumulativeGasUsed.mul(receipt.effectiveGasPrice)
+    const senderBalanceAfter = await ethers.provider.getBalance(sender.address)
+    expect(senderBalanceAfter).to.be.equal(senderBalanceBefore.sub(txFee).add(l1Fee))
+    expect(await ethers.provider.getBalance(recipient)).to.be.equal(aliceWithdrawAmount)
+  })
+
   it('should transfer funds to multisig in case of L1 deposit fail', async function () {
     const { tornadoPool, token, omniBridge, multisig } = await loadFixture(fixture)
     const aliceKeypair = new Keypair() // contains private and public keys

test/tree.test.js

@@ -77,7 +77,7 @@ describe('MerkleTreeWithHistory', function () {
 
       merkleTreeWithHistory.insert(toFixedHex(678), toFixedHex(876))
       tree.bulkInsert([678, 876])
-      expect(tree.root()).to.be.be.equal(await merkleTreeWithHistory.getLastRoot())
+      expect(tree.root()._hex).to.be.be.equal(await merkleTreeWithHistory.getLastRoot())
     })
 
     it('hasher gas', async () => {

test/utils.js

@@ -6,7 +6,7 @@ function encodeDataForBridge({ proof, extData }) {
   return abi.encode(
     [
       'tuple(bytes proof,bytes32 root,bytes32[] inputNullifiers,bytes32[2] outputCommitments,uint256 publicAmount,bytes32 extDataHash)',
-      'tuple(address recipient,int256 extAmount,address relayer,uint256 fee,bytes encryptedOutput1,bytes encryptedOutput2,bool isL1Withdrawal)',
+      'tuple(address recipient,int256 extAmount,address relayer,uint256 fee,bytes encryptedOutput1,bytes encryptedOutput2,bool isL1Withdrawal,uint256 l1Fee)',
     ],
     [proof, extData],
   )

yarn.lock

@@ -736,11 +736,6 @@
   resolved "https://registry.yarnpkg.com/@openzeppelin/contracts-upgradeable/-/contracts-upgradeable-3.4.2.tgz#2c2a1b0fa748235a1f495b6489349776365c51b3"
   integrity sha512-mDlBS17ymb2wpaLcrqRYdnBAmP1EwqhOXMvqWk2c5Q1N1pm5TkiCtXM9Xzznh4bYsQBq0aIWEkFFE2+iLSN1Tw==
 
-"@openzeppelin/contracts@3.2.2-solc-0.7":
-  version "3.2.2-solc-0.7"
-  resolved "https://registry.yarnpkg.com/@openzeppelin/contracts/-/contracts-3.2.2-solc-0.7.tgz#8ab169da64438d59f47ca285f1a10efe2f9ba19e"
-  integrity sha512-vFV53E4pvfsAEjzL9Um2VX9MEuXyq7Hyd9JjnP77AGsrEPxkJaYS06zZIVyhAt3rXTM6QGdW0C282Zv7fM93AA==
-
 "@openzeppelin@git+https://github.com/tornadocash/openzeppelin-contracts.git#6e46aa6946a7f215e7604169ddf46e1aebea850f":
   version "3.4.1-solc-0.7-2"
   resolved "git+https://github.com/tornadocash/openzeppelin-contracts.git#6e46aa6946a7f215e7604169ddf46e1aebea850f"
@@ -6918,11 +6913,10 @@ oboe@2.1.5:
   dependencies:
     http-https "^1.0.0"
 
-"omnibridge@git+https://github.com/peppersec/omnibridge.git#aa3a970c29752a4da5f3fc7ccf0733783c1acf0b":
+"omnibridge@git+https://github.com/peppersec/omnibridge.git#1f0baaa34bbfdc8f2ddb37c0554ad7d964a96803":
   version "1.1.0"
-  resolved "git+https://github.com/peppersec/omnibridge.git#aa3a970c29752a4da5f3fc7ccf0733783c1acf0b"
+  resolved "git+https://github.com/peppersec/omnibridge.git#1f0baaa34bbfdc8f2ddb37c0554ad7d964a96803"
   dependencies:
-    "@openzeppelin/contracts" "3.2.2-solc-0.7"
     axios "^0.21.0"
     bignumber.js "^9.0.1"
     dotenv "^8.2.0"