Add restoreKeybackup to CryptoApi. (#4476)

* First draft of moving out restoreKeyBackup out of MatrixClient

* Deprecate `restoreKeyBackup*` in `MatrixClient`

* Move types

* Handle only the room keys response

* Renaming and refactor `keysCountInBatch` & `getTotalKeyCount`

* Fix `importRoomKeysAsJson` tsdoc

* Fix typo

* Move `backupDecryptor.free()``

* Comment and simplify a bit `handleDecryptionOfAFullBackup`

* Fix decryption crash by moving`backupDecryptor.free`

* Use new api in `megolm-backup.spec.ts`

* Add tests to get recovery key from secret storage

* Add doc to `KeyBackupRestoreOpts` & `KeyBackupRestoreResult`

* Add doc to `restoreKeyBackupWithKey`

* Add doc to `backup.ts`

* Apply comment suggestions

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

* - Decryption key is recovered from the cache in `RustCrypto.restoreKeyBackup`
- Add `CryptoApi.getSecretStorageBackupPrivateKey` to get the decryption key from the secret storage.

* Add `CryptoApi.restoreKeyBackup` to `ImportRoomKeyProgressData` doc.

* Add deprecated symbol to all the `restoreKeyBackup*` overrides.

* Update tests

* Move `RustBackupManager.getTotalKeyCount` to `backup#calculateKeyCountInKeyBackup`

* Fix `RustBackupManager.restoreKeyBackup` tsdoc

* Move `backupDecryptor.free` in rust crypto.

* Move `handleDecryptionOfAFullBackup` in `importKeyBackup`

* Rename `calculateKeyCountInKeyBackup` to `countKeystInBackup`

* Fix `passphrase` typo

* Rename `backupInfoVersion` to `backupVersion`

* Complete restoreKeyBackup* methods documentation

* Add `loadSessionBackupPrivateKeyFromSecretStorage`

* Remove useless intermediary result variable.

* Check that decryption key matchs key backup info in `loadSessionBackupPrivateKeyFromSecretStorage`

* Get backup info from a specific version

* Fix typo in `countKeysInBackup`

* Improve documentation and naming

* Use `RustSdkCryptoJs.BackupDecryptionKey` as `decryptionKeyMatchesKeyBackupInfo` parameter.

* Call directly `olmMachine.getBackupKeys` in `restoreKeyBackup`

* Last review changes

---------

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

src/client.ts

@@ -3703,6 +3703,8 @@ export class MatrixClient extends TypedEventEmitter<EmittedEvents, ClientEventHa
      * @param opts - Optional params such as callbacks
      * @returns Status of restoration with `total` and `imported`
      * key counts.
+     *
+     * @deprecated Prefer {@link CryptoApi.restoreKeyBackupWithPassphrase | `CryptoApi.restoreKeyBackupWithPassphrase`}.
      */
     public async restoreKeyBackupWithPassword(
         password: string,
@@ -3711,6 +3713,9 @@ export class MatrixClient extends TypedEventEmitter<EmittedEvents, ClientEventHa
         backupInfo: IKeyBackupInfo,
         opts: IKeyBackupRestoreOpts,
     ): Promise<IKeyBackupRestoreResult>;
+    /**
+     * @deprecated Prefer {@link CryptoApi.restoreKeyBackupWithPassphrase | `CryptoApi.restoreKeyBackupWithPassphrase`}.
+     */
     public async restoreKeyBackupWithPassword(
         password: string,
         targetRoomId: string,
@@ -3718,6 +3723,9 @@ export class MatrixClient extends TypedEventEmitter<EmittedEvents, ClientEventHa
         backupInfo: IKeyBackupInfo,
         opts: IKeyBackupRestoreOpts,
     ): Promise<IKeyBackupRestoreResult>;
+    /**
+     * @deprecated Prefer {@link CryptoApi.restoreKeyBackupWithPassphrase | `CryptoApi.restoreKeyBackupWithPassphrase`}.
+     */
     public async restoreKeyBackupWithPassword(
         password: string,
         targetRoomId: string,
@@ -3725,6 +3733,9 @@ export class MatrixClient extends TypedEventEmitter<EmittedEvents, ClientEventHa
         backupInfo: IKeyBackupInfo,
         opts: IKeyBackupRestoreOpts,
     ): Promise<IKeyBackupRestoreResult>;
+    /**
+     * @deprecated Prefer {@link CryptoApi.restoreKeyBackupWithPassphrase | `CryptoApi.restoreKeyBackupWithPassphrase`}.
+     */
     public async restoreKeyBackupWithPassword(
         password: string,
         targetRoomId: string | undefined,
@@ -3748,6 +3759,8 @@ export class MatrixClient extends TypedEventEmitter<EmittedEvents, ClientEventHa
      * @param opts - Optional params such as callbacks
      * @returns Status of restoration with `total` and `imported`
      * key counts.
+     *
+     * @deprecated Prefer {@link CryptoApi.restoreKeyBackup | `CryptoApi.restoreKeyBackup`}.
      */
     public async restoreKeyBackupWithSecretStorage(
         backupInfo: IKeyBackupInfo,
@@ -3785,6 +3798,8 @@ export class MatrixClient extends TypedEventEmitter<EmittedEvents, ClientEventHa
 
      * @returns Status of restoration with `total` and `imported`
      * key counts.
+     *
+     * @deprecated Prefer {@link CryptoApi.restoreKeyBackup | `CryptoApi.restoreKeyBackup`}.
      */
     public restoreKeyBackupWithRecoveryKey(
         recoveryKey: string,
@@ -3793,6 +3808,9 @@ export class MatrixClient extends TypedEventEmitter<EmittedEvents, ClientEventHa
         backupInfo: IKeyBackupInfo,
         opts?: IKeyBackupRestoreOpts,
     ): Promise<IKeyBackupRestoreResult>;
+    /**
+     * @deprecated Prefer {@link CryptoApi.restoreKeyBackup | `CryptoApi.restoreKeyBackup`}.
+     */
     public restoreKeyBackupWithRecoveryKey(
         recoveryKey: string,
         targetRoomId: string,
@@ -3800,6 +3818,9 @@ export class MatrixClient extends TypedEventEmitter<EmittedEvents, ClientEventHa
         backupInfo: IKeyBackupInfo,
         opts?: IKeyBackupRestoreOpts,
     ): Promise<IKeyBackupRestoreResult>;
+    /**
+     * @deprecated Prefer {@link CryptoApi.restoreKeyBackup | `CryptoApi.restoreKeyBackup`}.
+     */
     public restoreKeyBackupWithRecoveryKey(
         recoveryKey: string,
         targetRoomId: string,
@@ -3807,6 +3828,9 @@ export class MatrixClient extends TypedEventEmitter<EmittedEvents, ClientEventHa
         backupInfo: IKeyBackupInfo,
         opts?: IKeyBackupRestoreOpts,
     ): Promise<IKeyBackupRestoreResult>;
+    /**
+     * @deprecated Prefer {@link CryptoApi.restoreKeyBackup | `CryptoApi.restoreKeyBackup`}.
+     */
     public restoreKeyBackupWithRecoveryKey(
         recoveryKey: string,
         targetRoomId: string | undefined,
@@ -3818,24 +3842,42 @@ export class MatrixClient extends TypedEventEmitter<EmittedEvents, ClientEventHa
         return this.restoreKeyBackup(privKey, targetRoomId!, targetSessionId!, backupInfo, opts);
     }
 
+    /**
+     * Restore from an existing key backup via a private key stored locally
+     * @param targetRoomId
+     * @param targetSessionId
+     * @param backupInfo
+     * @param opts
+     *
+     * @deprecated Prefer {@link CryptoApi.restoreKeyBackup | `CryptoApi.restoreKeyBackup`}.
+     */
     public async restoreKeyBackupWithCache(
         targetRoomId: undefined,
         targetSessionId: undefined,
         backupInfo: IKeyBackupInfo,
         opts?: IKeyBackupRestoreOpts,
     ): Promise<IKeyBackupRestoreResult>;
+    /**
+     * @deprecated Prefer {@link CryptoApi.restoreKeyBackup | `CryptoApi.restoreKeyBackup`}.
+     */
     public async restoreKeyBackupWithCache(
         targetRoomId: string,
         targetSessionId: undefined,
         backupInfo: IKeyBackupInfo,
         opts?: IKeyBackupRestoreOpts,
     ): Promise<IKeyBackupRestoreResult>;
+    /**
+     * @deprecated Prefer {@link CryptoApi.restoreKeyBackup | `CryptoApi.restoreKeyBackup`}.
+     */
     public async restoreKeyBackupWithCache(
         targetRoomId: string,
         targetSessionId: string,
         backupInfo: IKeyBackupInfo,
         opts?: IKeyBackupRestoreOpts,
     ): Promise<IKeyBackupRestoreResult>;
+    /**
+     * @deprecated Prefer {@link CryptoApi.restoreKeyBackup | `CryptoApi.restoreKeyBackup`}.
+     */
     public async restoreKeyBackupWithCache(
         targetRoomId: string | undefined,
         targetSessionId: string | undefined,

src/crypto-api/index.ts

@@ -22,7 +22,13 @@ import { DeviceMap } from "../models/device.ts";
 import { UIAuthCallback } from "../interactive-auth.ts";
 import { PassphraseInfo, SecretStorageCallbacks, SecretStorageKeyDescription } from "../secret-storage.ts";
 import { VerificationRequest } from "./verification.ts";
-import { BackupTrustInfo, KeyBackupCheck, KeyBackupInfo } from "./keybackup.ts";
+import {
+    BackupTrustInfo,
+    KeyBackupCheck,
+    KeyBackupInfo,
+    KeyBackupRestoreOpts,
+    KeyBackupRestoreResult,
+} from "./keybackup.ts";
 import { ISignatures } from "../@types/signed.ts";
 import { MatrixEvent } from "../models/event.ts";
 
@@ -502,6 +508,18 @@ export interface CryptoApi {
      */
     storeSessionBackupPrivateKey(key: Uint8Array, version: string): Promise<void>;
 
+    /**
+     * Attempt to fetch the backup decryption key from secret storage.
+     *
+     * If the key is found in secret storage, checks it against the latest backup on the server;
+     * if they match, stores the key in the crypto store by calling {@link storeSessionBackupPrivateKey},
+     * which enables automatic restore of individual keys when an Unable-to-decrypt error is encountered.
+     *
+     * If we are unable to fetch the key from secret storage, there is no backup on the server, or the key
+     * does not match, throws an exception.
+     */
+    loadSessionBackupPrivateKeyFromSecretStorage(): Promise<void>;
+
     /**
      * Get the current status of key backup.
      *
@@ -545,6 +563,36 @@ export interface CryptoApi {
      */
     deleteKeyBackupVersion(version: string): Promise<void>;
 
+    /**
+     * Download and restore the full key backup from the homeserver.
+     *
+     * Before calling this method, a decryption key, and the backup version to restore,
+     * must have been saved in the crypto store. This happens in one of the following ways:
+     *
+     * - When a new backup version is created with {@link CryptoApi.resetKeyBackup}, a new key is created and cached.
+     * - The key can be loaded from secret storage with {@link CryptoApi.loadSessionBackupPrivateKeyFromSecretStorage}.
+     * - The key can be received from another device via secret sharing, typically as part of the interactive verification flow.
+     * - The key and backup version can also be set explicitly via {@link CryptoApi.storeSessionBackupPrivateKey},
+     *   though this is not expected to be a common operation.
+     *
+     * Warning: the full key backup may be quite large, so this operation may take several hours to complete.
+     * Use of {@link KeyBackupRestoreOpts.progressCallback} is recommended.
+     *
+     * @param opts
+     */
+    restoreKeyBackup(opts?: KeyBackupRestoreOpts): Promise<KeyBackupRestoreResult>;
+
+    /**
+     * Restores a key backup using a passphrase.
+     * The decoded key (derived from the passphrase) is stored locally by calling {@link CryptoApi#storeSessionBackupPrivateKey}.
+     *
+     * @param passphrase - The passphrase to use to restore the key backup.
+     * @param opts
+     *
+     * @deprecated Deriving a backup key from a passphrase is not part of the matrix spec. Instead, a random key is generated and stored/shared via 4S.
+     */
+    restoreKeyBackupWithPassphrase(passphrase: string, opts?: KeyBackupRestoreOpts): Promise<KeyBackupRestoreResult>;
+
     ///////////////////////////////////////////////////////////////////////////////////////////////////////////////////
     //
     // Dehydrated devices
@@ -886,8 +934,8 @@ export class DeviceVerificationStatus {
 
 /**
  * Room key import progress report.
- * Used when calling {@link CryptoApi#importRoomKeys} or
- * {@link CryptoApi#importRoomKeysAsJson} as the parameter of
+ * 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 {

src/crypto-api/keybackup.ts

@@ -16,6 +16,7 @@ limitations under the License.
 
 import { ISigned } from "../@types/signed.ts";
 import { AESEncryptedSecretStoragePayload } from "../@types/AESEncryptedSecretStoragePayload.ts";
+import { ImportRoomKeyProgressData } from "./index.ts";
 
 export interface Curve25519AuthData {
     public_key: string;
@@ -87,3 +88,28 @@ export interface KeyBackupSession<T = Curve25519SessionData | AESEncryptedSecret
 export interface KeyBackupRoomSessions {
     [sessionId: string]: KeyBackupSession;
 }
+
+/**
+ * Extra parameters for {@link CryptoApi.restoreKeyBackup} and {@link CryptoApi.restoreKeyBackupWithPassphrase}.
+ */
+export interface KeyBackupRestoreOpts {
+    /**
+     * A callback which, if defined, will be called periodically to report ongoing progress of the backup restore process.
+     * @param progress
+     */
+    progressCallback?: (progress: ImportRoomKeyProgressData) => void;
+}
+
+/**
+ * The result of {@link CryptoApi.restoreKeyBackup}.
+ */
+export interface KeyBackupRestoreResult {
+    /**
+     * The total number of keys that were found in the backup.
+     */
+    total: number;
+    /**
+     * The number of keys that were imported.
+     */
+    imported: number;
+}

src/crypto/index.ts

@@ -102,6 +102,8 @@ import {
     OwnDeviceKeys,
     CryptoEvent as CryptoApiCryptoEvent,
     CryptoEventHandlerMap as CryptoApiCryptoEventHandlerMap,
+    KeyBackupRestoreResult,
+    KeyBackupRestoreOpts,
 } from "../crypto-api/index.ts";
 import { Device, DeviceMap } from "../models/device.ts";
 import { deviceInfoToDevice } from "./device-converter.ts";
@@ -1306,6 +1308,13 @@ export class Crypto extends TypedEventEmitter<CryptoEvent, CryptoEventHandlerMap
         });
     }
 
+    /**
+     * Implementation of {@link Crypto.loadSessionBackupPrivateKeyFromSecretStorage}.
+     */
+    public loadSessionBackupPrivateKeyFromSecretStorage(): Promise<void> {
+        throw new Error("Not implmeented");
+    }
+
     /**
      * Get the current status of key backup.
      *
@@ -4308,6 +4317,23 @@ export class Crypto extends TypedEventEmitter<CryptoEvent, CryptoEventHandlerMap
     public async startDehydration(createNewKey?: boolean): Promise<void> {
         throw new Error("Not implemented");
     }
+
+    /**
+     * Stub function -- restoreKeyBackup is not implemented here, so throw error
+     */
+    public restoreKeyBackup(opts: KeyBackupRestoreOpts): Promise<KeyBackupRestoreResult> {
+        throw new Error("Not implemented");
+    }
+
+    /**
+     * Stub function -- restoreKeyBackupWithPassphrase is not implemented here, so throw error
+     */
+    public restoreKeyBackupWithPassphrase(
+        passphrase: string,
+        opts: KeyBackupRestoreOpts,
+    ): Promise<KeyBackupRestoreResult> {
+        throw new Error("Not implemented");
+    }
 }
 
 /**

src/rust-crypto/backup.ts

@@ -24,6 +24,8 @@ import {
     KeyBackupInfo,
     KeyBackupSession,
     Curve25519SessionData,
+    KeyBackupRestoreOpts,
+    KeyBackupRestoreResult,
     KeyBackupRoomSessions,
 } from "../crypto-api/keybackup.ts";
 import { logger } from "../logger.ts";
@@ -218,7 +220,7 @@ export class RustBackupManager extends TypedEventEmitter<RustBackupCryptoEvents,
     /**
      * Import a list of room keys previously exported by exportRoomKeysAsJson
      *
-     * @param keys - a JSON string encoding a list of session export objects,
+     * @param jsonKeys - a JSON string encoding a list of session export objects,
      *    each of which is an IMegolmSessionData
      * @param opts - options object
      * @returns a promise which resolves once the keys have been imported
@@ -495,20 +497,19 @@ export class RustBackupManager extends TypedEventEmitter<RustBackupCryptoEvents,
      */
     private keysCountInBatch(batch: RustSdkCryptoJs.KeysBackupRequest): number {
         const parsedBody: KeyBackup = JSON.parse(batch.body);
-        let count = 0;
-        for (const { sessions } of Object.values(parsedBody.rooms)) {
-            count += Object.keys(sessions).length;
-        }
-        return count;
+        return countKeysInBackup(parsedBody);
     }
 
     /**
-     * Get information about the current key backup from the server
+     * Get information about a key backup from the server
+     * - If version is provided, get information about that backup version.
+     * - If no version is provided, get information about the latest backup.
      *
+     * @param version - The version of the backup to get information about.
      * @returns Information object from API or null if there is no active backup.
      */
-    private async requestKeyBackupVersion(): Promise<KeyBackupInfo | null> {
-        return await requestKeyBackupVersion(this.http);
+    public async requestKeyBackupVersion(version?: string): Promise<KeyBackupInfo | null> {
+        return await requestKeyBackupVersion(this.http, version);
     }
 
     /**
@@ -591,6 +592,148 @@ export class RustBackupManager extends TypedEventEmitter<RustBackupCryptoEvents,
     public createBackupDecryptor(decryptionKey: RustSdkCryptoJs.BackupDecryptionKey): BackupDecryptor {
         return new RustBackupDecryptor(decryptionKey);
     }
+
+    /**
+     * Restore a key backup.
+     *
+     * @param backupVersion - The version of the backup to restore.
+     * @param backupDecryptor - The backup decryptor to use to decrypt the keys.
+     * @param opts - Options for the restore.
+     * @returns The total number of keys and the total imported.
+     */
+    public async restoreKeyBackup(
+        backupVersion: string,
+        backupDecryptor: BackupDecryptor,
+        opts?: KeyBackupRestoreOpts,
+    ): Promise<KeyBackupRestoreResult> {
+        const keyBackup = await this.downloadKeyBackup(backupVersion);
+        opts?.progressCallback?.({
+            stage: "load_keys",
+        });
+
+        return this.importKeyBackup(keyBackup, backupVersion, backupDecryptor, opts);
+    }
+
+    /**
+     * Call `/room_keys/keys` to download the key backup (room keys) for the given backup version.
+     * https://spec.matrix.org/v1.12/client-server-api/#get_matrixclientv3room_keyskeys
+     *
+     * @param backupVersion
+     * @returns The key backup response.
+     */
+    private downloadKeyBackup(backupVersion: string): Promise<KeyBackup> {
+        return this.http.authedRequest<KeyBackup>(
+            Method.Get,
+            "/room_keys/keys",
+            { version: backupVersion },
+            undefined,
+            {
+                prefix: ClientPrefix.V3,
+            },
+        );
+    }
+
+    /**
+     * Import the room keys from a `/room_keys/keys` call.
+     * Calls `opts.progressCallback` with the progress of the import.
+     *
+     * @param keyBackup - The response from the server containing the keys to import.
+     * @param backupVersion - The version of the backup info.
+     * @param backupDecryptor - The backup decryptor to use to decrypt the keys.
+     * @param opts - Options for the import.
+     *
+     * @returns The total number of keys and the total imported.
+     *
+     * @private
+     */
+    private async importKeyBackup(
+        keyBackup: KeyBackup,
+        backupVersion: string,
+        backupDecryptor: BackupDecryptor,
+        opts?: KeyBackupRestoreOpts,
+    ): Promise<KeyBackupRestoreResult> {
+        // We have a full backup here, it can get quite big, so we need to decrypt and import it in chunks.
+
+        const CHUNK_SIZE = 200;
+        // Get the total count as a first pass
+        const totalKeyCount = countKeysInBackup(keyBackup);
+        let totalImported = 0;
+        let totalFailures = 0;
+
+        /**
+         * This method is called when we have enough chunks to decrypt.
+         * It will decrypt the chunks and try to import the room keys.
+         * @param roomChunks
+         */
+        const handleChunkCallback = async (roomChunks: Map<string, KeyBackupRoomSessions>): Promise<void> => {
+            const currentChunk: IMegolmSessionData[] = [];
+            for (const roomId of roomChunks.keys()) {
+                // Decrypt the sessions for the given room
+                const decryptedSessions = await backupDecryptor.decryptSessions(roomChunks.get(roomId)!);
+                // Add the decrypted sessions to the current chunk
+                decryptedSessions.forEach((session) => {
+                    // We set the room_id for each session
+                    session.room_id = roomId;
+                    currentChunk.push(session);
+                });
+            }
+
+            // We have a chunk of decrypted keys: import them
+            try {
+                await this.importBackedUpRoomKeys(currentChunk, backupVersion);
+                totalImported += currentChunk.length;
+            } catch (e) {
+                totalFailures += currentChunk.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);
+            }
+
+            opts?.progressCallback?.({
+                total: totalKeyCount,
+                successes: totalImported,
+                stage: "load_keys",
+                failures: totalFailures,
+            });
+        };
+
+        let groupChunkCount = 0;
+        let chunkGroupByRoom: Map<string, KeyBackupRoomSessions> = new Map();
+
+        // Iterate over the rooms and sessions to group them in chunks
+        // And we call the handleChunkCallback when we have enough chunks to decrypt
+        for (const [roomId, roomData] of Object.entries(keyBackup.rooms)) {
+            // If there are no sessions for the room, skip it
+            if (!roomData.sessions) continue;
+
+            // Initialize a new chunk group for the current room
+            chunkGroupByRoom.set(roomId, {});
+
+            for (const [sessionId, session] of Object.entries(roomData.sessions)) {
+                // We set previously the chunk group for the current room, so we can safely get it
+                const sessionsForRoom = chunkGroupByRoom.get(roomId)!;
+                sessionsForRoom[sessionId] = session;
+                groupChunkCount += 1;
+                // If we have enough chunks to decrypt, call the block callback
+                if (groupChunkCount >= CHUNK_SIZE) {
+                    // We have enough chunks to decrypt
+                    await handleChunkCallback(chunkGroupByRoom);
+                    // Reset the chunk group
+                    chunkGroupByRoom = new Map();
+                    // There might be remaining keys for that room, so add back an entry for the current room.
+                    chunkGroupByRoom.set(roomId, {});
+                    groupChunkCount = 0;
+                }
+            }
+        }
+
+        // Handle remaining chunk if needed
+        if (groupChunkCount > 0) {
+            await handleChunkCallback(chunkGroupByRoom);
+        }
+
+        return { total: totalKeyCount, imported: totalImported };
+    }
 }
 
 /**
@@ -657,11 +800,26 @@ export class RustBackupDecryptor implements BackupDecryptor {
     }
 }
 
+/**
+ * Fetch a key backup info from the server.
+ *
+ * If `version` is provided, calls `GET /room_keys/version/$version` and gets the backup info for that version.
+ * See https://spec.matrix.org/v1.12/client-server-api/#get_matrixclientv3room_keysversionversion.
+ *
+ * If not, calls `GET /room_keys/version` and gets the latest backup info.
+ * See https://spec.matrix.org/v1.12/client-server-api/#get_matrixclientv3room_keysversion
+ *
+ * @param http
+ * @param version - the specific version of the backup info to fetch
+ * @returns The key backup info or null if there is no backup.
+ */
 export async function requestKeyBackupVersion(
     http: MatrixHttpApi<IHttpOpts & { onlyData: true }>,
+    version?: string,
 ): Promise<KeyBackupInfo | null> {
     try {
-        return await http.authedRequest<KeyBackupInfo>(Method.Get, "/room_keys/version", undefined, undefined, {
+        const path = version ? encodeUri("/room_keys/version/$version", { $version: version }) : "/room_keys/version";
+        return await http.authedRequest<KeyBackupInfo>(Method.Get, path, undefined, undefined, {
             prefix: ClientPrefix.V3,
         });
     } catch (e) {
@@ -673,6 +831,34 @@ export async function requestKeyBackupVersion(
     }
 }
 
+/**
+ * Checks if the provided decryption key matches the public key of the key backup info.
+ *
+ * @param decryptionKey - The decryption key to check.
+ * @param keyBackupInfo - The key backup info to check against.
+ * @returns `true` if the decryption key matches the key backup info, `false` otherwise.
+ */
+export function decryptionKeyMatchesKeyBackupInfo(
+    decryptionKey: RustSdkCryptoJs.BackupDecryptionKey,
+    keyBackupInfo: KeyBackupInfo,
+): boolean {
+    const authData = <Curve25519AuthData>keyBackupInfo.auth_data;
+    return authData.public_key === decryptionKey.megolmV1PublicKey.publicKeyBase64;
+}
+
+/**
+ * Counts the total number of keys present in a key backup.
+ * @param keyBackup - The key backup to count the keys from.
+ * @returns The total number of keys in the backup.
+ */
+function countKeysInBackup(keyBackup: KeyBackup): number {
+    let count = 0;
+    for (const { sessions } of Object.values(keyBackup.rooms)) {
+        count += Object.keys(sessions).length;
+    }
+    return count;
+}
+
 export type RustBackupCryptoEvents =
     | CryptoEvent.KeyBackupStatus
     | CryptoEvent.KeyBackupSessionsRemaining

src/rust-crypto/rust-crypto.ts

@@ -46,7 +46,6 @@ import {
     CrossSigningStatus,
     CryptoApi,
     CryptoCallbacks,
-    Curve25519AuthData,
     DecryptionFailureCode,
     DeviceVerificationStatus,
     EventEncryptionInfo,
@@ -66,6 +65,8 @@ import {
     DeviceIsolationModeKind,
     CryptoEvent,
     CryptoEventHandlerMap,
+    KeyBackupRestoreOpts,
+    KeyBackupRestoreResult,
 } from "../crypto-api/index.ts";
 import { deviceKeysToDeviceMap, rustDeviceToJsDevice } from "./device-converter.ts";
 import { IDownloadKeyResult, IQueryKeysRequest } from "../client.ts";
@@ -76,16 +77,17 @@ import { secretStorageCanAccessSecrets, secretStorageContainsCrossSigningKeys }
 import { isVerificationEvent, RustVerificationRequest, verificationMethodIdentifierToMethod } from "./verification.ts";
 import { EventType, MsgType } from "../@types/event.ts";
 import { TypedEventEmitter } from "../models/typed-event-emitter.ts";
-import { RustBackupManager } from "./backup.ts";
+import { decryptionKeyMatchesKeyBackupInfo, RustBackupManager } from "./backup.ts";
 import { TypedReEmitter } from "../ReEmitter.ts";
 import { randomString } from "../randomstring.ts";
 import { ClientStoppedError } from "../errors.ts";
 import { ISignatures } from "../@types/signed.ts";
-import { encodeBase64 } from "../base64.ts";
+import { decodeBase64, encodeBase64 } from "../base64.ts";
 import { OutgoingRequestsManager } from "./OutgoingRequestsManager.ts";
 import { PerSessionKeyBackupDownloader } from "./PerSessionKeyBackupDownloader.ts";
 import { DehydratedDeviceManager } from "./DehydratedDeviceManager.ts";
 import { VerificationMethod } from "../types.ts";
+import { keyFromAuthData } from "../common-crypto/key-passphrase.ts";
 
 const ALL_VERIFICATION_METHODS = [
     VerificationMethod.Sas,
@@ -337,9 +339,7 @@ export class RustCrypto extends TypedEventEmitter<RustCryptoEvents, CryptoEventH
         }
 
         const backupDecryptionKey = RustSdkCryptoJs.BackupDecryptionKey.fromBase64(encodeBase64(privKey));
-
-        const authData = <Curve25519AuthData>backupInfo.auth_data;
-        if (authData.public_key != backupDecryptionKey.megolmV1PublicKey.publicKeyBase64) {
+        if (!decryptionKeyMatchesKeyBackupInfo(backupDecryptionKey, backupInfo)) {
             throw new Error(`getBackupDecryptor: key backup on server does not match the decryption key`);
         }
 
@@ -1202,6 +1202,28 @@ export class RustCrypto extends TypedEventEmitter<RustCryptoEvents, CryptoEventH
         );
     }
 
+    /**
+     * Implementation of {@link CryptoApi#loadSessionBackupPrivateKeyFromSecretStorage}.
+     */
+    public async loadSessionBackupPrivateKeyFromSecretStorage(): Promise<void> {
+        const backupKey = await this.secretStorage.get("m.megolm_backup.v1");
+        if (!backupKey) {
+            throw new Error("loadSessionBackupPrivateKeyFromSecretStorage: missing decryption key in secret storage");
+        }
+
+        const keyBackupInfo = await this.backupManager.getServerBackupInfo();
+        if (!keyBackupInfo || !keyBackupInfo.version) {
+            throw new Error("loadSessionBackupPrivateKeyFromSecretStorage: unable to get backup version");
+        }
+
+        const backupDecryptionKey = RustSdkCryptoJs.BackupDecryptionKey.fromBase64(backupKey);
+        if (!decryptionKeyMatchesKeyBackupInfo(backupDecryptionKey, keyBackupInfo)) {
+            throw new Error("loadSessionBackupPrivateKeyFromSecretStorage: decryption key does not match backup info");
+        }
+
+        await this.backupManager.saveBackupDecryptionKey(backupDecryptionKey, keyBackupInfo.version);
+    }
+
     /**
      * Get the current status of key backup.
      *
@@ -1280,6 +1302,53 @@ export class RustCrypto extends TypedEventEmitter<RustCryptoEvents, CryptoEventH
         obj.signatures = Object.fromEntries(sigs.entries());
     }
 
+    /**
+     * Implementation of {@link CryptoApi#restoreKeyBackupWithPassphrase}.
+     */
+    public async restoreKeyBackupWithPassphrase(
+        passphrase: string,
+        opts?: KeyBackupRestoreOpts,
+    ): Promise<KeyBackupRestoreResult> {
+        const backupInfo = await this.backupManager.getServerBackupInfo();
+        if (!backupInfo?.version) {
+            throw new Error("No backup info available");
+        }
+
+        const privateKey = await keyFromAuthData(backupInfo.auth_data, passphrase);
+
+        // Cache the key
+        await this.storeSessionBackupPrivateKey(privateKey, backupInfo.version);
+        return this.restoreKeyBackup(opts);
+    }
+
+    /**
+     * Implementation of {@link CryptoApi#restoreKeyBackup}.
+     */
+    public async restoreKeyBackup(opts?: KeyBackupRestoreOpts): Promise<KeyBackupRestoreResult> {
+        // Get the decryption key from the crypto store
+        const backupKeys: RustSdkCryptoJs.BackupKeys = await this.olmMachine.getBackupKeys();
+        const { decryptionKey, backupVersion } = backupKeys;
+        if (!decryptionKey || !backupVersion) throw new Error("No decryption key found in crypto store");
+
+        const decodedDecryptionKey = decodeBase64(decryptionKey.toBase64());
+
+        const backupInfo = await this.backupManager.requestKeyBackupVersion(backupVersion);
+        if (!backupInfo) throw new Error(`Backup version to restore ${backupVersion} not found on server`);
+
+        const backupDecryptor = await this.getBackupDecryptor(backupInfo, decodedDecryptionKey);
+
+        try {
+            opts?.progressCallback?.({
+                stage: "fetch",
+            });
+
+            return await this.backupManager.restoreKeyBackup(backupVersion, backupDecryptor, opts);
+        } finally {
+            // Free to avoid to keep in memory the decryption key stored in it. To avoid to exposing it to an attacker.
+            backupDecryptor.free();
+        }
+    }
+
     /**
      * Implementation of {@link CryptoApi#isDehydrationSupported}.
      */