WIP to make megolm session in memory once you start sending

deciced against merging this though as it increases the chance of
corrupting the outbound megolm session in case you have multiple tabs open on the same session

we don't officially support that usecase, and even try to automatically log the user out, but
I'm still not 100% sure if I'm comfortable with introducing more breakage if this does happen (no service worker, ...)
so parking this work here for now. I started working on this as part of sending out megolm keys when you start typing.

src/matrix/Session.js

@@ -123,7 +123,6 @@ export class Session {
             account: this._e2eeAccount,
             pickleKey: PICKLE_KEY,
             olm: this._olm,
-            storage: this._storage,
             now: this._platform.clock.now,
             ownDeviceId: this._sessionInfo.deviceId,
         });

src/matrix/e2ee/RoomEncryption.js

@@ -36,6 +36,7 @@ export class RoomEncryption {
         this._deviceTracker = deviceTracker;
         this._olmEncryption = olmEncryption;
         this._megolmEncryption = megolmEncryption;
+        this._megolmRoomEncryption = null;
         this._megolmDecryption = megolmDecryption;
         // content of the m.room.encryption event
         this._encryptionParams = encryptionParams;
@@ -253,6 +254,8 @@ export class RoomEncryption {
             this._storage.storeNames.outboundGroupSessions,
             this._storage.storeNames.inboundGroupSessions,
         ]);
