feat: method register and transact

contracts/TornadoPool.sol

@@ -36,6 +36,23 @@ contract TornadoPool {
     bytes encryptedOutput2;
   }
 
+  struct Proof {
+    bytes proof;
+    bytes32 root;
+    bytes32 newRoot;
+    bytes32[] inputNullifiers;
+    bytes32[2] outputCommitments;
+    uint256 outPathIndices;
+    uint256 extAmount;
+    uint256 fee;
+    bytes32 extDataHash;
+  }
+
+  struct Register {
+    bytes pubKey;
+    bytes account;
+  }
+
   event NewCommitment(bytes32 commitment, uint256 index, bytes encryptedOutput);
   event NewNullifier(bytes32 nullifier);
   event PublicKey(address indexed owner, bytes key);
@@ -56,37 +73,23 @@ contract TornadoPool {
     currentRoot = _currentRoot;
   }
 
-  function transaction(
+  function transaction(Proof calldata _args, ExtData calldata _extData) public payable {
-    bytes calldata _proof,
+    require(currentRoot == _args.root, "Invalid merkle root");
-    bytes32 _root,
+    for (uint256 i = 0; i < _args.inputNullifiers.length; i++) {
-    bytes32 _newRoot,
+      require(!isSpent(_args.inputNullifiers[i]), "Input is already spent");
-    bytes32[] calldata _inputNullifiers,
-    bytes32[2] calldata _outputCommitments,
-    uint256 _outPathIndices,
-    uint256 _extAmount,
-    uint256 _fee,
-    ExtData calldata _extData,
-    bytes32 _extDataHash
-  ) external payable {
-    require(currentRoot == _root, "Invalid merkle root");
-    for (uint256 i = 0; i < _inputNullifiers.length; i++) {
-      require(!isSpent(_inputNullifiers[i]), "Input is already spent");
     }
-    require(uint256(_extDataHash) == uint256(keccak256(abi.encode(_extData))) % FIELD_SIZE, "Incorrect external data hash");
+    require(uint256(_args.extDataHash) == uint256(keccak256(abi.encode(_extData))) % FIELD_SIZE, "Incorrect external data hash");
-    require(_outPathIndices == currentCommitmentIndex >> 1, "Invalid merkle tree insert position");
+    require(_args.outPathIndices == currentCommitmentIndex >> 1, "Invalid merkle tree insert position");
-    require(
+    require(verifyProof(_args), "Invalid transaction proof");
-      verifyProof(_proof, _root, _newRoot, _inputNullifiers, _outputCommitments, _outPathIndices, _extAmount, _fee, _extDataHash),
-      "Invalid transaction proof"
-    );
 
-    currentRoot = _newRoot;
+    currentRoot = _args.newRoot;
-    for (uint256 i = 0; i < _inputNullifiers.length; i++) {
+    for (uint256 i = 0; i < _args.inputNullifiers.length; i++) {
-      nullifierHashes[_inputNullifiers[i]] = true;
+      nullifierHashes[_args.inputNullifiers[i]] = true;
     }
 
-    int256 extAmount = calculateExternalAmount(_extAmount);
+    int256 extAmount = calculateExternalAmount(_args.extAmount);
     if (extAmount > 0) {
-      require(msg.value == uint256(extAmount), "Incorrect amount of ETH sent on deposit");
+      require(msg.value == uint256(_args.extAmount), "Incorrect amount of ETH sent on deposit");
     } else if (extAmount < 0) {
       require(msg.value == 0, "Sent ETH amount should be 0 for withdrawal");
       require(_extData.recipient != address(0), "Can't withdraw to zero address");
@@ -95,14 +98,14 @@ contract TornadoPool {
       require(msg.value == 0, "Sent ETH amount should be 0 for transaction");
     }
 
-    if (_fee > 0) {
+    if (_args.fee > 0) {
-      _extData.relayer.transfer(_fee);
+      _extData.relayer.transfer(_args.fee);
     }
 
-    emit NewCommitment(_outputCommitments[0], currentCommitmentIndex++, _extData.encryptedOutput1);
+    emit NewCommitment(_args.outputCommitments[0], currentCommitmentIndex++, _extData.encryptedOutput1);
-    emit NewCommitment(_outputCommitments[1], currentCommitmentIndex++, _extData.encryptedOutput2);
+    emit NewCommitment(_args.outputCommitments[1], currentCommitmentIndex++, _extData.encryptedOutput2);
-    for (uint256 i = 0; i < _inputNullifiers.length; i++) {
+    for (uint256 i = 0; i < _args.inputNullifiers.length; i++) {
-      emit NewNullifier(_inputNullifiers[i]);
+      emit NewNullifier(_args.inputNullifiers[i]);
     }
   }
 
@@ -123,63 +126,53 @@ contract TornadoPool {
     return nullifierHashes[_nullifierHash];
   }
 
-  function verifyProof(
+  function verifyProof(Proof calldata _args) public view returns (bool) {
-    bytes memory _proof,
+    if (_args.inputNullifiers.length == 2) {
-    bytes32 _root,
-    bytes32 _newRoot,
-    bytes32[] memory _inputNullifiers,
-    bytes32[2] memory _outputCommitments,
-    uint256 _outPathIndices,
-    uint256 _extAmount,
-    uint256 _fee,
-    bytes32 _extDataHash
-  ) public view returns (bool) {
-    if (_inputNullifiers.length == 2) {
       return
         verifier2.verifyProof(
-          _proof,
+          _args.proof,
           [
-            uint256(_root),
+            uint256(_args.root),
-            uint256(_newRoot),
+            uint256(_args.newRoot),
-            _extAmount,
+            _args.extAmount,
-            _fee,
+            _args.fee,
-            uint256(_extDataHash),
+            uint256(_args.extDataHash),
-            uint256(_inputNullifiers[0]),
+            uint256(_args.inputNullifiers[0]),
-            uint256(_inputNullifiers[1]),
+            uint256(_args.inputNullifiers[1]),
-            uint256(_outputCommitments[0]),
+            uint256(_args.outputCommitments[0]),
-            uint256(_outputCommitments[1]),
+            uint256(_args.outputCommitments[1]),
-            _outPathIndices
+            _args.outPathIndices
           ]
         );
-    } else if (_inputNullifiers.length == 16) {
+    } else if (_args.inputNullifiers.length == 16) {
       return
         verifier16.verifyProof(
-          _proof,
+          _args.proof,
           [
-            uint256(_root),
+            uint256(_args.root),
-            uint256(_newRoot),
+            uint256(_args.newRoot),
-            _extAmount,
+            _args.extAmount,
-            _fee,
+            _args.fee,
-            uint256(_extDataHash),
+            uint256(_args.extDataHash),
-            uint256(_inputNullifiers[0]),
+            uint256(_args.inputNullifiers[0]),
-            uint256(_inputNullifiers[1]),
+            uint256(_args.inputNullifiers[1]),
-            uint256(_inputNullifiers[2]),
+            uint256(_args.inputNullifiers[2]),
-            uint256(_inputNullifiers[3]),
+            uint256(_args.inputNullifiers[3]),
-            uint256(_inputNullifiers[4]),
+            uint256(_args.inputNullifiers[4]),
-            uint256(_inputNullifiers[5]),
+            uint256(_args.inputNullifiers[5]),
-            uint256(_inputNullifiers[6]),
+            uint256(_args.inputNullifiers[6]),
-            uint256(_inputNullifiers[7]),
+            uint256(_args.inputNullifiers[7]),
-            uint256(_inputNullifiers[8]),
+            uint256(_args.inputNullifiers[8]),
-            uint256(_inputNullifiers[9]),
+            uint256(_args.inputNullifiers[9]),
-            uint256(_inputNullifiers[10]),
+            uint256(_args.inputNullifiers[10]),
-            uint256(_inputNullifiers[11]),
+            uint256(_args.inputNullifiers[11]),
-            uint256(_inputNullifiers[12]),
+            uint256(_args.inputNullifiers[12]),
-            uint256(_inputNullifiers[13]),
+            uint256(_args.inputNullifiers[13]),
-            uint256(_inputNullifiers[14]),
+            uint256(_args.inputNullifiers[14]),
-            uint256(_inputNullifiers[15]),
+            uint256(_args.inputNullifiers[15]),
-            uint256(_outputCommitments[0]),
+            uint256(_args.outputCommitments[0]),
-            uint256(_outputCommitments[1]),
+            uint256(_args.outputCommitments[1]),
-            _outPathIndices
+            _args.outPathIndices
           ]
         );
     } else {
@@ -187,8 +180,17 @@ contract TornadoPool {
     }
   }
 
-  function register(bytes calldata _pubKey, bytes calldata _account) external {
+  function register(Register calldata args) public {
-    emit PublicKey(msg.sender, _pubKey);
+    emit PublicKey(msg.sender, args.pubKey);
-    emit EncryptedAccount(msg.sender, _account);
+    emit EncryptedAccount(msg.sender, args.account);
+  }
+
+  function registerAndTransact(
+    Register calldata _registerArgs,
+    Proof calldata _proofArgs,
+    ExtData calldata _extData
+  ) external payable {
+    register(_registerArgs);
+    transaction(_proofArgs, _extData);
   }
 }

src/index.js

@@ -80,26 +80,33 @@ async function getProof({ inputs, outputs, tree, extAmount, fee, recipient, rela
 
   const proof = await prove(input, `./artifacts/circuits/transaction${inputs.length}`)
 
-  const args = [
+  const args = {
-    toFixedHex(input.root),
+    proof,
-    toFixedHex(input.newRoot),
+    root: toFixedHex(input.root),
-    inputs.map((x) => toFixedHex(x.getNullifier())),
+    newRoot: toFixedHex(input.newRoot),
-    outputs.map((x) => toFixedHex(x.getCommitment())),
+    inputNullifiers: inputs.map((x) => toFixedHex(x.getNullifier())),
-    toFixedHex(outputIndex >> outputBatchBits),
+    outputCommitments: outputs.map((x) => toFixedHex(x.getCommitment())),
-    toFixedHex(extAmount),
+    outPathIndices: toFixedHex(outputIndex >> outputBatchBits),
-    toFixedHex(fee),
+    extAmount: toFixedHex(extAmount),
-    extData,
+    fee: toFixedHex(fee),
-    toFixedHex(extDataHash),
+    extDataHash: toFixedHex(extDataHash),
-  ]
+  }
   // console.log('Solidity args', args)
 
   return {
-    proof,
+    extData,
     args,
   }
 }
 
-async function transaction({ tornadoPool, inputs = [], outputs = [], fee = 0, recipient = 0, relayer = 0 }) {
+async function prepareTransaction({
+  tornadoPool,
+  inputs = [],
+  outputs = [],
+  fee = 0,
+  recipient = 0,
+  relayer = 0,
+}) {
   if (inputs.length > 16 || outputs.length > 2) {
     throw new Error('Incorrect inputs/outputs count')
   }
@@ -119,7 +126,7 @@ async function transaction({ tornadoPool, inputs = [], outputs = [], fee = 0, re
     extAmount = FIELD_SIZE.add(extAmount)
   }
 
-  const { proof, args } = await getProof({
+  const { args, extData } = await getProof({
     inputs,
     outputs,
     tree: await buildMerkleTree({ tornadoPool }),
@@ -129,11 +136,42 @@ async function transaction({ tornadoPool, inputs = [], outputs = [], fee = 0, re
     relayer,
   })
 
-  const receipt = await tornadoPool.transaction(proof, ...args, {
+  return {
+    args,
+    extData,
+    amount,
+  }
+}
+
+async function transaction({ tornadoPool, ...rest }) {
+  const { args, extData, amount } = await prepareTransaction({
+    tornadoPool,
+    ...rest,
+  })
+
+  const receipt = await tornadoPool.transaction(args, extData, {
     value: amount,
     gasLimit: 1e6,
   })
   await receipt.wait()
 }
 
-module.exports = { transaction }
+async function registerAndTransact({ tornadoPool, packedPrivateKeyData, poolAddress, ...rest }) {
+  const { args, extData, amount } = await prepareTransaction({
+    tornadoPool,
+    ...rest,
+  })
+
+  const params = {
+    pubKey: poolAddress,
+    account: packedPrivateKeyData,
+  }
+
+  const receipt = await tornadoPool.registerAndTransact(params, args, extData, {
+    value: amount,
+    gasLimit: 2e6,
+  })
+  await receipt.wait()
+}
+
+module.exports = { transaction, registerAndTransact }

src/keypair.js

@@ -101,4 +101,8 @@ class Keypair {
   }
 }
 
-module.exports = Keypair
+module.exports = {
+  Keypair,
+  packEncryptedMessage,
+  unpackEncryptedMessage,
+}

src/utxo.js

@@ -1,7 +1,7 @@
 const { ethers } = require('hardhat')
 const { BigNumber } = ethers
 const { randomBN, poseidonHash, toBuffer } = require('./utils')
-const Keypair = require('./keypair')
+const { Keypair } = require('./keypair')
 
 class Utxo {
   /** Initialize a new UTXO - unspent transaction output or input. Note, a full TX consists of 2/16 inputs and 2 outputs

test/full.test.js

@@ -8,14 +8,16 @@ const Utxo = require('../src/utxo')
 const MERKLE_TREE_HEIGHT = 5
 const MerkleTree = require('fixed-merkle-tree')
 
-const { transaction } = require('../src/index')
+const { transaction, registerAndTransact } = require('../src/index')
-const Keypair = require('../src/keypair')
+const { Keypair } = require('../src/keypair')
 
 describe('TornadoPool', () => {
-  let snapshotId, tornadoPool
+  let snapshotId, tornadoPool, sender
 
   /* prettier-ignore */
   before(async function () {
+    ;[sender] = await ethers.getSigners()
+
     const Verifier2 = await ethers.getContractFactory('Verifier2')
     const verifier2 = await Verifier2.deploy()
     await verifier2.deployed()
@@ -42,6 +44,64 @@ describe('TornadoPool', () => {
     expect(result).to.be.deep.equal(data)
   })
 
+  it('should register and deposit', async function () {
+    // Alice deposits into tornado pool
+    const aliceDepositAmount = 1e7
+    const aliceDepositUtxo = new Utxo({ amount: aliceDepositAmount })
+
+    const backupAccount = new Keypair()
+
+    const bufferPrivateKey = Buffer.from(aliceDepositUtxo.keypair.privkey)
+    const packedPrivateKeyData = backupAccount.encrypt(bufferPrivateKey)
+
+    tornadoPool = tornadoPool.connect(sender)
+    await registerAndTransact({
+      tornadoPool,
+      packedPrivateKeyData,
+      outputs: [aliceDepositUtxo],
+      poolAddress: aliceDepositUtxo.keypair.address(),
+    })
+
+    const filter = tornadoPool.filters.NewCommitment()
+    const fromBlock = await ethers.provider.getBlock()
+    const events = await tornadoPool.queryFilter(filter, fromBlock.number)
+
+    let aliceReceiveUtxo
+    try {
+      aliceReceiveUtxo = Utxo.decrypt(
+        aliceDepositUtxo.keypair,
+        events[0].args.encryptedOutput,
+        events[0].args.index,
+      )
+    } catch (e) {
+      // we try to decrypt another output here because it shuffles outputs before sending to blockchain
+      aliceReceiveUtxo = Utxo.decrypt(
+        aliceDepositUtxo.keypair,
+        events[1].args.encryptedOutput,
+        events[1].args.index,
+      )
+    }
+    expect(aliceReceiveUtxo.amount).to.be.equal(aliceDepositAmount)
+
+    const filterRegister = tornadoPool.filters.PublicKey(sender.address)
+    const filterFromBlock = await ethers.provider.getBlock()
+    const registerEvents = await tornadoPool.queryFilter(filterRegister, filterFromBlock.number)
+
+    const [registerEvent] = registerEvents.sort((a, b) => a.blockNumber - b.blockNumber).slice(-1)
+
+    expect(registerEvent.args.key).to.be.equal(aliceDepositUtxo.keypair.address())
+
+    const accountFilter = tornadoPool.filters.EncryptedAccount(sender.address)
+    const accountFromBlock = await ethers.provider.getBlock()
+    const accountEvents = await tornadoPool.queryFilter(accountFilter, accountFromBlock.number)
+
+    const [accountEvent] = accountEvents.sort((a, b) => a.blockNumber - b.blockNumber).slice(-1)
+
+    const privateKey = backupAccount.decrypt(accountEvent.args.account)
+
+    expect(bufferPrivateKey.toString('hex')).to.be.equal(privateKey.toString('hex'))
+  })
+
   it('should deposit, transact and withdraw', async function () {
     // Alice deposits into tornado pool
     const aliceDepositAmount = 1e7

yarn.lock

@@ -4230,10 +4230,10 @@ find-yarn-workspace-root@^2.0.0:
   dependencies:
     micromatch "^4.0.2"
 
-fixed-merkle-tree@^0.5.0:
+fixed-merkle-tree@^0.5.1:
-  version "0.5.0"
+  version "0.5.1"
-  resolved "https://registry.yarnpkg.com/fixed-merkle-tree/-/fixed-merkle-tree-0.5.0.tgz#401cdcf3d670c1e18bc7d3a8e81322eb1b27c1d1"
+  resolved "https://registry.npmjs.org/fixed-merkle-tree/-/fixed-merkle-tree-0.5.1.tgz#770bbf64b174e88b1133841721b79ea99a63d0b5"
-  integrity sha512-egOy12EzVATX3Ru2/SLtnWprVpy/sbPCt/MbeG3ANB28jykWLEYj7EjinFnOxtsgR3gTHU6xYXX53yMn/bZqyw==
+  integrity sha512-9daba9QU5tWa1scXpWusMFVJgYbSIPJUXdgn0JaaXSK1m4xSjzfTjdT3bYuo/x1xeLWrYj5i21E8q+vxyP7TCw==
   dependencies:
     circomlib "git+https://github.com/tornadocash/circomlib.git#5beb6aee94923052faeecea40135d45b6ce6172c"
     snarkjs "git+https://github.com/tornadocash/snarkjs.git#869181cfaf7526fe8972073d31655493a04326d5"
@@ -7847,20 +7847,9 @@ snapdragon@^0.8.1:
     source-map-resolve "^0.5.0"
     use "^3.1.0"
 
-"snarkjs@git+https://github.com/tornadocash/snarkjs.git#869181cfaf7526fe8972073d31655493a04326d5":
+"snarkjs@git+https://github.com/tornadocash/snarkjs.git#616c2d30699f28c8f3ab737b877402ccbb604cfe":
-  version "0.1.20"
-  resolved "git+https://github.com/tornadocash/snarkjs.git#869181cfaf7526fe8972073d31655493a04326d5"
-  dependencies:
-    big-integer "^1.6.43"
-    chai "^4.2.0"
-    escape-string-regexp "^1.0.5"
-    eslint "^5.16.0"
-    keccak "^2.0.0"
-    yargs "^12.0.5"
-
-"snarkjs@git+https://github.com/tornadocash/snarkjs.git#c103e3bf10e95e2e9bbf0f7952ed13812f8e39d3":
   version "0.4.5"
-  resolved "git+https://github.com/tornadocash/snarkjs.git#c103e3bf10e95e2e9bbf0f7952ed13812f8e39d3"
+  resolved "git+https://github.com/tornadocash/snarkjs.git#616c2d30699f28c8f3ab737b877402ccbb604cfe"
   dependencies:
     "@iden3/binfileutils" "0.0.8"
     blake2b-wasm "https://github.com/jbaylina/blake2b-wasm.git"
@@ -7873,6 +7862,17 @@ snapdragon@^0.8.1:
     r1csfile "0.0.32"
     readline "^1.3.0"
 
+"snarkjs@git+https://github.com/tornadocash/snarkjs.git#869181cfaf7526fe8972073d31655493a04326d5":
+  version "0.1.20"
+  resolved "git+https://github.com/tornadocash/snarkjs.git#869181cfaf7526fe8972073d31655493a04326d5"
+  dependencies:
+    big-integer "^1.6.43"
+    chai "^4.2.0"
+    escape-string-regexp "^1.0.5"
+    eslint "^5.16.0"
+    keccak "^2.0.0"
+    yargs "^12.0.5"
+
 solc@0.7.3:
   version "0.7.3"
   resolved "https://registry.yarnpkg.com/solc/-/solc-0.7.3.tgz#04646961bd867a744f63d2b4e3c0701ffdc7d78a"