Feat v2/get all balances (#228)

* feat: add services to get evm and cosmos balances

* add constants, types, and dependencies

* feat: add methods to get token balances

* fix: use a different rpc url for each chain and catch errors

* refactor: replace ethers-multicall with ethers-multical-provider

* chore: remove unused dependency

* refactor: move services to evm and cosmos handlers

* fix yarn lock

* remove log

---------

Co-authored-by: jmd3v <juan.villarraza@gmail.com>

package.json

@@ -33,6 +33,7 @@
     "axios": "^1.5.0",
     "cosmjs-types": "^0.8.0",
     "ethers": "6.7.1",
+    "ethers-multicall-provider": "^5.0.0",
     "lodash": "^4.17.21"
   },
   "resolutions": {

src/constants/index.ts

@@ -2,3 +2,28 @@ export const uint256MaxValue =
   "115792089237316195423570985008687907853269984665640564039457584007913129639935";
 
 export const nativeTokenConstant = "0xEeeeeEeeeEeEeeEeEeEeeEEEeeeeEeeeeeeeEEeE";
+export const NATIVE_EVM_TOKEN_ADDRESS =
+  "0xeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee";
+export const MULTICALL_ADDRESS = "0xcA11bde05977b3631167028862bE2a173976CA11";
+export const multicallAbi = [
+  {
+    inputs: [
+      {
+        internalType: "address",
+        name: "account",
+        type: "address"
+      }
+    ],
+    name: "getEthBalance",
+    outputs: [
+      {
+        internalType: "uint256",
+        name: "",
+        type: "uint256"
+      }
+    ],
+    stateMutability: "view",
+    type: "function"
+  }
+];
+export const CHAINS_WITHOUT_MULTICALL = [314, 3141]; // Filecoin, & Filecoin testnet

src/handlers/cosmos/index.ts

@@ -1,5 +1,5 @@
-import { toUtf8 } from "@cosmjs/encoding";
-import { calculateFee, Coin, GasPrice } from "@cosmjs/stargate";
+import { fromBech32, toBech32, toUtf8 } from "@cosmjs/encoding";
+import { calculateFee, Coin, GasPrice, StargateClient } from "@cosmjs/stargate";
 import { MsgExecuteContract } from "cosmjs-types/cosmwasm/wasm/v1/tx";
 
 import {
@@ -9,7 +9,11 @@ import {
   CosmosMsg,
   IBC_TRANSFER_TYPE,
   WasmHookMsg,
-  WASM_TYPE
+  WASM_TYPE,
+  CosmosBalance,
+  CosmosChain,
+  CosmosAddress,
+  ChainType
 } from "../../types";
 import { TxRaw } from "cosmjs-types/cosmos/tx/v1beta1/tx";
 
@@ -115,4 +119,58 @@ export class CosmosHandler {
       ""
     );
   }
+
+  async getBalances({
+    addresses,
+    cosmosChains
+  }: {
+    addresses: CosmosAddress[];
+    cosmosChains: CosmosChain[];
+  }): Promise<CosmosBalance[]> {
+    const cosmosBalances: CosmosBalance[] = [];
+
+    for (const chain of cosmosChains) {
+      if (chain.chainType !== ChainType.COSMOS) continue;
+
+      const addressData = addresses.find(
+        address => address.coinType === chain.coinType
+      );
+
+      if (!addressData) continue;
+
+      const cosmosAddress = this.deriveCosmosAddress(
+        chain.bech32Config.bech32PrefixAccAddr,
+        addressData.address
+      );
+
+      try {
+        const client = await StargateClient.connect(chain.rpc);
+
+        const balances = (await client.getAllBalances(cosmosAddress)) ?? [];
+
+        if (balances.length === 0) continue;
+
+        balances.forEach(balance => {
+          const { amount, denom } = balance;
+
+          cosmosBalances.push({
+            balance: amount,
+            denom,
+            chainId: String(chain.chainId),
+            decimals:
+              chain.currencies.find(currency => currency.coinDenom === denom)
+                ?.coinDecimals ?? 6
+          });
+        });
+      } catch (error) {
+        //
+      }
+    }
+
+    return cosmosBalances;
+  }
+
+  deriveCosmosAddress(chainPrefix: string, address: string): string {
+    return toBech32(chainPrefix, fromBech32(address).data);
+  }
 }

