From 9badb2f3fe66a3c5a5737a0cf79a45090cd5f10e Mon Sep 17 00:00:00 2001
From: Richard Ramos <info@richardramos.me>
Date: Wed, 15 Aug 2018 13:53:09 -0400
Subject: [PATCH] Simplified getBalance()

---
 gas-relayer/src/strategy/BaseStrategy.js     | 5 ++---
 gas-relayer/src/strategy/IdentityStrategy.js | 3 +--
 gas-relayer/src/strategy/SNTStrategy.js      | 4 ++--
 3 files changed, 5 insertions(+), 7 deletions(-)

diff --git a/gas-relayer/src/strategy/BaseStrategy.js b/gas-relayer/src/strategy/BaseStrategy.js
index 4f4083d..b00d9a1 100644
--- a/gas-relayer/src/strategy/BaseStrategy.js
+++ b/gas-relayer/src/strategy/BaseStrategy.js
@@ -10,14 +10,13 @@ class BaseStrategy {
         this.config = config;
     }
 
-    async getBalance(address, token, message, gasToken){
+    async getBalance(address, token){
         // Determining balances of token used
-        // TODO: probably token and gasToken can be unified
         if(token.symbol == "ETH"){
             return new this.web3.utils.BN(await this.web3.eth.getBalance(address));
         } else {
             const Token = new this.web3.eth.Contract(erc20ABI.abi);
-            Token.options.address = gasToken;
+            Token.options.address = token.address;
             return new this.web3.utils.BN(await Token.methods.balanceOf(address).call());  
         }
     }
diff --git a/gas-relayer/src/strategy/IdentityStrategy.js b/gas-relayer/src/strategy/IdentityStrategy.js
index e8b4f1b..749079e 100644
--- a/gas-relayer/src/strategy/IdentityStrategy.js
+++ b/gas-relayer/src/strategy/IdentityStrategy.js
@@ -42,8 +42,7 @@ class IdentityStrategy extends Strategy {
         }
 
         // gasPrice * limit calculation
-        const gasToken = params('_gasToken');
-        const balance = await this.getBalance(message.input.address, token, message, gasToken);
+        const balance = await this.getBalance(message.input.address, token);
         if(balance.lt(this.web3.utils.toBN(gasPrice.mul(gasLimit)))) {
             return {success: false, message: "Identity has not enough tokens for gasPrice*gasLimit"};
         }
diff --git a/gas-relayer/src/strategy/SNTStrategy.js b/gas-relayer/src/strategy/SNTStrategy.js
index ce3db2b..d7345ad 100644
--- a/gas-relayer/src/strategy/SNTStrategy.js
+++ b/gas-relayer/src/strategy/SNTStrategy.js
@@ -13,6 +13,8 @@ class SNTStrategy extends Strategy {
         const token = this.settings.getTokenBySymbol("SNT");
         if(token == undefined) return {success: false, message: "Token not allowed"};
         
+        const balance = await this.getBalance(message.input.wallet, token);
+
         if(message.input.functionName == TransferSNT){
             const estimatedGas = await this.web3.eth.estimateGas({
                 data: message.input.payload,
@@ -21,7 +23,6 @@ class SNTStrategy extends Strategy {
             });
 
             const gas = this.web3.utils.toBN(estimatedGas);
-            const balance = await this.getBalance(message.input.wallet, token, message, token.address);
             const value = this.web3.utils.toBN(params('_amount'));
             const requiredGas = value.add(gas);
 
@@ -37,7 +38,6 @@ class SNTStrategy extends Strategy {
                 if(exc.message.indexOf("revert") > -1) return {success: false, message: "Transaction will revert"};
             }
 
-            const balance = await this.getBalance(message.input.wallet, token, message, token.address);
             if(balance.lt(estimatedGas)){
                 return {success: false, message: "Address has not enough balance to execute the transaction (" + estimatedGas.toString() + ")"};
             }