fix: unify asset decoding for Namada and Cosmos (#1235)

This fixes cases where Namada addresses are decoded to the wrong chain's
asset. Also:

- Add NAM address to registry and remove special handling for Namada
- Store addresses for all assets, replacing IBC-specific address entry
- Remove unused chain lookup in mapCoinsToAssets

Author: emccorson

apps/namadillo/src/App/Ibc/ShieldAllAssetList.tsx

@@ -4,15 +4,15 @@ import { TokenCurrency } from "App/Common/TokenCurrency";
 import BigNumber from "bignumber.js";
 import clsx from "clsx";
 import { getAssetImageUrl } from "integrations/utils";
-import { AssetWithBalanceAndIbcInfo } from "types";
+import { AddressWithAssetAndBalance } from "types";
 
-export type SelectableAssetWithBalanceAndIbcInfo =
-  AssetWithBalanceAndIbcInfo & {
+export type SelectableAddressWithAssetAndBalance =
+  AddressWithAssetAndBalance & {
     checked: boolean;
   };
 
 type ShieldAllAssetListProps = {
-  assets: SelectableAssetWithBalanceAndIbcInfo[];
+  assets: SelectableAddressWithAssetAndBalance[];
   onToggleAsset: (asset: Asset) => void;
 };
 
@@ -24,7 +24,7 @@ export const ShieldAllAssetList = ({
     <ul className="max-h-[200px] dark-scrollbar -mr-2">
       {assets.map(
         (
-          assetWithBalance: SelectableAssetWithBalanceAndIbcInfo,
+          assetWithBalance: SelectableAddressWithAssetAndBalance,
           idx: number
         ) => {
           const image = getAssetImageUrl(assetWithBalance.asset);

apps/namadillo/src/App/Ibc/ShieldAllPanel.tsx

@@ -11,12 +11,12 @@ import { TransferTransactionFee } from "App/Transfer/TransferTransactionFee";
 import { getTransactionFee } from "integrations/utils";
 import { useEffect, useMemo, useState } from "react";
 import {
-  AssetWithBalanceAndIbcInfo,
+  AddressWithAssetAndBalance,
   ChainRegistryEntry,
   WalletProvider,
 } from "types";
 import {
-  SelectableAssetWithBalanceAndIbcInfo,
+  SelectableAddressWithAssetAndBalance,
   ShieldAllAssetList,
 } from "./ShieldAllAssetList";
 import { ShieldAllContainer } from "./ShieldAllContainer";
@@ -26,7 +26,7 @@ type ShieldAllPanelProps = {
   wallet: WalletProvider;
   walletAddress: string;
   isLoading: boolean;
-  assetList: AssetWithBalanceAndIbcInfo[];
+  assetList: AddressWithAssetAndBalance[];
   onShieldAll: (assets: Asset[]) => void;
 };
 
@@ -39,7 +39,7 @@ export const ShieldAllPanel = ({
   onShieldAll,
 }: ShieldAllPanelProps): JSX.Element => {
   const [selectableAssets, setSelectableAssets] = useState<
-    SelectableAssetWithBalanceAndIbcInfo[]
+    SelectableAddressWithAssetAndBalance[]
   >([]);
 
   useEffect(() => {

apps/namadillo/src/atoms/balance/atoms.ts

@@ -4,21 +4,26 @@ import {
   defaultAccountAtom,
   transparentBalanceAtom,
 } from "atoms/accounts/atoms";
-import { nativeTokenAddressAtom, tokenAddressesAtom } from "atoms/chain";
+import {
+  chainParametersAtom,
+  nativeTokenAddressAtom,
+  tokenAddressesAtom,
+} from "atoms/chain";
 import { shouldUpdateBalanceAtom } from "atoms/etc";
 import { availableAssetsAtom } from "atoms/integrations";
 import { queryDependentFn } from "atoms/utils";
 import BigNumber from "bignumber.js";
 import { atomWithQuery } from "jotai-tanstack-query";
 import { namadaAsset } from "registry/namadaAsset";
-import { unknownAsset } from "registry/unknownAsset";
-import { toDisplayAmount } from "utils";
-import { findAssetByToken } from "./functions";
+import { AddressWithAssetAndBalance } from "types";
+import {
+  findAssetByToken,
+  mapNamadaAddressesToAssets,
+  mapNamadaAssetsToTokenBalances,
+} from "./functions";
 import { fetchCoinPrices, fetchShieldedBalance } from "./services";
 
-export type TokenBalance = {
-  asset: Asset;
-  balance: BigNumber;
+export type TokenBalance = AddressWithAssetAndBalance & {
   dollar?: BigNumber;
 };
 
@@ -148,75 +153,88 @@ export const tokenPricesAtom = atomWithQuery((get) => {
   };
 });
 
-export const shieldedTokensAtom = atomWithQuery<TokenBalance[]>((get) => {
-  const shieldedBalanceQuery = get(shieldedBalanceAtom);
-  const assetByAddressQuery = get(assetByAddressAtom);
-  const tokenPricesQuery = get(tokenPricesAtom);
+export const namadaShieldedAssetsAtom = atomWithQuery((get) => {
+  const shieldedBalances = get(shieldedBalanceAtom);
+  const tokenAddresses = get(tokenAddressesAtom);
+  const chainParameters = get(chainParametersAtom);
 
   return {
     queryKey: [
-      "shielded-tokens",
-      shieldedBalanceQuery.data,
-      assetByAddressQuery.data,
-      tokenPricesQuery.data,
+      "namada-shielded-assets",
+      shieldedBalances.data,
+      tokenAddresses.data,
+      chainParameters.data!.chainId,
     ],
     ...queryDependentFn(
       async () =>
-        shieldedBalanceQuery.data?.map(({ address, amount }) =>
-          formatTokenBalance(
-            address,
-            amount,
-            assetByAddressQuery.data,
-            tokenPricesQuery.data
-          )
-        ) ?? [],
-      [shieldedBalanceQuery, tokenPricesQuery, assetByAddressQuery]
+        await mapNamadaAddressesToAssets(
+          shieldedBalances.data!,
+          tokenAddresses.data!,
+          chainParameters.data!.chainId
+        ),
+      [shieldedBalances, tokenAddresses, chainParameters]
     ),
   };
 });
 
-export const transparentTokensAtom = atomWithQuery<TokenBalance[]>((get) => {
-  const transparentBalanceQuery = get(transparentBalanceAtom);
-  const assetByAddressQuery = get(assetByAddressAtom);
-  const tokenPricesQuery = get(tokenPricesAtom);
+export const namadaTransparentAssetsAtom = atomWithQuery((get) => {
+  const transparentBalances = get(transparentBalanceAtom);
+  const tokenAddresses = get(tokenAddressesAtom);
+  const chainParameters = get(chainParametersAtom);
 
   return {
     queryKey: [
-      "transparent-tokens",
-      transparentBalanceQuery.data,
-      assetByAddressQuery.data,
-      tokenPricesQuery.data,
+      "namada-transparent-assets",
+      transparentBalances.data,
+      tokenAddresses.data,
+      chainParameters.data!.chainId,
     ],
     ...queryDependentFn(
       async () =>
-        transparentBalanceQuery.data?.map(({ address, amount }) =>
-          formatTokenBalance(
-            address,
-            amount,
-            assetByAddressQuery.data,
-            tokenPricesQuery.data
-          )
-        ) ?? [],
-      [transparentBalanceQuery, tokenPricesQuery, assetByAddressQuery]
+        await mapNamadaAddressesToAssets(
+          transparentBalances.data!,
+          tokenAddresses.data!,
+          chainParameters.data!.chainId
+        ),
+      [transparentBalances, tokenAddresses, chainParameters]
     ),
   };
 });
 
-const formatTokenBalance = (
-  address: string,
-  amount: BigNumber,
-  assetsByAddress?: Record<string, Asset>,
-  tokenPrices?: Record<string, BigNumber>
-): TokenBalance => {
-  const asset = assetsByAddress?.[address] ?? unknownAsset;
-  const balance = toDisplayAmount(asset, amount);
+export const shieldedTokensAtom = atomWithQuery<TokenBalance[]>((get) => {
+  const shieldedAssets = get(namadaShieldedAssetsAtom);
+  const tokenPrices = get(tokenPricesAtom);
+
+  return {
+    queryKey: ["shielded-tokens", shieldedAssets.data, tokenPrices.data],
+    ...queryDependentFn(
+      () =>
+        Promise.resolve(
+          mapNamadaAssetsToTokenBalances(
+            shieldedAssets.data!,
+            tokenPrices.data!
+          )
+        ),
+      [shieldedAssets, tokenPrices]
+    ),
+  };
+});
 
-  const tokenPrice = tokenPrices?.[address];
-  const dollar = tokenPrice ? balance.multipliedBy(tokenPrice) : undefined;
+export const transparentTokensAtom = atomWithQuery<TokenBalance[]>((get) => {
+  const transparentAssets = get(namadaTransparentAssetsAtom);
+  const tokenPrices = get(tokenPricesAtom);
 
   return {
-    asset,
-    balance,
-    dollar,
+    queryKey: ["transparent-tokens", transparentAssets.data, tokenPrices.data],
+    ...queryDependentFn(
+      () =>
+        Promise.resolve(
+          mapNamadaAssetsToTokenBalances(
+            transparentAssets.data!,
+            tokenPrices.data!
+          )
+        ),
+      [transparentAssets, tokenPrices]
+    ),
   };
-};
+});

apps/namadillo/src/atoms/balance/functions.ts

@@ -1,7 +1,10 @@
 import { IbcToken, NativeToken } from "@anomaorg/namada-indexer-client";
 import { Asset, AssetList } from "@chain-registry/types";
+import { mapCoinsToAssets } from "atoms/integrations";
 import BigNumber from "bignumber.js";
+import { DenomTrace } from "cosmjs-types/ibc/applications/transfer/v1/transfer";
 import { namadaAsset } from "registry/namadaAsset";
+import { AddressWithAssetAndBalanceMap } from "types";
 import { TokenBalance } from "./atoms";
 
 // TODO upgrade this function to be as smart as possible
@@ -56,3 +59,59 @@ export const getTotalDollar = (list?: TokenBalance[]): BigNumber | undefined =>
 export const getTotalNam = (list?: TokenBalance[]): BigNumber =>
   list?.find((i) => i.asset.base === namadaAsset.base)?.balance ??
   new BigNumber(0);
+
+const tnamAddressToDenomTrace = (
+  address: string,
+  tokenAddresses: (NativeToken | IbcToken)[]
+): DenomTrace | undefined => {
+  const token = tokenAddresses.find((entry) => entry.address === address);
+  const trace = token && "trace" in token ? token.trace : undefined;
+
+  // If no trace, the token is NAM, but return undefined because we only care
+  // about IBC tokens here
+  if (typeof trace === "undefined") {
+    return undefined;
+  }
+
+  const separatorIndex = trace.lastIndexOf("/");
+
+  if (separatorIndex === -1) {
+    return undefined;
+  }
+
+  return {
+    path: trace.substring(0, separatorIndex),
+    baseDenom: trace.substring(separatorIndex + 1),
+  };
+};
+
+export const mapNamadaAddressesToAssets = async (
+  balances: { address: string; amount: BigNumber }[],
+  tokenAddresses: (NativeToken | IbcToken)[],
+  chainName: string
+): Promise<AddressWithAssetAndBalanceMap> => {
+  const coins = balances.map(({ address, amount }) => ({
+    denom: address,
+    amount: amount.toString(), // TODO: don't convert back to string
+  }));
+
+  return await mapCoinsToAssets(coins, chainName, (tnamAddress) =>
+    Promise.resolve(tnamAddressToDenomTrace(tnamAddress, tokenAddresses))
+  );
+};
+
+export const mapNamadaAssetsToTokenBalances = (
+  assets: AddressWithAssetAndBalanceMap,
+  tokenPrices: Record<string, BigNumber>
+): TokenBalance[] => {
+  return Object.values(assets).map((assetEntry) => {
+    const tokenPrice = tokenPrices[assetEntry.address];
+    const dollar =
+      tokenPrice ? assetEntry.balance.multipliedBy(tokenPrice) : undefined;
+
+    return {
+      ...assetEntry,
+      dollar,
+    };
+  });
+};

apps/namadillo/src/atoms/integrations/atoms.ts

@@ -6,7 +6,11 @@ import { atom } from "jotai";
 import { atomWithMutation, atomWithQuery } from "jotai-tanstack-query";
 import { atomFamily, atomWithStorage } from "jotai/utils";
 import { ChainId, ChainRegistryEntry } from "types";
-import { getKnownChains, mapCoinsToAssets } from "./functions";
+import {
+  getKnownChains,
+  ibcAddressToDenomTrace,
+  mapCoinsToAssets,
+} from "./functions";
 import {
   IbcTransferParams,
   queryAndStoreRpc,
@@ -60,7 +64,11 @@ export const assetBalanceAtomFamily = atomFamily(
       ...queryDependentFn(async () => {
         return await queryAndStoreRpc(chain!, async (rpc: string) => {
           const assetsBalances = await queryAssetBalances(walletAddress!, rpc);
-          return await mapCoinsToAssets(assetsBalances, chain!.chain_name, rpc);
+          return await mapCoinsToAssets(
+            assetsBalances,
+            chain!.chain_name,
+            ibcAddressToDenomTrace(rpc)
+          );
         });
       }, [!!walletAddress, !!chain]),
     }));