src/handlers/evm/index.ts

@@ -5,12 +5,14 @@ import {
   EvmWallet,
   ExecuteRoute,
   RouteParamsPopulated,
+  Token,
+  TokenBalance,
   TransactionRequest,
   TransactionResponse,
   WalletV6
 } from "../../types";
 
-import { uint256MaxValue } from "../../constants";
+import { CHAINS_WITHOUT_MULTICALL, uint256MaxValue } from "../../constants";
 import { Utils } from "./utils";
 
 const ethersAdapter = new EthersAdapter();
@@ -238,4 +240,71 @@ export class EvmHandler extends Utils {
       ...gasData
     });
   }
+
+  async getBalances(
+    evmTokens: Token[],
+    userAddress: string,
+    chainRpcUrls: {
+      [chainId: string]: string;
+    }
+  ): Promise<TokenBalance[]> {
+    try {
+      // Some tokens don't support multicall, so we need to fetch them with Promise.all
+      // TODO: Once we support multicall on all chains, we can remove this split
+      const splittedTokensByMultiCallSupport = evmTokens.reduce(
+        (acc, token) => {
+          if (CHAINS_WITHOUT_MULTICALL.includes(Number(token.chainId))) {
+            acc[0].push(token);
+          } else {
+            acc[1].push(token);
+          }
+          return acc;
+        },
+        [[], []] as Token[][]
+      );
+
+      const tokensNotSupportingMulticall = splittedTokensByMultiCallSupport[0];
+      const tokensSupportingMulticall = splittedTokensByMultiCallSupport[1];
+
+      const tokensByChainId = tokensSupportingMulticall.reduce(
+        (groupedTokens, token) => {
+          if (!groupedTokens[token.chainId]) {
+            groupedTokens[token.chainId] = [];
+          }
+
+          groupedTokens[token.chainId].push(token);
+
+          return groupedTokens;
+        },
+        {} as Record<string, Token[]>
+      );
+
+      const tokensMulticall: TokenBalance[] = [];
+
+      for (const chainId in tokensByChainId) {
+        const tokens = tokensByChainId[chainId];
+        const rpcUrl = chainRpcUrls[chainId];
+
+        if (!rpcUrl) continue;
+
+        const tokensBalances = await this.getTokensBalanceSupportingMultiCall(
+          tokens,
+          rpcUrl,
+          userAddress
+        );
+
+        tokensMulticall.push(...tokensBalances);
+      }
+
+      const tokensNotMultiCall = await this.getTokensBalanceWithoutMultiCall(
+        tokensNotSupportingMulticall,
+        userAddress,
+        chainRpcUrls
+      );
+
+      return [...tokensMulticall, ...tokensNotMultiCall];
+    } catch (error) {
+      return [];
+    }
+  }
 }

src/handlers/evm/utils.ts

@@ -1,6 +1,19 @@
-import { ChainData, SquidData } from "@0xsquid/squid-types";
+import { ChainData, SquidData, Token } from "@0xsquid/squid-types";
 
