From 5d12aef6db5f9d576c687e1189bdc7e2f952e4b9 Mon Sep 17 00:00:00 2001
From: Bruno Windels <brunow@element.io>
Date: Fri, 6 Nov 2020 10:32:37 +0100
Subject: [PATCH] support pre-sharing room keys in room encryption

---
 src/matrix/e2ee/RoomEncryption.js    | 25 ++++++++++--
 src/matrix/e2ee/megolm/Encryption.js | 61 +++++++++++++++++++---------
 2 files changed, 63 insertions(+), 23 deletions(-)

diff --git a/src/matrix/e2ee/RoomEncryption.js b/src/matrix/e2ee/RoomEncryption.js
index 90c63b4a..7d60a475 100644
--- a/src/matrix/e2ee/RoomEncryption.js
+++ b/src/matrix/e2ee/RoomEncryption.js
@@ -244,7 +244,26 @@ export class RoomEncryption {
         }
         
         return matches;
-        
+    }
+
+    /** shares the encryption key for the next message if needed */
+    async ensureNextMessageEncryptionKeyIsShared(hsApi) {
+        const txn = this._storage.readWriteTxn([
+            this._storage.storeNames.operations,
+            this._storage.storeNames.outboundGroupSessions,
+            this._storage.storeNames.inboundGroupSessions,
+        ]);
+        let roomKeyMessage;
+        try {
+            roomKeyMessage = await this._megolmEncryption.ensureOutboundSession(this._room.id, this._encryptionParams, txn);
+        } catch (err) {
+            txn.abort();
+            throw err;
+        }
+        // will complete the txn
+        if (roomKeyMessage) {
+            await this._shareNewRoomKey(roomKeyMessage, hsApi, txn);
+        }
     }
 
     async encrypt(type, content, hsApi) {
@@ -269,12 +288,12 @@ export class RoomEncryption {
         return false;
     }
 
-    async _shareNewRoomKey(roomKeyMessage, hsApi) {
+    async _shareNewRoomKey(roomKeyMessage, hsApi, txn = null) {
         const devices = await this._deviceTracker.devicesForTrackedRoom(this._room.id, hsApi);
         const userIds = Array.from(devices.reduce((set, device) => set.add(device.userId), new Set()));
 
         // store operation for room key share, in case we don't finish here
-        const writeOpTxn = this._storage.readWriteTxn([this._storage.storeNames.operations]);
+        const writeOpTxn = txn || this._storage.readWriteTxn([this._storage.storeNames.operations]);
         let operationId;
         try {
             operationId = this._writeRoomKeyShareOperation(roomKeyMessage, userIds, writeOpTxn);
diff --git a/src/matrix/e2ee/megolm/Encryption.js b/src/matrix/e2ee/megolm/Encryption.js
index a0769ba1..a1257199 100644
--- a/src/matrix/e2ee/megolm/Encryption.js
+++ b/src/matrix/e2ee/megolm/Encryption.js
@@ -43,6 +43,45 @@ export class Encryption {
         }
     }
 
+    async ensureOutboundSession(roomId, encryptionParams, txn) {
+        let session = new this._olm.OutboundGroupSession();
+        try {
+            let sessionEntry = await txn.outboundGroupSessions.get(roomId);
+            const roomKeyMessage = this._readOrCreateSession(session, sessionEntry, roomId, encryptionParams, txn);
+            if (roomKeyMessage) {
+                this._writeSession(sessionEntry, session, roomId, txn);
+                return roomKeyMessage;
+            }
+        } finally {
+            session.free();
+        }
+    }
+
+    _readOrCreateSession(session, sessionEntry, roomId, encryptionParams, txn) {
+        if (sessionEntry) {
+            session.unpickle(this._pickleKey, sessionEntry.session);
+        }
+        if (!sessionEntry || this._needsToRotate(session, sessionEntry.createdAt, encryptionParams)) {
+            // in the case of rotating, recreate a session as we already unpickled into it
+            if (sessionEntry) {
+                session.free();
+                session = new this._olm.OutboundGroupSession();
+            }
+            session.create();
+            const roomKeyMessage = this._createRoomKeyMessage(session, roomId);
+            this._storeAsInboundSession(session, roomId, txn);
+            return roomKeyMessage;
+        }
+    }
+
+    _writeSession(sessionEntry, session, roomId, txn) {
+        txn.outboundGroupSessions.set({
+            roomId,
+            session: session.pickle(this._pickleKey),
+            createdAt: sessionEntry?.createdAt || this._now(),
+        });
+    }
+
     /**
      * Encrypts a message with megolm
      * @param  {string} roomId           
@@ -61,28 +100,10 @@ export class Encryption {
             let roomKeyMessage;
             let encryptedContent;
             try {
-                // TODO: we could consider keeping the session in memory for the current room
                 let sessionEntry = await txn.outboundGroupSessions.get(roomId);
-                if (sessionEntry) {
-                    session.unpickle(this._pickleKey, sessionEntry.session);
-                }
-                if (!sessionEntry || this._needsToRotate(session, sessionEntry.createdAt, encryptionParams)) {
-                    // in the case of rotating, recreate a session as we already unpickled into it
-                    if (sessionEntry) {
-                        session.free();
-                        session = new this._olm.OutboundGroupSession();
-                    }
-                    session.create();
-                    roomKeyMessage = this._createRoomKeyMessage(session, roomId);
-                    this._storeAsInboundSession(session, roomId, txn);
-                    // TODO: we could tell the Decryption here that we have a new session so it can add it to its cache
-                }
+                roomKeyMessage = this._readOrCreateSession(session, sessionEntry, roomId, encryptionParams, txn);
                 encryptedContent = this._encryptContent(roomId, session, type, content);
-                txn.outboundGroupSessions.set({
-                    roomId,
-                    session: session.pickle(this._pickleKey),
-                    createdAt: sessionEntry?.createdAt || this._now(),
-                });
+                this._writeSession(sessionEntry, session, roomId, txn);
 
             } catch (err) {
                 txn.abort();