split up megolm decryption so it can happen in multiple steps,see README

src/matrix/e2ee/megolm/Decryption.js

@@ -15,102 +15,101 @@ limitations under the License.
 */
 
 import {DecryptionError} from "../common.js";
-import {DecryptionResult} from "../DecryptionResult.js";
+import {groupBy} from "../../../utils/groupBy.js";
 
-const CACHE_MAX_SIZE = 10;
+import {SessionInfo} from "./decryption/SessionInfo.js";
+import {DecryptionPreparation} from "./decryption/DecryptionPreparation.js";
+import {SessionDecryption} from "./decryption/SessionDecryption.js";
+import {SessionCache} from "./decryption/SessionCache.js";
+
+function getSenderKey(event) {
+    return event.content?.["sender_key"];
+}
+
+function getSessionId(event) {
+    return event.content?.["session_id"];
+}
+
+function getCiphertext(event) {
+    return event.content?.ciphertext;
+}
 
 export class Decryption {
     constructor({pickleKey, olm}) {
         this._pickleKey = pickleKey;
         this._olm = olm;
+        // this._worker = new MessageHandler(new Worker("worker-2580578233.js"));
     }
 
-    createSessionCache() {
-        return new SessionCache();
+    createSessionCache(fallback) {
+        return new SessionCache(fallback);
     }
 
     /**
-     * [decrypt description]
+     * Reads all the state from storage to be able to decrypt the given events.
+     * Decryption can then happen outside of a storage transaction.
      * @param  {[type]} roomId       [description]
-     * @param  {[type]} event        [description]
+     * @param  {[type]} events        [description]
      * @param  {[type]} sessionCache [description]
      * @param  {[type]} txn          [description]
-     * @return {DecryptionResult?} the decrypted event result, or undefined if the session id is not known.
+     * @return {DecryptionPreparation}
      */
-    async decrypt(roomId, event, sessionCache, txn) {
-        const senderKey = event.content?.["sender_key"];
-        const sessionId = event.content?.["session_id"];
-        const ciphertext = event.content?.ciphertext;
+    async prepareDecryptAll(roomId, events, sessionCache, txn) {
+        const errors = new Map();
+        const validEvents = [];
 
-        if (
-            typeof senderKey !== "string" ||
-            typeof sessionId !== "string" ||
-            typeof ciphertext !== "string"
-        ) {
-            throw new DecryptionError("MEGOLM_INVALID_EVENT", event);
+        for (const event of events) {
+            const isValid = typeof getSenderKey(event) === "string" &&
+                            typeof getSessionId(event) === "string" &&
+                            typeof getCiphertext(event) === "string";
+            if (isValid) {
+                validEvents.push(event);
+            } else {
+                errors.set(event.event_id, new DecryptionError("MEGOLM_INVALID_EVENT", event))
+            }
         }
 
-        let session;
-        let claimedKeys;
-        const cacheEntry = sessionCache.get(roomId, senderKey, sessionId);
-        if (cacheEntry) {
-            session = cacheEntry.session;
-            claimedKeys = cacheEntry.claimedKeys;
+        const eventsBySession = groupBy(validEvents, event => {
+            return `${getSenderKey(event)}|${getSessionId(event)}`;
+        });
+
+        const sessionDecryptions = [];
+
+        await Promise.all(Array.from(eventsBySession.values()).map(async eventsForSession => {
+            const first = eventsForSession[0];
+            const senderKey = getSenderKey(first);
+            const sessionId = getSessionId(first);
+            const sessionInfo = await this._getSessionInfo(roomId, senderKey, sessionId, sessionCache, txn);
+            if (!sessionInfo) {
+                for (const event of eventsForSession) {
+                    errors.set(event.event_id, new DecryptionError("MEGOLM_NO_SESSION", event));
+                }
             } else {
+                sessionDecryptions.push(new SessionDecryption(sessionInfo, eventsForSession));
+            }
+        }));
+
+        return new DecryptionPreparation(roomId, sessionDecryptions, errors);
+    }
+
+    async _getSessionInfo(roomId, senderKey, sessionId, sessionCache, txn) {
+        let sessionInfo;
+        sessionInfo = sessionCache.get(roomId, senderKey, sessionId);
+        if (!sessionInfo) {
             const sessionEntry = await txn.inboundGroupSessions.get(roomId, senderKey, sessionId);
             if (sessionEntry) {
-                session = new this._olm.InboundGroupSession();
+                let session = new this._olm.InboundGroupSession();
                 try {
                     session.unpickle(this._pickleKey, sessionEntry.session);
+                    sessionInfo = new SessionInfo(roomId, senderKey, session, sessionEntry.claimedKeys);
                 } catch (err) {
                     session.free();
                     throw err;
                 }
-                claimedKeys = sessionEntry.claimedKeys;
-                sessionCache.add(roomId, senderKey, session, claimedKeys);
+                sessionCache.add(sessionInfo);
             }
         }
-        if (!session) {
-            return;
-        }
-        const {plaintext, message_index: messageIndex} = session.decrypt(ciphertext);
-        let payload;
-        try {
-            payload = JSON.parse(plaintext);
-        } catch (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});
-        }
-        await this._handleReplayAttack(roomId, sessionId, messageIndex, event, txn);
-        return new DecryptionResult(payload, senderKey, claimedKeys);
-    }
-
-    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);
-        if (decryption && decryption.eventId !== eventId) {
-            // 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_REPLAYED_INDEX", event, {
-                messageIndex,
-                badEventId,
-                otherEventId: decryption.eventId
-            });
-        }
-        if (!decryption) {
-            txn.groupSessionDecryptions.set({
-                roomId,
-                sessionId,
-                messageIndex,
-                eventId,
-                timestamp
-            });
-        }
+        return sessionInfo;
     }
 
     /**
@@ -165,55 +164,3 @@ export class Decryption {
     }
 }
 
-class SessionCache {
-    constructor() {
-        this._sessions = [];
-    }
-
-    /**
-     * @type {CacheEntry}
-     * @property {InboundGroupSession} session the unpickled session
-     * @property {Object} claimedKeys an object with the claimed ed25519 key
-     *
-     * 
-     * @param  {string} roomId
-     * @param  {string} senderKey
-     * @param  {string} sessionId
-     * @return {CacheEntry?}
-     */
-    get(roomId, senderKey, sessionId) {
-        const idx = this._sessions.findIndex(s => {
-            return s.roomId === roomId &&
-                s.senderKey === senderKey &&
-                sessionId === s.session.session_id();
-        });
-        if (idx !== -1) {
-            const entry = this._sessions[idx];
-            // move to top
-            if (idx > 0) {
-                this._sessions.splice(idx, 1);
-                this._sessions.unshift(entry);
-            }
-            return entry;
-        }
-    }
-
-    add(roomId, senderKey, session, claimedKeys) {
-        // add new at top
-        this._sessions.unshift({roomId, senderKey, session, claimedKeys});
-        if (this._sessions.length > CACHE_MAX_SIZE) {
-            // free sessions we're about to remove
-            for (let i = CACHE_MAX_SIZE; i < this._sessions.length; i += 1) {
-                this._sessions[i].session.free();
-            }
-            this._sessions = this._sessions.slice(0, CACHE_MAX_SIZE);
-        }
-    }
-
-    dispose() {
-        for (const entry of this._sessions) {
-            entry.session.free();
-        }
-
-    }
-}