-import { OverrideParams, Contract, GasData, RpcProvider } from "../../types";
+import {
+  OverrideParams,
+  Contract,
+  GasData,
+  RpcProvider,
+  TokenBalance
+} from "../../types";
+import { MulticallWrapper } from "ethers-multicall-provider";
+import { Provider, ethers } from "ethers";
+import {
+  multicallAbi,
+  MULTICALL_ADDRESS,
+  NATIVE_EVM_TOKEN_ADDRESS
+} from "../../constants";
 
 export class Utils {
   async validateNativeBalance({
@@ -108,4 +121,134 @@ export class Utils {
 
     return overrides ? { ...gasParams, ...overrides } : (gasParams as GasData);
   };
+
+  async getTokensBalanceSupportingMultiCall(
+    tokens: Token[],
+    chainRpcUrl: string,
+    userAddress?: string
+  ): Promise<TokenBalance[]> {
+    if (!userAddress) return [];
+
+    const multicallProvider = MulticallWrapper.wrap(
+      // eslint-disable-next-line @typescript-eslint/ban-ts-comment
+      // @ts-ignore
+      new ethers.JsonRpcProvider(chainRpcUrl)
+    );
+
+    const tokenBalances: Promise<TokenBalance>[] = tokens.map(token => {
+      const isNativeToken =
+        token.address.toLowerCase() === NATIVE_EVM_TOKEN_ADDRESS.toLowerCase();
+
+      const contract = new ethers.Contract(
+        isNativeToken ? MULTICALL_ADDRESS : token.address,
+        isNativeToken
+          ? multicallAbi
+          : [
+              {
+                name: "balanceOf",
+                type: "function",
+                inputs: [{ name: "_owner", type: "address" }],
+                outputs: [{ name: "balance", type: "uint256" }],
+                stateMutability: "view"
+              }
+            ],
+        multicallProvider as unknown as Provider
+      );
+
+      const getTokenData = async () => {
+        const balanceInWei = await contract[
+          isNativeToken ? "getEthBalance" : "balanceOf"
+        ](userAddress);
+
+        return {
+          balance: balanceInWei.toString(),
+          symbol: token.symbol,
+          address: token.address,
+          decimals: token.decimals,
+          chainId: token.chainId
+        };
+      };
+
+      return getTokenData();
+    });
+
+    try {
+      return Promise.all(tokenBalances);
+    } catch (error) {
+      return [];
+    }
+  }
+
+  async getTokensBalanceWithoutMultiCall(
+    tokens: Token[],
+    userAddress: string,
+    rpcUrlsPerChain: {
+      [chainId: string]: string;
+    }
+  ): Promise<TokenBalance[]> {
+    const balances: (TokenBalance | null)[] = await Promise.all(
+      tokens.map(async t => {
+        let balance: TokenBalance | null;
+        try {
+          if (t.address === NATIVE_EVM_TOKEN_ADDRESS) {
+            balance = await this.fetchBalance({
+              token: t,
+              userAddress,
+              rpcUrl: rpcUrlsPerChain[t.chainId]
+            });
+          } else {
+            balance = await this.fetchBalance({
+              token: t,
+              userAddress,
+              rpcUrl: rpcUrlsPerChain[t.chainId]
+            });
+          }
+
+          return balance;
+        } catch (error) {
+          return null;
+        }
+      })
+    );
+
+    // filter out null values
+    return balances.filter(Boolean) as TokenBalance[];
+  }
+
+  async fetchBalance({
+    token,
+    userAddress,
+    rpcUrl
+  }: {
+    token: Token;
+    userAddress: string;
+    rpcUrl: string;
+  }): Promise<TokenBalance | null> {
+    try {
+      const provider = new ethers.JsonRpcProvider(rpcUrl);
+
+      const tokenAbi = ["function balanceOf(address) view returns (uint256)"];
+      const tokenContract = new ethers.Contract(
+        token.address ?? "",
+        tokenAbi,
+        provider
+      );
+
+      const balance = (await tokenContract.balanceOf(userAddress)) ?? "0";
+
+      if (!token) return null;
+
+      const { decimals, symbol, address } = token;
+
+      return {
+        address,
+        // balance in wei
+        balance: parseInt(balance, 16).toString(),
+        decimals,
+        symbol
+      };
+    } catch (error) {
+      return null;
+    }
+  }
 }

src/index.ts

@@ -4,7 +4,11 @@ import {
   RouteResponse,
   StatusResponse,
   EvmWallet,
-  Token
+  Token,
+  TokenBalance,
+  CosmosAddress,
+  CosmosChain,
+  CosmosBalance
 } from "./types";
 
 import HttpAdapter from "./adapter/HttpAdapter";
@@ -301,4 +305,122 @@ export class Squid extends TokensChains {
 
     return fromAmountPlusSlippage.toString();
   }
+
+  public async getEvmBalances({
+    userAddress,
+    chains
+  }: {
+    userAddress: string;
+    chains: (string | number)[];
+  }): Promise<TokenBalance[]> {
+    // remove invalid and duplicate chains and convert to number
+    const filteredChains = new Set(chains.map(Number).filter(c => !isNaN(c)));
+    const chainRpcUrls = this.chains.reduce(
+      (acc, chain) => ({
+        ...acc,
+        [chain.chainId]: chain.rpc
+      }),
+      {}
+    );
+
+    return this.handlers.evm.getBalances(
+      this.tokens.filter(t => filteredChains.has(Number(t.chainId))),
+      userAddress,
+      chainRpcUrls
+    );
+  }
+
+  public async getCosmosBalances({
+    addresses,
+    chainIds = []
+  }: {
+    addresses: CosmosAddress[];
+    chainIds?: (string | number)[];
+  }) {
+    const cosmosChains = this.chains.filter(c =>
+      c.chainType === ChainType.COSMOS &&
+      // if chainIds is not provided, return all cosmos chains
+      chainIds.length === 0
+        ? true
+        : // else return only chains that are in chainIds
+          chainIds?.includes(c.chainId)
+    ) as CosmosChain[];
+    return this.handlers.cosmos.getBalances({
+      addresses,
+      cosmosChains
+    });
+  }
+
+  public async getAllBalances({
+    chainIds,
+    cosmosAddresses,
+    evmAddress
+  }: {
+    chainIds?: (string | number)[];
+    cosmosAddresses?: CosmosAddress[];
+    evmAddress?: string;
+  }): Promise<{
+    cosmosBalances?: CosmosBalance[];
+    evmBalances?: TokenBalance[];
+  }> {
+    if (!chainIds) {
+      // fetch balances for all chains compatible with provided addresses
+      const evmBalances = evmAddress
+        ? await this.getEvmBalances({
+            chains: this.tokens.map(t => String(t.chainId)),
+            userAddress: evmAddress
+          })
+        : [];
+
+      const cosmosBalances = cosmosAddresses
+        ? await this.getCosmosBalances({
+            addresses: cosmosAddresses
+          })
+        : [];
+
+      return {
+        evmBalances,
+        cosmosBalances
+      };
+    }
+
+    const normalizedChainIds = chainIds.map(String);
+
+    // fetch balances for provided chains
+    const [evmChainIds, cosmosChainIds] = this.chains.reduce(
+      (cosmosAndEvmChains, chain) => {
+        if (!normalizedChainIds.includes(String(chain.chainId))) {
+          return cosmosAndEvmChains;
+        }
+
+        if (chain.chainType === ChainType.COSMOS) {
+          cosmosAndEvmChains[1].push(chain.chainId);
+        } else {
+          cosmosAndEvmChains[0].push(chain.chainId);
+        }
+        return cosmosAndEvmChains;
+      },
+
+      [[], []] as [(string | number)[], (string | number)[]]
+    );
+
+    const evmBalances = evmAddress
+      ? await this.getEvmBalances({
+          chains: evmChainIds,
+          userAddress: evmAddress
+        })
+      : [];
+
+    const cosmosBalances = cosmosAddresses
+      ? await this.getCosmosBalances({
+          addresses: cosmosAddresses,
+          chainIds: cosmosChainIds
+        })
+      : [];
+
+    return {
+      evmBalances,
+      cosmosBalances
+    };
+  }
 }

src/types/index.ts

@@ -105,3 +105,23 @@ export type StatusResponse = ApiBasicResponse & {
   squidTransactionStatus?: string;
 };
 // END STATUS TYPE
+
+export type TokenBalance = {
+  symbol: string;
+  address: string;
+  decimals: number;
+  balance: string;
+};
+
+export type CosmosAddress = {
+  coinType: number;
+  chainId: string;
+  address: string;
+};
+
+export type CosmosBalance = {
+  decimals: number;
+  balance: string;
+  denom: string;
+  chainId: string;
+};