From 4b62e0a2ce7ea6a88168b96039e3847b90c11b80 Mon Sep 17 00:00:00 2001
From: Bruno Windels <brunow@element.io>
Date: Mon, 15 Mar 2021 13:38:27 +0100
Subject: [PATCH] extract retry event ids for key before overwriting key on key
 backup

---
 src/matrix/e2ee/RoomEncryption.js            | 12 +++++---
 src/matrix/e2ee/megolm/decryption/RoomKey.js | 10 ++++++
 src/matrix/room/Room.js                      | 32 +++++++++++---------
 3 files changed, 34 insertions(+), 20 deletions(-)

diff --git a/src/matrix/e2ee/RoomEncryption.js b/src/matrix/e2ee/RoomEncryption.js
index aaddd174..3659c447 100644
--- a/src/matrix/e2ee/RoomEncryption.js
+++ b/src/matrix/e2ee/RoomEncryption.js
@@ -210,11 +210,14 @@ export class RoomEncryption {
                 }
                 let roomKey = this._megolmDecryption.roomKeyFromBackup(this._room.id, sessionId, session);
                 if (roomKey) {
-                    let keyIsBestOne = false;
+                    let retryEventIds;
                     try {
                         const txn = await this._storage.readWriteTxn([this._storage.storeNames.inboundGroupSessions]);
                         try {
-                            keyIsBestOne = await this._megolmDecryption.writeRoomKey(roomKey, txn);
+                            const keyIsBestOne = await this._megolmDecryption.writeRoomKey(roomKey, txn);
+                            if (keyIsBestOne) {
+                                retryEventIds = roomKey.eventIds;
+                            }
                         } catch (err) {
                             txn.abort();
                             throw err;
@@ -225,9 +228,8 @@ export class RoomEncryption {
                         // this is just clearing the internal sessionInfo
                         roomKey.dispose();
                     }
-                    if (keyIsBestOne) {
-                        // wrote the key, meaning we didn't have a better one, go ahead and reattempt decryption
-                        await this._room.notifyRoomKey(roomKey);
+                    if (retryEventIds?.length) {
+                        await this._room.retryDecryption(retryEventIds);
                     }
                 }
             } else if (session?.algorithm) {
diff --git a/src/matrix/e2ee/megolm/decryption/RoomKey.js b/src/matrix/e2ee/megolm/decryption/RoomKey.js
index dc961627..6fcd738b 100644
--- a/src/matrix/e2ee/megolm/decryption/RoomKey.js
+++ b/src/matrix/e2ee/megolm/decryption/RoomKey.js
@@ -4,6 +4,7 @@ export class BaseRoomKey {
     constructor() {
         this._sessionInfo = null;
         this._isBetter = null;
+        this._eventIds = null;
     }
 
     async createSessionInfo(olm, pickleKey, txn) {
@@ -44,6 +45,11 @@ export class BaseRoomKey {
                 existingSession.free();
             }
         }
+        // store the event ids that can be decrypted with this key
+        // before we overwrite them if called from `write`.
+        if (existingSessionEntry?.eventIds) {
+            this._eventIds = existingSessionEntry.eventIds;
+        }
         return isBetter;
     }
 
@@ -71,6 +77,10 @@ export class BaseRoomKey {
         return false;
     }
 
+    get eventIds() {
+        return this._eventIds;
+    }
+
     dispose() {
         if (this._sessionInfo) {
             this._sessionInfo.release();
diff --git a/src/matrix/room/Room.js b/src/matrix/room/Room.js
index 671db2a7..f5dd38de 100644
--- a/src/matrix/room/Room.js
+++ b/src/matrix/room/Room.js
@@ -55,27 +55,24 @@ export class Room extends EventEmitter {
         this._observedEvents = null;
     }
 
-    async _getRetryDecryptEntriesForKey(roomKey, roomEncryption, txn) {
-        const retryEventIds = await roomEncryption.getEventIdsForMissingKey(roomKey, txn);
+    async _eventIdsToEntries(eventIds, txn) {
         const retryEntries = [];
-        if (retryEventIds) {
-            for (const eventId of retryEventIds) {
-                const storageEntry = await txn.timelineEvents.getByEventId(this._roomId, eventId);
-                if (storageEntry) {
-                    retryEntries.push(new EventEntry(storageEntry, this._fragmentIdComparer));
-                }
+        await Promise.all(eventIds.map(async eventId => {
+            const storageEntry = await txn.timelineEvents.getByEventId(this._roomId, eventId);
+            if (storageEntry) {
+                retryEntries.push(new EventEntry(storageEntry, this._fragmentIdComparer));
             }
-        }
+        }));
         return retryEntries;
     }
 
     /**
-     * Used for keys received from other sources than sync, like key backup.
+     * Used for retrying decryption from other sources than sync, like key backup.
      * @internal
-     * @param  {RoomKey} roomKey
+     * @param  {Array<string>} eventIds
      * @return {Promise}
      */
-    async notifyRoomKey(roomKey) {
+    async retryDecryption(eventIds) {
         if (!this._roomEncryption) {
             return;
         }
@@ -83,7 +80,7 @@ export class Room extends EventEmitter {
             this._storage.storeNames.timelineEvents,
             this._storage.storeNames.inboundGroupSessions,
         ]);
-        const retryEntries = await this._getRetryDecryptEntriesForKey(roomKey, this._roomEncryption, txn);
+        const retryEntries = await this._eventIdsToEntries(eventIds, txn);
         if (retryEntries.length) {
             const decryptRequest = this._decryptEntries(DecryptionSource.Retry, retryEntries, txn);
             // this will close txn while awaiting decryption
@@ -157,8 +154,13 @@ export class Room extends EventEmitter {
     }
 
     async _getSyncRetryDecryptEntries(newKeys, roomEncryption, txn) {
-        const entriesPerKey = await Promise.all(newKeys.map(key => this._getRetryDecryptEntriesForKey(key, roomEncryption, txn)));
-        let retryEntries = entriesPerKey.reduce((allEntries, entries) => allEntries.concat(entries), []);
+        const entriesPerKey = await Promise.all(newKeys.map(async key => {
+            const retryEventIds = await roomEncryption.getEventIdsForMissingKey(key, txn);
+            if (retryEventIds) {
+                return this._eventIdsToEntries(retryEventIds, txn);
+            }
+        }));
+        let retryEntries = entriesPerKey.reduce((allEntries, entries) => entries ? allEntries.concat(entries) : allEntries, []);
         // If we have the timeline open, see if there are more entries for the new keys
         // as we only store missing session information for synced events, not backfilled.
         // We want to decrypt all events we can though if the user is looking