Functions to cache session backups key automatically

spec/unit/crypto/backup.spec.js

@@ -541,5 +541,31 @@ describe("MegolmBackup", function() {
                 expect(res.clearEvent.content).toEqual('testytest');
             });
         });
+
+        it('has working cache functions', async function() {
+            const key = Uint8Array.from([1, 2, 3, 4, 5, 6, 7, 8]);
+            await client._crypto.storeSessionBackupPrivateKey(key);
+            const result = await client._crypto.getSessionBackupPrivateKey();
+            expect(result).toEqual(key);
+        });
+
+        it('caches session backup keys as it encounters them', async function() {
+            const cachedNull = await client._crypto.getSessionBackupPrivateKey();
+            expect(cachedNull).toBeNull();
+            client._http.authedRequest = function() {
+                return Promise.resolve(KEY_BACKUP_DATA);
+            };
+            await new Promise((resolve) => {
+                client.restoreKeyBackupWithRecoveryKey(
+                    "EsTc LW2K PGiF wKEA 3As5 g5c4 BXwk qeeJ ZJV8 Q9fu gUMN UE4d",
+                    ROOM_ID,
+                    SESSION_ID,
+                    BACKUP_INFO,
+                    { cacheCompleteCallback: resolve },
+                );
+            });
+            const cachedKey = await client._crypto.getSessionBackupPrivateKey();
+            expect(cachedKey).not.toBeNull();
+        });
     });
 });

src/client.js

@@ -1745,15 +1745,16 @@ MatrixClient.RESTORE_BACKUP_ERROR_BAD_KEY = 'RESTORE_BACKUP_ERROR_BAD_KEY';
  * @param {string} [targetSessionId] Session ID to target a specific session.
  * Restores all sessions if omitted.
  * @param {object} backupInfo Backup metadata from `checkKeyBackup`
+ * @param {object} opts Optional params such as callbacks
  * @return {Promise<object>} Status of restoration with `total` and `imported`
  * key counts.
  */
 MatrixClient.prototype.restoreKeyBackupWithPassword = async function(
-    password, targetRoomId, targetSessionId, backupInfo,
+    password, targetRoomId, targetSessionId, backupInfo, opts,
 ) {
     const privKey = await keyFromAuthData(backupInfo.auth_data, password);
     return this._restoreKeyBackup(
-        privKey, targetRoomId, targetSessionId, backupInfo,
+        privKey, targetRoomId, targetSessionId, backupInfo, opts,
     );
 };
 
