feat: enhance account migration process and ensure complete before lo…

…ading keys

packages/core/src/store/accountMigrations.ts

@@ -12,15 +12,46 @@ export default class AccountMigrations extends BaseStore<KeyringJson> implements
   // EXTENSION_PREFIX configured in webpack.shared.cjs
   public async migrateUnPrefixedAccounts (): Promise<void> {
     const accountsStore = new AccountsStore();
+    const migrationPromises: Promise<void>[] = [];
 
     const migrateAccount = async (key: string, value: KeyringJson): Promise<void> => {
-      if (key.startsWith('account:')) {
-        await accountsStore.set(key, value);
+      let existingValue: KeyringJson | null = null;
+
+      await accountsStore.get(key, (value) => {
+        existingValue = value;
+      });
+
+      // If the account already exists, remove the old storage key
+      if (existingValue) {
+        console.log(`Account with key ${key} already migrated, removing the old storage key`);
+        await super.remove(key);
+
+        return;
+      }
+
+      // migrate the account to the new storage key
+      await accountsStore.set(key, value);
+
+      // Verify the migration was successful
+      await accountsStore.get(key, (value) => {
+        existingValue = value;
+      });
+
+      if (existingValue) {
+        // Migration successful, remove the old storage key
+        await super.remove(key);
+        console.log(`Migration successful for key ${key}`);
+      } else {
+        console.error(`Failed to verify migration for key ${key}`);
       }
     };
 
     await super.all((key, value): void => {
-      migrateAccount(key, value).catch(console.error);
+      if (key.startsWith('account:')) {
+        migrationPromises.push(migrateAccount(key, value).catch(console.error));
+      }
     });
+
+    await Promise.all(migrationPromises);
   }
 }

packages/extension/src/background-init.ts

@@ -1,13 +1,32 @@
 /* global chrome */
 
+import { AccountsStore } from '@polkadot/extension-base/stores';
+import { keyring } from '@polkadot/ui-keyring';
+import { cryptoWaitReady } from '@polkadot/util-crypto';
+
+import { fatalErrorHandler } from '@polymeshassociation/extension-core/utils';
+
 import { checkForUpdateAndMigrate } from './migrations';
 
+// Ensure crypto is ready and load all keyring data
+async function initializeCryptoAndKeyring () {
+  await cryptoWaitReady();
+  console.log('crypto initialized');
+
+  keyring.loadAll({ store: new AccountsStore(), type: 'sr25519' });
+  console.log('initialization completed');
+}
+
 /*
  This is placed in a separate file to ensure that the listener is registered immediately,
  avoiding any delays caused by other imports or initializations. When placed in background.ts
  it was not being registered before the event was emitted.
 */
 
 chrome.runtime.onInstalled.addListener((details) => {
-  checkForUpdateAndMigrate(details).catch(console.error);
+  checkForUpdateAndMigrate(details)
+    .then(() => {
+      return initializeCryptoAndKeyring();
+    })
+    .catch(fatalErrorHandler);
 });

packages/extension/src/background.ts

@@ -8,14 +8,10 @@ import type { PolyRequestSignatures, PolyTransportRequestMessage } from '@polyme
 
 import { withErrorLog } from '@polkadot/extension-base/background';
 import { PORT_CONTENT, PORT_EXTENSION } from '@polkadot/extension-base/defaults';
-import { AccountsStore } from '@polkadot/extension-base/stores';
-import { keyring } from '@polkadot/ui-keyring';
 import { assert } from '@polkadot/util';
-import { cryptoWaitReady } from '@polkadot/util-crypto';
 
 import subscribePolymesh, { accountsSynchronizer } from '@polymeshassociation/extension-core';
 import handlers from '@polymeshassociation/extension-core/background/handlers';
-import { fatalErrorHandler } from '@polymeshassociation/extension-core/utils';
 
 // setup the notification (same a FF default background, white text)
 withErrorLog(() => chrome.action.setBadgeBackgroundColor({ color: '#d90000' }));
@@ -114,15 +110,3 @@ chrome.tabs.onActivated.addListener(() => {
 chrome.tabs.onRemoved.addListener(() => {
   getActiveTabs();
 });
-
-// initial setup
-cryptoWaitReady()
-  .then((): void => {
-    console.log('crypto initialized');
-
-    // load all the keyring data
-    keyring.loadAll({ store: new AccountsStore(), type: 'sr25519' });
-
-    console.log('initialization completed');
-  })
-  .catch(fatalErrorHandler);

packages/extension/src/migrations.ts

@@ -20,7 +20,7 @@ function isVersionEarlierOrEqual (version: string, compareTo: string): boolean {
   return true;
 }
 
-// Run the migration applicable to version 2.0.2 or earlier
+// Run the migration applicable to version 2.2.0 or earlier
 async function runMigrations () {
   const migrate = new AccountMigrations();
 
@@ -29,12 +29,12 @@ async function runMigrations () {
 }
 
 // Check for version update and perform migration if needed
-export async function checkForUpdateAndMigrate (details: chrome.runtime.InstalledDetails) {
+export async function checkForUpdateAndMigrate (details: chrome.runtime.InstalledDetails): Promise<void> {
   if (details.reason === chrome.runtime.OnInstalledReason.UPDATE) {
     const previousVersion = details.previousVersion;
 
     // Migrate Account Prefixes. Runs when no lastVersion or a previous version of 2.0.2 or earlier
-    if (previousVersion && isVersionEarlierOrEqual(previousVersion, '2.0.2')) {
+    if (!previousVersion || isVersionEarlierOrEqual(previousVersion, '2.2.0')) {
       await runMigrations();
     }
   }