first draft of retrying decryption when receiving room keys

src/matrix/DeviceMessageHandler.js

@@ -55,6 +55,8 @@ export class DeviceMessageHandler {
     _applyDecryptChanges(rooms, {roomKeys}) {
         const roomKeysByRoom = groupBy(roomKeys, s => s.roomId);
         for (const [roomId, roomKeys] of roomKeysByRoom) {
+            const room = rooms.get(roomId);
+            room?.notifyRoomKeys(roomKeys);
         }
     }
 

src/matrix/e2ee/RoomEncryption.js

@@ -30,6 +30,8 @@ export class RoomEncryption {
 
         this._megolmBackfillCache = this._megolmDecryption.createSessionCache();
         this._megolmSyncCache = this._megolmDecryption.createSessionCache();
+        // not `event_id`, but an internal event id passed in to the decrypt methods
+        this._eventIdsByMissingSession = new Map();
     }
 
     notifyTimelineClosed() {
@@ -45,21 +47,56 @@ export class RoomEncryption {
     async decryptNewSyncEvent(id, event, txn) {
         const payload = await this._megolmDecryption.decryptNewEvent(
             this._room.id, event, this._megolmSyncCache, txn);
+        if (!payload) {
+            this._addMissingSessionEvent(id, event);
+        }
         return payload;
     }
 
     async decryptNewGapEvent(id, event, txn) {
         const payload = await this._megolmDecryption.decryptNewEvent(
             this._room.id, event, this._megolmBackfillCache, txn);
+        if (!payload) {
+            this._addMissingSessionEvent(id, event);
+        }
         return payload;
     }
 
     async decryptStoredEvent(id, event, txn) {
         const payload = await this._megolmDecryption.decryptStoredEvent(
             this._room.id, event, this._megolmBackfillCache, txn);
+        if (!payload) {
+            this._addMissingSessionEvent(id, event);
+        }
         return payload;
     }
 
+    _addMissingSessionEvent(id, event) {
+        const senderKey = event.content?.["sender_key"];
+        const sessionId = event.content?.["session_id"];
+        const key = `${senderKey}|${sessionId}`;
+        let eventIds = this._eventIdsByMissingSession.get(key);
+        if (!eventIds) {
+            eventIds = new Set();
+            this._eventIdsByMissingSession.set(key, eventIds);
+        }
+        eventIds.add(id);
+    }
+
+    applyRoomKeys(roomKeys) {
+        // retry decryption with the new sessions
+        const idsToRetry = [];
+        for (const roomKey of roomKeys) {
+            const key = `${roomKey.senderKey}|${roomKey.sessionId}`;
+            const idsForSession = this._eventIdsByMissingSession.get(key);
+            if (idsForSession) {
+                this._eventIdsByMissingSession.delete(key);
+                idsToRetry.push(...Array.from(idsForSession));
+            }
+        }
+        return idsToRetry;
+    }
+
     async encrypt(type, content, hsApi) {
         const megolmResult = await this._megolmEncryption.encrypt(this._room.id, type, content, this._encryptionParams);
         // share the new megolm session if needed

src/matrix/room/Room.js

@@ -45,6 +45,15 @@ export class Room extends EventEmitter {
         this._roomEncryption = null;
 	}
 
+    notifyRoomKeys(roomKeys) {
+        if (this._roomEncryption) {
+            const internalIdsToRetry = this._roomEncryption.applyRoomKeys(roomKeys);
+            if (this._timeline) {
+
+            }
+        }
+    }
+
     async _decryptSyncEntries(entries, txn) {
         await Promise.all(entries.map(async e => {
             if (e.eventType === "m.room.encrypted") {