Emit events to indicate migration progress

spec/integ/crypto/rust-crypto.spec.ts

@@ -18,7 +18,7 @@ import "fake-indexeddb/auto";
 import { IDBFactory } from "fake-indexeddb";
 import fetchMock from "fetch-mock-jest";
 
-import { createClient, IndexedDBCryptoStore } from "../../../src";
+import { createClient, CryptoEvent, IndexedDBCryptoStore } from "../../../src";
 import { populateStore } from "../../test-utils/test_indexeddb_cryptostore_dump";
 
 jest.setTimeout(15000);
@@ -124,6 +124,9 @@ describe("MatrixClient.initRustCrypto", () => {
             pickleKey: "+1k2Ppd7HIisUY824v7JtV3/oEE4yX0TqtmNPyhaD7o",
         });
 
+        const progressListener = jest.fn();
+        matrixClient.addListener(CryptoEvent.LegacyCryptoStoreMigrationProgress, progressListener);
+
         await matrixClient.initRustCrypto();
 
         // Do some basic checks on the imported data
@@ -132,7 +135,25 @@ describe("MatrixClient.initRustCrypto", () => {
         expect(deviceKeys.ed25519).toEqual("qK70DEqIXq7T+UU3v/al47Ab4JkMEBLpNrTBMbS5rrw");
 
         expect(await matrixClient.getCrypto()!.getActiveSessionBackupVersion()).toEqual("7");
-    });
+
+        // check the progress callback
+        expect(progressListener.mock.calls.length).toBeGreaterThan(50);
+
+        // The first call should have progress == 0
+        const [firstProgress, totalSteps] = progressListener.mock.calls[0];
+        expect(totalSteps).toBeGreaterThan(3000);
+        expect(firstProgress).toEqual(0);
+
+        for (let i = 1; i < progressListener.mock.calls.length - 1; i++) {
+            const [progress, total] = progressListener.mock.calls[i];
+            expect(total).toEqual(totalSteps);
+            expect(progress).toBeGreaterThan(progressListener.mock.calls[i - 1][0]);
+            expect(progress).toBeLessThanOrEqual(totalSteps);
+        }
+
+        // The final call should have progress == total == -1
+        expect(progressListener).toHaveBeenLastCalledWith(-1, -1);
+    }, 60000);
 });
 
 describe("MatrixClient.clearStores", () => {

spec/unit/rust-crypto/rust-crypto.spec.ts

@@ -193,6 +193,10 @@ describe("initRustCrypto", () => {
 
             fetchMock.get("path:/_matrix/client/v3/room_keys/version", { version: "45" });
 
+            function legacyMigrationProgressListener(progress: number, total: number): void {
+                console.log(`migrated ${progress} of ${total}`);
+            }
+
             await initRustCrypto({
                 logger,
                 http: makeMatrixHttpApi(),
@@ -204,6 +208,7 @@ describe("initRustCrypto", () => {
                 storePassphrase: "storePassphrase",
                 legacyCryptoStore: legacyStore,
                 legacyPickleKey: PICKLE_KEY,
+                legacyMigrationProgressListener,
             });
 
             // Check that the migration functions were correctly called

src/client.ts

@@ -962,7 +962,8 @@ type CryptoEvents =
     | CryptoEvent.KeysChanged
     | CryptoEvent.Warning
     | CryptoEvent.DevicesUpdated
-    | CryptoEvent.WillUpdateDevices;
+    | CryptoEvent.WillUpdateDevices
+    | CryptoEvent.LegacyCryptoStoreMigrationProgress;
 
 type MatrixEventEvents = MatrixEventEvent.Decrypted | MatrixEventEvent.Replaced | MatrixEventEvent.VisibilityChange;
 
@@ -2330,6 +2331,9 @@ export class MatrixClient extends TypedEventEmitter<EmittedEvents, ClientEventHa
             storePassphrase: this.pickleKey,
             legacyCryptoStore: this.cryptoStore,
             legacyPickleKey: this.pickleKey ?? "DEFAULT_KEY",
+            legacyMigrationProgressListener: (progress, total) => {
+                this.emit(CryptoEvent.LegacyCryptoStoreMigrationProgress, progress, total);
+            },
         });
 
         rustCrypto.setSupportedVerificationMethods(this.verificationMethods);

src/crypto/index.ts

@@ -259,6 +259,15 @@ export enum CryptoEvent {
     WillUpdateDevices = "crypto.willUpdateDevices",
     DevicesUpdated = "crypto.devicesUpdated",
     KeysChanged = "crossSigning.keysChanged",
+
+    /**
+     * Fires when data is being migrated from legacy crypto to rust crypto.
+     *
+     * The payload is a pair `(progress, total)`, where `progress` is the number of steps completed so far, and
+     * `total` is the total number of steps. When migration is complete, a final instance of the event is emitted, with
+     * `progress === total === -1`.
+     */
+    LegacyCryptoStoreMigrationProgress = "crypto.legacyCryptoStoreMigrationProgress",
 }
 
 export type CryptoEventHandlerMap = {
@@ -368,6 +377,8 @@ export type CryptoEventHandlerMap = {
      */
     [CryptoEvent.DevicesUpdated]: (users: string[], initialFetch: boolean) => void;
     [CryptoEvent.UserCrossSigningUpdated]: (userId: string) => void;
+
+    [CryptoEvent.LegacyCryptoStoreMigrationProgress]: (progress: number, total: number) => void;
 };
 
 export class Crypto extends TypedEventEmitter<CryptoEvent, CryptoEventHandlerMap> implements CryptoBackend {

src/rust-crypto/index.ts

@@ -72,6 +72,13 @@ export async function initRustCrypto(args: {
 
     /** The pickle key for `legacyCryptoStore` */
     legacyPickleKey?: string;
+
+    /**
+     * A callback which will receive progress updates on migration from `legacyCryptoStore`.
+     *
+     * Called with (-1, -1) to mark the end of migration.
+     */
+    legacyMigrationProgressListener?: (progress: number, total: number) => void;
 }): Promise<RustCrypto> {
     const { logger } = args;
 

src/rust-crypto/libolm_migration.ts

@@ -17,11 +17,11 @@ limitations under the License.
 import { Logger } from "../logger";
 import { CryptoStore, MigrationState, SecretStorePrivateKeys } from "../crypto/store/base";
 import * as RustSdkCryptoJs from "../../../matrix-rust-sdk-crypto-wasm";
+import { PickledInboundGroupSession, PickledSession } from "../../../matrix-rust-sdk-crypto-wasm";
 import { IndexedDBCryptoStore } from "../crypto/store/indexeddb-crypto-store";
 import { decryptAES, IEncryptedPayload } from "../crypto/aes";
 import { IHttpOpts, MatrixHttpApi } from "../http-api";
 import { requestKeyBackupVersion } from "./backup";
-import { PickledInboundGroupSession, PickledSession } from "../../../matrix-rust-sdk-crypto-wasm";
 
 /**
  * Determine if any data needs migrating from the legacy store, and do so.
@@ -52,6 +52,13 @@ export async function migrateFromLegacyCrypto(args: {
 
     /** Rust crypto store to migrate data into. */
     storeHandle: RustSdkCryptoJs.StoreHandle;
+
+    /**
+     * A callback which will receive progress updates on migration from `legacyStore`.
+     *
+     * Called with (-1, -1) to mark the end of migration.
+     */
+    legacyMigrationProgressListener?: (progress: number, total: number) => void;
 }): Promise<void> {
     const { logger, legacyStore } = args;
 
@@ -74,6 +81,17 @@ export async function migrateFromLegacyCrypto(args: {
         return;
     }
 
+    const nOlmSessions = await countOlmSessions(logger, legacyStore);
+    const nMegolmSessions = await countMegolmSessions(logger, legacyStore);
+    const totalSteps = 1 + nOlmSessions + nMegolmSessions;
+
+    let stepsDone = 0;
+    function onProgress(steps: number): void {
+        stepsDone += steps;
+        args.legacyMigrationProgressListener?.(stepsDone, totalSteps);
+    }
+    onProgress(0);
+
     const pickleKey = new TextEncoder().encode(args.legacyPickleKey);
 
     if (migrationState === MigrationState.NOT_STARTED) {
@@ -83,23 +101,26 @@ export async function migrateFromLegacyCrypto(args: {
         migrationState = MigrationState.INITIAL_DATA_MIGRATED;
         await legacyStore.setMigrationState(migrationState);
     }
+    onProgress(1);
 
     if (migrationState === MigrationState.INITIAL_DATA_MIGRATED) {
         logger.info("Migrating data from legacy crypto store. Step 2: olm sessions");
-        await migrateOlmSessions(logger, legacyStore, pickleKey, args.storeHandle);
+        await migrateOlmSessions(logger, legacyStore, pickleKey, args.storeHandle, onProgress);
 
         migrationState = MigrationState.OLM_SESSIONS_MIGRATED;
         await legacyStore.setMigrationState(migrationState);
     }
 
     if (migrationState === MigrationState.OLM_SESSIONS_MIGRATED) {
         logger.info("Migrating data from legacy crypto store. Step 3: megolm sessions");
-        await migrateMegolmSessions(logger, legacyStore, pickleKey, args.storeHandle);
+        await migrateMegolmSessions(logger, legacyStore, pickleKey, args.storeHandle, onProgress);
 
         migrationState = MigrationState.MEGOLM_SESSIONS_MIGRATED;
         await legacyStore.setMigrationState(migrationState);
     }
 
+    // Migration is done.
+    args.legacyMigrationProgressListener?.(-1, -1);
     logger.info("Migration from legacy crypto store complete");
 }
 
@@ -147,11 +168,26 @@ async function migrateBaseData(
     await RustSdkCryptoJs.Migration.migrateBaseData(migrationData, pickleKey, storeHandle);
 }
 
+async function countOlmSessions(logger: Logger, legacyStore: CryptoStore): Promise<number> {
+    logger.debug("Counting olm sessions to be migrated");
+    let nSessions: number;
+    await legacyStore.doTxn("readonly", [IndexedDBCryptoStore.STORE_SESSIONS], (txn) =>
+        legacyStore.countEndToEndSessions(txn, (n) => (nSessions = n)),
+    );
+    return nSessions!;
+}
+
+async function countMegolmSessions(logger: Logger, legacyStore: CryptoStore): Promise<number> {
+    logger.debug("Counting megolm sessions to be migrated");
+    return await legacyStore.countEndToEndInboundGroupSessions();
+}
+
 async function migrateOlmSessions(
     logger: Logger,
     legacyStore: CryptoStore,
     pickleKey: Uint8Array,
     storeHandle: RustSdkCryptoJs.StoreHandle,
+    onBatchDone: (batchSize: number) => void,
 ): Promise<void> {
     // eslint-disable-next-line no-constant-condition
     while (true) {
@@ -170,6 +206,7 @@ async function migrateOlmSessions(
 
         await RustSdkCryptoJs.Migration.migrateOlmSessions(migrationData, pickleKey, storeHandle);
         await legacyStore.deleteEndToEndSessionsBatch(batch);
+        onBatchDone(batch.length);
     }
 }
 
@@ -178,6 +215,7 @@ async function migrateMegolmSessions(
     legacyStore: CryptoStore,
     pickleKey: Uint8Array,
     storeHandle: RustSdkCryptoJs.StoreHandle,
+    onBatchDone: (batchSize: number) => void,
 ): Promise<void> {
     // eslint-disable-next-line no-constant-condition
     while (true) {
@@ -204,6 +242,7 @@ async function migrateMegolmSessions(
 
         await RustSdkCryptoJs.Migration.migrateMegolmSessions(migrationData, pickleKey, storeHandle);
         await legacyStore.deleteEndToEndInboundGroupSessionsBatch(batch);
+        onBatchDone(batch.length);
     }
 }