src/matrix/e2ee/megolm/decryption/DecryptionChanges.js

@@ -0,0 +1,78 @@
+/*
+Copyright 2020 The Matrix.org Foundation C.I.C.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+import {DecryptionError} from "../../common.js";
+
+export class DecryptionChanges {
+    constructor(roomId, results, errors, replayEntries) {
+        this._roomId = roomId;
+        this._results = results;
+        this._errors = errors;
+        this._replayEntries = replayEntries;
+    }
+
+    /**
+     * @type MegolmBatchDecryptionResult
+     * @property {Map<string, DecryptionResult>} results a map of event id to decryption result
+     * @property {Map<string, Error>} errors event id -> errors
+     * 
+     * Handle replay attack detection, and return result
+     * @param  {[type]} txn [description]
+     * @return {MegolmBatchDecryptionResult}
+     */
+    async write(txn) {
+        await Promise.all(this._replayEntries.map(async replayEntry => {
+            try {
+                this._handleReplayAttack(this._roomId, replayEntry, txn);
+            } catch (err) {
+                this._errors.set(replayEntry.eventId, err);
+            }
+        }));
+        return {
+            results: this._results,
+            errors: this._errors
+        };
+    }
+
+    async _handleReplayAttack(roomId, replayEntry, txn) {
+        const {messageIndex, sessionId, eventId, timestamp} = replayEntry;
+        const decryption = await txn.groupSessionDecryptions.get(roomId, sessionId, messageIndex);
+
+        if (decryption && decryption.eventId !== eventId) {
+            // the one with the newest timestamp should be the attack
+            const decryptedEventIsBad = decryption.timestamp < timestamp;
+            const badEventId = decryptedEventIsBad ? eventId : decryption.eventId;
+            // discard result
+            this._results.delete(eventId);
+
+            throw new DecryptionError("MEGOLM_REPLAYED_INDEX", event, {
+                messageIndex,
+                badEventId,
+                otherEventId: decryption.eventId
+            });
+        }
+
+        if (!decryption) {
+            txn.groupSessionDecryptions.set({
+                roomId,
+                sessionId,
+                messageIndex,
+                eventId,
+                timestamp
+            });
+        }
+    }
+}

