From d7f4998f5281dd3b371184df756f9f8c570a1132 Mon Sep 17 00:00:00 2001
From: andyk <kyzia.ru@gmail.com>
Date: Fri, 6 Nov 2020 18:01:19 +0300
Subject: [PATCH]  fix getUserWalletBalances of WalletBalanceProvider

---
 contracts/misc/WalletBalanceProvider.sol | 24 ++++++++----------------
 1 file changed, 8 insertions(+), 16 deletions(-)

diff --git a/contracts/misc/WalletBalanceProvider.sol b/contracts/misc/WalletBalanceProvider.sol
index df9569c8..9153e3c6 100644
--- a/contracts/misc/WalletBalanceProvider.sol
+++ b/contracts/misc/WalletBalanceProvider.sol
@@ -46,12 +46,13 @@ contract WalletBalanceProvider {
       - return 0 on non-contract address
     **/
   function balanceOf(address user, address token) public view returns (uint256) {
-    // check if token is actually a contract
-    if (token.isContract()) {
+    if (token == MOCK_ETH_ADDRESS) {
+      return user.balance; // ETH balance
+      // check if token is actually a contract
+    } else if (token.isContract()) {
       return IERC20(token).balanceOf(user);
-    } else {
-      return 0;
     }
+    revert('INVALID_TOKEN');
   }
 
   /**
@@ -69,16 +70,7 @@ contract WalletBalanceProvider {
 
     for (uint256 i = 0; i < users.length; i++) {
       for (uint256 j = 0; j < tokens.length; j++) {
-        uint256 offset = i * tokens.length;
-        if (tokens[j] == MOCK_ETH_ADDRESS) {
-          balances[offset + j] = users[i].balance; // ETH balance
-        } else {
-          if (!tokens[j].isContract()) {
-            revert('INVALID_TOKEN');
-          } else {
-            balances[offset + j] = balanceOf(users[i], tokens[j]);
-          }
-        }
+        balances[i * tokens.length + j] = balanceOf(users[i], tokens[j]);
       }
     }
 
@@ -113,9 +105,9 @@ contract WalletBalanceProvider {
         balances[j] = 0;
         continue;
       }
-      balances[j] = balanceOf(user, reserves[j]);
+      balances[j] = balanceOf(user, reservesWithEth[j]);
     }
 
-    return (reserves, balances);
+    return (reservesWithEth, balances);
   }
 }