@@ -1766,15 +1767,16 @@ MatrixClient.prototype.restoreKeyBackupWithPassword = async function(
  * Restores all rooms if omitted.
  * @param {string} [targetSessionId] Session ID to target a specific session.
  * Restores all sessions if omitted.
+ * @param {object} opts Optional params such as callbacks
  * @return {Promise<object>} Status of restoration with `total` and `imported`
  * key counts.
  */
 MatrixClient.prototype.restoreKeyBackupWithSecretStorage = async function(
-    backupInfo, targetRoomId, targetSessionId,
+    backupInfo, targetRoomId, targetSessionId, opts,
 ) {
     const privKey = decodeBase64(await this.getSecret("m.megolm_backup.v1"));
     return this._restoreKeyBackup(
-        privKey, targetRoomId, targetSessionId, backupInfo,
+        privKey, targetRoomId, targetSessionId, backupInfo, opts,
     );
 };
 
@@ -1787,20 +1789,47 @@ MatrixClient.prototype.restoreKeyBackupWithSecretStorage = async function(
  * @param {string} [targetSessionId] Session ID to target a specific session.
  * Restores all sessions if omitted.
  * @param {object} backupInfo Backup metadata from `checkKeyBackup`
+ * @param {object} opts Optional params such as callbacks
+
  * @return {Promise<object>} Status of restoration with `total` and `imported`
  * key counts.
  */
 MatrixClient.prototype.restoreKeyBackupWithRecoveryKey = function(
-    recoveryKey, targetRoomId, targetSessionId, backupInfo,
+    recoveryKey, targetRoomId, targetSessionId, backupInfo, opts,
 ) {
     const privKey = decodeRecoveryKey(recoveryKey);
     return this._restoreKeyBackup(
-        privKey, targetRoomId, targetSessionId, backupInfo,
+        privKey, targetRoomId, targetSessionId, backupInfo, opts,
+    );
+};
+
+/**
+ * Restore from an existing key backup using a cached key, or fail
+ *
+ * @param {string} [targetRoomId] Room ID to target a specific room.
+ * Restores all rooms if omitted.
+ * @param {string} [targetSessionId] Session ID to target a specific session.
+ * Restores all sessions if omitted.
+ * @param {object} backupInfo Backup metadata from `checkKeyBackup`
+ * @param {object} opts Optional params such as callbacks
+ * @return {Promise<object>} Status of restoration with `total` and `imported`
+ * key counts.
+ */
+MatrixClient.prototype.restoreKeyBackupWithCache = async function(
+    targetRoomId, targetSessionId, backupInfo, opts,
+) {
+    const privKey = await this._crypto.getSessionBackupPrivateKey();
+    if (!privKey) {
+        return Promise.reject(new Error("Couldn't get key"));
+    }
+    return this._restoreKeyBackup(
+        privKey, targetRoomId, targetSessionId, backupInfo, opts,
     );
 };
 
 MatrixClient.prototype._restoreKeyBackup = function(
     privKey, targetRoomId, targetSessionId, backupInfo,
+    { cacheCompleteCallback }={}, // For sequencing during tests
 ) {
     if (this._crypto === null) {
         throw new Error("End-to-end encryption disabled");
@@ -1828,6 +1857,13 @@ MatrixClient.prototype._restoreKeyBackup = function(
         return Promise.reject({errcode: MatrixClient.RESTORE_BACKUP_ERROR_BAD_KEY});
     }
 
+    // Cache the key, if possible.
+    // This is async.
+    this._crypto.storeSessionBackupPrivateKey(privKey)
+    .catch((e) => {
+        console.warn("Error caching session backup key:", e);
+    }).then(cacheCompleteCallback);
+
     return this._http.authedRequest(
         undefined, "GET", path.path, path.queryData, undefined,
         {prefix: PREFIX_UNSTABLE},

src/crypto/index.js

@@ -583,6 +583,41 @@ Crypto.prototype.checkSecretStoragePrivateKey = function(privateKey, expectedPub
     }
 };
 
+/**
+ * Fetches the backup private key, if cached
+ * @returns {Promise} the key, if any, or null
+ */
+Crypto.prototype.getSessionBackupPrivateKey = async function() {
+    return new Promise((resolve) => {
+        this._cryptoStore.doTxn(
+            'readonly',
+            [IndexedDBCryptoStore.STORE_ACCOUNT],
+            (txn) => {
+                this._cryptoStore.getSecretStorePrivateKey(
+                    txn,
+                    resolve,
+                    "m.megolm_backup.v1",
+                );
+            },
+        );
+    });
+};
+
+/**
+ * Stores the session backup key to the cache
+ * @param {Uint8Array} key the private key
+ * @returns {Promise} so you can catch failures
+ */
+Crypto.prototype.storeSessionBackupPrivateKey = async function(key) {
+    return this._cryptoStore.doTxn(
+        'readwrite',
+        [IndexedDBCryptoStore.STORE_ACCOUNT],
+        (txn) => {
+            this._cryptoStore.storeSecretStorePrivateKey(txn, "m.megolm_backup.v1", key);
+        },
+    );
+};
+
 /**
  * Checks that a given cross-signing private key matches a given public key.
  * This can be used by the getCrossSigningKey callback to verify that the