src/matrix/e2ee/megolm/decryption/DecryptionPreparation.js

@@ -0,0 +1,52 @@
+/*
+Copyright 2020 The Matrix.org Foundation C.I.C.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+import {DecryptionChanges} from "./DecryptionChanges.js";
+import {mergeMap} from "../../../../utils/mergeMap.js";
+
+/**
+ * Class that contains all the state loaded from storage to decrypt the given events
+ */
+export class DecryptionPreparation {
+    constructor(roomId, sessionDecryptions, errors) {
+        this._roomId = roomId;
+        this._sessionDecryptions = sessionDecryptions;
+        this._initialErrors = errors;
+    }
+
+    async decrypt() {
+        try {
+            const errors = this._initialErrors;
+            const results = new Map();
+            const replayEntries = [];
+            await Promise.all(this._sessionDecryptions.map(async sessionDecryption => {
+                const sessionResult = await sessionDecryption.decryptAll();
+                mergeMap(sessionResult.errors, errors);
+                mergeMap(sessionResult.results, results);
+                replayEntries.push(...sessionResult.replayEntries);
+            }));
+            return new DecryptionChanges(this._roomId, results, errors, replayEntries);
+        } finally {
+            this.dispose();
+        }
+    }
+
+    dispose() {
+        for (const sd of this._sessionDecryptions) {
+            sd.dispose();
+        }
+    }
+}

src/matrix/e2ee/megolm/decryption/README.md

@@ -0,0 +1,6 @@
+Lots of classes here. The complexity comes from needing to offload decryption to a webworker, mainly for IE11. We can't keep a idb transaction open while waiting for the response from the worker, so need to batch decryption of multiple events and do decryption in multiple steps:
+
+ 1. Read all used inbound sessions for the batch of events, requires a read txn. This happens in `Decryption`. Sessions are loaded into `SessionInfo` objects, which are also kept in a `SessionCache` to prevent having to read and unpickle them all the time.
+ 2. Actually decrypt. No txn can stay open during this step, as it can be offloaded to a worker and is thus async. This happens in `DecryptionPreparation`, which delegates to `SessionDecryption` per session.
+ 3. Read and write for the replay detection, requires a read/write txn. This happens in `DecryptionChanges`
+ 4. Return the decrypted entries, and errors if any

src/matrix/e2ee/megolm/decryption/ReplayDetectionEntry.js

@@ -0,0 +1,24 @@
+/*
+Copyright 2020 The Matrix.org Foundation C.I.C.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+export class ReplayDetectionEntry {
+    constructor(sessionId, messageIndex, event) {
+        this.sessionId = sessionId;
+        this.messageIndex = messageIndex;
+        this.eventId = event.event_id;
+        this.timestamp = event.origin_server_ts;
+    }
+}

src/matrix/e2ee/megolm/decryption/SessionCache.js

@@ -0,0 +1,68 @@
+/*
+Copyright 2020 The Matrix.org Foundation C.I.C.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+const CACHE_MAX_SIZE = 10;
+
+/**
+ * Cache of unpickled inbound megolm session.
+ */
+export class SessionCache {
+    constructor() {
+        this._sessions = [];
+    }
+
+    /**
+     * @param  {string} roomId
+     * @param  {string} senderKey
+     * @param  {string} sessionId
+     * @return {SessionInfo?}
+     */
+    get(roomId, senderKey, sessionId) {
+        const idx = this._sessions.findIndex(s => {
+            return s.roomId === roomId &&
+                s.senderKey === senderKey &&
+                sessionId === s.session.session_id();
+        });
+        if (idx !== -1) {
+            const sessionInfo = this._sessions[idx];
+            // move to top
+            if (idx > 0) {
+                this._sessions.splice(idx, 1);
+                this._sessions.unshift(sessionInfo);
+            }
+            return sessionInfo;
+        }
+    }
+
+    add(sessionInfo) {
+        sessionInfo.retain();
+        // add new at top
+        this._sessions.unshift(sessionInfo);
+        if (this._sessions.length > CACHE_MAX_SIZE) {
+            // free sessions we're about to remove
+            for (let i = CACHE_MAX_SIZE; i < this._sessions.length; i += 1) {
+                this._sessions[i].release();
+            }
+            this._sessions = this._sessions.slice(0, CACHE_MAX_SIZE);
+        }
+    }
+
+    dispose() {
+        for (const sessionInfo of this._sessions) {
+            sessionInfo.release();
+        }
+    }
+}

