From 8ab4b33d915df0e6d54e0ea4bde9c9e1c39b6ed0 Mon Sep 17 00:00:00 2001
From: poma <semenov.roma@gmail.com>
Date: Mon, 16 Aug 2021 19:53:18 +0300
Subject: [PATCH] refactor publicAmount

---
 README.md                 |  1 +
 contracts/TornadoPool.sol | 29 ++++++++++-------------------
 src/index.js              |  5 +----
 src/utils.js              | 20 +++++++++++++-------
 4 files changed, 25 insertions(+), 30 deletions(-)

diff --git a/README.md b/README.md
index 5c87310..5107c5d 100644
--- a/README.md
+++ b/README.md
@@ -9,4 +9,5 @@ yarn test
 ```
 
 TODO
+
 1. deposit from mainnet to the pool on optimism in one tx
diff --git a/contracts/TornadoPool.sol b/contracts/TornadoPool.sol
index 01f589c..bd68b75 100644
--- a/contracts/TornadoPool.sol
+++ b/contracts/TornadoPool.sol
@@ -22,6 +22,7 @@ interface IVerifier {
 contract TornadoPool {
   uint256 public constant FIELD_SIZE = 21888242871839275222246405745257275088548364400416034343698204186575808495617;
   uint256 public constant MAX_EXT_AMOUNT = 2**248 - 1;
+  uint256 public constant MAX_FEE = 2**248;
 
   mapping(bytes32 => bool) public nullifierHashes;
   bytes32 public currentRoot;
@@ -31,7 +32,7 @@ contract TornadoPool {
 
   struct ExtData {
     address payable recipient;
-    uint256 extAmount;
+    int256 extAmount;
     address payable relayer;
     uint256 fee;
     bytes encryptedOutput1;
@@ -82,8 +83,11 @@ 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 < 2**248, "Invalid fee");
-    require((_args.publicAmount + _extData.extAmount) % FIELD_SIZE == _extData.fee % FIELD_SIZE, "Invalid public amount");
+    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(verifyProof(_args), "Invalid transaction proof");
 
     currentRoot = _args.newRoot;
@@ -92,13 +96,12 @@ contract TornadoPool {
       nullifierHashes[_args.inputNullifiers[i]] = true;
     }
 
-    int256 extAmount = calculateExternalAmount(_extData.extAmount);
-    if (extAmount > 0) {
+    if (_extData.extAmount > 0) {
       require(msg.value == uint256(_extData.extAmount), "Incorrect amount of ETH sent on deposit");
-    } else if (extAmount < 0) {
+    } else if (_extData.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");
-      _extData.recipient.transfer(uint256(-extAmount));
+      _extData.recipient.transfer(uint256(-_extData.extAmount));
     } else {
       require(msg.value == 0, "Sent ETH amount should be 0 for transaction");
     }
@@ -114,18 +117,6 @@ contract TornadoPool {
     }
   }
 
-  function calculateExternalAmount(uint256 _extAmount) public pure returns (int256) {
-    // -MAX_EXT_AMOUNT < extAmount < MAX_EXT_AMOUNT
-    if (_extAmount < MAX_EXT_AMOUNT) {
-      return int256(_extAmount);
-    } else if (_extAmount > FIELD_SIZE - MAX_EXT_AMOUNT) {
-      // FIELD_SIZE - MAX_EXT_AMOUNT < _extAmount < FIELD_SIZE
-      return -(int256(FIELD_SIZE) - int256(_extAmount));
-    } else {
-      revert("Invalid extAmount value");
-    }
-  }
-
   /** @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 24f592b..8f8231f 100644
--- a/src/index.js
+++ b/src/index.js
@@ -121,10 +121,7 @@ async function prepareTransaction({
     .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 // extAmount will be positive for a deposit, zero for a transact and negative for withdraw
-  if (extAmount < 0) {
-    extAmount = FIELD_SIZE.add(extAmount)
-  }
+  const amount = extAmount > 0 ? extAmount : 0
 
   const { args, extData } = await getProof({
     inputs,
diff --git a/src/utils.js b/src/utils.js
index f237031..87b9e5f 100644
--- a/src/utils.js
+++ b/src/utils.js
@@ -18,7 +18,7 @@ function getExtDataHash({ recipient, extAmount, relayer, fee, encryptedOutput1,
 
   const encodedData = abi.encode(
     [
-      'tuple(address recipient,uint256 extAmount,address relayer,uint256 fee,bytes encryptedOutput1,bytes encryptedOutput2)',
+      'tuple(address recipient,int256 extAmount,address relayer,uint256 fee,bytes encryptedOutput1,bytes encryptedOutput2)',
     ],
     [
       {
@@ -36,12 +36,18 @@ function getExtDataHash({ recipient, extAmount, relayer, fee, encryptedOutput1,
 }
 
 /** BigNumber to hex string of specified length */
-const toFixedHex = (number, length = 32) =>
-  '0x' +
-  (number instanceof Buffer
-    ? number.toString('hex')
-    : BigNumber.from(number).toHexString().slice(2)
-  ).padStart(length * 2, '0')
+function toFixedHex(number, length = 32) {
+  let result =
+    '0x' +
+    (number instanceof Buffer
+      ? number.toString('hex')
+      : BigNumber.from(number).toHexString().replace('0x', '')
+    ).padStart(length * 2, '0')
+  if (result.indexOf('-') > -1) {
+    result = '-' + result.replace('-', '')
+  }
+  return result
+}
 
 /** Convert value into buffer of specified byte length */
 const toBuffer = (value, length) =>