From bc1eef23443dcad1e9f0e2ffdaee43fbbca50073 Mon Sep 17 00:00:00 2001
From: Shreerang Kale <shreerangkale@gmail.com>
Date: Fri, 9 Aug 2024 17:00:08 +0530
Subject: [PATCH] Wait for session to load before redirecting

---
 src/context/WalletConnectContext.tsx | 14 +++++++++++++-
 src/pages/LandingPage.tsx            |  2 +-
 src/pages/SignWithNitroKey.tsx       |  8 ++++++--
 src/pages/Validator.tsx              |  8 ++++++--
 4 files changed, 26 insertions(+), 6 deletions(-)

diff --git a/src/context/WalletConnectContext.tsx b/src/context/WalletConnectContext.tsx
index 18b21ba..f4ba36d 100644
--- a/src/context/WalletConnectContext.tsx
+++ b/src/context/WalletConnectContext.tsx
@@ -22,6 +22,7 @@ assert(PROJECT_ID, "Wallet connect project id not provided");
 interface ContextValue {
   connect: () => Promise<void>;
   session: SessionTypes.Struct | null;
+  isSessionLoading: boolean;
   signClient: SignClient | undefined;
   checkPersistedState: (client: SignClient) => Promise<void>;
   disconnect: () => Promise<void>;
@@ -30,6 +31,7 @@ interface ContextValue {
 const walletConnectContext = createContext<ContextValue>({
   connect: () => Promise.resolve(),
   session: null,
+  isSessionLoading: true,
   signClient: undefined,
   checkPersistedState: () => Promise.resolve(),
   disconnect: () => Promise.resolve(),
@@ -47,6 +49,7 @@ export const WalletConnectProvider = ({
 }) => {
   const [signClient, setSignClient] = useState<SignClient>();
   const [session, setSession] = useState<SessionTypes.Struct | null>(null);
+  const [isSessionLoading, setIsSessionLoading] = useState(true);
   const isSignClientInitializing = useRef<boolean>(false);
   const navigate = useNavigate();
 
@@ -59,6 +62,7 @@ export const WalletConnectProvider = ({
     }
 
     setSession(null);
+    setIsSessionLoading(false);
   }, [signClient, session]);
 
   const checkPersistedState = useCallback(async (client: SignClient) => {
@@ -67,6 +71,8 @@ export const WalletConnectProvider = ({
       const session = client.session.get(client.session.keys[lastKeyIndex]);
       setSession(session);
     }
+
+    setIsSessionLoading(false);
   }, []);
 
   const subscribeToEvents = useCallback(
@@ -76,10 +82,12 @@ export const WalletConnectProvider = ({
         const currentSession = client.session.get(topic);
         const updatedSession = { ...currentSession, namespaces };
         setSession(updatedSession);
+        setIsSessionLoading(false);
       });
 
       client.on("session_delete", () => {
         setSession(null);
+        setIsSessionLoading(false);
         navigate("/");
       });
     },
@@ -133,12 +141,14 @@ export const WalletConnectProvider = ({
       try {
         const session = await approval();
         setSession(session);
+        setIsSessionLoading(false);
       } catch (error) {
         enqueueSnackbar("User rejected pairing request", { variant: "error" });
         console.log(error);
       }
       web3Modal.closeModal();
     }
+
   };
 
   useEffect(() => {
@@ -152,6 +162,7 @@ export const WalletConnectProvider = ({
       value={{
         connect,
         session,
+        isSessionLoading,
         signClient,
         checkPersistedState,
         disconnect,
@@ -163,12 +174,13 @@ export const WalletConnectProvider = ({
 };
 
 export const useWalletConnectContext = () => {
-  const { connect, session, signClient, checkPersistedState, disconnect } =
+  const { connect, session, signClient, isSessionLoading, checkPersistedState, disconnect } =
     useContext(walletConnectContext);
 
   return {
     connect,
     session,
+    isSessionLoading,
     signClient,
     checkPersistedState,
     disconnect,
diff --git a/src/pages/LandingPage.tsx b/src/pages/LandingPage.tsx
index 4ec78e6..4fa53df 100644
--- a/src/pages/LandingPage.tsx
+++ b/src/pages/LandingPage.tsx
@@ -37,7 +37,7 @@ const LandingPage = () => {
         </Typography>
       </Box>
       <TermsAndConditionsBox height="43vh" onLoad={()=>{setIsDisabled(false);}} />
-      <Box mt={2} display="flex" justifyContent="center">
+      <Box m={2} display="flex" justifyContent="center">
         <Button variant="contained" color="primary" onClick={handleAccept} disabled={isDisabled}>
           Accept
         </Button>
diff --git a/src/pages/SignWithNitroKey.tsx b/src/pages/SignWithNitroKey.tsx
index b03b431..740090b 100644
--- a/src/pages/SignWithNitroKey.tsx
+++ b/src/pages/SignWithNitroKey.tsx
@@ -17,19 +17,23 @@ import { ENABLE_KYC, HASHED_SUBSCRIBER_ID_KEY } from "../constants";
 
 const SignWithNitroKey = () => {
 
-  const { session, signClient } =
+  const { session, signClient, isSessionLoading } =
     useWalletConnectContext();
 
   const navigate = useNavigate();
   const location = useLocation();
 
   useEffect(() => {
+    if (isSessionLoading) {
+      return;
+    }
+
     if (!session) {
       navigate("/connect-wallet?redirectTo=sign-with-nitro-key", {
         state: location.state,
       });
     }
-  }, [session, navigate, location.state]);
+  }, [session, isSessionLoading, navigate, location.state]);
 
   const [ethAddress, setEthAddress] = useState("");
   const [ethSignature, setEthSignature] = useState("");
diff --git a/src/pages/Validator.tsx b/src/pages/Validator.tsx
index 1977e46..81c8c99 100644
--- a/src/pages/Validator.tsx
+++ b/src/pages/Validator.tsx
@@ -13,7 +13,7 @@ import { useWalletConnectContext } from '../context/WalletConnectContext';
 import { Participant } from '../types';
 
 const Validator = () => {
-  const { session, signClient } = useWalletConnectContext();
+  const { session, signClient, isSessionLoading } = useWalletConnectContext();
   const navigate = useNavigate();
 
   const [cosmosAddress, setCosmosAddress] = useState('');
@@ -24,10 +24,14 @@ const Validator = () => {
   const [isError, setIsError] = useState(false);
 
   useEffect(() => {
+    if (isSessionLoading) {
+      return;
+    }
+
     if (!session) {
       navigate("/connect-wallet?redirectTo=validator");
     }
-  }, [session, navigate]);
+  }, [session, navigate, isSessionLoading]);
 
   useEffect(() => {
     if (!cosmosAddress) {