From ecdcb545fc47c2ce6c087bff903c27a1d48b9bad Mon Sep 17 00:00:00 2001
From: smart_ex <sergei.smart@gmail.com>
Date: Mon, 18 Apr 2022 17:18:57 +1000
Subject: [PATCH] fix isKnownContract validation

---
 src/modules/validator.js |  2 +-
 src/utils.js             | 24 +++++++++++++++++-------
 src/worker.js            |  4 +---
 test/validator.js        |  8 +++++++-
 4 files changed, 26 insertions(+), 12 deletions(-)

diff --git a/src/modules/validator.js b/src/modules/validator.js
index 577dc24..864f36e 100644
--- a/src/modules/validator.js
+++ b/src/modules/validator.js
@@ -19,7 +19,7 @@ ajv.addKeyword('isAddress', {
 ajv.addKeyword('isKnownContract', {
   validate: (schema, data) => {
     try {
-      return getInstance(data) !== null
+      return !!getInstance(data)
     } catch (e) {
       return false
     }
diff --git a/src/utils.js b/src/utils.js
index 6845745..614c52d 100644
--- a/src/utils.js
+++ b/src/utils.js
@@ -1,16 +1,26 @@
 const { instances } = require('./config')
 const { toChecksumAddress, BN } = require('web3-utils')
 
+const addressMap = new Map()
+for (const [currency, { instanceAddress, symbol, decimals }] of Object.entries(instances)) {
+  Object.entries(instanceAddress).forEach(([amount, address]) =>
+    addressMap.set(address, {
+      currency,
+      amount,
+      symbol,
+      decimals,
+    }),
+  )
+}
+
 function getInstance(address) {
   address = toChecksumAddress(address)
-  for (const currency of Object.keys(instances)) {
-    for (const amount of Object.keys(instances[currency].instanceAddress)) {
-      if (instances[currency].instanceAddress[amount] === address) {
-        return { currency, amount }
-      }
-    }
+  const key = toChecksumAddress(address)
+  if (addressMap.has(key)) {
+    return addressMap.get(key)
+  } else {
+    throw new Error('Unknown contact address')
   }
-  return null
 }
 
 function setSafeInterval(func, interval) {
diff --git a/src/worker.js b/src/worker.js
index 03ab235..4708a25 100644
--- a/src/worker.js
+++ b/src/worker.js
@@ -10,7 +10,6 @@ const {
   netId,
   gasPrices,
   gasLimits,
-  instances,
   privateKey,
   proxyLight,
   httpRpcUrl,
@@ -74,8 +73,7 @@ function getGasLimit() {
 }
 
 async function checkTornadoFee({ args, contract }) {
-  const { currency, amount } = getInstance(contract)
-  const { decimals } = instances[currency]
+  const { amount, decimals } = getInstance(contract)
   const fee = toBN(args[4])
 
   const { fast } = await getGasPrices()
diff --git a/test/validator.js b/test/validator.js
index 7a1b25e..6464178 100644
--- a/test/validator.js
+++ b/test/validator.js
@@ -15,7 +15,13 @@ describe('Validator', () => {
         '.proof should match pattern "^0x[a-fA-F0-9]{512}$"',
       )
     })
-
+    it('should throw if unknown contract', () => {
+      const malformedData = { ...withdrawData }
+      malformedData.contract = '0xf17f52151ebef6c7334fad080c5704d77216b732'
+      getTornadoWithdrawInputError(malformedData).should.be.equal(
+        '.contract should pass "isKnownContract" keyword validation',
+      )
+    })
     it('should throw something is missing', () => {
       const malformedData = { ...withdrawData }
       delete malformedData.proof