+                await this._deviceTracker.trackRoom(this._room);
+
         let roomKeyMessage;
         try {
             roomKeyMessage = await this._megolmEncryption.ensureOutboundSession(this._room.id, this._encryptionParams, txn);
@@ -268,10 +271,16 @@ export class RoomEncryption {
 
     async encrypt(type, content, hsApi) {
         await this._deviceTracker.trackRoom(this._room);
-        const megolmResult = await this._megolmEncryption.encrypt(this._room.id, type, content, this._encryptionParams);
+        const txn = this._storage.readWriteTxn([
+            this._storage.storeNames.operations,
+            this._storage.storeNames.outboundGroupSessions,
+            this._storage.storeNames.inboundGroupSessions,
+        ]);
+        await this._ensureMegolmEncryption(txn);
+        const megolmResult = this._megolmRoomEncryption.encrypt(type, content, txn);
         if (megolmResult.roomKeyMessage) {
             // TODO: should we await this??
-            this._shareNewRoomKey(megolmResult.roomKeyMessage, hsApi);
+            this._shareNewRoomKey(megolmResult.roomKeyMessage, hsApi, txn);
         }
         return {
             type: ENCRYPTED_TYPE,
@@ -288,12 +297,11 @@ export class RoomEncryption {
         return false;
     }
 
-    async _shareNewRoomKey(roomKeyMessage, hsApi, txn = null) {
+    async _shareNewRoomKey(roomKeyMessage, hsApi, writeOpTxn) {
         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 = txn || this._storage.readWriteTxn([this._storage.storeNames.operations]);
         let operationId;
         try {
             operationId = this._writeRoomKeyShareOperation(roomKeyMessage, userIds, writeOpTxn);
@@ -322,8 +330,8 @@ export class RoomEncryption {
 
     async _addShareRoomKeyOperationForNewMembers(memberChangesArray, txn) {
         const userIds = memberChangesArray.filter(m => m.hasJoined).map(m => m.userId);
-        const roomKeyMessage = await this._megolmEncryption.createRoomKeyMessage(
-            this._room.id, txn);
+        await this._ensureMegolmEncryption(txn);
+        const roomKeyMessage = this._megolmRoomEncryption.createRoomKeyMessage(txn);
         if (roomKeyMessage) {
             this._writeRoomKeyShareOperation(roomKeyMessage, userIds, txn);
         }
@@ -394,10 +402,23 @@ export class RoomEncryption {
         await hsApi.sendToDevice(type, payload, txnId).response();
     }
 
+    async _ensureMegolmEncryption(txn) {
+        if (!this._megolmRoomEncryption) {
+            this._megolmRoomEncryption = await this._megolmEncryption.openRoomEncryption(
+                this._room.id,
+                this._encryptionParams,
+                txn
+            );
+        }
+    }
+
     dispose() {
         this._disposed = true;
         this._megolmBackfillCache.dispose();
         this._megolmSyncCache.dispose();
+        if (this._megolmRoomEncryption) {
+            this._megolmRoomEncryption.dispose();
+        }
     }
 }
 

src/matrix/e2ee/megolm/Encryption.js

@@ -17,164 +17,186 @@ limitations under the License.
 import {MEGOLM_ALGORITHM} from "../common.js";
 
 export class Encryption {
-    constructor({pickleKey, olm, account, storage, now, ownDeviceId}) {
+    constructor({pickleKey, olm, account, now, ownDeviceId}) {
         this._pickleKey = pickleKey;
         this._olm = olm;
         this._account = account;
-        this._storage = storage;
         this._now = now;
         this._ownDeviceId = ownDeviceId;
     }
 
-    discardOutboundSession(roomId, txn) {
-        txn.outboundGroupSessions.remove(roomId);
-    }
-
-    async createRoomKeyMessage(roomId, txn) {
-        let sessionEntry = await txn.outboundGroupSessions.get(roomId);
+    async openRoomEncryption(roomId, encryptionParams, txn) {
+        const sessionEntry = await txn.outboundGroupSessions.get(roomId);
+        let session = null;
         if (sessionEntry) {
-            const session = new this._olm.OutboundGroupSession();
-            try {
-                session.unpickle(this._pickleKey, sessionEntry.session);
-                return this._createRoomKeyMessage(session, roomId);
-            } finally {
-                session.free();
-            }
-        }
-    }
-
-    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.unpickle(this._pickleKey, sessionEntry.session);
         }
-            session.create();
-            const roomKeyMessage = this._createRoomKeyMessage(session, roomId);
-            this._storeAsInboundSession(session, roomId, txn);
-            return roomKeyMessage;
+        return new RoomEncryption({
+            pickleKey: this._pickleKey,
+            olm: this._olm,
+            account: this._account,
+            now: this._now,
+            ownDeviceId: this._ownDeviceId,
+            sessionEntry,
+            session,
+            roomId,
+            encryptionParams
+        });
     }
+}
+
+export class RoomEncryption {
+    constructor({pickleKey, olm, account, now, roomId, encryptionParams, sessionEntry, session, ownDeviceId}) {
+        this._pickleKey = pickleKey;
+        this._olm = olm;
+        this._account = account;
+        this._now = now;
+        this._roomId = roomId;
+        this._encryptionParams = encryptionParams;
+        this._ownDeviceId = ownDeviceId;
+        this._sessionEntry = sessionEntry;
+        this._session = session;
     }
 
-    _writeSession(sessionEntry, session, roomId, txn) {
-        txn.outboundGroupSessions.set({
-            roomId,
-            session: session.pickle(this._pickleKey),
-            createdAt: sessionEntry?.createdAt || this._now(),
-        });
+    /**
+     * Discards the outbound session, if any.
+     * @param  {Transaction} txn         a storage transaction with readwrite access to outboundGroupSessions and inboundGroupSessions stores
+     */
+    discardOutboundSession(txn) {
+        txn.outboundGroupSessions.remove(this._roomId);
+        if (this._session) {
+            this._session.free();
+        }
+        this._session = null;
+        this._sessionEntry = null;
+    }
+
+    /**
+     * Creates an outbound session if non exists already
+     * @param  {Transaction} txn         a storage transaction with readwrite access to outboundGroupSessions and inboundGroupSessions stores
+     * @return {boolean} true if a session has been created. Call `createRoomKeyMessage` to share the new session.
+     */
+    ensureOutboundSession(txn) {
+        if (this._readOrCreateSession(txn)) {
+            this._writeSession(txn);
+            return true;
+        }
+        return false;
     }
 
     /**
      * Encrypts a message with megolm
-     * @param  {string} roomId           
      * @param  {string} type             event type to encrypt
      * @param  {string} content          content to encrypt
-     * @param  {object} encryptionParams the content of the m.room.encryption event
+     * @param  {Transaction} txn         a storage transaction with readwrite access to outboundGroupSessions and inboundGroupSessions stores
      * @return {Promise<EncryptionResult>}
      */
-    async encrypt(roomId, type, content, encryptionParams) {
-        let session = new this._olm.OutboundGroupSession();
-        try {
-            const txn = this._storage.readWriteTxn([
-                this._storage.storeNames.inboundGroupSessions,
-                this._storage.storeNames.outboundGroupSessions,
-            ]);
+    encrypt(type, content, txn) {
         let roomKeyMessage;
-            let encryptedContent;
-            try {
-                let sessionEntry = await txn.outboundGroupSessions.get(roomId);
-                roomKeyMessage = this._readOrCreateSession(session, sessionEntry, roomId, encryptionParams, txn);
-                encryptedContent = this._encryptContent(roomId, session, type, content);
-                this._writeSession(sessionEntry, session, roomId, txn);
-
-            } catch (err) {
-                txn.abort();
-                throw err;
+        if (this._readOrCreateSession(txn)) {
+            // important to create the room key message before encrypting
+            // so the message index isn't advanced yet
+            roomKeyMessage = this.createRoomKeyMessage();
         }
-            await txn.complete();
+        const encryptedContent = this._encryptContent(type, content);
+        this._writeSession(txn);
         return new EncryptionResult(encryptedContent, roomKeyMessage);
-        } finally {
-            if (session) {
-                session.free();
-            }
-        }
     }
 
-    _needsToRotate(session, createdAt, encryptionParams) {
+    needsNewSession() {
+        if (!this._session) {
+            return true;
+        }
         let rotationPeriodMs = 604800000; // default
-        if (Number.isSafeInteger(encryptionParams?.rotation_period_ms)) {
-            rotationPeriodMs = encryptionParams?.rotation_period_ms;
+        if (Number.isSafeInteger(this._encryptionParams?.rotation_period_ms)) {
+            rotationPeriodMs = this._encryptionParams?.rotation_period_ms;
         }
         let rotationPeriodMsgs = 100; // default
-        if (Number.isSafeInteger(encryptionParams?.rotation_period_msgs)) {
-            rotationPeriodMsgs = encryptionParams?.rotation_period_msgs;
+        if (Number.isSafeInteger(this._encryptionParams?.rotation_period_msgs)) {
+            rotationPeriodMsgs = this._encryptionParams?.rotation_period_msgs;
+        }
+        // assume this is a new session if sessionEntry hasn't been created/written yet
+        if (this._sessionEntry && this._now() > (this._sessionEntry.createdAt + rotationPeriodMs)) {
+            return true;
+        }
+        if (this._session.message_index() >= rotationPeriodMsgs) {
+            return true;
+        }
+        return false;
     }
 
-        if (this._now() > (createdAt + rotationPeriodMs)) {
-            return true;
+    createRoomKeyMessage() {
+        if (!this._session) {
+            return;
         }
-        if (session.message_index() >= rotationPeriodMsgs) {
-            return true;
+        return {
+            room_id: this._roomId,
+            session_id: this._session.session_id(),
+            session_key: this._session.session_key(),
+            algorithm: MEGOLM_ALGORITHM,
+            // chain_index is ignored by element-web if not all clients
+            // but let's send it anyway, as element-web does so
+            chain_index: this._session.message_index()
         }
     }
 
-    _encryptContent(roomId, session, type, content) {
+    dispose() {
+        if (this._session) {
+            this._session.free();
+        }
+    }
+
+    _encryptContent(type, content) {
         const plaintext = JSON.stringify({
-            room_id: roomId,
+            room_id: this._roomId,
             type,
             content
         });
-        const ciphertext = session.encrypt(plaintext);
+        const ciphertext = this._session.encrypt(plaintext);
 
         const encryptedContent = {
             algorithm: MEGOLM_ALGORITHM,
             sender_key: this._account.identityKeys.curve25519,
             ciphertext,
-            session_id: session.session_id(),
+            session_id: this._session.session_id(),
             device_id: this._ownDeviceId
         };
 
         return encryptedContent;
     }
 
-    _createRoomKeyMessage(session, roomId) {
-        return {
-            room_id: roomId,
-            session_id: session.session_id(),
-            session_key: session.session_key(),
-            algorithm: MEGOLM_ALGORITHM,
-            // chain_index is ignored by element-web if not all clients
-            // but let's send it anyway, as element-web does so
-            chain_index: session.message_index()
+
+    _readOrCreateSession(txn) {
+        if (this.needsNewSession()) {
+            if (this._session) {
+                this._session.free();
+                this._session = new this._olm.OutboundGroupSession();
             }
+            this._session.create();
+            this._storeAsInboundSession(txn);
+            return true;
+        }
+        return false;
     }
 
-    _storeAsInboundSession(outboundSession, roomId, txn) {
+    _writeSession(txn) {
+        this._sessionEntry = {
+            roomId: this._roomId,
+            session: this._session.pickle(this._pickleKey),
+            createdAt: this._sessionEntry?.createdAt || this._now(),
+        };
+        txn.outboundGroupSessions.set(this._sessionEntry);
+    }
+
+    _storeAsInboundSession(txn) {
         const {identityKeys} = this._account;
         const claimedKeys = {ed25519: identityKeys.ed25519};
         const session = new this._olm.InboundGroupSession();
         try {
-            session.create(outboundSession.session_key());
+            session.create(this._session.session_key());
             const sessionEntry = {
-                roomId,
+                roomId: this._roomId,
                 senderKey: identityKeys.curve25519,
                 sessionId: session.session_id(),
                 session: session.pickle(this._pickleKey),
@@ -187,7 +209,6 @@ export class Encryption {
         }
     }
 }
-
 /**
  * @property {object?} roomKeyMessage  if encrypting this message
  *                                     created a new outbound session,