Persist wallet in localstorage (#1447)

Author: jeesunikim

src/components/WalletKit/ConnectWallet.tsx

@@ -1,4 +1,4 @@
-import { useContext, useState } from "react";
+import { useContext, useEffect, useState } from "react";
 import { Button, Modal, Text } from "@stellar/design-system";
 import { ISupportedWallet } from "@creit.tech/stellar-wallets-kit";
 import { useStore } from "@/store/useStore";
@@ -7,6 +7,7 @@ import { useAccountInfo } from "@/query/useAccountInfo";
 
 import { shortenStellarAddress } from "@/helpers/shortenStellarAddress";
 import { getNetworkHeaders } from "@/helpers/getNetworkHeaders";
+import { localStorageSavedWallet } from "@/helpers/localStorageSavedWallet";
 
 import { ConnectedModal } from "@/components/WalletKit/ConnectedModal";
 import { WalletKitContext } from "@/components/WalletKit/WalletKitContextProvider";
@@ -15,56 +16,109 @@ import { trackEvent, TrackingEvent } from "@/metrics/tracking";
 
 export const ConnectWallet = () => {
   const { network, walletKit, updateWalletKit } = useStore();
+  const [connected, setConnected] = useState<boolean>(false);
   const [isModalVisible, setShowModal] = useState(false);
   const [errorMessageOnConnect, setErrorMessageOnConnect] = useState("");
   const walletKitInstance = useContext(WalletKitContext);
+  const savedWallet = localStorageSavedWallet.get();
 
   const { data: accountInfo, refetch: fetchAccountInfo } = useAccountInfo({
     publicKey: walletKit?.publicKey || "",
     horizonUrl: network?.horizonUrl || "",
     headers: network ? getNetworkHeaders(network, "horizon") : {},
   });
 
-  const resetWalletKit = () => {
+  const disconnect = () => {
     updateWalletKit({
       publicKey: undefined,
       walletType: undefined,
     });
 
     setShowModal(false);
+    setConnected(false);
+    localStorageSavedWallet.remove();
   };
 
+  useEffect(() => {
+    let t: NodeJS.Timeout;
+
+    if (
+      !connected &&
+      !!savedWallet?.id &&
+      ![undefined, "false", "wallet_connect"].includes(savedWallet?.id) &&
+      savedWallet.network.id === network.id
+    ) {
+      t = setTimeout(() => {
+        walletKitInstance.walletKit?.setWallet(savedWallet.id);
+        handleSetWalletAddress();
+        clearTimeout(t);
+      }, 750);
+    }
+
+    return () => {
+      clearTimeout(t);
+    };
+  }, [savedWallet?.id, connected, walletKitInstance]);
+
+  async function handleSetWalletAddress(): Promise<boolean> {
+    try {
+      const addressResult = await walletKitInstance.walletKit?.getAddress();
+
+      if (!addressResult?.address) {
+        return false;
+      }
+      const publicKey = addressResult.address;
+
+      if (!publicKey) {
+        return false;
+      }
+
+      updateWalletKit({
+        publicKey,
+        walletType: savedWallet?.id,
+      });
+      setConnected(true);
+
+      return true;
+      // eslint-disable-next-line @typescript-eslint/no-unused-vars
+    } catch (e) {
+      return false;
+    }
+  }
+
   const connectWallet = async () => {
     try {
       await walletKitInstance.walletKit?.openModal({
         onWalletSelected: async (option: ISupportedWallet) => {
-          try {
-            walletKitInstance.walletKit?.setWallet(option.id);
-            const addressResult =
-              await walletKitInstance.walletKit?.getAddress();
-
-            if (addressResult?.address) {
-              updateWalletKit({
-                publicKey: addressResult.address,
-                walletType: option.id,
-              });
-
-              trackEvent(TrackingEvent.WALLET_KIT_SELECTED, {
-                walletType: option.id,
-              });
-            }
-          } catch (e: unknown) {
-            // Ledger sends a message with the error code, so we need to check for that
-            const errorMessage =
-              (e as { message?: string })?.message || "Unknown error occurred";
+          walletKitInstance.walletKit?.setWallet(option.id);
+          const isWalletConnected = await handleSetWalletAddress();
+
+          if (!isWalletConnected) {
+            const errorMessage = "Unable to load wallet information";
             setErrorMessageOnConnect(errorMessage);
-            resetWalletKit();
+            disconnect();
+            return;
           }
+
+          localStorageSavedWallet.set({
+            id: option.id,
+            network: {
+              id: network.id,
+              label: network.label,
+            },
+          });
+
+          trackEvent(TrackingEvent.WALLET_KIT_SELECTED, {
+            walletType: option.id,
+          });
         },
       });
       // eslint-disable-next-line @typescript-eslint/no-unused-vars
     } catch (e) {
-      resetWalletKit();
+      const errorMessage =
+        (e as { message?: string })?.message || "Unknown error occurred";
+      setErrorMessageOnConnect(errorMessage);
+      disconnect();
     }
   };
 
@@ -82,7 +136,7 @@ export const ConnectWallet = () => {
         isVisible={isModalVisible}
         showModal={setShowModal}
         publicKey={walletKit?.publicKey || ""}
-        onDisconnect={resetWalletKit}
+        onDisconnect={disconnect}
         balance={
           xlmBalance?.balance
             ? `${xlmBalance.balance} XLM`

src/components/WalletKit/WalletKitContextProvider.tsx

@@ -1,6 +1,6 @@
 "use client";
 
-import { createContext, useEffect, useMemo } from "react";
+import { createContext, useEffect, useMemo, useState } from "react";
 import { useStore } from "@/store/useStore";
 
 import {
@@ -17,9 +17,12 @@ import { LedgerModule } from "@creit.tech/stellar-wallets-kit/modules/ledger.mod
 
 import { getWalletKitNetwork } from "@/helpers/getWalletKitNetwork";
 import { localStorageSavedTheme } from "@/helpers/localStorageSavedTheme";
+import { localStorageSavedWallet } from "@/helpers/localStorageSavedWallet";
+import { SavedWallet } from "@/types/types";
 
 type WalletKitProps = {
   walletKit?: StellarWalletsKit;
+  walletId?: string;
 };
 
 export const WalletKitContext = createContext<WalletKitProps>({
@@ -32,6 +35,7 @@ export const WalletKitContextProvider = ({
   children: React.ReactNode;
 }) => {
   const { network, theme, setTheme } = useStore();
+  const [savedWallet, setSavedWallet] = useState<SavedWallet | null>(null);
   const networkType = getWalletKitNetwork(network.id);
 
   useEffect(() => {
@@ -44,6 +48,15 @@ export const WalletKitContextProvider = ({
     // eslint-disable-next-line react-hooks/exhaustive-deps
   }, []);
 
+  useEffect(() => {
+    const savedWallet = localStorageSavedWallet.get();
+
+    if (savedWallet && savedWallet.network.id === network.id) {
+      setSavedWallet(savedWallet);
+    }
+    // eslint-disable-next-line react-hooks/exhaustive-deps
+  }, []);
+
   const walletKitInstance = useMemo(() => {
     const isDarkTheme = theme === "sds-theme-dark";
 
@@ -111,7 +124,12 @@ export const WalletKitContextProvider = ({
   }, [networkType, theme]);
 
   return (
-    <WalletKitContext.Provider value={{ walletKit: walletKitInstance }}>
+    <WalletKitContext.Provider
+      value={{
+        walletKit: walletKitInstance,
+        walletId: savedWallet?.id,
+      }}
+    >
       {children}
     </WalletKitContext.Provider>
   );

src/constants/settings.ts

@@ -16,6 +16,7 @@ export const LOCAL_STORAGE_SAVED_CONTRACTS = "stellar_lab_saved_contract_ids";
 export const LOCAL_STORAGE_SAVED_THEME = "stellarTheme:Laboratory";
 export const LOCAL_STORAGE_SAVED_EXPLORER_TRANSACTIONS =
   "stellar_lab_saved_explorer_transactions";
+export const LOCAL_STORAGE_SAVED_WALLET = "stellar_lab_wallet";
 
 export const XDR_TYPE_TRANSACTION_ENVELOPE = "TransactionEnvelope";
 

src/helpers/localStorageSavedWallet.ts

@@ -0,0 +1,20 @@
+import { LOCAL_STORAGE_SAVED_WALLET } from "@/constants/settings";
+import { SavedWallet } from "@/types/types";
+
+export const localStorageSavedWallet = {
+  get: () => {
+    const savedWalletString = localStorage.getItem(LOCAL_STORAGE_SAVED_WALLET);
+    return savedWalletString
+      ? (JSON.parse(savedWalletString) as SavedWallet)
+      : null;
+  },
+  set: (savedWallet: SavedWallet) => {
+    return localStorage.setItem(
+      LOCAL_STORAGE_SAVED_WALLET,
+      JSON.stringify(savedWallet),
+    );
+  },
+  remove: () => {
+    return localStorage.removeItem(LOCAL_STORAGE_SAVED_WALLET);
+  },
+};

src/types/types.ts

@@ -396,6 +396,11 @@ export interface SavedContract extends LocalStorageSavedItem {
   shareableUrl: string;
 }
 
+export interface SavedWallet {
+  id: string;
+  network: LocalStorageSavedNetwork;
+}
+
 export type SavedTransactionPage = "build" | "sign" | "simulate" | "submit";
 
 export interface SavedEndpointHorizon extends LocalStorageSavedItem {