fix(native): fallback to recreate device share if private key recovery 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/wallet/retrieval.ts

@@ -7,8 +7,12 @@ import {
 import type { Account } from "../../../../interfaces/wallet.js";
 import { privateKeyToAccount } from "../../../../private-key.js";
 import type { ClientScopedStorage } from "../../../core/authentication/client-scoped-storage.js";
-import type { SetUpWalletRpcReturnType } from "../../../core/authentication/types.js";
+import type {
+  AuthStoredTokenWithCookieReturnType,
+  SetUpWalletRpcReturnType,
+} from "../../../core/authentication/types.js";
 import { getUserShares } from "../api/fetchers.js";
+import { getRecoveryCode } from "../auth/middleware.js";
 import {
   DEVICE_SHARE_MISSING_MESSAGE,
   ROUTE_GET_USER_SHARES,
@@ -24,8 +28,9 @@ import { decryptShareWeb } from "./encryption.js";
 export async function getExistingUserAccount(args: {
   client: ThirdwebClient;
   storage: ClientScopedStorage;
+  storedToken?: AuthStoredTokenWithCookieReturnType["storedToken"];
 }) {
-  const { client, storage } = args;
+  const { client, storage, storedToken } = args;
   const { authShare, deviceShare } = await getShares({
     client,
     authShare: { toRetrieve: true },
@@ -36,6 +41,8 @@ export async function getExistingUserAccount(args: {
   return getAccountFromShares({
     client,
     shares: [authShare, deviceShare],
+    storage,
+    storedToken,
   });
 }
 
@@ -56,11 +63,34 @@ async function getWalletPrivateKeyFromShares(shares: string[]) {
 async function getAccountFromShares(args: {
   client: ThirdwebClient;
   shares: string[];
+  storage: ClientScopedStorage;
+  storedToken?: AuthStoredTokenWithCookieReturnType["storedToken"];
 }): Promise<Account> {
-  const { client, shares } = args;
+  const { client, shares, storage, storedToken } = args;
+  const privateKey = await (async () => {
+    try {
+      return await getWalletPrivateKeyFromShares(shares);
+    } catch (e) {
+      // If the private key reconstruction fails, try to reset the device share
+      if (storedToken) {
+        await setUpShareForNewDevice({
+          client,
+          recoveryCode: await getRecoveryCode({
+            storedToken,
+            client,
+            storage,
+          }),
+          storage,
+        });
+        return await getWalletPrivateKeyFromShares(shares);
+      }
+      throw e;
+    }
+  })();
+
   return privateKeyToAccount({
     client,
-    privateKey: await getWalletPrivateKeyFromShares(shares),
+    privateKey,
   });
 }
 
@@ -179,6 +209,7 @@ async function getShares<
 async function getAccountAddressFromShares(args: {
   client: ThirdwebClient;
   shares: string[];
+  storage: ClientScopedStorage;
 }) {
   const wallet = await getAccountFromShares(args);
   return wallet.address;
@@ -205,6 +236,7 @@ export async function setUpShareForNewDevice({
   const walletAddress = await getAccountAddressFromShares({
     client,
     shares: [recoveryShare, authShare],
+    storage,
   });
 
   const maybeDeviceShare = await storeShares({

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

@@ -3,6 +3,7 @@ import type { Account } from "../../../../interfaces/wallet.js";
 import type { ClientScopedStorage } from "../../../core/authentication/client-scoped-storage.js";
 import type {
   AuthResultAndRecoveryCode,
+  AuthStoredTokenWithCookieReturnType,
   GetUser,
 } from "../../../core/authentication/types.js";
 import type { IWebWallet } from "../../../core/wallet/web-wallet.js";
@@ -14,6 +15,7 @@ import { getExistingUserAccount } from "./retrieval.js";
 export class ShardedWallet implements IWebWallet {
   private client: ThirdwebClient;
   private storage: ClientScopedStorage;
+  private storedToken?: AuthStoredTokenWithCookieReturnType["storedToken"];
 
   constructor(args: {
     client: ThirdwebClient;
@@ -24,6 +26,7 @@ export class ShardedWallet implements IWebWallet {
   }
 
   async postWalletSetUp(authResult: AuthResultAndRecoveryCode): Promise<void> {
+    this.storedToken = authResult.storedToken;
     await postAuth({
       storedToken: authResult.storedToken,
       client: this.client,
@@ -67,6 +70,7 @@ export class ShardedWallet implements IWebWallet {
     return getExistingUserAccount({
       client: this.client,
       storage: this.storage,
+      storedToken: this.storedToken,
     });
   }
 }