From 9643de889ba3c6968fc2e777e50f778fdf569b57 Mon Sep 17 00:00:00 2001
From: Gary Wang <gary@garywang.net>
Date: Wed, 5 Aug 2020 21:40:42 -0700
Subject: [PATCH] Handle token accounts with invalid mints

---
 src/utils/tokens/instructions.js |  4 +++
 src/utils/wallet.js              | 54 +++++++++++++++++++++++++-------
 2 files changed, 46 insertions(+), 12 deletions(-)

diff --git a/src/utils/tokens/instructions.js b/src/utils/tokens/instructions.js
index c64ffbe..364259b 100644
--- a/src/utils/tokens/instructions.js
+++ b/src/utils/tokens/instructions.js
@@ -5,6 +5,10 @@ export const TOKEN_PROGRAM_ID = new PublicKey(
   'TokenSVp5gheXUvJ6jGWGeCsgPKgnE3YgdGKRVCMY9o',
 );
 
+export const WRAPPED_SOL_MINT = new PublicKey(
+  'So11111111111111111111111111111111111111111',
+);
+
 const LAYOUT = BufferLayout.union(BufferLayout.u8('instruction'));
 LAYOUT.addVariant(
   0,
diff --git a/src/utils/wallet.js b/src/utils/wallet.js
index eeca7ba..935817a 100644
--- a/src/utils/wallet.js
+++ b/src/utils/wallet.js
@@ -12,7 +12,7 @@ import {
   getOwnedTokenAccounts,
   transferTokens,
 } from './tokens';
-import { TOKEN_PROGRAM_ID } from './tokens/instructions';
+import { TOKEN_PROGRAM_ID, WRAPPED_SOL_MINT } from './tokens/instructions';
 import {
   ACCOUNT_LAYOUT,
   parseMintData,
@@ -142,30 +142,60 @@ export function useBalanceInfo(publicKey) {
   let [mintInfo, mintInfoLoaded] = useAccountInfo(mint);
   let { name, symbol } = useTokenName(mint);
 
-  if (accountInfoLoaded && mint && mintInfoLoaded) {
-    let { decimals } = parseMintData(mintInfo.data);
+  if (!accountInfoLoaded) {
+    return null;
+  }
+
+  if (mint && mint.equals(WRAPPED_SOL_MINT)) {
     return {
       amount,
-      decimals,
+      decimals: 9,
       mint,
       owner,
-      tokenName: name,
-      tokenSymbol: symbol,
-      initialized: true,
+      tokenName: 'Wrapped SOL',
+      tokenSymbol: 'SOL',
+      valid: true,
     };
-  } else if (accountInfoLoaded && !mint) {
+  }
+
+  if (mint && mintInfoLoaded) {
+    try {
+      let { decimals } = parseMintData(mintInfo.data);
+      return {
+        amount,
+        decimals,
+        mint,
+        owner,
+        tokenName: name,
+        tokenSymbol: symbol,
+        valid: true,
+      };
+    } catch (e) {
+      return {
+        amount,
+        decimals: 0,
+        mint,
+        owner,
+        tokenName: 'Invalid',
+        tokenSymbol: 'INVALID',
+        valid: false,
+      };
+    }
+  }
+
+  if (!mint) {
     return {
       amount: accountInfo?.lamports ?? 0,
       decimals: 9,
       mint: null,
       owner: publicKey,
-      tokenName: 'Solana',
+      tokenName: 'SOL',
       tokenSymbol: 'SOL',
-      initialized: false,
+      valid: true,
     };
-  } else {
-    return null;
   }
+
+  return null;
 }
 
 export function useWalletSelector() {