refactor transaction and encryption functions

src/index.js

@@ -1,6 +1,7 @@
 /* eslint-disable no-console */
 const MerkleTree = require('fixed-merkle-tree')
 const { ethers } = require('hardhat')
+const { BigNumber } = ethers
 const { toFixedHex, poseidonHash2, getExtDataHash, FIELD_SIZE, packEncryptedMessage } = require('./utils')
 const Utxo = require('./utxo')
 
@@ -52,12 +53,8 @@ async function getProof({ inputs, outputs, tree, extAmount, fee, recipient, rela
   const extData = {
     recipient: toFixedHex(recipient, 20),
     relayer: toFixedHex(relayer, 20),
-    encryptedOutput1: packEncryptedMessage(
+    encryptedOutput1: outputs[0].encrypt(),
-      outputs[0].keypair.encrypt({ blinding: outputs[0].blinding, amount: outputs[0].amount }),
+    encryptedOutput2: outputs[1].encrypt(),
-    ),
-    encryptedOutput2: packEncryptedMessage(
-      outputs[1].keypair.encrypt({ blinding: outputs[1].blinding, amount: outputs[1].amount }),
-    ),
   }
 
   const extDataHash = getExtDataHash(extData)
@@ -108,89 +105,41 @@ async function getProof({ inputs, outputs, tree, extAmount, fee, recipient, rela
   }
 }
 
-async function deposit({ tornadoPool, utxo }) {
+async function transaction({ tornadoPool, inputs = [], outputs = [], fee = 0, recipient = 0, relayer = 0 }) {
-  const inputs = [new Utxo(), new Utxo()]
+  if (inputs.length > 16 || outputs.length > 2) {
-  const outputs = [utxo, new Utxo()]
+    throw new Error('Incorrect inputs/outputs count')
+  }
+  while(inputs.length !== 2 && inputs.length < 16) {
+    inputs.push(new Utxo())
+  }
+  while(outputs.length < 2) {
+    outputs.push(new Utxo())
+  }
+
+  let extAmount = BigNumber.from(fee)
+    .add(outputs.reduce((sum, x) => sum.add(x.amount), BigNumber.from(0)))
+    .sub(inputs.reduce((sum, x) => sum.add(x.amount), BigNumber.from(0)))
+  const amount = extAmount > 0 ? extAmount : 0
+  if (extAmount < 0) {
+    extAmount = FIELD_SIZE.add(extAmount)
+  }
 
   const { proof, args } = await getProof({
     inputs,
     outputs,
     tree: await buildMerkleTree({ tornadoPool }),
-    extAmount: utxo.amount,
+    extAmount,
-    fee: 0,
+    fee,
-    recipient: 0,
+    recipient,
-    relayer: 0,
+    relayer,
-  })
-
-  console.log('Sending deposit transaction...')
-  const receipt = await tornadoPool.transaction(proof, ...args, {
-    value: utxo.amount,
-    gasLimit: 1e6,
-  })
-  console.log(`Receipt ${receipt.hash}`)
-  return outputs[0]
-}
-
-async function merge({ tornadoPool }) {
-  const amount = 1e6
-  const inputs = new Array(16).fill(0).map((_) => new Utxo())
-  const outputs = [new Utxo({ amount }), new Utxo()]
-
-  const { proof, args } = await getProof({
-    inputs,
-    outputs,
-    tree: await buildMerkleTree({ tornadoPool }),
-    extAmount: amount,
-    fee: 0,
-    recipient: 0,
-    relayer: 0,
   })
 
+  console.log('Sending transaction...')
   const receipt = await tornadoPool.transaction(proof, ...args, {
     value: amount,
     gasLimit: 1e6,
   })
   console.log(`Receipt ${receipt.hash}`)
-  return outputs[0]
 }
 
-async function transact({ tornadoPool, input, output }) {
+module.exports = { transaction }
-  const inputs = [input, new Utxo()]
-  const outputs = [output, new Utxo()]
-
-  const { proof, args } = await getProof({
-    inputs,
-    outputs,
-    tree: await buildMerkleTree({ tornadoPool }),
-    extAmount: 0,
-    fee: 0,
-    recipient: 0,
-    relayer: 0,
-  })
-
-  console.log('Sending transfer transaction...')
-  const receipt = await tornadoPool.transaction(proof, ...args, { gasLimit: 1e6 })
-  console.log(`Receipt ${receipt.hash}`)
-  return outputs[0]
-}
-
-async function withdraw({ tornadoPool, input, change, recipient }) {
-  const inputs = [input, new Utxo()]
-  const outputs = [change, new Utxo()]
-
-  const { proof, args } = await getProof({
-    inputs,
-    outputs,
-    tree: await buildMerkleTree({ tornadoPool }),
-    extAmount: FIELD_SIZE.sub(input.amount.sub(change.amount)),
-    fee: 0,
-    recipient,
-    relayer: 0,
-  })
-
-  console.log('Sending withdraw transaction...')
-  const receipt = await tornadoPool.transaction(proof, ...args, { gasLimit: 1e6 })
-  console.log(`Receipt ${receipt.hash}`)
-}
-
-module.exports = { deposit, withdraw, transact, merge }

src/keypair.js

@@ -1,8 +1,37 @@
 const { encrypt, decrypt, getEncryptionPublicKey } = require('eth-sig-util')
 const { ethers } = require('hardhat')
 const { BigNumber } = ethers
-const { randomBN, poseidonHash, toFixedHex } = require('./utils')
+const { poseidonHash, toFixedHex } = require('./utils')
-const BNjs = require('bn.js')
+
+function packEncryptedMessage(encryptedMessage) {
+  const nonceBuf = Buffer.from(encryptedMessage.nonce, 'base64')
+  const ephemPublicKeyBuf = Buffer.from(encryptedMessage.ephemPublicKey, 'base64')
+  const ciphertextBuf = Buffer.from(encryptedMessage.ciphertext, 'base64')
+  const messageBuff = Buffer.concat([
+    Buffer.alloc(24 - nonceBuf.length),
+    nonceBuf,
+    Buffer.alloc(32 - ephemPublicKeyBuf.length),
+    ephemPublicKeyBuf,
+    ciphertextBuf,
+  ])
+  return '0x' + messageBuff.toString('hex')
+}
+
+function unpackEncryptedMessage(encryptedMessage) {
+  if (encryptedMessage.slice(0, 2) === '0x') {
+    encryptedMessage = encryptedMessage.slice(2)
+  }
+  const messageBuff = Buffer.from(encryptedMessage, 'hex')
+  const nonceBuf = messageBuff.slice(0, 24)
+  const ephemPublicKeyBuf = messageBuff.slice(24, 56)
+  const ciphertextBuf = messageBuff.slice(56)
+  return {
+    version: 'x25519-xsalsa20-poly1305',
+    nonce: nonceBuf.toString('base64'),
+    ephemPublicKey: ephemPublicKeyBuf.toString('base64'),
+    ciphertext: ciphertextBuf.toString('base64'),
+  }
+}
 
 class Keypair {
   constructor(privkey = ethers.Wallet.createRandom().privateKey) {
@@ -33,21 +62,12 @@ class Keypair {
     })
   }
 
-  encrypt({ blinding, amount }) {
+  encrypt(bytes) {
-    const bytes = Buffer.concat([
+    return packEncryptedMessage(encrypt(this.encryptionKey, { data: bytes.toString('base64') }, 'x25519-xsalsa20-poly1305'))
-      new BNjs(blinding.toString()).toBuffer('be', 31),
-      new BNjs(amount.toString()).toBuffer('be', 31),
-    ])
-    return encrypt(this.encryptionKey, { data: bytes.toString('base64') }, 'x25519-xsalsa20-poly1305')
   }
 
   decrypt(data) {
-    const decryptedMessage = decrypt(data, this.privkey.slice(2))
+    return Buffer.from(decrypt(unpackEncryptedMessage(data), this.privkey.slice(2)), 'base64')
-    const buf = Buffer.from(decryptedMessage, 'base64')
-    return {
-      blinding: BigNumber.from('0x' + buf.slice(0, 31).toString('hex')),
-      amount: BigNumber.from('0x' + buf.slice(31, 62).toString('hex')),
-    }
   }
 }
 

src/utils.js

@@ -55,36 +55,6 @@ async function revertSnapshot(id) {
   await ethers.provider.send('evm_revert', [id])
 }
 
-function packEncryptedMessage(encryptedMessage) {
-  const nonceBuf = Buffer.from(encryptedMessage.nonce, 'base64')
-  const ephemPublicKeyBuf = Buffer.from(encryptedMessage.ephemPublicKey, 'base64')
-  const ciphertextBuf = Buffer.from(encryptedMessage.ciphertext, 'base64')
-  const messageBuff = Buffer.concat([
-    Buffer.alloc(24 - nonceBuf.length),
-    nonceBuf,
-    Buffer.alloc(32 - ephemPublicKeyBuf.length),
-    ephemPublicKeyBuf,
-    ciphertextBuf,
-  ])
-  return '0x' + messageBuff.toString('hex')
-}
-
-function unpackEncryptedMessage(encryptedMessage) {
-  if (encryptedMessage.slice(0, 2) === '0x') {
-    encryptedMessage = encryptedMessage.slice(2)
-  }
-  const messageBuff = Buffer.from(encryptedMessage, 'hex')
-  const nonceBuf = messageBuff.slice(0, 24)
-  const ephemPublicKeyBuf = messageBuff.slice(24, 56)
-  const ciphertextBuf = messageBuff.slice(56)
-  return {
-    version: 'x25519-xsalsa20-poly1305',
-    nonce: nonceBuf.toString('base64'),
-    ephemPublicKey: ephemPublicKeyBuf.toString('base64'),
-    ciphertext: ciphertextBuf.toString('base64'),
-  }
-}
-
 module.exports = {
   FIELD_SIZE,
   randomBN,
@@ -95,6 +65,4 @@ module.exports = {
   getExtDataHash,
   takeSnapshot,
   revertSnapshot,
-  packEncryptedMessage,
-  unpackEncryptedMessage,
 }

src/utxo.js

@@ -27,6 +27,28 @@ class Utxo {
     }
     return this._nullifier
   }
+
+  encrypt() {
+    const blindingBuf = Buffer.from(this.blinding.toHexString().slice(2), 'hex')
+    const amountBuf = Buffer.from(this.amount.toHexString().slice(2), 'hex')
+    const bytes = Buffer.concat([
+      Buffer.alloc(31 - blindingBuf.length),
+      blindingBuf,
+      Buffer.alloc(31 - amountBuf.length),
+      amountBuf,
+    ])
+    return this.keypair.encrypt(bytes)
+  }
+
+  static decrypt(keypair, data, index) {
+    const buf = keypair.decrypt(data)
+    return new Utxo({
+      blinding: BigNumber.from('0x' + buf.slice(0, 31).toString('hex')),
+      amount: BigNumber.from('0x' + buf.slice(31, 62).toString('hex')),
+      keypair,
+      index,
+    })
+  }
 }
 
 module.exports = Utxo

test/full.test.js

@@ -16,7 +16,7 @@ const Utxo = require('../src/utxo')
 const MERKLE_TREE_HEIGHT = 5
 const MerkleTree = require('fixed-merkle-tree')
 
-const { deposit, transact, withdraw, merge } = require('../src/index')
+const { transaction } = require('../src/index')
 const Keypair = require('../src/keypair')
 
 describe('TornadoPool', () => {
@@ -41,89 +41,50 @@ describe('TornadoPool', () => {
     snapshotId = await takeSnapshot()
   })
 
-  it('encryp -> pack -> unpack -> decrypt should work', () => {
+  it('encrypt -> decrypt should work', () => {
-    const blinding = 3
+    const data = Buffer.from([0xff, 0xaa, 0x00, 0x01])
-    const amount = 5
     const keypair = new Keypair()
 
-    const cyphertext = keypair.encrypt({ blinding, amount })
+    const ciphertext = keypair.encrypt(data)
-
+    const result = keypair.decrypt(ciphertext)
-    const packedMessage = packEncryptedMessage(cyphertext)
+    expect(result).to.be.deep.equal(data)
-
-    const unpackedMessage = unpackEncryptedMessage(packedMessage)
-
-    const result = keypair.decrypt(unpackedMessage)
-
-    expect(result.blinding).to.be.equal(blinding)
-    expect(result.amount).to.be.equal(amount)
   })
 
   it('should deposit, transact and withdraw', async function () {
-    /// deposit phase
     // Alice deposits into tornado pool
-    const amount = BigNumber.from('10000000')
+    const aliceDepositAmount = 1e7
-    const alicePrivateKey = ethers.Wallet.createRandom().privateKey // the private key we use for snarks and encryption, not for transactions
+    const aliceDepositUtxo = new Utxo({ amount: aliceDepositAmount })
-    const aliceKeypair = new Keypair(alicePrivateKey)
+    await transaction({ tornadoPool, outputs: [aliceDepositUtxo] })
 
-    const depositInput = new Utxo({ amount, keypair: aliceKeypair })
+    // Bob gives Alice address to send some eth inside the shielded pool
-    await deposit({ tornadoPool, utxo: depositInput })
+    const bobKeypair = new Keypair()
-
-    // getting account data from chain to verify that Alice has an Input to spend now
-    const filter = tornadoPool.filters.NewCommitment()
-    let events = await tornadoPool.queryFilter(filter)
-    let unpackedMessage = unpackEncryptedMessage(events[0].args.encryptedOutput)
-    let decryptedMessage = aliceKeypair.decrypt(unpackedMessage)
-    let aliceInputIndex = events[0].args.index
-    expect(decryptedMessage.amount).to.be.equal(amount)
-    expect(decryptedMessage.blinding).to.be.equal(depositInput.blinding)
-
-    /// transact phase.
-    // Bob gives Alice address to send some eth inside pool
-    const bobPrivateKey = ethers.Wallet.createRandom().privateKey
-    const bobKeypair = new Keypair(bobPrivateKey)
     const bobAddress = bobKeypair.address()
 
-    // but alice does not have Bob's privkey so let's build keypair without it
+    // Alice sends some funds to Bob
-    const bobKeypairForEncryption = Keypair.fromString(bobAddress)
+    const bobSendAmount = 3e6
+    const bobSendUtxo = new Utxo({ amount: bobSendAmount, keypair: Keypair.fromString(bobAddress) })
+    const aliceChangeUtxo = new Utxo({ amount: aliceDepositAmount - bobSendAmount, keypair: aliceDepositUtxo.keypair })
+    await transaction({ tornadoPool, inputs: [aliceDepositUtxo], outputs: [bobSendUtxo, aliceChangeUtxo] })
 
-    // let's build input for the shielded transaction
+    // Bob parses chain to detect incoming funds
-    const aliceInput = new Utxo({
+    const filter = tornadoPool.filters.NewCommitment()
-      amount,
-      blinding: depositInput.blinding,
-      index: aliceInputIndex,
-      keypair: aliceKeypair,
-    })
-    const bobInput = new Utxo({ amount, keypair: bobKeypairForEncryption })
-
-    await transact({ tornadoPool, input: aliceInput, output: bobInput })
-
-    // getting account data from chain to verify that Bob has an Input to spend now
     const fromBlock = await ethers.provider.getBlock()
-    events = await tornadoPool.queryFilter(filter, fromBlock.number)
+    const events = await tornadoPool.queryFilter(filter, fromBlock.number)
-    const bobInputIndex = events[0].args.index
+    const bobReceiveUtxo = Utxo.decrypt(bobKeypair, events[0].args.encryptedOutput, events[0].args.index)
-    unpackedMessage = unpackEncryptedMessage(events[0].args.encryptedOutput)
+    expect(bobReceiveUtxo.amount).to.be.equal(bobSendAmount)
-    decryptedMessage = bobKeypair.decrypt(unpackedMessage)
-    expect(decryptedMessage.amount).to.be.equal(amount)
-    expect(decryptedMessage.blinding).to.be.equal(bobInput.blinding)
 
-    /// withdraw phase
+    // Bob withdraws part of his funds from the shielded pool
-    // now Bob wants to exit the pool using a half of its funds
+    const bobWithdrawAmount = 2e6
-    const bobInputForWithdraw = new Utxo({
+    const bobEthAddress = '0xDeaD00000000000000000000000000000000BEEf'
-      amount,
+    const bobChangeUtxo = new Utxo({ amount: bobSendAmount - bobWithdrawAmount, keypair: bobKeypair })
-      blinding: bobInput.blinding,
+    await transaction({ tornadoPool, inputs: [bobReceiveUtxo], outputs: [bobChangeUtxo], recipient: bobEthAddress })
-      index: bobInputIndex,
-      keypair: bobKeypair,
-    })
-    const bobChange = new Utxo({ amount: amount.div(2), keypair: bobKeypair })
-    const recipient = '0xc2Ba33d4c0d2A92fb4f1a07C273c5d21E688Eb48'
-    await withdraw({ tornadoPool, input: bobInputForWithdraw, change: bobChange, recipient })
 
-    const bal = await ethers.provider.getBalance(recipient)
+    const bobBalance = await ethers.provider.getBalance(bobEthAddress)
-    expect(bal).to.be.gt(0)
+    expect(bobBalance).to.be.equal(bobWithdrawAmount)
   })
 
   it('should work with 16 inputs', async function () {
-    const utxo1 = await merge({ tornadoPool })
+    await transaction({ tornadoPool, inputs: [new Utxo(), new Utxo(), new Utxo()] })
   })
 
   afterEach(async () => {