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/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

@@ -13,7 +13,7 @@ import {
   DEVICE_SHARE_MISSING_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";
 
@@ -58,9 +58,27 @@ 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 private key reconstructed from shares"
+      ) {
+        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,
   });
 }
 
@@ -179,6 +197,7 @@ async function getShares<
 async function getAccountAddressFromShares(args: {
   client: ThirdwebClient;
   shares: string[];
+  storage: ClientScopedStorage;
 }) {
   const wallet = await getAccountFromShares(args);
   return wallet.address;
@@ -205,6 +224,7 @@ export async function setUpShareForNewDevice({
   const walletAddress = await getAccountAddressFromShares({
     client,
     shares: [recoveryShare, authShare],
+    storage,
   });
 
   const maybeDeviceShare = await storeShares({