From fe9245dd04afa6123df15e7c341bca6b44b0d939 Mon Sep 17 00:00:00 2001
From: Bruno Windels <brunow@matrix.org>
Date: Fri, 4 Sep 2020 12:10:12 +0200
Subject: [PATCH] first draft of retrying decryption when receiving room keys

---
 src/matrix/DeviceMessageHandler.js |  2 ++
 src/matrix/e2ee/RoomEncryption.js  | 37 ++++++++++++++++++++++++++++++
 src/matrix/room/Room.js            |  9 ++++++++
 3 files changed, 48 insertions(+)

diff --git a/src/matrix/DeviceMessageHandler.js b/src/matrix/DeviceMessageHandler.js
index 51a2378c..d8698127 100644
--- a/src/matrix/DeviceMessageHandler.js
+++ b/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);
         }
     }
 
diff --git a/src/matrix/e2ee/RoomEncryption.js b/src/matrix/e2ee/RoomEncryption.js
index f5e91913..caec39ce 100644
--- a/src/matrix/e2ee/RoomEncryption.js
+++ b/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
diff --git a/src/matrix/room/Room.js b/src/matrix/room/Room.js
index 14579b73..493febdb 100644
--- a/src/matrix/room/Room.js
+++ b/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") {