Fix progress reporting for key import

spec/integ/crypto/megolm-backup.spec.ts

@@ -46,6 +46,7 @@ import { IKeyBackup } from "../../../src/crypto/backup";
 import { flushPromises } from "../../test-utils/flushPromises";
 import { defer, IDeferred } from "../../../src/utils";
 import { DecryptionFailureCode } from "../../../src/crypto-api";
+import { ImportRoomKeysOpts } from "../../../src/crypto-api";
 
 const ROOM_ID = testData.TEST_ROOM_ID;
 
@@ -311,6 +312,7 @@ describe.each(Object.entries(CRYPTO_BACKENDS))("megolm-keys backup (%s)", (backe
 
     describe("recover from backup", () => {
         let aliceCrypto: Crypto.CryptoApi;
+        let importMockImpl: jest.Mock;
 
         beforeEach(async () => {
             fetchMock.get("path:/_matrix/client/v3/room_keys/version", testData.SIGNED_BACKUP_DATA);
@@ -322,6 +324,22 @@ describe.each(Object.entries(CRYPTO_BACKENDS))("megolm-keys backup (%s)", (backe
             // tell Alice to trust the dummy device that signed the backup
             await waitForDeviceList();
             await aliceCrypto.setDeviceVerified(testData.TEST_USER_ID, testData.TEST_DEVICE_ID);
+
+            importMockImpl = jest
+                .fn()
+                .mockImplementation((keys: IMegolmSessionData[], version: String, opts?: ImportRoomKeysOpts) => {
+                    // need to report progress
+                    if (opts?.progressCallback) {
+                        opts.progressCallback({
+                            stage: "load_keys",
+                            successes: keys.length,
+                            failures: 0,
+                            total: keys.length,
+                        });
+                    }
+                });
+            // @ts-ignore - mock a private method for testing purpose
+            aliceCrypto.importBackedUpRoomKeys = importMockImpl;
         });
 
         it("can restore from backup (Curve25519 version)", async function () {
@@ -397,10 +415,6 @@ describe.each(Object.entries(CRYPTO_BACKENDS))("megolm-keys backup (%s)", (backe
         }
 
         it("Should import full backup in chunks", async function () {
-            const importMockImpl = jest.fn();
-            // @ts-ignore - mock a private method for testing purpose
-            aliceCrypto.importBackedUpRoomKeys = importMockImpl;
-
             // We need several rooms with several sessions to test chunking
             const { response, expectedTotal } = createBackupDownloadResponse([45, 300, 345, 12, 130]);
 
@@ -459,7 +473,7 @@ describe.each(Object.entries(CRYPTO_BACKENDS))("megolm-keys backup (%s)", (backe
                     throw new Error("test error");
                 })
                 // Ok for other chunks
-                .mockResolvedValue(undefined);
+                .mockImplementation(importMockImpl);
 
             const { response, expectedTotal } = createBackupDownloadResponse([100, 300]);
 
@@ -498,9 +512,6 @@ describe.each(Object.entries(CRYPTO_BACKENDS))("megolm-keys backup (%s)", (backe
         });
 
         it("Should continue if some keys fails to decrypt", async function () {
-            // @ts-ignore - mock a private method for testing purpose
-            aliceCrypto.importBackedUpRoomKeys = jest.fn();
-
             const decryptionFailureCount = 2;
 
             const mockDecryptor = {
@@ -540,6 +551,85 @@ describe.each(Object.entries(CRYPTO_BACKENDS))("megolm-keys backup (%s)", (backe
             expect(result.imported).toStrictEqual(expectedTotal - decryptionFailureCount);
         });
 
+        it("Should report failures when decryption works but import fails", async function () {
+            // @ts-ignore - mock a private method for testing purpose
+            aliceCrypto.importBackedUpRoomKeys = jest
+                .fn()
+                .mockImplementationOnce((keys: IMegolmSessionData[], version: String, opts?: ImportRoomKeysOpts) => {
+                    // report 10 failures to import
+                    opts!.progressCallback!({
+                        stage: "load_keys",
+                        successes: 20,
+                        failures: 10,
+                        total: 30,
+                    });
+                    return Promise.resolve();
+                })
+                // Ok for other chunks
+                .mockResolvedValue(importMockImpl);
+
+            const { response, expectedTotal } = createBackupDownloadResponse([30]);
+
+            fetchMock.get("express:/_matrix/client/v3/room_keys/keys", response);
+
+            const check = await aliceCrypto.checkKeyBackupAndEnable();
+
+            const progressCallback = jest.fn();
+            const result = await aliceClient.restoreKeyBackupWithRecoveryKey(
+                testData.BACKUP_DECRYPTION_KEY_BASE58,
+                undefined,
+                undefined,
+                check!.backupInfo!,
+                {
+                    progressCallback,
+                },
+            );
+
+            expect(result.total).toStrictEqual(expectedTotal);
+            // A chunk failed to import
+            expect(result.imported).toStrictEqual(20);
+        });
+
+        it("Should report failures when decryption works but import fails - per room variant", async function () {
+            // @ts-ignore - mock a private method for testing purpose
+            aliceCrypto.importBackedUpRoomKeys = jest
+                .fn()
+                .mockImplementationOnce((keys: IMegolmSessionData[], version: String, opts?: ImportRoomKeysOpts) => {
+                    // report 10 failures to import
+                    opts!.progressCallback!({
+                        stage: "load_keys",
+                        successes: 20,
+                        failures: 10,
+                        total: 30,
+                    });
+                    return Promise.resolve();
+                })
+                // Ok for other chunks
+                .mockResolvedValue(importMockImpl);
+
+            const { response, expectedTotal } = createBackupDownloadResponse([30]);
+            const roomId = Object.keys(response.rooms)[0];
+
+            fetchMock.get(`express:/_matrix/client/v3/room_keys/keys/${roomId}`, response.rooms[roomId]);
+
+            const check = await aliceCrypto.checkKeyBackupAndEnable();
+
+            const progressCallback = jest.fn();
+            const result = await aliceClient.restoreKeyBackupWithRecoveryKey(
+                testData.BACKUP_DECRYPTION_KEY_BASE58,
+                roomId,
+                undefined,
+                check!.backupInfo!,
+                {
+                    progressCallback,
+                },
+            );
+
+            expect(result.total).toStrictEqual(expectedTotal);
+            // A chunk failed to import
+            expect(result.imported).toStrictEqual(20);
+        });
+
         it("recover specific session from backup", async function () {
             fetchMock.get(
                 "express:/_matrix/client/v3/room_keys/keys/:room_id/:session_id",

src/client.ts

@@ -212,7 +212,13 @@ import { LocalNotificationSettings } from "./@types/local_notifications";
 import { buildFeatureSupportMap, Feature, ServerSupport } from "./feature";
 import { BackupDecryptor, CryptoBackend } from "./common-crypto/CryptoBackend";
 import { RUST_SDK_STORE_PREFIX } from "./rust-crypto/constants";
-import { BootstrapCrossSigningOpts, CrossSigningKeyInfo, CryptoApi, ImportRoomKeysOpts } from "./crypto-api";
+import {
+    BootstrapCrossSigningOpts,
+    CrossSigningKeyInfo,
+    CryptoApi,
+    ImportRoomKeyProgressData,
+    ImportRoomKeysOpts,
+} from "./crypto-api";
 import { DeviceInfoMap } from "./crypto/DeviceList";
 import {
     AddSecretStorageKeyOpts,
@@ -3916,11 +3922,18 @@ export class MatrixClient extends TypedEventEmitter<EmittedEvents, ClientEventHa
                     async (chunk) => {
                         // We have a chunk of decrypted keys: import them
                         try {
-                            const backupVersion = backupInfo.version!;
+                            let success = 0;
+                            let failures = 0;
+                            const partialProgress = (stage: ImportRoomKeyProgressData): void => {
+                                success = stage.successes ?? 0;
+                                failures = stage.failures ?? 0;
+                            };
                             await this.cryptoBackend!.importBackedUpRoomKeys(chunk, backupVersion, {
                                 untrusted,
+                                progressCallback: partialProgress,
                             });
-                            totalImported += chunk.length;
+                            totalImported += success;
+                            totalFailures += failures;
                         } catch (e) {
                             totalFailures += chunk.length;
                             // We failed to import some keys, but we should still try to import the rest?
@@ -3947,11 +3960,25 @@ export class MatrixClient extends TypedEventEmitter<EmittedEvents, ClientEventHa
                 for (const k of keys) {
                     k.room_id = targetRoomId!;
                 }
-                await this.cryptoBackend.importBackedUpRoomKeys(keys, backupVersion, {
-                    progressCallback,
-                    untrusted,
-                });
-                totalImported = keys.length;
+                try {
+                    let success = 0;
+                    let failures = 0;
+                    const partialProgress = (stage: ImportRoomKeyProgressData): void => {
+                        success = stage.successes ?? 0;
+                        failures = stage.failures ?? 0;
+                    };
+                    await this.cryptoBackend!.importBackedUpRoomKeys(keys, backupVersion, {
+                        untrusted,
+                        progressCallback: partialProgress,
+                    });
+                    totalImported += success;
+                    totalFailures += failures;
+                } catch (e) {
+                    totalFailures += keys.length;
+                    // We failed to import some keys, but we should still try to import the rest?
+                    // Log the error and continue
+                    logger.error("Error importing keys from backup", e);
+                }
             } else {
                 totalKeyCount = 1;
                 try {
@@ -3967,6 +3994,7 @@ export class MatrixClient extends TypedEventEmitter<EmittedEvents, ClientEventHa
                     });
                     totalImported = 1;
                 } catch (e) {
+                    totalFailures = 1;
                     this.logger.debug("Failed to decrypt megolm session from backup", e);
                 }
             }

src/rust-crypto/backup.ts

@@ -252,7 +252,7 @@ export class RustBackupManager extends TypedEventEmitter<RustBackupCryptoEvents,
             }
             keysByRoom.get(roomId)!.set(key.session_id, key);
         }
-        await this.olmMachine.importBackedUpRoomKeys(
+        const result: RustSdkCryptoJs.RoomKeyImportResult = await this.olmMachine.importBackedUpRoomKeys(
             keysByRoom,
             (progress: BigInt, total: BigInt, failures: BigInt): void => {
                 const importOpt: ImportRoomKeyProgressData = {
@@ -265,6 +265,17 @@ export class RustBackupManager extends TypedEventEmitter<RustBackupCryptoEvents,
             },
             backupVersion,
         );
+        // call the progress callback one last time with the final state
+        if (opts?.progressCallback) {
+            // We use total count here and not imported count.
+            // Imported count could be 0 if all the keys were already imported.
+            opts.progressCallback({
+                total: result.totalCount,
+                successes: result.totalCount,
+                stage: "load_keys",
+                failures: keys.length - result.totalCount,
+            });
+        }
     }
 
     private keyBackupCheckInProgress: Promise<KeyBackupCheck | null> | null = null;