fix: remove gaia call when unlocking wallet

closes #1877

package.json

@@ -104,6 +104,8 @@
     "argon2-browser": "1.18.0",
     "assert": "2.0.0",
     "bignumber.js": "9.0.1",
+    "bip32": "2.0.6",
+    "bitcoinjs-lib": "5.2.0",
     "bn.js": "5.2.0",
     "c32check": "1.1.3",
     "capsize": "2.0.0",

src/background/vault.ts

@@ -15,6 +15,9 @@ import { DEFAULT_PASSWORD } from '@common/types';
 import { InternalMethods } from '@common/message-types';
 import { logger } from '@common/logger';
 import { getHasSetPassword, hasSetPasswordIdentifier } from '@common/storage';
+import { getDecryptedWalletDetails } from '@background/wallet/unlock-wallet';
+import { saveWalletConfigLocally } from '@common/wallet/wallet-config-helper';
+import { setToLocalstorageIfDefined } from '@common/storage';
 
 // In-memory (background) wallet instance
 export interface InMemoryVault {
@@ -42,18 +45,10 @@ let inMemoryVault: InMemoryVault = {
   salt: localStorage.getItem(saltIdentifier) || undefined,
 };
 
-function persistOptional(storageKey: string, value?: string) {
-  if (value) {
-    localStorage.setItem(storageKey, value);
-  } else {
-    localStorage.removeItem(storageKey);
-  }
-}
-
 export async function vaultMessageHandler(message: VaultActions) {
   inMemoryVault = await vaultReducer(message);
-  persistOptional(encryptedKeyIdentifier, inMemoryVault.encryptedSecretKey);
-  persistOptional(saltIdentifier, inMemoryVault.salt);
+  setToLocalstorageIfDefined(encryptedKeyIdentifier, inMemoryVault.encryptedSecretKey);
+  setToLocalstorageIfDefined(saltIdentifier, inMemoryVault.salt);
   localStorage.setItem(hasSetPasswordIdentifier, JSON.stringify(inMemoryVault.hasSetPassword));
   return inMemoryVault;
 }
@@ -71,7 +66,8 @@ async function storeSeed(secretKey: string, password?: string): Promise<InMemory
     currentAccountIndex: 0,
     hasSetPassword,
   };
-  // This method is called on `unlockWallet`.
+
+  // This method is sometimes called on `unlockWallet` (see unlockWallet below)
   // `restoreWalletAccounts` is reliant on external resources.
   // If this method fails, we return a single wallet instance,
   // the root wallet.
@@ -80,8 +76,10 @@ async function storeSeed(secretKey: string, password?: string): Promise<InMemory
       wallet: generatedWallet,
       gaiaHubUrl: gaiaUrl,
     });
+
     return { ...inMemoryVault, ...keyInfo, wallet: _wallet };
   } catch (error) {
+    logger.error('Failed to restore accounts', error);
     return { ...inMemoryVault, ...keyInfo, wallet: generatedWallet };
   }
 }
