Fix(extension): [LW-8723] make deletion of wallet more robust (#674)

* refactor(extension): make wallet cleanup more robust

* fix(extension): refresh wallet after deleting to reset all stores

* refactor(extension): refactor cleanup of storage when deleting wallet

* fix(extension): fix keys whitelist

* fix(extension): fix more keys

* fix(extension): fix wallet manager unit tests

* test(extension): fix background service unit test

* fix(extension): key iteration in clearLocalStorage

---------

Co-authored-by: przemyslaw.wlodek <przem.wlodek.github@gmail.com>

Author: refi93

apps/browser-extension-wallet/src/components/ResetDataError/ResetDataError.tsx

@@ -22,14 +22,17 @@ export const ResetDataError = ({
   buttonLabel
 }: ResetDataErrorProps): React.ReactElement => {
   const { deleteWallet } = useWalletManager();
-  const { walletManagerUi } = useWalletStore();
+  const { walletManagerUi, setDeletingWallet } = useWalletStore();
   const { theme } = useTheme();
   const backgroundService = useBackgroundServiceAPIContext();
 
   const Layout = appMode === 'browser' ? WalletSetupLayout : React.Fragment;
 
   const resetData = async () => {
-    if (walletManagerUi) await deleteWallet();
+    if (walletManagerUi) {
+      setDeletingWallet(true);
+      await deleteWallet();
+    }
     window.localStorage.clear();
     window.localStorage.setItem('mode', theme.name);
     await backgroundService.resetStorage();

apps/browser-extension-wallet/src/features/unlock-wallet/components/UnlockWalletContainer.tsx

@@ -19,7 +19,7 @@ export interface UnlockWalletContainerProps {
 export const UnlockWalletContainer = ({ validateMnemonic }: UnlockWalletContainerProps): React.ReactElement => {
   const analytics = useAnalyticsContext();
   const { unlockWallet, lockWallet, deleteWallet } = useWalletManager();
-  const { environmentName, setKeyAgentData } = useWalletStore();
+  const { environmentName, setKeyAgentData, setDeletingWallet } = useWalletStore();
   const backgroundService = useBackgroundServiceAPIContext();
 
   const [isVerifyingPassword, setIsVerifyingPassword] = useState(false);
@@ -75,8 +75,10 @@ export const UnlockWalletContainer = ({ validateMnemonic }: UnlockWalletContaine
   const onForgotPasswordClick = async (): Promise<void> => {
     await analytics.sendEventToPostHog(PostHogAction.UnlockWalletForgotPasswordProceedClick);
     saveValueInLocalStorage({ key: 'isForgotPasswordFlow', value: true });
+    setDeletingWallet(true);
     await deleteWallet(true);
     await backgroundService.handleOpenBrowser({ section: BrowserViewSections.FORGOT_PASSWORD });
+    setDeletingWallet(false);
   };
 
   useKeyboardShortcut(['Enter'], onUnlock);

apps/browser-extension-wallet/src/hooks/__tests__/useWalletManager.test.tsx

@@ -162,7 +162,7 @@ describe('Testing useWalletManager hook', () => {
 
       await lockWallet();
       expect(lockWallet).toBeDefined();
-      expect(clearBackgroundStorage).toBeCalledWith(['keyAgentsByChain']);
+      expect(clearBackgroundStorage).toBeCalledWith({ keys: ['keyAgentsByChain'] });
       expect(deleteFromLocalStorage).toBeCalledWith('keyAgentData');
       expect(setCardanoWallet).toBeCalledWith();
       expect(setKeyAgentData).toBeCalledWith();
@@ -728,8 +728,8 @@ describe('Testing useWalletManager hook', () => {
       const shutdownWalletMocked = jest.fn();
       mockShutdownWallet.mockImplementation(shutdownWalletMocked);
 
-      const deleteFromLocalStorage = jest.fn();
-      jest.spyOn(localStorage, 'deleteFromLocalStorage').mockImplementation(deleteFromLocalStorage);
+      const clearLocalStorage = jest.fn();
+      jest.spyOn(localStorage, 'clearLocalStorage').mockImplementation(clearLocalStorage);
 
       const resetWalletLock = jest.fn();
       const setCardanoWallet = jest.fn();
@@ -757,12 +757,20 @@ describe('Testing useWalletManager hook', () => {
       expect(deleteWallet).toBeDefined();
       await deleteWallet();
       expect(shutdownWalletMocked).toBeCalledWith(walletManagerUi);
-      expect(deleteFromLocalStorage.mock.calls[0]).toEqual(['appSettings']);
-      expect(deleteFromLocalStorage.mock.calls[1]).toEqual(['showDappBetaModal']);
-      expect(deleteFromLocalStorage.mock.calls[2]).toEqual(['lastStaking']);
-      expect(deleteFromLocalStorage.mock.calls[3]).toEqual(['userInfo']);
-      expect(deleteFromLocalStorage.mock.calls[4]).toEqual(['keyAgentData']);
-      expect(clearBackgroundStorage).toBeCalledWith(['message', 'mnemonic', 'keyAgentsByChain']);
+      expect(clearLocalStorage).toBeCalledWith({
+        except: [
+          'currency',
+          'lock',
+          'mode',
+          'hideBalance',
+          'analyticsAccepted',
+          'isForgotPasswordFlow',
+          'multidelegationFirstVisit'
+        ]
+      });
+      expect(clearBackgroundStorage).toBeCalledWith({
+        except: ['fiatPrices', 'userId', 'usePersistentUserId', 'experimentsConfiguration']
+      });
       expect(resetWalletLock).toBeCalledWith();
       expect(setCardanoWallet).toBeCalledWith();
       expect(setKeyAgentData).toBeCalledWith();

apps/browser-extension-wallet/src/hooks/useWalletManager.ts

@@ -7,11 +7,17 @@ import { useCardanoWalletManagerContext } from '@providers/CardanoWalletManager'
 import { useAppSettingsContext } from '@providers/AppSettings';
 import { useBackgroundServiceAPIContext } from '@providers/BackgroundServiceAPI';
 import { AddressBookSchema, addressBookSchema, NftFoldersSchema, nftFoldersSchema, useDbState } from '@src/lib/storage';
-import { deleteFromLocalStorage, getValueFromLocalStorage, saveValueInLocalStorage } from '@src/utils/local-storage';
+import {
+  deleteFromLocalStorage,
+  getValueFromLocalStorage,
+  clearLocalStorage,
+  saveValueInLocalStorage
+} from '@src/utils/local-storage';
 import { config } from '@src/config';
 import { getWalletFromStorage } from '@src/utils/get-wallet-from-storage';
 import { getUserIdService } from '@providers/AnalyticsProvider/getUserIdService';
 import { ENHANCED_ANALYTICS_OPT_IN_STATUS_LS_KEY } from '@providers/AnalyticsProvider/matomo/config';
+import { ILocalStorage } from '@src/types';
 
 const { AVAILABLE_CHAINS, CHAIN } = config();
 
@@ -145,7 +151,7 @@ export const useWalletManager = (): UseWalletManager => {
   const lockWallet = useCallback(async (): Promise<void> => {
     if (!walletLock) return;
     // Deletes key agent data from storage and clears states
-    await backgroundService.clearBackgroundStorage(['keyAgentsByChain']);
+    await backgroundService.clearBackgroundStorage({ keys: ['keyAgentsByChain'] });
     deleteFromLocalStorage('keyAgentData');
 
     setKeyAgentData();
@@ -296,35 +302,49 @@ export const useWalletManager = (): UseWalletManager => {
   const deleteWallet = useCallback(
     async (isForgotPasswordFlow = false): Promise<void> => {
       await Wallet.shutdownWallet(walletManagerUi);
-      deleteFromLocalStorage('appSettings');
-      deleteFromLocalStorage('showDappBetaModal');
-      deleteFromLocalStorage('lastStaking');
-      deleteFromLocalStorage('userInfo');
-      deleteFromLocalStorage('keyAgentData');
-      await backgroundService.clearBackgroundStorage(['message', 'mnemonic', 'keyAgentsByChain']);
       setKeyAgentData();
       resetWalletLock();
       setCardanoWallet();
+
+      await backgroundService.clearBackgroundStorage({
+        except: ['fiatPrices', 'userId', 'usePersistentUserId', 'experimentsConfiguration']
+      });
+
+      const commonLocalStorageKeysToKeep: (keyof ILocalStorage)[] = [
+        'currency',
+        'lock',
+        'mode',
+        'hideBalance',
+        'analyticsAccepted',
+        'isForgotPasswordFlow',
+        'multidelegationFirstVisit'
+      ];
+
       setCurrentChain(CHAIN);
 
-      if (!isForgotPasswordFlow) {
-        deleteFromLocalStorage('wallet');
-        deleteFromLocalStorage(ENHANCED_ANALYTICS_OPT_IN_STATUS_LS_KEY);
-        await userIdService.clearId();
-        clearAddressBook();
-        clearNftsFolders();
+      if (isForgotPasswordFlow) {
+        const additionalKeysToKeep: (keyof ILocalStorage)[] = ['wallet', ENHANCED_ANALYTICS_OPT_IN_STATUS_LS_KEY];
+        clearLocalStorage({
+          except: [...commonLocalStorageKeysToKeep, ...additionalKeysToKeep]
+        });
+        return;
       }
+
+      clearLocalStorage({ except: commonLocalStorageKeysToKeep });
+      await userIdService.clearId();
+      clearAddressBook();
+      clearNftsFolders();
     },
     [
       walletManagerUi,
-      backgroundService,
-      userIdService,
-      clearAddressBook,
-      clearNftsFolders,
       setKeyAgentData,
       resetWalletLock,
       setCardanoWallet,
-      setCurrentChain
+      backgroundService,
+      setCurrentChain,
+      userIdService,
+      clearAddressBook,
+      clearNftsFolders
     ]
   );
 

apps/browser-extension-wallet/src/lib/scripts/background/services/userIdService.test.ts

@@ -112,7 +112,7 @@ describe('userIdService', () => {
       await userIdService.clearId();
       const newUserId = await userIdService.getUserId(1);
 
-      expect(clearStorageMock).toHaveBeenCalledWith(expect.arrayContaining(['userId', 'usePersistentUserId']));
+      expect(clearStorageMock).toHaveBeenCalledWith({ keys: ['userId', 'usePersistentUserId'] });
       expect(previousUserId).not.toEqual(newUserId);
       const subscription = userIdService.userTrackingType$.subscribe((trackingType) => {
         expect(trackingType).toEqual(UserTrackingType.Basic);

apps/browser-extension-wallet/src/lib/scripts/background/services/userIdService.ts

@@ -104,7 +104,7 @@ export class UserIdService implements UserIdServiceInterface {
     this.walletBasedUserId = undefined;
     this.userTrackingType$.next(UserTrackingType.Basic);
     this.clearSessionTimeout();
-    await this.clearStorage(['userId', 'usePersistentUserId']);
+    await this.clearStorage({ keys: ['userId', 'usePersistentUserId'] });
   }
 
   async makePersistent(): Promise<void> {

apps/browser-extension-wallet/src/lib/scripts/background/util.ts

@@ -37,18 +37,33 @@ export const getADAPriceFromBackgroundStorage = async (): Promise<BackgroundStor
 /**
  * Deletes the specified `keys` from the background storage.
  *
- * If no `keys` are passed then **ALL** of it is cleared.
+ * If no `options` are passed then **ALL** of it is cleared.
  *
- * @param keys Optional. List of keys to delete from storage
+ * @param options Optional. List of keys to either delete or remove from storage
  */
-export const clearBackgroundStorage = async (keys?: BackgroundStorageKeys[]): Promise<void> => {
-  if (!keys?.length) {
+type ClearBackgroundStorageOptions =
+  | {
+      keys: BackgroundStorageKeys[];
+      except?: never;
+    }
+  | {
+      keys?: never;
+      except: BackgroundStorageKeys[];
+    };
+export const clearBackgroundStorage = async (options?: ClearBackgroundStorageOptions): Promise<void> => {
+  if (!options) {
     await webStorage.local.remove('BACKGROUND_STORAGE');
     return;
   }
   const backgroundStorage = await getBackgroundStorage();
-
-  for (const key of keys) delete backgroundStorage?.[key];
+  for (const key in backgroundStorage) {
+    if (options.keys && options.keys.includes(key as BackgroundStorageKeys)) {
+      delete backgroundStorage[key as BackgroundStorageKeys];
+    }
+    if (options.except && options.except.includes(key as BackgroundStorageKeys)) {
+      delete backgroundStorage[key as BackgroundStorageKeys];
+    }
+  }
   await webStorage.local.set({ BACKGROUND_STORAGE: backgroundStorage ?? {} });
 };
 

apps/browser-extension-wallet/src/lib/scripts/migrations/versions/v0_6_0.ts

@@ -129,7 +129,7 @@ export const v0_6_0: Migration = {
           removeItemFromLocalStorage('wallet_tmp');
           removeItemFromLocalStorage('keyAgentData_tmp');
           removeItemFromLocalStorage('lock_tmp');
-          await clearBackgroundStorage(['keyAgentsByChain_tmp' as BackgroundStorageKeys]);
+          await clearBackgroundStorage({ keys: ['keyAgentsByChain_tmp' as BackgroundStorageKeys] });
           throw error;
         }
       },
@@ -214,13 +214,13 @@ export const v0_6_0: Migration = {
         if (tmpBackgroundStorage?.keyAgentsByChain_tmp) {
           await setBackgroundStorage({ keyAgentsByChain: tmpBackgroundStorage.keyAgentsByChain_tmp });
         }
-        await clearBackgroundStorage(['walletName' as BackgroundStorageKeys]);
+        await clearBackgroundStorage({ keys: ['walletName' as BackgroundStorageKeys] });
         // Delete temporary storage
         removeItemFromLocalStorage('appSettings_tmp');
         removeItemFromLocalStorage('wallet_tmp');
         removeItemFromLocalStorage('keyAgentData_tmp');
         removeItemFromLocalStorage('lock_tmp');
-        await clearBackgroundStorage(['keyAgentsByChain_tmp' as BackgroundStorageKeys]);
+        await clearBackgroundStorage({ keys: ['keyAgentsByChain_tmp' as BackgroundStorageKeys] });
       },
       rollback: async () => {
         console.info(`Rollback migrated data for ${MIGRATION_VERSION} upgrade`);
@@ -229,15 +229,15 @@ export const v0_6_0: Migration = {
         if (oldWalletInfo) setItemInLocalStorage('wallet', oldWalletInfo);
         if (oldLock) setItemInLocalStorage('lock', oldLock);
         removeItemFromLocalStorage('keyAgentData');
-        await clearBackgroundStorage(['keyAgentsByChain']);
+        await clearBackgroundStorage({ keys: ['keyAgentsByChain'] });
         if (backgroundStorage?.walletName)
           await setBackgroundStorage({ walletName: backgroundStorage.walletName } as BackgroundStorage);
         // Delete any temporary storage that may have been created
         removeItemFromLocalStorage('appSettings_tmp');
         removeItemFromLocalStorage('wallet_tmp');
         removeItemFromLocalStorage('keyAgentData_tmp');
         removeItemFromLocalStorage('lock_tmp');
-        await clearBackgroundStorage(['keyAgentsByChain_tmp' as BackgroundStorageKeys]);
+        await clearBackgroundStorage({ keys: ['keyAgentsByChain_tmp' as BackgroundStorageKeys] });
       }
     };
   }

apps/browser-extension-wallet/src/lib/scripts/types/background-service.ts

@@ -1,7 +1,8 @@
 import { BehaviorSubject, Subject } from 'rxjs';
 import { themes } from '@providers/ThemeProvider';
-import { BackgroundStorage, BackgroundStorageKeys, MigrationState } from './storage';
+import { BackgroundStorage, MigrationState } from './storage';
 import { CoinPrices } from './prices';
+import type { clearBackgroundStorage } from '../background/util';
 
 export enum BaseChannels {
   BACKGROUND_ACTIONS = 'background-actions'
@@ -60,14 +61,7 @@ export type BackgroundService = {
   handleChangeTheme: (data: ChangeThemeData) => void;
   setBackgroundStorage: (data: BackgroundStorage) => Promise<void>;
   getBackgroundStorage: () => Promise<BackgroundStorage>;
-  /**
-   * Deletes the specified `keys` from the background storage.
-   *
-   * If no `keys` are passed then **ALL** of it is cleared.
-   *
-   * @param keys Optional. List of keys to delete from storage
-   */
-  clearBackgroundStorage: (keys?: BackgroundStorageKeys[]) => Promise<void>;
+  clearBackgroundStorage: typeof clearBackgroundStorage;
   getWalletPassword: () => Uint8Array;
   setWalletPassword: (password?: Uint8Array) => void;
   resetStorage: () => Promise<void>;

apps/browser-extension-wallet/src/stores/slices/wallet-info-slice.ts

@@ -28,5 +28,6 @@ export const walletInfoSlice: SliceCreator<WalletInfoSlice & BlockchainProviderS
     set({ currentChain: Wallet.Cardano.ChainIds[chain], environmentName: chain });
     get().setBlockchainProvider(chain);
   },
-  getKeyAgentType: () => get()?.cardanoWallet?.keyAgent.serializableData.__typename
+  getKeyAgentType: () => get()?.cardanoWallet?.keyAgent.serializableData.__typename,
+  setDeletingWallet: (deletingWallet: boolean) => set({ deletingWallet })
 });