From 17fc249fa88ef4b0cb2039e1639f6a3cbbba7b6f Mon Sep 17 00:00:00 2001
From: Bruno Windels <brunow@matrix.org>
Date: Thu, 17 Sep 2020 14:20:15 +0200
Subject: [PATCH] integrate session backup with room encryption and megolm
 decryption

---
 src/matrix/e2ee/RoomEncryption.js    | 47 ++++++++++++++
 src/matrix/e2ee/megolm/Decryption.js | 93 +++++++++++++++++++---------
 2 files changed, 110 insertions(+), 30 deletions(-)

diff --git a/src/matrix/e2ee/RoomEncryption.js b/src/matrix/e2ee/RoomEncryption.js
index b5b56ec2..f337875f 100644
--- a/src/matrix/e2ee/RoomEncryption.js
+++ b/src/matrix/e2ee/RoomEncryption.js
@@ -37,6 +37,13 @@ export class RoomEncryption {
         this._eventIdsByMissingSession = new Map();
         this._senderDeviceCache = new Map();
         this._storage = storage;
+        this._sessionBackup = null;
+    }
+
+    setSessionBackup(sessionBackup) {
+        this._sessionBackup = sessionBackup;
+        // TODO: query session backup for all missing sessions so far
+        // can we query multiple? no, only for sessionId, all for room, or all
     }
 
     notifyTimelineClosed() {
@@ -125,13 +132,53 @@ export class RoomEncryption {
         const sessionId = event.content?.["session_id"];
         const key = `${senderKey}|${sessionId}`;
         let eventIds = this._eventIdsByMissingSession.get(key);
+        // new missing session
         if (!eventIds) {
+            this._requestMissingSessionFromBackup(sessionId).catch(err => {
+                console.error(`Could not get session ${sessionId} from backup`, err);
+            });
             eventIds = new Set();
             this._eventIdsByMissingSession.set(key, eventIds);
         }
         eventIds.add(event.event_id);
     }
 
+    async _requestMissingSessionFromBackup(sessionId) {
+        if (!this._sessionBackup) {
+            // somehow prompt for passphrase here
+            return;
+        }
+        const session = await this._sessionBackup.getSession(this._room.id, sessionId);
+        if (session?.algorithm === MEGOLM_ALGORITHM) {
+            const txn = await this._storage.readWriteTxn([this._storage.storeNames.inboundGroupSessions]);
+            let roomKey;
+            try {
+                roomKey = await this._megolmDecryption.addRoomKeyFromBackup(
+                    this._room.id, sessionId, session, txn);
+            } catch (err) {
+                txn.abort();
+                throw err;
+            }
+            await txn.complete();
+
+            if (roomKey) {
+                // this will call into applyRoomKeys below
+                await this._room.notifyRoomKeys([roomKey]);
+            }
+        } else if (session?.algorithm) {
+            console.info(`Backed-up session of unknown algorithm: ${session.algorithm}`);
+        }
+    }
+
+    /**
+     * @type {RoomKeyDescription}
+     * @property {RoomKeyDescription} senderKey the curve25519 key of the sender
+     * @property {RoomKeyDescription} sessionId
+     * 
+     * 
+     * @param  {Array<RoomKeyDescription>} roomKeys
+     * @return {Array<string>} the event ids that should be retried to decrypt
+     */
     applyRoomKeys(roomKeys) {
         // retry decryption with the new sessions
         const retryEventIds = [];
diff --git a/src/matrix/e2ee/megolm/Decryption.js b/src/matrix/e2ee/megolm/Decryption.js
index 4d756dcb..b53a345e 100644
--- a/src/matrix/e2ee/megolm/Decryption.js
+++ b/src/matrix/e2ee/megolm/Decryption.js
@@ -138,40 +138,73 @@ export class Decryption {
                 return;
             }
 
-            const session = new this._olm.InboundGroupSession();
-            try {
-                session.create(sessionKey);
-
-                let incomingSessionIsBetter = true;
-                const existingSessionEntry = await txn.inboundGroupSessions.get(roomId, senderKey, sessionId);
-                if (existingSessionEntry) {
-                    const existingSession = new this._olm.InboundGroupSession();
-                    try {
-                        existingSession.unpickle(this._pickleKey, existingSessionEntry.session);
-                        incomingSessionIsBetter = session.first_known_index() < existingSession.first_known_index();
-                    } finally {
-                        existingSession.free();
-                    }
-                }
-
-                if (incomingSessionIsBetter) {
-                    const sessionEntry = {
-                        roomId,
-                        senderKey,
-                        sessionId,
-                        session: session.pickle(this._pickleKey),
-                        claimedKeys: {ed25519: claimedEd25519Key},
-                    };
-                    txn.inboundGroupSessions.set(sessionEntry);
-                    newSessions.push(sessionEntry);
-                }
-            } finally {
-                session.free();
+            const sessionEntry = await this._writeInboundSession(
+                roomId, senderKey, claimedEd25519Key, sessionId, sessionKey, txn);
+            if (sessionEntry) {
+                newSessions.push(sessionEntry);
             }
-
         }
         // this will be passed to the Room in notifyRoomKeys
         return newSessions;
     }
+
+    /*
+    sessionInfo is a response from key backup and has the following keys:
+        algorithm
+        forwarding_curve25519_key_chain
+        sender_claimed_keys
+        sender_key
+        session_key
+     */
+    async addRoomKeyFromBackup(roomId, sessionId, sessionInfo, txn) {
+        const sessionKey = sessionInfo["session_key"];
+        const senderKey = sessionInfo["sender_key"];
+        const claimedEd25519Key = sessionInfo["sender_claimed_keys"]?.["ed25519"];
+
+        if (
+            typeof roomId !== "string" || 
+            typeof sessionId !== "string" || 
+            typeof senderKey !== "string" ||
+            typeof sessionKey !== "string" ||
+            typeof claimedEd25519Key !== "string"
+        ) {
+            return;
+        }
+        return await this._writeInboundSession(
+            roomId, senderKey, claimedEd25519Key, sessionId, sessionKey, txn);
+    }
+
+    async _writeInboundSession(roomId, senderKey, claimedEd25519Key, sessionId, sessionKey, txn) {
+        const session = new this._olm.InboundGroupSession();
+        try {
+            session.create(sessionKey);
+
+            let incomingSessionIsBetter = true;
+            const existingSessionEntry = await txn.inboundGroupSessions.get(roomId, senderKey, sessionId);
+            if (existingSessionEntry) {
+                const existingSession = new this._olm.InboundGroupSession();
+                try {
+                    existingSession.unpickle(this._pickleKey, existingSessionEntry.session);
+                    incomingSessionIsBetter = session.first_known_index() < existingSession.first_known_index();
+                } finally {
+                    existingSession.free();
+                }
+            }
+
+            if (incomingSessionIsBetter) {
+                const sessionEntry = {
+                    roomId,
+                    senderKey,
+                    sessionId,
+                    session: session.pickle(this._pickleKey),
+                    claimedKeys: {ed25519: claimedEd25519Key},
+                };
+                txn.inboundGroupSessions.set(sessionEntry);
+                return sessionEntry;
+            }
+        } finally {
+            session.free();
+        }
+    }
 }