Report backup key import progress on start and improve types (#4711)

* report key import progress on start and improve types

* fix lint

* add documentation for exported types

* link `ImportRoomKeyProgressData` type in `ImportRoomKeyStage`

Co-authored-by: Richard van der Hoff <1389908+richvdh@users.noreply.github.com>

* link `ImportRoomKeyFetchProgress` in fetch stage doc

Co-authored-by: Richard van der Hoff <1389908+richvdh@users.noreply.github.com>

* link `ImportRoomKeyLoadProgress` in load_keys stage

Co-authored-by: Richard van der Hoff <1389908+richvdh@users.noreply.github.com>

* link `ImportRoomKeyProgressData` in `ImportRoomKeyFetchProgress` type doc

Co-authored-by: Richard van der Hoff <1389908+richvdh@users.noreply.github.com>

* link `ImportRoomKeyProgressData` in `ImportRoomKeyLoadProgress` type doc

Co-authored-by: Richard van der Hoff <1389908+richvdh@users.noreply.github.com>

* link `ImportRoomKeyStage.Fetch`

Co-authored-by: Richard van der Hoff <1389908+richvdh@users.noreply.github.com>

* convert `ImportRoomKeyStage.LoadKeys` to link in `ImportRoomKeyLoadProgress` stage doc

Co-authored-by: Richard van der Hoff <1389908+richvdh@users.noreply.github.com>

* remove whitespace

* improve `ImportRoomKeyStage.Fetch` stage doc

Co-authored-by: Richard van der Hoff <1389908+richvdh@users.noreply.github.com>

* improve `ImportRoomKeyStage.LoadKeys ` stage doc

Co-authored-by: Richard van der Hoff <1389908+richvdh@users.noreply.github.com>

---------

Co-authored-by: Richard van der Hoff <1389908+richvdh@users.noreply.github.com>

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

@@ -500,8 +500,10 @@ describe("RustCrypto", () => {
             const someRoomKeys = testData.MEGOLM_SESSION_DATA_ARRAY;
             let importTotal = 0;
             const opt: ImportRoomKeysOpts = {
-                progressCallback: (stage) => {
+                progressCallback: (progress) => {
-                    importTotal = stage.total ?? 0;
+                    if (progress.stage === "load_keys") {
+                        importTotal = progress.total;
+                    }
                 },
             };
             await rustCrypto.importRoomKeys(someRoomKeys, opt);
@@ -523,8 +525,10 @@ describe("RustCrypto", () => {
             const someRoomKeys = testData.MEGOLM_SESSION_DATA_ARRAY;
             let importTotal = 0;
             const opt: ImportRoomKeysOpts = {
-                progressCallback: (stage) => {
+                progressCallback: (progress) => {
-                    importTotal = stage.total ?? 0;
+                    if (progress.stage === "load_keys") {
+                        importTotal = progress.total;
+                    }
                 },
             };
             await rustCrypto.importRoomKeysAsJson(JSON.stringify(someRoomKeys), opt);

src/crypto-api/index.ts

@@ -1000,18 +1000,72 @@ export class DeviceVerificationStatus {
     }
 }
 
+/**
+ * Enum representing the different stages of importing room keys.
+ *
+ * This is the type of the `stage` property of {@link ImportRoomKeyProgressData}.
+ */
+export enum ImportRoomKeyStage {
+    /**
+     * The stage where room keys are being fetched.
+     *
+     * @see {@link ImportRoomKeyFetchProgress}.
+     */
+    Fetch = "fetch",
+    /**
+     * The stage where room keys are being loaded.
+     *
+     * @see {@link ImportRoomKeyLoadProgress}.
+     */
+    LoadKeys = "load_keys",
+}
+
+/**
+ * Type representing the progress during the 'fetch' stage of the room key import process.
+ *
+ * @see {@link ImportRoomKeyProgressData}.
+ */
+export type ImportRoomKeyFetchProgress = {
+    /**
+     * The current stage of the import process.
+     */
+    stage: ImportRoomKeyStage.Fetch;
+};
+
+/**
+ * Type representing the progress during the 'load_keys' stage of the room key import process.
+ *
+ * @see {@link ImportRoomKeyProgressData}.
+ */
+export type ImportRoomKeyLoadProgress = {
+    /**
+     * The current stage of the import process.
+     */
+    stage: ImportRoomKeyStage.LoadKeys;
+
+    /**
+     * The number of successfully loaded room keys so far.
+     */
+    successes: number;
+
+    /**
+     * The number of room keys that failed to load so far.
+     */
+    failures: number;
+
+    /**
+     * The total number of room keys being loaded.
+     */
+    total: number;
+};
+
 /**
  * Room key import progress report.
  * Used when calling {@link CryptoApi#importRoomKeys},
  * {@link CryptoApi#importRoomKeysAsJson} or {@link CryptoApi#restoreKeyBackup} as the parameter of
  * the progressCallback. Used to display feedback.
  */
-export interface ImportRoomKeyProgressData {
+export type ImportRoomKeyProgressData = ImportRoomKeyFetchProgress | ImportRoomKeyLoadProgress;
-    stage: string; // TODO: Enum
-    successes?: number;
-    failures?: number;
-    total?: number;
-}
 
 /**
  * Options object for {@link CryptoApi#importRoomKeys} and

src/rust-crypto/backup.ts

@@ -35,7 +35,12 @@ import { encodeUri, logDuration } from "../utils.ts";
 import { type OutgoingRequestProcessor } from "./OutgoingRequestProcessor.ts";
 import { sleep } from "../utils.ts";
 import { type BackupDecryptor } from "../common-crypto/CryptoBackend.ts";
-import { type ImportRoomKeyProgressData, type ImportRoomKeysOpts, CryptoEvent } from "../crypto-api/index.ts";
+import {
+    type ImportRoomKeyProgressData,
+    type ImportRoomKeysOpts,
+    CryptoEvent,
+    ImportRoomKeyStage,
+} from "../crypto-api/index.ts";
 import { type AESEncryptedSecretStoragePayload } from "../@types/AESEncryptedSecretStoragePayload.ts";
 import { type IMegolmSessionData } from "../@types/crypto.ts";
 
@@ -235,7 +240,7 @@ export class RustBackupManager extends TypedEventEmitter<RustBackupCryptoEvents,
             const importOpt: ImportRoomKeyProgressData = {
                 total: Number(total),
                 successes: Number(progress),
-                stage: "load_keys",
+                stage: ImportRoomKeyStage.LoadKeys,
                 failures: 0,
             };
             opts?.progressCallback?.(importOpt);
@@ -264,7 +269,7 @@ export class RustBackupManager extends TypedEventEmitter<RustBackupCryptoEvents,
                 const importOpt: ImportRoomKeyProgressData = {
                     total: Number(total),
                     successes: Number(progress),
-                    stage: "load_keys",
+                    stage: ImportRoomKeyStage.LoadKeys,
                     failures: Number(failures),
                 };
                 opts?.progressCallback?.(importOpt);
@@ -619,9 +624,6 @@ export class RustBackupManager extends TypedEventEmitter<RustBackupCryptoEvents,
         opts?: KeyBackupRestoreOpts,
     ): Promise<KeyBackupRestoreResult> {
         const keyBackup = await this.downloadKeyBackup(backupVersion);
-        opts?.progressCallback?.({
-            stage: "load_keys",
-        });
 
         return this.importKeyBackup(keyBackup, backupVersion, backupDecryptor, opts);
     }
@@ -672,6 +674,13 @@ export class RustBackupManager extends TypedEventEmitter<RustBackupCryptoEvents,
         let totalImported = 0;
         let totalFailures = 0;
 
+        opts?.progressCallback?.({
+            total: totalKeyCount,
+            successes: totalImported,
+            stage: ImportRoomKeyStage.LoadKeys,
+            failures: totalFailures,
+        });
+
         /**
          * This method is called when we have enough chunks to decrypt.
          * It will decrypt the chunks and try to import the room keys.
@@ -704,7 +713,7 @@ export class RustBackupManager extends TypedEventEmitter<RustBackupCryptoEvents,
             opts?.progressCallback?.({
                 total: totalKeyCount,
                 successes: totalImported,
-                stage: "load_keys",
+                stage: ImportRoomKeyStage.LoadKeys,
                 failures: totalFailures,
             });
         };

src/rust-crypto/rust-crypto.ts

@@ -67,6 +67,7 @@ import {
     type KeyBackupRestoreOpts,
     type KeyBackupRestoreResult,
     type StartDehydrationOpts,
+    ImportRoomKeyStage,
 } from "../crypto-api/index.ts";
 import { deviceKeysToDeviceMap, rustDeviceToJsDevice } from "./device-converter.ts";
 import { type IDownloadKeyResult, type IQueryKeysRequest } from "../client.ts";
@@ -1340,7 +1341,7 @@ export class RustCrypto extends TypedEventEmitter<RustCryptoEvents, CryptoEventH
 
         try {
             opts?.progressCallback?.({
-                stage: "fetch",
+                stage: ImportRoomKeyStage.Fetch,
             });
 
             return await this.backupManager.restoreKeyBackup(backupVersion, backupDecryptor, opts);