implement decryption retrying and decrypting of gap/load entries

turns out we do have to always check for replay attacks because
failing to decrypt doesn't prevent an item from being stored,
so if you reload and then load you might be decrypting it
for the first time

src/matrix/Sync.js

@@ -133,6 +133,8 @@ export class Sync {
             storeNames.timelineFragments,
             storeNames.pendingEvents,
             storeNames.userIdentities,
+            storeNames.inboundGroupSessions,
+            storeNames.groupSessionDecryptions,
         ]);
         const roomChanges = [];
         let sessionChanges;

src/matrix/e2ee/RoomEncryption.js

@@ -14,6 +14,7 @@ See the License for the specific language governing permissions and
 limitations under the License.
 */
 
+import {MEGOLM_ALGORITHM} from "./common.js";
 import {groupBy} from "../../utils/groupBy.js";
 import {makeTxnId} from "../common.js";
 
@@ -44,57 +45,43 @@ export class RoomEncryption {
         return await this._deviceTracker.writeMemberChanges(this._room, memberChanges, txn);
     }
 
-    async decryptNewSyncEvent(id, event, txn) {
-        const payload = await this._megolmDecryption.decryptNewEvent(
-            this._room.id, event, this._megolmSyncCache, txn);
+    async decrypt(event, isSync, retryData, txn) {
+        if (event.content?.algorithm !== MEGOLM_ALGORITHM) {
+            throw new Error("Unsupported algorithm: " + event.content?.algorithm);
+        }
+        let sessionCache = isSync ? this._megolmSyncCache : this._megolmBackfillCache;
+        const payload = await this._megolmDecryption.decrypt(
+            this._room.id, event, sessionCache, txn);
         if (!payload) {
-            this._addMissingSessionEvent(id, event);
+            this._addMissingSessionEvent(event, isSync, retryData);
         }
         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) {
+    _addMissingSessionEvent(event, isSync, data) {
         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();
+            eventIds = new Map();
             this._eventIdsByMissingSession.set(key, eventIds);
         }
-        eventIds.add(id);
+        eventIds.set(event.event_id, {data, isSync});
     }
 
     applyRoomKeys(roomKeys) {
         // retry decryption with the new sessions
-        const idsToRetry = [];
+        const retryEntries = [];
         for (const roomKey of roomKeys) {
             const key = `${roomKey.senderKey}|${roomKey.sessionId}`;
-            const idsForSession = this._eventIdsByMissingSession.get(key);
-            if (idsForSession) {
+            const entriesForSession = this._eventIdsByMissingSession.get(key);
+            if (entriesForSession) {
                 this._eventIdsByMissingSession.delete(key);
-                idsToRetry.push(...Array.from(idsForSession));
+                retryEntries.push(...entriesForSession.values());
             }
         }
-        return idsToRetry;
+        return retryEntries;
     }
 
     async encrypt(type, content, hsApi) {

src/matrix/e2ee/megolm/Decryption.js

@@ -28,19 +28,7 @@ export class Decryption {
         return new SessionCache();
     }
 
-    async decryptNewEvent(roomId, event, sessionCache, txn) {
-        const {payload, messageIndex} = this._decrypt(roomId, event, sessionCache, txn);
-        const sessionId = event.content?.["session_id"];
-        this._handleReplayAttacks(roomId, sessionId, messageIndex, event, txn);
-        return payload;
-    }
-
-    async decryptStoredEvent(roomId, event, sessionCache, txn) {
-        const {payload} = this._decrypt(roomId, event, sessionCache, txn);
-        return payload;
-    }
-
-    async _decrypt(roomId, event, sessionCache, txn) {
+    async decrypt(roomId, event, sessionCache, txn) {
         const senderKey = event.content?.["sender_key"];
         const sessionId = event.content?.["session_id"];
         const ciphertext = event.content?.ciphertext;
@@ -75,16 +63,18 @@ export class Decryption {
         try {
             payload = JSON.parse(plaintext);
         } catch (err) {
-            throw new DecryptionError("NOT_JSON", event, {plaintext, err});
+            throw new DecryptionError("PLAINTEXT_NOT_JSON", event, {plaintext, err});
         }
         if (payload.room_id !== roomId) {
             throw new DecryptionError("MEGOLM_WRONG_ROOM", event,
                 {encryptedRoomId: payload.room_id, eventRoomId: roomId});
         }
-        return {payload, messageIndex};
+        await this._handleReplayAttack(roomId, sessionId, messageIndex, event, txn);
+        // TODO: verify event came from said senderKey
+        return payload;
     }
 
-    async _handleReplayAttacks(roomId, sessionId, messageIndex, event, txn) {
+    async _handleReplayAttack(roomId, sessionId, messageIndex, event, txn) {
         const eventId = event.event_id;
         const timestamp = event.origin_server_ts;
         const decryption = await txn.groupSessionDecryptions.get(roomId, sessionId, messageIndex);
@@ -92,7 +82,7 @@ export class Decryption {
             // the one with the newest timestamp should be the attack
             const decryptedEventIsBad = decryption.timestamp < timestamp;
             const badEventId = decryptedEventIsBad ? eventId : decryption.eventId;
-            throw new DecryptionError("MEGOLM_REPLAY_ATTACK", event, {badEventId, otherEventId: decryption.eventId});
+            throw new DecryptionError("MEGOLM_REPLAYED_INDEX", event, {badEventId, otherEventId: decryption.eventId});
         }
         if (!decryption) {
             txn.groupSessionDecryptions.set({

src/matrix/room/Room.js

@@ -25,6 +25,8 @@ import {WrappedError} from "../error.js"
 import {fetchOrLoadMembers} from "./members/load.js";
 import {MemberList} from "./members/MemberList.js";
 import {Heroes} from "./members/Heroes.js";
+import {EventEntry} from "./timeline/entries/EventEntry.js";
+import {EventKey} from "./timeline/EventKey.js";
 
 export class Room extends EventEmitter {
 	constructor({roomId, storage, hsApi, emitCollectionChange, sendScheduler, pendingEvents, user, createRoomEncryption}) {
@@ -45,29 +47,75 @@ export class Room extends EventEmitter {
         this._roomEncryption = null;
 	}
 
-    notifyRoomKeys(roomKeys) {
+    async notifyRoomKeys(roomKeys) {
         if (this._roomEncryption) {
-            const internalIdsToRetry = this._roomEncryption.applyRoomKeys(roomKeys);
-            if (this._timeline) {
-
+            // array of {data, source}
+            let retryEntries = this._roomEncryption.applyRoomKeys(roomKeys);
+            let decryptedEntries = [];
+            if (retryEntries.length) {
+                // groupSessionDecryptions can be written, the other stores not
+                const txn = await this._storage.readWriteTxn([
+                    this._storage.storeNames.timelineEvents,
+                    this._storage.storeNames.inboundGroupSessions,
+                    this._storage.storeNames.groupSessionDecryptions,
+                ]);
+                try {
+                    for (const retryEntry of retryEntries) {
+                        const {data: eventKey} = retryEntry;
+                        let entry = this._timeline?.findEntry(eventKey);
+                        if (!entry) {
+                            const storageEntry = await txn.timelineEvents.get(this._roomId, eventKey.fragmentId, eventKey.entryIndex);
+                            if (storageEntry) {
+                                entry = new EventEntry(storageEntry, this._fragmentIdComparer);
+                            }
+                        }
+                        if (entry) {
+                            entry = await this._decryptEntry(entry, txn, retryEntry.isSync);
+                            decryptedEntries.push(entry);
+                        }
+                    }
+                } catch (err) {
+                    txn.abort();
+                    throw err;
+                }
+                await txn.complete();
             }
+            if (this._timeline) {
+                // only adds if already present
+                this._timeline.replaceEntries(decryptedEntries);
+            }
+            // pass decryptedEntries to roomSummary
         }
     }
 
-    async _decryptSyncEntries(entries, txn) {
-        await Promise.all(entries.map(async e => {
-            if (e.eventType === "m.room.encrypted") {
-                try {
-                    const decryptedEvent = await this._roomEncryption.decryptNewSyncEvent(e.internalId, e.event, txn);
-                    if (decryptedEvent) {
-                        e.replaceWithDecrypted(decryptedEvent);
-                    }
-                } catch (err) {
-                    e.setDecryptionError(err);
+    _enableEncryption(encryptionParams) {
+        this._roomEncryption = this._createRoomEncryption(this, encryptionParams);
+        if (this._roomEncryption) {
+            this._sendQueue.enableEncryption(this._roomEncryption);
+            this._timeline.enableEncryption(this._decryptEntries.bind(this));
+        }
+    }
+
+    async _decryptEntry(entry, txn, isSync) {
+        if (entry.eventType === "m.room.encrypted") {
+            try {
+                const {fragmentId, entryIndex} = entry;
+                const key = new EventKey(fragmentId, entryIndex);
+                const decryptedEvent = await this._roomEncryption.decrypt(
+                    entry.event, isSync, key, txn);
+                if (decryptedEvent) {
+                    entry.replaceWithDecrypted(decryptedEvent);
                 }
+            } catch (err) {
+                console.warn("event decryption error", err, entry.event);
+                entry.setDecryptionError(err);
             }
-        }));
-        return entries;
+        }
+        return entry;
+    }
+
+    async _decryptEntries(entries, txn, isSync = false) {
+        return await Promise.all(entries.map(async e => this._decryptEntry(e, txn, isSync)));
     }
 
     /** @package */
@@ -83,7 +131,7 @@ export class Room extends EventEmitter {
         // decrypt if applicable
         let entries = encryptedEntries;
         if (this._roomEncryption) {
-            entries = await this._decryptSyncEntries(encryptedEntries, txn);
+            entries = await this._decryptEntries(encryptedEntries, txn, true);
         }
         // fetch new members while we have txn open,
         // but don't make any in-memory changes yet
@@ -116,12 +164,8 @@ export class Room extends EventEmitter {
     /** @package */
     afterSync({summaryChanges, newTimelineEntries, newLiveKey, removedPendingEvents, memberChanges, heroChanges}) {
         this._syncWriter.afterSync(newLiveKey);
-        // encryption got enabled
         if (!this._summary.encryption && summaryChanges.encryption && !this._roomEncryption) {
-            this._roomEncryption = this._createRoomEncryption(this, summaryChanges.encryption);
-            if (this._roomEncryption) {
-                this._sendQueue.enableEncryption(this._roomEncryption);
-            }
+            this._enableEncryption(summaryChanges.encryption);
         }
         if (memberChanges.size) {
             if (this._changedMembersDuringSync) {
@@ -170,10 +214,7 @@ export class Room extends EventEmitter {
         try {
             this._summary.load(summary);
             if (this._summary.encryption) {
-                this._roomEncryption = this._createRoomEncryption(this, this._summary.encryption);
-                if (this._roomEncryption) {
-                    this._sendQueue.enableEncryption(this._roomEncryption);
-                }
+                this._enableEncryption(this._summary.encryption);
             }
             // need to load members for name?
             if (this._summary.needsHeroes) {
@@ -231,11 +272,18 @@ export class Room extends EventEmitter {
             }
         }).response();
 
-        const txn = await this._storage.readWriteTxn([
+        let stores = [
             this._storage.storeNames.pendingEvents,
             this._storage.storeNames.timelineEvents,
             this._storage.storeNames.timelineFragments,
-        ]);
+        ];
+        if (this._roomEncryption) {
+            stores = stores.concat([
+                this._storage.storeNames.inboundGroupSessions,
+                this._storage.storeNames.groupSessionDecryptions,
+            ]);
+        }
+        const txn = await this._storage.readWriteTxn(stores);
         let removedPendingEvents;
         let gapResult;
         try {
@@ -245,9 +293,12 @@ export class Room extends EventEmitter {
             const gapWriter = new GapWriter({
                 roomId: this._roomId,
                 storage: this._storage,
-                fragmentIdComparer: this._fragmentIdComparer
+                fragmentIdComparer: this._fragmentIdComparer,
             });
             gapResult = await gapWriter.writeFragmentFill(fragmentEntry, response, txn);
+            if (this._roomEncryption) {
+                gapResult.entries = await this._decryptEntries(gapResult.entries, false, txn);
+            }
         } catch (err) {
             txn.abort();
             throw err;
@@ -378,6 +429,9 @@ export class Room extends EventEmitter {
             },
             user: this._user,
         });
+        if (this._roomEncryption) {
+            this._timeline.enableEncryption(this._decryptEntries.bind(this));
+        }
         await this._timeline.load();
         return this._timeline;
     }

src/matrix/room/timeline/Timeline.js

@@ -18,6 +18,7 @@ import {SortedArray, MappedList, ConcatList} from "../../../observable/index.js"
 import {Direction} from "./Direction.js";
 import {TimelineReader} from "./persistence/TimelineReader.js";
 import {PendingEventEntry} from "./entries/PendingEventEntry.js";
+import {EventEntry} from "./entries/EventEntry.js";
 
 export class Timeline {
     constructor({roomId, storage, closeCallback, fragmentIdComparer, pendingEvents, user}) {
@@ -45,6 +46,27 @@ export class Timeline {
         this._remoteEntries.setManySorted(entries);
     }
 
+    findEntry(eventKey) {
+        // a storage event entry has a fragmentId and eventIndex property, used for sorting,
+        // just like an EventKey, so this will work, but perhaps a bit brittle.
+        const entry = new EventEntry(eventKey, this._fragmentIdComparer);
+        try {
+            const idx = this._remoteEntries.indexOf(entry);
+            if (idx !== -1) {
+                return this._remoteEntries.get(idx);
+            }
+        } catch (err) {
+            // fragmentIdComparer threw, ignore
+            return;
+        }
+    }
+
+    replaceEntries(entries) {
+        for (const entry of entries) {
+            this._remoteEntries.replace(entry);
+        }
+    }
+
     // TODO: should we rather have generic methods for
     // - adding new entries
     // - updating existing entries (redaction, relations)
@@ -84,4 +106,8 @@ export class Timeline {
             this._closeCallback = null;
         }
     }
+
+    enableEncryption(decryptEntries) {
+        this._timelineReader.enableEncryption(decryptEntries);
+    }
 }

src/observable/list/SortedArray.js

@@ -41,6 +41,22 @@ export class SortedArray extends BaseObservableList {
         }
     }
 
+    replace(item) {
+        const idx = this.indexOf(item);
+        if (idx !== -1) {
+            this._items[idx] = item;
+        }
+    }
+
+    indexOf(item) {
+        const idx = sortedIndex(this._items, item, this._comparator);
+        if (idx < this._items.length && this._comparator(this._items[idx], item) === 0) {
+            return idx;
+        } else {
+            return -1;
+        }
+    }
+
     set(item, updateParams = null) {
         const idx = sortedIndex(this._items, item, this._comparator);
         if (idx >= this._items.length || this._comparator(this._items[idx], item) !== 0) {