fix(native): remove device share to reset state if private key reconstruction fails

Author: gregfromstl

packages/thirdweb/src/wallets/in-app/native/helpers/auth/middleware.ts

@@ -81,7 +81,7 @@ export async function postAuth({
   return storedToken;
 }
 
-async function getRecoveryCode(args: {
+export async function getRecoveryCode(args: {
   storedToken: AuthStoredTokenWithCookieReturnType["storedToken"];
   client: ThirdwebClient;
   storage: ClientScopedStorage;

packages/thirdweb/src/wallets/in-app/native/helpers/constants.ts

@@ -6,6 +6,8 @@ export const AUTH_SHARE_INDEX = AUTH_SHARE_ID - 1;
 const DEVICE_SHARE_ID = 1;
 export const DEVICE_SHARE_INDEX = DEVICE_SHARE_ID - 1;
 export const DEVICE_SHARE_MISSING_MESSAGE = "Missing device share.";
+export const INVALID_DEVICE_SHARE_MESSAGE =
+  "Invalid private key reconstructed from shares";
 
 const RECOVERY_SHARE_ID = 2;
 export const RECOVERY_SHARE_INDEX = RECOVERY_SHARE_ID - 1;

packages/thirdweb/src/wallets/in-app/native/helpers/storage/local.ts

@@ -112,6 +112,22 @@ export async function setDeviceShare({
   return deviceShare;
 }
 
+export async function removeDeviceShare({
+  clientId,
+}: {
+  clientId: string;
+}): Promise<void> {
+  const userDetails = await getWalletUserDetails(clientId);
+
+  if (!userDetails) {
+    throw new Error("Missing wallet user ID");
+  }
+
+  const name = DEVICE_SHARE_LOCAL_STORAGE_NAME(clientId, userDetails.userId);
+  await removeItemInAsyncStorage(name);
+  return;
+}
+
 export async function getDeviceShare(clientId: string) {
   const cachedWalletUserId = await getWalletUserDetails(clientId);
   if (!cachedWalletUserId) {

packages/thirdweb/src/wallets/in-app/native/helpers/wallet/retrieval.ts

@@ -11,9 +11,10 @@ import type { SetUpWalletRpcReturnType } from "../../../core/authentication/type
 import { getUserShares } from "../api/fetchers.js";
 import {
   DEVICE_SHARE_MISSING_MESSAGE,
+  INVALID_DEVICE_SHARE_MESSAGE,
   ROUTE_GET_USER_SHARES,
 } from "../constants.js";
-import { getDeviceShare } from "../storage/local.js";
+import { getDeviceShare, removeDeviceShare } from "../storage/local.js";
 import { storeShares } from "./creation.js";
 import { decryptShareWeb } from "./encryption.js";
 
@@ -47,7 +48,7 @@ async function getWalletPrivateKeyFromShares(shares: string[]) {
   }
   const prefixPrivateKey = hexToString(privateKeyHex as Hex);
   if (!prefixPrivateKey.startsWith("thirdweb_")) {
-    throw new Error("Invalid private key reconstructed from shares");
+    throw new Error(INVALID_DEVICE_SHARE_MESSAGE);
   }
   const privateKey = prefixPrivateKey.replace("thirdweb_", "");
   return privateKey;
@@ -58,9 +59,24 @@ async function getAccountFromShares(args: {
   shares: string[];
 }): Promise<Account> {
   const { client, shares } = args;
+  const privateKey = await (async () => {
+    try {
+      return await getWalletPrivateKeyFromShares(shares);
+    } catch (e) {
+      // If the private key reconstruction fails, try to reset the device share and prompt the user to try again
+      // This can happen if a user's account has been migrated or otherwise modified in the backend to use a new wallet. In that case, we need to reset their device state to get a new share
+      if (e instanceof Error && e.message === INVALID_DEVICE_SHARE_MESSAGE) {
+        await removeDeviceShare({ clientId: client.clientId });
+        throw new Error("Invalid device state, please try again.");
+      }
+      // Otherwise this is a legitimate error, throw it
+      throw e;
+    }
+  })();
+
   return privateKeyToAccount({
     client,
-    privateKey: await getWalletPrivateKeyFromShares(shares),
+    privateKey,
   });
 }
 
@@ -84,13 +100,13 @@ async function getShares<
   client: ThirdwebClient;
   authShare: { toRetrieve: A };
   recoveryShare: R extends true
-    ? {
-        toRetrieve: R;
-        recoveryCode: string;
-      }
-    : {
-        toRetrieve: R;
-      };
+  ? {
+    toRetrieve: R;
+    recoveryCode: string;
+  }
+  : {
+    toRetrieve: R;
+  };
   deviceShare: { toRetrieve: D };
   storage: ClientScopedStorage;
 }): Promise<{
@@ -147,7 +163,7 @@ async function getShares<
         }
         // if we get here, decryption was successful, so we stop trying
         break;
-      } catch {}
+      } catch { }
     }
     if (!recoverShareToReturn) {
       throw new Error("Invalid recovery code.");
@@ -179,6 +195,7 @@ async function getShares<
 async function getAccountAddressFromShares(args: {
   client: ThirdwebClient;
   shares: string[];
+  storage: ClientScopedStorage;
 }) {
   const wallet = await getAccountFromShares(args);
   return wallet.address;
@@ -205,6 +222,7 @@ export async function setUpShareForNewDevice({
   const walletAddress = await getAccountAddressFromShares({
     client,
     shares: [recoveryShare, authShare],
+    storage,
   });
 
   const maybeDeviceShare = await storeShares({