From 339c62a0fcf23be73731d47651b4b60439b7d71c Mon Sep 17 00:00:00 2001
From: Evan Salzbrenner <evan.salzbrenner@gmail.com>
Date: Mon, 14 Oct 2024 22:56:50 +0200
Subject: [PATCH] WalletLink: use getRootKeyForWallet to check if wallet is
 linked to any root (#1106)

User A links wallet X.
User B tries to link wallet X.

Previously, this client check would pass b/c it's comparing X to root
key B, but X is linked to A, and then the transaction would go on to
revert.

Now the client check will fail b/c X is already linked, so no
transaction attempt will be made.
---
 packages/web3/src/v3/WalletLink.ts | 37 +++++++++++++++---------------
 1 file changed, 18 insertions(+), 19 deletions(-)

diff --git a/packages/web3/src/v3/WalletLink.ts b/packages/web3/src/v3/WalletLink.ts
index 284ff4e04..9745388fa 100644
--- a/packages/web3/src/v3/WalletLink.ts
+++ b/packages/web3/src/v3/WalletLink.ts
@@ -25,31 +25,27 @@ export class WalletLink {
         }
     }
 
-    private async assertNotAlreadyLinked(rootKey: ethers.Signer, wallet: ethers.Signer | Address) {
-        const rootKeyAddress = await rootKey.getAddress()
+    public async isLinked(walletAddress: string): Promise<boolean> {
+        const rootKeyAddress = await this.walletLinkShim.read.getRootKeyForWallet(walletAddress)
+
+        return rootKeyAddress !== INVALID_ADDRESS
+    }
+
+    private async assertNotLinked(wallet: ethers.Signer | Address) {
         const walletAddress = typeof wallet === 'string' ? wallet : await wallet.getAddress()
-        const isLinkedAlready = await this.walletLinkShim.read.checkIfLinked(
-            rootKeyAddress,
-            walletAddress,
-        )
 
-        if (isLinkedAlready) {
+        if (await this.isLinked(walletAddress)) {
             throw new WalletAlreadyLinkedError()
         }
 
-        return { rootKeyAddress, walletAddress }
+        return { walletAddress }
     }
 
-    private async assertAlreadyLinked(rootKey: ethers.Signer, walletAddress: string) {
-        const rootKeyAddress = await rootKey.getAddress()
-        const isLinkedAlready = await this.walletLinkShim.read.checkIfLinked(
-            rootKeyAddress,
-            walletAddress,
-        )
-        if (!isLinkedAlready) {
+    private async assertLinked(walletAddress: string) {
+        if (!(await this.isLinked(walletAddress))) {
             throw new WalletNotLinkedError()
         }
-        return { rootKeyAddress, walletAddress }
+        return { walletAddress }
     }
 
     private generateRootKeySignatureForWallet({
@@ -111,7 +107,8 @@ export class WalletLink {
         rootKey: ethers.Signer,
         wallet: ethers.Signer | Address,
     ) {
-        const { rootKeyAddress, walletAddress } = await this.assertNotAlreadyLinked(rootKey, wallet)
+        const { walletAddress } = await this.assertNotLinked(wallet)
+        const rootKeyAddress = await rootKey.getAddress()
 
         const nonce = await this.walletLinkShim.read.getLatestNonceForRootKey(rootKeyAddress)
         const rootKeySignature = await this.generateRootKeySignatureForCallerData({
@@ -134,7 +131,8 @@ export class WalletLink {
         rootKey: ethers.Signer,
         wallet: ethers.Signer,
     ) {
-        const { rootKeyAddress, walletAddress } = await this.assertNotAlreadyLinked(rootKey, wallet)
+        const { walletAddress } = await this.assertNotLinked(wallet)
+        const rootKeyAddress = await rootKey.getAddress()
 
         const nonce = await this.walletLinkShim.read.getLatestNonceForRootKey(rootKeyAddress)
 
@@ -259,7 +257,8 @@ export class WalletLink {
     }
 
     private async generateRemoveLinkData(rootKey: ethers.Signer, walletAddress: string) {
-        const { rootKeyAddress } = await this.assertAlreadyLinked(rootKey, walletAddress)
+        await this.assertLinked(walletAddress)
+        const rootKeyAddress = await rootKey.getAddress()
         const nonce = await this.walletLinkShim.read.getLatestNonceForRootKey(rootKeyAddress)
         const { domain, types, value } = createEip712LinkedWalletdData({
             domain: this.eip712Domain,