From 544b840ede3c9cc0f2b7352b77e77050a0c15fee Mon Sep 17 00:00:00 2001
From: tomiir <rocchitomas@gmail.com>
Date: Fri, 2 Feb 2024 03:55:16 -0600
Subject: [PATCH] fix: smart account creation on session requests (#442)

---
 .../src/components/SmartAccountCard.tsx       | 19 +++++-----
 .../src/hooks/useSmartAccount.ts              |  6 ++--
 .../src/lib/SmartAccountLib.ts                | 17 ++++++---
 .../src/utils/EIP155RequestHandlerUtil.ts     | 25 ++++++-------
 .../src/views/SessionProposalModal.tsx        | 36 ++++++++++++-------
 5 files changed, 61 insertions(+), 42 deletions(-)

diff --git a/advanced/wallets/react-wallet-v2/src/components/SmartAccountCard.tsx b/advanced/wallets/react-wallet-v2/src/components/SmartAccountCard.tsx
index d2e9a3c..dc71677 100644
--- a/advanced/wallets/react-wallet-v2/src/components/SmartAccountCard.tsx
+++ b/advanced/wallets/react-wallet-v2/src/components/SmartAccountCard.tsx
@@ -100,7 +100,6 @@ export default function SmartAccountCard({
       >
         {activeChainId === chainId ? `✅` : `🔄`}
       </Button>
-
       {smartAccountAddress ? (
         <>
           <Text h5 css={{ marginTop: 20 }}>
@@ -109,15 +108,17 @@ export default function SmartAccountCard({
           <Text small css={{ marginTop: 5 }}>
             {smartAccountAddress}
           </Text>
-          <Button
-            size="md"
-            css={{ marginTop: 20, width: '100%' }}
-            onClick={sendTestTransaction}
-            disabled={!isActiveChain || loading}
-          >
-            {loading ? <Loading size="sm" /> : 'Send Test TX'}
-          </Button>
         </>
+      ) : null}
+      {isDeployed && smartAccountAddress ? (
+        <Button
+          size="md"
+          css={{ marginTop: 20, width: '100%' }}
+          onClick={sendTestTransaction}
+          disabled={!isActiveChain || loading}
+        >
+          {loading ? <Loading size="sm" /> : 'Send Test TX'}
+        </Button>
       ) : (
         <>
           <Button
diff --git a/advanced/wallets/react-wallet-v2/src/hooks/useSmartAccount.ts b/advanced/wallets/react-wallet-v2/src/hooks/useSmartAccount.ts
index 38fce46..f8e4d08 100644
--- a/advanced/wallets/react-wallet-v2/src/hooks/useSmartAccount.ts
+++ b/advanced/wallets/react-wallet-v2/src/hooks/useSmartAccount.ts
@@ -58,9 +58,9 @@ export default function useSmartAccount(signerPrivateKey: Hex, chain: Chain) {
     }, [signerPrivateKey, smartAccountSponsorshipEnabled, chain])
 
     useEffect(() => {
-        client?.checkIfSmartAccountDeployed()
-            .then((deployed: boolean) => {
-                setIsDeployed(deployed)
+        client?.init()
+            .then(() => {
+                setIsDeployed(client?.isDeployed)
                 setAddress(client?.address)
             })
     }, [client, chain])
diff --git a/advanced/wallets/react-wallet-v2/src/lib/SmartAccountLib.ts b/advanced/wallets/react-wallet-v2/src/lib/SmartAccountLib.ts
index 960f3a9..31dcf9b 100644
--- a/advanced/wallets/react-wallet-v2/src/lib/SmartAccountLib.ts
+++ b/advanced/wallets/react-wallet-v2/src/lib/SmartAccountLib.ts
@@ -162,8 +162,13 @@ export class SmartAccountLib {
   }
 
   // -- Public ------------------------------------------------------------------
-  public getAccount = async () =>
-    privateKeyToSafeSmartAccount(this.publicClient, {
+  public async init() {
+    await this.checkIfSmartAccountDeployed()
+    this.address = (await this.getAccount()).address
+  }
+
+  public getAccount = async () => {
+    const account = await privateKeyToSafeSmartAccount(this.publicClient, {
       privateKey: this.#signerPrivateKey,
       safeVersion: '1.4.1', // simple version
       entryPoint: ENTRYPOINT_ADDRESSES[this.chain.name], // global entrypoint
@@ -178,6 +183,10 @@ export class SmartAccountLib {
         }
       ]
     })
+    this.address = account.address
+    
+    return account;
+  }
 
   static isSmartAccount = async (address: Address, chain: Chain) => {
     const client = createPublicClient({
@@ -263,9 +272,7 @@ export class SmartAccountLib {
     this.isDeployed = Boolean(bytecode)
     
     console.log(`Smart Account Deployed: ${this.isDeployed}`)
-    if (this.isDeployed) {
-      this.address = smartAccountClient.account.address
-    }
+  
     return this.isDeployed;
 }
 
diff --git a/advanced/wallets/react-wallet-v2/src/utils/EIP155RequestHandlerUtil.ts b/advanced/wallets/react-wallet-v2/src/utils/EIP155RequestHandlerUtil.ts
index 6cc20b5..93c494e 100644
--- a/advanced/wallets/react-wallet-v2/src/utils/EIP155RequestHandlerUtil.ts
+++ b/advanced/wallets/react-wallet-v2/src/utils/EIP155RequestHandlerUtil.ts
@@ -12,33 +12,34 @@ import { SignClientTypes } from '@walletconnect/types'
 import { getSdkError } from '@walletconnect/utils'
 import { providers } from 'ethers'
 import { Hex } from 'viem'
-import { allowedChains } from './SmartAccountUtils'
+import { Chain, allowedChains } from './SmartAccountUtils'
 type RequestEventArgs = Omit<SignClientTypes.EventArguments['session_request'], 'verifyContext'>
 
 
 const getWallet = async (params: any) => {
-  const requestParams: Array<any> = params?.request?.params || []
+  console.log('get wallet params', params)
+  const chainId = params?.chainId?.split(':')[1]
+  console.log('chain id', chainId)
   const eoaWallet = eip155Wallets[getWalletAddressFromParams(eip155Addresses, params)]
   if (eoaWallet) {
     return eoaWallet
   }
 
-  const deployedSmartAccounts = await Promise.all(Object.values(eip155Wallets).map(async (wallet) => {
+  const smartAccountEnabledChain = allowedChains.find((chain) => chain.id.toString() === chainId) as Chain
+  console.log('smart account enabled chain', smartAccountEnabledChain)
+  const smartAccounts = await Promise.all(Object.values(eip155Wallets).map(async (wallet) => {
     const smartAccount = new SmartAccountLib({
       privateKey: wallet.getPrivateKey() as Hex,
-      chain: allowedChains[0], // TODO: FIX FOR MULTI NETWORK
+      chain: smartAccountEnabledChain,
       sponsored: true, // TODO: Sponsor for now but should be dynamic according to SettingsStore
     })
-    const isDeployed = await smartAccount.checkIfSmartAccountDeployed()
-    if (isDeployed) {
-      return smartAccount
-    }
-    return null
+    await smartAccount.init()
+    return smartAccount
   }));
-  const validSmartAccounts = deployedSmartAccounts.filter(Boolean) as Array<SmartAccountLib>
-  const smartAccountAddress = getWalletAddressFromParams(validSmartAccounts.map(acc => acc.address!), params)
 
-  return validSmartAccounts.find((smartAccount) => smartAccount?.address === smartAccountAddress) as SmartAccountLib
+  const smartAccountAddress = getWalletAddressFromParams(smartAccounts.map(acc => acc.address!), params)
+
+  return smartAccounts.find((smartAccount) => smartAccount?.address === smartAccountAddress) as SmartAccountLib
 }
 
 
diff --git a/advanced/wallets/react-wallet-v2/src/views/SessionProposalModal.tsx b/advanced/wallets/react-wallet-v2/src/views/SessionProposalModal.tsx
index e29b6ea..9f04afe 100644
--- a/advanced/wallets/react-wallet-v2/src/views/SessionProposalModal.tsx
+++ b/advanced/wallets/react-wallet-v2/src/views/SessionProposalModal.tsx
@@ -36,7 +36,7 @@ import { Hex } from 'viem'
 import ChainSmartAddressMini from '@/components/ChainSmartAddressMini'
 import { useSnapshot } from 'valtio'
 import SettingsStore from '@/store/SettingsStore'
-import { allowedChains } from '@/utils/SmartAccountUtils'
+import { Chain, allowedChains } from '@/utils/SmartAccountUtils'
 
 const StyledText = styled(Text, {
   fontWeight: 400
@@ -236,18 +236,28 @@ export default function SessionProposalModal() {
       })
 
       // TODO: improve for multi network
-      console.log('Checking if SA is deployed', eip155Wallets[eip155Addresses[0]])
-      const chainId = namespaces['eip155'].chains?.[0]
-      const smartAccountClient = new SmartAccountLib({
-        privateKey: eip155Wallets[eip155Addresses[0]].getPrivateKey() as Hex,
-        chain: allowedChains.find(chain => chain.id.toString() === chainId)!,
-        sponsored: smartAccountSponsorshipEnabled,
-      })
-      const isDeployed = await smartAccountClient.checkIfSmartAccountDeployed()
-      console.log('isDeployed', isDeployed)
-
-      if (isDeployed) {
-        namespaces.eip155.accounts = [...namespaces.eip155.accounts, `eip155:5:${smartAccountClient.address}`]
+      console.log('namespaces', namespaces['eip155'])
+      const namespaceChains = namespaces['eip155'].chains?.map((c: string) => c.split(':')[1])
+      const smartAccountEnabledChains: Chain[] = allowedChains.filter(chain => namespaceChains?.includes(chain.id.toString()))
+      // We find a request for a chain that is enabled for smart account
+      if (smartAccountEnabledChains.length) {
+        const signerAddress = namespaces['eip155'].accounts[0].split(':')[2]
+        const wallet = eip155Wallets[signerAddress]
+        const chain = smartAccountEnabledChains[0]
+        if (wallet) {
+          const smartAccountClient = new SmartAccountLib({
+            privateKey: wallet.getPrivateKey() as Hex,
+            chain,
+            sponsored: smartAccountSponsorshipEnabled,
+          })
+          await smartAccountClient.init()
+          const isDeployed = await smartAccountClient.checkIfSmartAccountDeployed()
+          console.log('isDeployed', isDeployed, smartAccountClient.address)
+    
+          if (isDeployed) {
+            namespaces.eip155.accounts = [...namespaces.eip155.accounts, `eip155:${chain.id}:${smartAccountClient.address}`]
+          }
+        }
       }
 
       console.log('approving namespaces:', namespaces)