bridge_ui: unified balance handling

Change-Id: I5d6dd49c7f05fbc7d1f3a579d14c8c0786e63aac

bridge_ui/src/components/KeyAndBalance.tsx

@@ -1,39 +1,20 @@
 import { Typography } from "@material-ui/core";
-import { useEthereumProvider } from "../contexts/EthereumProviderContext";
-import { useSolanaWallet } from "../contexts/SolanaWalletContext";
-import useEthereumBalance from "../hooks/useEthereumBalance";
-import useSolanaBalance from "../hooks/useSolanaBalance";
 import { ChainId, CHAIN_ID_ETH, CHAIN_ID_SOLANA } from "../utils/consts";
 import EthereumSignerKey from "./EthereumSignerKey";
 import SolanaWalletKey from "./SolanaWalletKey";
 
 function KeyAndBalance({
   chainId,
-  tokenAddress,
+  balance,
 }: {
   chainId: ChainId;
-  tokenAddress?: string;
+  balance?: string;
 }) {
-  // TODO: more generic way to get balance
-  const { provider, signerAddress } = useEthereumProvider();
-  const { uiAmountString: ethBalance } = useEthereumBalance(
-    tokenAddress,
-    signerAddress,
-    provider,
-    chainId === CHAIN_ID_ETH
-  );
-  const { wallet: solWallet } = useSolanaWallet();
-  const solPK = solWallet?.publicKey;
-  const { uiAmountString: solBalance } = useSolanaBalance(
-    tokenAddress,
-    solPK,
-    chainId === CHAIN_ID_SOLANA
-  );
   if (chainId === CHAIN_ID_ETH) {
     return (
       <>
         <EthereumSignerKey />
-        <Typography>{ethBalance}</Typography>
+        <Typography>{balance}</Typography>
       </>
     );
   }
@@ -41,7 +22,7 @@ function KeyAndBalance({
     return (
       <>
         <SolanaWalletKey />
-        <Typography>{solBalance}</Typography>
+        <Typography>{balance}</Typography>
       </>
     );
   }

bridge_ui/src/components/Transfer/Send.tsx

@@ -8,13 +8,12 @@ import { useCallback } from "react";
 import { useDispatch, useSelector } from "react-redux";
 import { useEthereumProvider } from "../../contexts/EthereumProviderContext";
 import { useSolanaWallet } from "../../contexts/SolanaWalletContext";
-import useEthereumBalance from "../../hooks/useEthereumBalance";
-import useSolanaBalance from "../../hooks/useSolanaBalance";
 import useWrappedAsset from "../../hooks/useWrappedAsset";
 import {
   selectAmount,
   selectSourceAsset,
   selectSourceChain,
+  selectSourceParsedTokenAccount,
   selectTargetChain,
 } from "../../store/selectors";
 import { setSignedVAAHex } from "../../store/transferSlice";
@@ -51,20 +50,12 @@ function Send() {
   const amount = useSelector(selectAmount);
   const targetChain = useSelector(selectTargetChain);
   const { provider, signer, signerAddress } = useEthereumProvider();
-  const { decimals: ethDecimals, uiAmountString: ethBalance } =
-    useEthereumBalance(
-      sourceAsset,
-      signerAddress,
-      provider,
-      sourceChain === CHAIN_ID_ETH
-    );
   const { wallet } = useSolanaWallet();
   const solPK = wallet?.publicKey;
-  const {
+  const sourceParsedTokenAccount = useSelector(selectSourceParsedTokenAccount);
-    tokenAccount: solTokenPK,
+  const tokenPK = sourceParsedTokenAccount?.publicKey;
-    decimals: solDecimals,
+  const decimals = sourceParsedTokenAccount?.decimals;
-    uiAmount: solBalance,
+  const uiAmountString = sourceParsedTokenAccount?.uiAmountString;
-  } = useSolanaBalance(sourceAsset, solPK, sourceChain === CHAIN_ID_SOLANA);
   const {
     isLoading: isCheckingWrapped,
     // isWrapped,
@@ -88,7 +79,8 @@ function Send() {
       }
       if (
         sourceChain === CHAIN_ID_SOLANA &&
-        attestFrom[sourceChain] === attestFromSolana
+        attestFrom[sourceChain] === attestFromSolana &&
+        decimals
       ) {
         //TODO: just for testing, this should eventually use the store to communicate between steps
         (async () => {
@@ -96,20 +88,21 @@ function Send() {
             wallet,
             solPK?.toString(),
             sourceAsset,
-            solDecimals
+            decimals
           );
           console.log("bytes in transfer", vaaBytes);
         })();
       }
     }
-  }, [sourceChain, provider, signer, wallet, solPK, sourceAsset, solDecimals]);
+  }, [sourceChain, provider, signer, wallet, solPK, sourceAsset, decimals]);
   // TODO: dynamically get "to" wallet
   const handleTransferClick = useCallback(() => {
     // TODO: more generic way of calling these
     if (transferFrom[sourceChain]) {
       if (
         sourceChain === CHAIN_ID_ETH &&
-        transferFrom[sourceChain] === transferFromEth
+        transferFrom[sourceChain] === transferFromEth &&
+        decimals
       ) {
         //TODO: just for testing, this should eventually use the store to communicate between steps
         (async () => {
@@ -117,7 +110,7 @@ function Send() {
             provider,
             signer,
             sourceAsset,
-            ethDecimals,
+            decimals,
             amount,
             targetChain,
             solPK?.toBytes()
@@ -128,17 +121,18 @@ function Send() {
       }
       if (
         sourceChain === CHAIN_ID_SOLANA &&
-        transferFrom[sourceChain] === transferFromSolana
+        transferFrom[sourceChain] === transferFromSolana &&
+        decimals
       ) {
         //TODO: just for testing, this should eventually use the store to communicate between steps
         (async () => {
           const vaaBytes = await transferFromSolana(
             wallet,
             solPK?.toString(),
-            solTokenPK?.toString(),
+            tokenPK,
             sourceAsset,
             amount,
-            solDecimals,
+            decimals,
             signerAddress,
             targetChain
           );
@@ -155,15 +149,15 @@ function Send() {
     signerAddress,
     wallet,
     solPK,
-    solTokenPK,
+    tokenPK,
     sourceAsset,
     amount,
-    ethDecimals,
+    decimals,
-    solDecimals,
     targetChain,
   ]);
   // update this as we develop, just setting expectations with the button state
-  const balance = Number(ethBalance) || solBalance;
+  const hasDecimals = decimals !== undefined;
+  const balance = Number(uiAmountString);
   const isAttestImplemented = !!attestFrom[sourceChain];
   const isTransferImplemented = !!transferFrom[sourceChain];
   const isProviderConnected = !!provider;
@@ -172,11 +166,13 @@ function Send() {
   const isAmountPositive = Number(amount) > 0; // TODO: this needs per-chain, bn parsing
   const isBalanceAtLeastAmount = balance >= Number(amount); // TODO: ditto
   const canAttemptAttest =
+    hasDecimals &&
     isAttestImplemented &&
     isProviderConnected &&
     isRecipientAvailable &&
     isAddressDefined;
   const canAttemptTransfer =
+    hasDecimals &&
     isTransferImplemented &&
     isProviderConnected &&
     isRecipientAvailable &&

bridge_ui/src/components/Transfer/Source.tsx

@@ -1,9 +1,11 @@
 import { Button, makeStyles, MenuItem, TextField } from "@material-ui/core";
 import { useCallback } from "react";
 import { useDispatch, useSelector } from "react-redux";
+import useGetBalanceEffect from "../../hooks/useGetBalanceEffect";
 import {
   selectAmount,
   selectSourceAsset,
+  selectSourceBalanceString,
   selectSourceChain,
 } from "../../store/selectors";
 import {
@@ -24,8 +26,10 @@ const useStyles = makeStyles((theme) => ({
 function Source() {
   const classes = useStyles();
   const dispatch = useDispatch();
+  useGetBalanceEffect();
   const sourceChain = useSelector(selectSourceChain);
   const sourceAsset = useSelector(selectSourceAsset);
+  const uiAmountString = useSelector(selectSourceBalanceString);
   const amount = useSelector(selectAmount);
   const handleSourceChange = useCallback(
     (event) => {
@@ -65,7 +69,7 @@ function Source() {
           </MenuItem>
         ))}
       </TextField>
-      <KeyAndBalance chainId={sourceChain} tokenAddress={sourceAsset} />
+      <KeyAndBalance chainId={sourceChain} balance={uiAmountString} />
       <TextField
         placeholder="Asset"
         fullWidth

bridge_ui/src/hooks/useEthereumBalance.ts

@@ -1,55 +0,0 @@
-import { ethers } from "ethers";
-import { formatUnits } from "ethers/lib/utils";
-import { useEffect, useState } from "react";
-import { TokenImplementation__factory } from "../ethers-contracts";
-
-// TODO: can this be shared with other balances
-export interface Balance {
-  decimals: number;
-  uiAmountString: string;
-}
-
-function createBalance(decimals: number, uiAmountString: string) {
-  return {
-    decimals,
-    uiAmountString,
-  };
-}
-
-function useEthereumBalance(
-  address: string | undefined,
-  ownerAddress: string | undefined,
-  provider: ethers.providers.Web3Provider | undefined,
-  shouldCalculate?: boolean
-) {
-  //TODO: should this check allowance too or subtract allowance?
-  const [balance, setBalance] = useState<Balance>(createBalance(0, ""));
-  useEffect(() => {
-    if (!address || !ownerAddress || !provider || !shouldCalculate) {
-      setBalance(createBalance(0, ""));
-      return;
-    }
-    let cancelled = false;
-    const token = TokenImplementation__factory.connect(address, provider);
-    token
-      .decimals()
-      .then((decimals) => {
-        token.balanceOf(ownerAddress).then((n) => {
-          if (!cancelled) {
-            setBalance(createBalance(decimals, formatUnits(n, decimals)));
-          }
-        });
-      })
-      .catch(() => {
-        if (!cancelled) {
-          setBalance(createBalance(0, ""));
-        }
-      });
-    return () => {
-      cancelled = true;
-    };
-  }, [address, ownerAddress, provider, shouldCalculate]);
-  return balance;
-}
-
-export default useEthereumBalance;

bridge_ui/src/hooks/useGetBalanceEffect.ts

@@ -0,0 +1,112 @@
+import { Connection, PublicKey } from "@solana/web3.js";
+import { formatUnits } from "ethers/lib/utils";
+import { useEffect } from "react";
+import { useDispatch, useSelector } from "react-redux";
+import { useEthereumProvider } from "../contexts/EthereumProviderContext";
+import { useSolanaWallet } from "../contexts/SolanaWalletContext";
+import { TokenImplementation__factory } from "../ethers-contracts";
+import { selectSourceAsset, selectSourceChain } from "../store/selectors";
+import { setSourceParsedTokenAccount } from "../store/transferSlice";
+import { CHAIN_ID_ETH, CHAIN_ID_SOLANA, SOLANA_HOST } from "../utils/consts";
+
+function createParsedTokenAccount(
+  publicKey: PublicKey | undefined,
+  amount: string,
+  decimals: number,
+  uiAmount: number,
+  uiAmountString: string
+) {
+  return {
+    publicKey: publicKey?.toString(),
+    amount,
+    decimals,
+    uiAmount,
+    uiAmountString,
+  };
+}
+
+function useGetBalanceEffect() {
+  const dispatch = useDispatch();
+  const sourceChain = useSelector(selectSourceChain);
+  const sourceAsset = useSelector(selectSourceAsset);
+  const { wallet } = useSolanaWallet();
+  const solPK = wallet?.publicKey;
+  const { provider, signerAddress } = useEthereumProvider();
+  useEffect(() => {
+    // TODO: loading state
+    dispatch(setSourceParsedTokenAccount(undefined));
+    if (!sourceAsset) {
+      return;
+    }
+    let cancelled = false;
+    if (sourceChain === CHAIN_ID_SOLANA && solPK) {
+      let mint;
+      try {
+        mint = new PublicKey(sourceAsset);
+      } catch (e) {
+        return;
+      }
+      const connection = new Connection(SOLANA_HOST, "finalized");
+      connection
+        .getParsedTokenAccountsByOwner(solPK, { mint })
+        .then(({ value }) => {
+          if (!cancelled) {
+            if (value.length) {
+              dispatch(
+                setSourceParsedTokenAccount(
+                  createParsedTokenAccount(
+                    value[0].pubkey,
+                    value[0].account.data.parsed?.info?.tokenAmount?.amount,
+                    value[0].account.data.parsed?.info?.tokenAmount?.decimals,
+                    value[0].account.data.parsed?.info?.tokenAmount?.uiAmount,
+                    value[0].account.data.parsed?.info?.tokenAmount
+                      ?.uiAmountString
+                  )
+                )
+              );
+            } else {
+              // TODO: error state
+            }
+          }
+        })
+        .catch(() => {
+          if (!cancelled) {
+            // TODO: error state
+          }
+        });
+    }
+    if (sourceChain === CHAIN_ID_ETH && provider && signerAddress) {
+      const token = TokenImplementation__factory.connect(sourceAsset, provider);
+      token
+        .decimals()
+        .then((decimals) => {
+          token.balanceOf(signerAddress).then((n) => {
+            if (!cancelled) {
+              dispatch(
+                setSourceParsedTokenAccount(
+                  // TODO: verify accuracy
+                  createParsedTokenAccount(
+                    undefined,
+                    n.toString(),
+                    decimals,
+                    Number(formatUnits(n, decimals)),
+                    formatUnits(n, decimals)
+                  )
+                )
+              );
+            }
+          });
+        })
+        .catch(() => {
+          if (!cancelled) {
+            // TODO: error state
+          }
+        });
+    }
+    return () => {
+      cancelled = true;
+    };
+  }, [dispatch, sourceChain, sourceAsset, solPK, provider, signerAddress]);
+}
+
+export default useGetBalanceEffect;

bridge_ui/src/hooks/useSolanaBalance.ts

@@ -1,83 +0,0 @@
-import { Connection, PublicKey } from "@solana/web3.js";
-import { useEffect, useState } from "react";
-import { SOLANA_HOST } from "../utils/consts";
-
-export interface Balance {
-  tokenAccount: PublicKey | undefined;
-  amount: string;
-  decimals: number;
-  uiAmount: number;
-  uiAmountString: string;
-}
-
-function createBalance(
-  tokenAccount: PublicKey | undefined,
-  amount: string,
-  decimals: number,
-  uiAmount: number,
-  uiAmountString: string
-) {
-  return {
-    tokenAccount,
-    amount,
-    decimals,
-    uiAmount,
-    uiAmountString,
-  };
-}
-
-function useSolanaBalance(
-  tokenAddress: string | undefined,
-  ownerAddress: PublicKey | null | undefined,
-  shouldCalculate?: boolean
-) {
-  //TODO: should connection happen in a context?
-  const [balance, setBalance] = useState<Balance>(
-    createBalance(undefined, "", 0, 0, "")
-  );
-  useEffect(() => {
-    if (!tokenAddress || !ownerAddress || !shouldCalculate) {
-      setBalance(createBalance(undefined, "", 0, 0, ""));
-      return;
-    }
-    let mint;
-    try {
-      mint = new PublicKey(tokenAddress);
-    } catch (e) {
-      setBalance(createBalance(undefined, "", 0, 0, ""));
-      return;
-    }
-    let cancelled = false;
-    const connection = new Connection(SOLANA_HOST, "finalized");
-    connection
-      .getParsedTokenAccountsByOwner(ownerAddress, { mint })
-      .then(({ value }) => {
-        if (!cancelled) {
-          if (value.length) {
-            setBalance(
-              createBalance(
-                value[0].pubkey,
-                value[0].account.data.parsed?.info?.tokenAmount?.amount,
-                value[0].account.data.parsed?.info?.tokenAmount?.decimals,
-                value[0].account.data.parsed?.info?.tokenAmount?.uiAmount,
-                value[0].account.data.parsed?.info?.tokenAmount?.uiAmountString
-              )
-            );
-          } else {
-            setBalance(createBalance(undefined, "0", 0, 0, "0"));
-          }
-        }
-      })
-      .catch(() => {
-        if (!cancelled) {
-          setBalance(createBalance(undefined, "", 0, 0, ""));
-        }
-      });
-    return () => {
-      cancelled = true;
-    };
-  }, [tokenAddress, ownerAddress, shouldCalculate]);
-  return balance;
-}
-
-export default useSolanaBalance;

bridge_ui/src/store/selectors.ts

@@ -5,6 +5,10 @@ export const selectSourceChain = (state: RootState) =>
   state.transfer.sourceChain;
 export const selectSourceAsset = (state: RootState) =>
   state.transfer.sourceAsset;
+export const selectSourceParsedTokenAccount = (state: RootState) =>
+  state.transfer.sourceParsedTokenAccount;
+export const selectSourceBalanceString = (state: RootState) =>
+  state.transfer.sourceParsedTokenAccount?.uiAmountString || "";
 export const selectAmount = (state: RootState) => state.transfer.amount;
 export const selectTargetChain = (state: RootState) =>
   state.transfer.targetChain;

bridge_ui/src/store/transferSlice.ts

@@ -11,10 +11,19 @@ const LAST_STEP = 3;
 
 type Steps = 0 | 1 | 2 | 3;
 
+export interface ParsedTokenAccount {
+  publicKey: string | undefined;
+  amount: string;
+  decimals: number;
+  uiAmount: number;
+  uiAmountString: string;
+}
+
 export interface TransferState {
   activeStep: Steps;
   sourceChain: ChainId;
   sourceAsset: string;
+  sourceParsedTokenAccount: ParsedTokenAccount | undefined;
   amount: string;
   targetChain: ChainId;
   signedVAAHex: string | undefined;
@@ -24,6 +33,7 @@ const initialState: TransferState = {
   activeStep: 0,
   sourceChain: CHAIN_ID_SOLANA,
   sourceAsset: SOL_TEST_TOKEN_ADDRESS,
+  sourceParsedTokenAccount: undefined,
   amount: "",
   targetChain: CHAIN_ID_ETH,
   signedVAAHex: undefined,
@@ -59,6 +69,12 @@ export const transferSlice = createSlice({
     setSourceAsset: (state, action: PayloadAction<string>) => {
       state.sourceAsset = action.payload;
     },
+    setSourceParsedTokenAccount: (
+      state,
+      action: PayloadAction<ParsedTokenAccount | undefined>
+    ) => {
+      state.sourceParsedTokenAccount = action.payload;
+    },
     setAmount: (state, action: PayloadAction<string>) => {
       state.amount = action.payload;
     },
@@ -90,6 +106,7 @@ export const {
   setStep,
   setSourceChain,
   setSourceAsset,
+  setSourceParsedTokenAccount,
   setAmount,
   setTargetChain,
   setSignedVAAHex,