From 042be187d10d331024ab9fa371b8284c449fa77e Mon Sep 17 00:00:00 2001
From: Alexey <peretsmobil@gmail.com>
Date: Mon, 16 Aug 2021 22:17:07 +0300
Subject: [PATCH] calc public amount

---
 circuits/transaction.circom |  2 +-
 contracts/TornadoPool.sol   | 15 +++++++++------
 src/index.js                |  2 +-
 test/full.test.js           |  8 ++++++++
 4 files changed, 19 insertions(+), 8 deletions(-)

diff --git a/circuits/transaction.circom b/circuits/transaction.circom
index bbd47d0..2009590 100644
--- a/circuits/transaction.circom
+++ b/circuits/transaction.circom
@@ -117,7 +117,7 @@ template Transaction(levels, nIns, nOuts, zeroLeaf) {
     }
 
     // verify amount invariant
-    sumIns === sumOuts + publicAmount;
+    sumIns + publicAmount === sumOuts;
 
     // Check merkle tree update with inserted transaction outputs
     component treeUpdater = TreeUpdater(levels, 1 /* log2(nOuts) */, zeroLeaf);
diff --git a/contracts/TornadoPool.sol b/contracts/TornadoPool.sol
index bd68b75..9f7bf8c 100644
--- a/contracts/TornadoPool.sol
+++ b/contracts/TornadoPool.sol
@@ -21,7 +21,7 @@ interface IVerifier {
 
 contract TornadoPool {
   uint256 public constant FIELD_SIZE = 21888242871839275222246405745257275088548364400416034343698204186575808495617;
-  uint256 public constant MAX_EXT_AMOUNT = 2**248 - 1;
+  int256 public constant MAX_EXT_AMOUNT = 2**248;
   uint256 public constant MAX_FEE = 2**248;
 
   mapping(bytes32 => bool) public nullifierHashes;
@@ -83,11 +83,7 @@ contract TornadoPool {
     require(uint256(_args.extDataHash) == uint256(keccak256(abi.encode(_extData))) % FIELD_SIZE, "Incorrect external data hash");
     uint256 cachedCommitmentIndex = currentCommitmentIndex;
     require(_args.outPathIndices == cachedCommitmentIndex >> 1, "Invalid merkle tree insert position");
-    require(_extData.fee < MAX_FEE, "Invalid fee");
-    require(_extData.extAmount < int256(MAX_EXT_AMOUNT) && _extData.extAmount > -int256(MAX_EXT_AMOUNT), "Invalid ext amount");
-    int256 _publicAmount = int256(_extData.fee) - _extData.extAmount;
-    uint256 publicAmount = (_publicAmount >= 0) ? uint256(_publicAmount) : FIELD_SIZE + uint256(_publicAmount);
-    require(_args.publicAmount == publicAmount, "Invalid public amount");
+    require(_args.publicAmount == calculatePublicAmount(_extData.extAmount, _extData.fee), "Invalid public amount");
     require(verifyProof(_args), "Invalid transaction proof");
 
     currentRoot = _args.newRoot;
@@ -117,6 +113,13 @@ contract TornadoPool {
     }
   }
 
+  function calculatePublicAmount(int256 _extAmount, uint256 _fee) public pure returns(uint256) {
+    require(_fee < MAX_FEE, "Invalid fee");
+    require(_extAmount > -MAX_EXT_AMOUNT && _extAmount < MAX_EXT_AMOUNT, "Invalid ext amount");
+    int256 publicAmount = _extAmount - int256(_fee);
+    return (publicAmount >= 0) ? uint256(publicAmount) : FIELD_SIZE - uint256(-publicAmount);
+  }
+
   /** @dev whether a note is already spent */
   function isSpent(bytes32 _nullifierHash) public view returns (bool) {
     return nullifierHashes[_nullifierHash];
diff --git a/src/index.js b/src/index.js
index 8f8231f..80e461e 100644
--- a/src/index.js
+++ b/src/index.js
@@ -61,7 +61,7 @@ async function getProof({ inputs, outputs, tree, extAmount, fee, recipient, rela
     newRoot: tree.root(),
     inputNullifier: inputs.map((x) => x.getNullifier()),
     outputCommitment: outputs.map((x) => x.getCommitment()),
-    publicAmount: BigNumber.from(fee).sub(extAmount).add(FIELD_SIZE).mod(FIELD_SIZE).toString(),
+    publicAmount: BigNumber.from(extAmount).sub(fee).add(FIELD_SIZE).mod(FIELD_SIZE).toString(),
     extDataHash,
 
     // data for 2 transaction inputs
diff --git a/test/full.test.js b/test/full.test.js
index 99518d2..9512d8a 100644
--- a/test/full.test.js
+++ b/test/full.test.js
@@ -44,6 +44,14 @@ describe('TornadoPool', () => {
     expect(result).to.be.deep.equal(data)
   })
 
+  it('constants check', async () => {
+    const maxFee = await tornadoPool.MAX_FEE()
+    const maxExtAmount = await tornadoPool.MAX_EXT_AMOUNT()
+    const fieldSize = await tornadoPool.FIELD_SIZE()
+
+    expect(maxExtAmount.add(maxFee)).to.be.lt(fieldSize)
+  })
+
   it('should register and deposit', async function () {
     // Alice deposits into tornado pool
     const aliceDepositAmount = 1e7