From e55aee9617cced9771e1d6c28a76e45e596fce6e Mon Sep 17 00:00:00 2001
From: Florian Duros <florian.duros@ormaz.fr>
Date: Wed, 6 Nov 2024 18:37:56 +0100
Subject: [PATCH] Add `loadSessionBackupPrivateKeyFromSecretStorage`

---
 spec/integ/crypto/megolm-backup.spec.ts |  8 ++---
 src/crypto-api/index.ts                 | 15 ++++-----
 src/crypto/index.ts                     | 14 ++++-----
 src/rust-crypto/rust-crypto.ts          | 41 +++++++++++++++----------
 4 files changed, 41 insertions(+), 37 deletions(-)

diff --git a/spec/integ/crypto/megolm-backup.spec.ts b/spec/integ/crypto/megolm-backup.spec.ts
index de693c4c8ae..94d98b4cef8 100644
--- a/spec/integ/crypto/megolm-backup.spec.ts
+++ b/spec/integ/crypto/megolm-backup.spec.ts
@@ -23,6 +23,7 @@ import {
     createClient,
     Crypto,
     CryptoEvent,
+    encodeBase64,
     ICreateClientOpts,
     IEvent,
     IMegolmSessionData,
@@ -621,11 +622,10 @@ describe.each(Object.entries(CRYPTO_BACKENDS))("megolm-keys backup (%s)", (backe
                 };
                 fetchMock.get("express:/_matrix/client/v3/room_keys/keys", fullBackup);
 
-                const check = await aliceCrypto.checkKeyBackupAndEnable();
-                const recoveryKey = await aliceCrypto.getSecretStorageBackupPrivateKey();
-                expect(recoveryKey).not.toBeNull();
+                await aliceCrypto.loadSessionBackupPrivateKeyFromSecretStorage();
+                const decryptionKey = await aliceCrypto.getSessionBackupPrivateKey();
+                expect(encodeBase64(decryptionKey!)).toStrictEqual(testData.BACKUP_DECRYPTION_KEY_BASE64);
 
-                await aliceCrypto.storeSessionBackupPrivateKey(recoveryKey!, check!.backupInfo!.version!);
                 const result = await aliceCrypto.restoreKeyBackup();
                 expect(result.imported).toStrictEqual(1);
             },
diff --git a/src/crypto-api/index.ts b/src/crypto-api/index.ts
index d62daba23ea..d9ab886c262 100644
--- a/src/crypto-api/index.ts
+++ b/src/crypto-api/index.ts
@@ -470,15 +470,6 @@ export interface CryptoApi {
     //
     ///////////////////////////////////////////////////////////////////////////////////////////////////////////////////
 
-    /**
-     * Fetch the backup decryption key we have saved in our secret storage.
-     *
-     * This can be used for gossiping the key to other devices.
-     *
-     * @returns the key, if any, or null
-     */
-    getSecretStorageBackupPrivateKey(): Promise<Uint8Array | null>;
-
     /**
      * Fetch the backup decryption key we have saved in our store.
      *
@@ -511,6 +502,12 @@ export interface CryptoApi {
      */
     storeSessionBackupPrivateKey(key: Uint8Array, version: string): Promise<void>;
 
+    /**
+     * Fetch the backup decryption key from the secret storage, fetch the backup info version.
+     * Store locally the key and the backup info version by calling {@link storeSessionBackupPrivateKey}.
+     */
+    loadSessionBackupPrivateKeyFromSecretStorage(): Promise<void>;
+
     /**
      * Get the current status of key backup.
      *
diff --git a/src/crypto/index.ts b/src/crypto/index.ts
index ab0245eb55b..c75589813e7 100644
--- a/src/crypto/index.ts
+++ b/src/crypto/index.ts
@@ -1308,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.
      *
@@ -4327,13 +4334,6 @@ export class Crypto extends TypedEventEmitter<CryptoEvent, CryptoEventHandlerMap
     ): Promise<KeyBackupRestoreResult> {
         throw new Error("Not implemented");
     }
-
-    /**
-     * Stub function -- getSecretStorageBackupPrivateKey is not implemented here, so throw error
-     */
-    public getSecretStorageBackupPrivateKey(): Promise<Uint8Array | null> {
-        throw new Error("Not implemented");
-    }
 }
 
 /**
diff --git a/src/rust-crypto/rust-crypto.ts b/src/rust-crypto/rust-crypto.ts
index 7fa360b60fc..95f2667449c 100644
--- a/src/rust-crypto/rust-crypto.ts
+++ b/src/rust-crypto/rust-crypto.ts
@@ -1171,15 +1171,6 @@ export class RustCrypto extends TypedEventEmitter<RustCryptoEvents, CryptoEventH
         }
     }
 
-    /**
-     * Implementation of {@link CryptoApi#getSecretStorageBackupPrivateKey}.
-     */
-    public async getSecretStorageBackupPrivateKey(): Promise<Uint8Array | null> {
-        const backupKey = await this.secretStorage.get("m.megolm_backup.v1");
-        if (!backupKey) return null;
-        return decodeBase64(backupKey);
-    }
-
     /**
      * Fetch the backup decryption key we have saved in our store.
      *
@@ -1214,6 +1205,24 @@ 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 decodedKey = decodeBase64(backupKey);
+        const keyBackupInfo = await this.backupManager.getServerBackupInfo();
+        if (!keyBackupInfo) {
+            throw new Error("loadSessionBackupPrivateKeyFromSecretStorage: unable to get backup version");
+        }
+
+        await this.storeSessionBackupPrivateKey(decodedKey, keyBackupInfo.version);
+    }
+
     /**
      * Get the current status of key backup.
      *
@@ -1323,19 +1332,17 @@ export class RustCrypto extends TypedEventEmitter<RustCryptoEvents, CryptoEventH
 
         const backupDecryptor = await this.getBackupDecryptor(backupInfo, privateKeyFromCache);
 
-        opts?.progressCallback?.({
-            stage: "fetch",
-        });
-
-        let result: KeyBackupRestoreResult;
         try {
-            result = await this.backupManager.restoreKeyBackup(backupInfo.version, backupDecryptor, opts);
+            opts?.progressCallback?.({
+                stage: "fetch",
+            });
+
+            const result = await this.backupManager.restoreKeyBackup(backupInfo.version, backupDecryptor, opts);
+            return result;
         } finally {
             // Free to avoid to keep in memory the decryption key stored in it. To avoid to exposing it to an attacker.
             backupDecryptor.free();
         }
-
-        return result;
     }
 
     /**