@@ -124,10 +122,13 @@ const vaultReducer = async (message: VaultActions): Promise<InMemoryVault> => {
       try {
         const updateConfig = async () => {
           const gaiaHubConfig = await createWalletGaiaConfig({ gaiaHubUrl: gaiaUrl, wallet });
-          await updateWalletConfig({
+          const walletConfig = await updateWalletConfig({
             wallet: newWallet,
             gaiaHubConfig,
           });
+          // The gaia wallet config is saved locally so we don't have
+          // to fetch it again from gaia on wallet unlock
+          saveWalletConfigLocally(walletConfig);
         };
         await updateConfig();
       } catch (e) {
@@ -164,6 +165,17 @@ const vaultReducer = async (message: VaultActions): Promise<InMemoryVault> => {
       if (!encryptedSecretKey) {
         throw new Error('Unable to unlock - logged out.');
       }
+      const vault = await getDecryptedWalletDetails(encryptedSecretKey, password, salt);
+      if (vault) {
+        return {
+          ...inMemoryVault,
+          ...vault,
+        };
+      }
+      // Since the user does not have the gaia wallet config saved locally, we use the legacy way
+      // i.e fetching it vis storeSeed. This can only happen when users have their wallet locked
+      // and then got the wallet upgraded. They won't have the config saved yet (this is done on account creation and login)
+      // This code path can be deleted after some months
       const decryptedData = await decryptMnemonic({
         encryptedSecretKey,
         password,

src/background/wallet/unlock-wallet.ts

@@ -0,0 +1,36 @@
+import { decryptMnemonic } from '@background/crypto/mnemonic-encryption';
+import { getWallet } from './wallet-utils';
+
+export async function getDecryptedWalletDetails(
+  encryptedSecretKey: string,
+  password: string,
+  salt: string | undefined
+) {
+  const hasSetPassword = password !== undefined;
+  const decryptedData = await decryptMnemonic({
+    encryptedSecretKey,
+    password,
+    salt,
+  });
+
+  const keyInfo = {
+    secretKey: decryptedData.secretKey,
+    encryptedSecretKey: encryptedSecretKey,
+    currentAccountIndex: 0,
+    hasSetPassword,
+  };
+
+  const wallet = await getWallet({
+    secretKey: decryptedData.secretKey,
+    salt: decryptedData.salt,
+    password,
+  });
+
+  const result = {
+    ...keyInfo,
+    wallet,
+    salt: decryptedData.salt,
+    encryptedSecretKey: decryptedData.encryptedSecretKey,
+  };
+  return result;
+}

src/background/wallet/wallet-utils.ts

@@ -0,0 +1,80 @@
+import { Account, deriveAccount, generateWallet, WalletConfig } from '@stacks/wallet-sdk';
+import { Wallet as SDKWallet } from '@stacks/wallet-sdk';
+import { mnemonicToSeed } from 'bip39';
+import { fromSeed } from 'bip32';
+import { BIP32Interface } from 'bitcoinjs-lib';
+
+function getSavedWalletConfig() {
+  const walletConfig = localStorage.getItem('walletConfig');
+  if (typeof walletConfig !== 'string') return;
+  try {
+    return JSON.parse(walletConfig) as WalletConfig;
+  } catch (e) {
+    return;
+  }
+}
+
+function accountsFromWalletConfig(
+  walletConfig: WalletConfig,
+  rootNode: BIP32Interface,
+  salt: string
+) {
+  return walletConfig.accounts.map((account, index) => {
+    const existingAccount = deriveAccount({
+      rootNode,
+      index,
+      salt,
+    });
+    return {
+      ...existingAccount,
+      username: account.username,
+    };
+  });
+}
+
+async function rootNodeFromSecretKey(secretKey: string) {
+  const rootPrivateKey = await mnemonicToSeed(secretKey);
+  return fromSeed(rootPrivateKey);
+}
+
+async function getSavedWalletAccounts({
+  secretKey,
+  walletConfig,
+  salt,
+}: {
+  secretKey: string;
+  walletConfig: WalletConfig;
+  salt: string;
+}): Promise<Account[]> {
+  // Restore from existing config
+  const rootNode = await rootNodeFromSecretKey(secretKey);
+  if (!walletConfig) return [] as Account[];
+  return accountsFromWalletConfig(walletConfig, rootNode, salt);
+}
+
+interface GetWalletParams {
+  secretKey: string;
+  salt: string;
+  password: string;
+}
+
+export async function getWallet(params: GetWalletParams): Promise<SDKWallet | undefined> {
+  const { secretKey, salt, password } = params;
+  const wallet = await generateWallet({
+    secretKey,
+    password,
+  });
+
+  const walletConfig = getSavedWalletConfig();
+  if (!walletConfig) return;
+
+  const accounts = await getSavedWalletAccounts({
+    secretKey,
+    walletConfig,
+    salt,
+  });
+
+  if (accounts.length === 0) return wallet;
+
+  return { ...wallet, accounts };
+}

src/common/storage.ts

@@ -71,3 +71,9 @@ export function getHasSetPassword() {
   }
   return false;
 }
+
+export function setToLocalstorageIfDefined(storageKey: string, value?: string) {
+  if (value) {
+    localStorage.setItem(storageKey, value);
+  }
+}

src/common/wallet/wallet-config-helper.ts

@@ -0,0 +1,26 @@
+import { setToLocalstorageIfDefined } from '@common/storage';
+import { WalletConfig } from '@stacks/wallet-sdk';
+import { createWalletGaiaConfig, getOrCreateWalletConfig, Wallet } from '@stacks/wallet-sdk';
+import { gaiaUrl } from '@common/constants';
+import { logger } from '@common/logger';
+
+export function saveWalletConfigLocally(walletConfig: WalletConfig) {
+  setToLocalstorageIfDefined('walletConfig', JSON.stringify(walletConfig));
+  return walletConfig;
+}
+
+export async function getWalletConfig(wallet: Wallet) {
+  try {
+    const gaiaHubConfig = await createWalletGaiaConfig({ gaiaHubUrl: gaiaUrl, wallet });
+    const walletConfig = await getOrCreateWalletConfig({
+      wallet,
+      gaiaHubConfig,
+      skipUpload: true,
+    });
+    saveWalletConfigLocally(walletConfig);
+    return walletConfig;
+  } catch (e) {
+    logger.error('useWalletConfig error', e);
+    return;
+  }
+}

src/pages/set-password.tsx

@@ -15,6 +15,7 @@ import { Body, Caption } from '@components/typography';
 import { Header } from '@components/header';
 import { RouteUrls } from '@routes/route-urls';
 import { OnboardingSelectors } from '@tests/integration/onboarding.selectors';
+import { getWalletConfig } from '@common/wallet/wallet-config-helper';
 
 interface SetPasswordProps {
   placeholder?: string;
@@ -35,6 +36,12 @@ export const SetPasswordPage = ({ placeholder }: SetPasswordProps) => {
     void analytics.page('view', '/set-password');
   }, [analytics]);
 
+  useEffect(() => {
+    // Proactively fetch the gaia wallet config
+    if (!wallet) return;
+    void getWalletConfig(wallet);
+  });
+
   const submit = useCallback(
     async (password: string) => {
       if (!wallet) throw 'Please log in before setting a password.';

yarn.lock

@@ -5691,7 +5691,7 @@ bitcoin-ops@^1.3.0, bitcoin-ops@^1.4.0:
   resolved "https://registry.yarnpkg.com/bitcoin-ops/-/bitcoin-ops-1.4.1.tgz#e45de620398e22fd4ca6023de43974ff42240278"
   integrity sha512-pef6gxZFztEhaE9RY9HmWVmiIHqCb2OyS4HPKkpc6CIiiOa3Qmuoylxc5P2EkU3w+5eTSifI9SEZC88idAIGow==
 
-bitcoinjs-lib@^5.1.10, bitcoinjs-lib@^5.1.6, bitcoinjs-lib@^5.2.0:
+bitcoinjs-lib@5.2.0, bitcoinjs-lib@^5.1.10, bitcoinjs-lib@^5.1.6, bitcoinjs-lib@^5.2.0:
   version "5.2.0"
   resolved "https://registry.yarnpkg.com/bitcoinjs-lib/-/bitcoinjs-lib-5.2.0.tgz#caf8b5efb04274ded1b67e0706960b93afb9d332"
   integrity sha512-5DcLxGUDejgNBYcieMIUfjORtUeNWl828VWLHJGVKZCb4zIS1oOySTUr0LGmcqJBQgTBz3bGbRQla4FgrdQEIQ==