src/matrix/e2ee/megolm/decryption/SessionDecryption.js

@@ -0,0 +1,75 @@
+/*
+Copyright 2020 The Matrix.org Foundation C.I.C.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+import {DecryptionResult} from "../../DecryptionResult.js";
+import {DecryptionError} from "../../common.js";
+import {ReplayDetectionEntry} from "./ReplayDetectionEntry.js";
+
+/**
+ * Does the actual decryption of all events for a given megolm session in a batch
+ */
+export class SessionDecryption {
+    constructor(sessionInfo, events) {
+        sessionInfo.retain();
+        this._sessionInfo = sessionInfo;
+        this._events = events;
+    }
+
+    async decryptAll() {
+        const replayEntries = [];
+        const results = new Map();
+        let errors;
+        const roomId = this._sessionInfo.roomId;
+
+        await Promise.all(this._events.map(async event => {
+            try {
+                const {session} = this._sessionInfo;
+                const ciphertext = event.content.ciphertext;
+                const {plaintext, message_index: messageIndex} = await this._decrypt(session, ciphertext);
+                let payload;
+                try {
+                    payload = JSON.parse(plaintext);
+                } catch (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});
+                }
+                replayEntries.push(new ReplayDetectionEntry(session.session_id(), messageIndex, event));
+                const result = new DecryptionResult(payload, this._sessionInfo.senderKey, this._sessionInfo.claimedKeys);
+                results.set(event.event_id, result);
+            } catch (err) {
+                if (!errors) {
+                    errors = new Map();
+                }
+                errors.set(event.event_id, err);
+            }
+        }));
+
+        return {results, errors, replayEntries};
+    }
+
+    async _decrypt(session, ciphertext) {
+        // const sessionKey = session.export_session(session.first_known_index());
+        // return this._worker.decrypt(sessionKey, ciphertext);
+        return session.decrypt(ciphertext);
+    }
+
+    dispose() {
+        this._sessionInfo.release();
+    }
+}

src/matrix/e2ee/megolm/decryption/SessionInfo.js

@@ -0,0 +1,44 @@
+/*
+Copyright 2020 The Matrix.org Foundation C.I.C.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+/**
+ * session loaded in memory with everything needed to create DecryptionResults
+ * and to store/retrieve it in the SessionCache
+ */
+export class SessionInfo {
+    constructor(roomId, senderKey, session, claimedKeys) {
+        this.roomId = roomId;
+        this.senderKey = senderKey;
+        this.session = session;
+        this.claimedKeys = claimedKeys;
+        this._refCounter = 0;
+    }
+
+    retain() {
+        this._refCounter += 1;
+    }
+
+    release() {
+        this._refCounter -= 1;
+        if (this._refCounter <= 0) {
+            this.dispose();
+        }
+    }
+
+    dispose() {
+        this.session.free();
+    }
+}

src/utils/mergeMap.js

@@ -0,0 +1,41 @@
+/*
+Copyright 2020 Bruno Windels <bruno@windels.cloud>
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+export function mergeMap(src, dst) {
+    if (src) {
+        for (const [key, value] of src.entries()) {
+            dst.set(key, value);
+        }
+    }
+}
+
+export function tests() {
+    return {
+        "mergeMap with src": assert => {
+            const src = new Map();
+            src.set(1, "a");
+            const dst = new Map();
+            dst.set(2, "b");
+            mergeMap(src, dst);
+            assert.equal(dst.get(1), "a");
+            assert.equal(dst.get(2), "b");
+            assert.equal(src.get(2), null);
+        },
+        "mergeMap without src doesn't fail": () => {
+            mergeMap(undefined, new Map());
+        }
+    }
+}