Merge pull request #191 from vector-im/bwindels/preshare-megolmsessions

Share megolm session once you start typing

src/domain/session/room/RoomViewModel.js

@@ -188,8 +188,14 @@ class ComposerViewModel extends ViewModel {
         return !this._isEmpty;
     }
 
-    setInput(text) {
+    async setInput(text) {
+        const wasEmpty = this._isEmpty;
         this._isEmpty = text.length === 0;
+        if (wasEmpty && !this._isEmpty) {
+            this._roomVM._room.ensureMessageKeyIsShared();
+        }
+        if (wasEmpty !== this._isEmpty) {
             this.emitChange("canSend");
         }
     }
+}

src/matrix/e2ee/Account.js

@@ -149,10 +149,14 @@ export class Account {
         }
     }
 
-    createOutboundOlmSession(theirIdentityKey, theirOneTimeKey) {
+    async createOutboundOlmSession(theirIdentityKey, theirOneTimeKey) {
         const newSession = new this._olm.Session();
         try {
+            if (this._olmWorker) {
+                await this._olmWorker.createOutboundOlmSession(this._account, newSession, theirIdentityKey, theirOneTimeKey);
+            } else {
                 newSession.create_outbound(this._account, theirIdentityKey, theirOneTimeKey);
+            }
             return newSession;
         } catch (err) {
             newSession.free();

src/matrix/e2ee/OlmWorker.js

@@ -37,6 +37,12 @@ export class OlmWorker {
         account.unpickle("", pickle);
     }
 
+    async createOutboundSession(account, newSession, theirIdentityKey, theirOneTimeKey) {
+        const accountPickle = account.pickle("");
+        const sessionPickle = await this._workerPool.send({type: "olm_create_outbound", accountPickle, theirIdentityKey, theirOneTimeKey}).response();
+        newSession.unpickle("", sessionPickle);
+    }
+
     dispose() {
         this._workerPool.dispose();
     }

src/matrix/e2ee/RoomEncryption.js

@@ -20,6 +20,10 @@ import {mergeMap} from "../../utils/mergeMap.js";
 import {makeTxnId} from "../common.js";
 
 const ENCRYPTED_TYPE = "m.room.encrypted";
+// how often ensureMessageKeyIsShared can check if it needs to
+// create a new outbound session
+// note that encrypt could still create a new session
+const MIN_PRESHARE_INTERVAL = 60 * 1000; // 1min
 
 function encodeMissingSessionKey(senderKey, sessionId) {
     return `${senderKey}|${sessionId}`;
@@ -55,6 +59,7 @@ export class RoomEncryption {
         this._clock = clock;
         this._disposed = false;
         this._isFlushingRoomKeyShares = false;
+        this._lastKeyPreShareTime = null;
     }
 
     async enableSessionBackup(sessionBackup) {
@@ -244,14 +249,23 @@ export class RoomEncryption {
         }
         
         return matches;
+    }
 
+    /** shares the encryption key for the next message if needed */
+    async ensureMessageKeyIsShared(hsApi) {
+        if (this._lastKeyPreShareTime?.measure() < MIN_PRESHARE_INTERVAL) {
+            return;
+        }
+        this._lastKeyPreShareTime = this._clock.createMeasure();
+        const roomKeyMessage = await this._megolmEncryption.ensureOutboundSession(this._room.id, this._encryptionParams);
+        if (roomKeyMessage) {
+            await this._shareNewRoomKey(roomKeyMessage, hsApi);
+        }
     }
 
     async encrypt(type, content, hsApi) {
-        await this._deviceTracker.trackRoom(this._room);
         const megolmResult = await this._megolmEncryption.encrypt(this._room.id, type, content, this._encryptionParams);
         if (megolmResult.roomKeyMessage) {
-            // TODO: should we await this??
             this._shareNewRoomKey(megolmResult.roomKeyMessage, hsApi);
         }
         return {
@@ -270,6 +284,7 @@ export class RoomEncryption {
     }
 
     async _shareNewRoomKey(roomKeyMessage, hsApi) {
+        await this._deviceTracker.trackRoom(this._room);
         const devices = await this._deviceTracker.devicesForTrackedRoom(this._room.id, hsApi);
         const userIds = Array.from(devices.reduce((set, device) => set.add(device.userId), new Set()));
 

src/matrix/e2ee/megolm/Encryption.js

@@ -43,6 +43,56 @@ export class Encryption {
         }
     }
 
+    async ensureOutboundSession(roomId, encryptionParams) {
+        let session = new this._olm.OutboundGroupSession();
+        try {
+            const txn = this._storage.readWriteTxn([
+                this._storage.storeNames.inboundGroupSessions,
+                this._storage.storeNames.outboundGroupSessions,
+            ]);
+            let roomKeyMessage;
+            try {
+                let sessionEntry = await txn.outboundGroupSessions.get(roomId);
+                roomKeyMessage = this._readOrCreateSession(session, sessionEntry, roomId, encryptionParams, txn);
+                if (roomKeyMessage) {
+                    this._writeSession(sessionEntry, session, roomId, txn);
+                }
+            } catch (err) {
+                txn.abort();
+                throw err;
+            }
+            await txn.complete();
+            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 +111,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) {
+                roomKeyMessage = this._readOrCreateSession(session, sessionEntry, roomId, encryptionParams, txn);
-                    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
-                }
                 encryptedContent = this._encryptContent(roomId, session, type, content);
-                txn.outboundGroupSessions.set({
+                this._writeSession(sessionEntry, session, roomId, txn);
-                    roomId,
-                    session: session.pickle(this._pickleKey),
-                    createdAt: sessionEntry?.createdAt || this._now(),
-                });
 
             } catch (err) {
                 txn.abort();

src/matrix/e2ee/olm/Encryption.js

@@ -154,7 +154,7 @@ export class Encryption {
         try {
             for (const target of newEncryptionTargets) {
                 const {device, oneTimeKey} = target;
-                target.session = this._account.createOutboundOlmSession(device.curve25519Key, oneTimeKey);
+                target.session = await this._account.createOutboundOlmSession(device.curve25519Key, oneTimeKey);
             }
             this._storeSessions(newEncryptionTargets, timestamp);
         } catch (err) {

src/matrix/room/Room.js

@@ -354,6 +354,10 @@ export class Room extends EventEmitter {
         return this._sendQueue.enqueueEvent(eventType, content);
     }
 
+    async ensureMessageKeyIsShared() {
+        return this._roomEncryption?.ensureMessageKeyIsShared(this._hsApi);
+    }
+
     /** @public */
     async loadMemberList() {
         if (this._memberList) {

src/platform/web/ui/css/main.css

@@ -44,3 +44,8 @@ body.hydrogen {
 .hidden {
     display: none !important;
 }
+
+/* hide clear buttons in IE */
+input::-ms-clear {
+    display: none;
+}

src/platform/web/worker/main.js

@@ -116,14 +116,13 @@ class MessageHandler {
 
     _megolmDecrypt(sessionKey, ciphertext) {
         return this._toMessage(() => {
-            let session;
+            const session = new this._olm.InboundGroupSession();
             try {
-                session = new this._olm.InboundGroupSession();
                 session.import_session(sessionKey);
                 // returns object with plaintext and message_index
                 return session.decrypt(ciphertext);
             } finally {
-                session?.free();
+                session.free();
             }
         });
     }
@@ -132,10 +131,29 @@ class MessageHandler {
         return this._toMessage(() => {
             this._feedRandomValues(randomValues);
             const account = new this._olm.Account();
+            try {
                 account.create();
                 account.generate_one_time_keys(otkAmount);
                 this._checkRandomValuesUsed();
                 return account.pickle("");
+            } finally {
+                account.free();
+            }
+        });
+    }
+
+    _olmCreateOutbound(accountPickle, theirIdentityKey, theirOneTimeKey) {
+        return this._toMessage(() => {
+            const account = new this._olm.Account();
+            const newSession = new this._olm.Session();
+            try {
+                account.unpickle("", accountPickle);
+                newSession.create_outbound(account, newSession, theirIdentityKey, theirOneTimeKey);
+                return newSession.pickle("");
+            } finally {
+                account.free();
+                newSession.free();
+            }
         });
     }
 
@@ -149,6 +167,8 @@ class MessageHandler {
             this._sendReply(message, this._megolmDecrypt(message.sessionKey, message.ciphertext));
         } else if (type === "olm_create_account_otks") {
             this._sendReply(message, this._olmCreateAccountAndOTKs(message.randomValues, message.otkAmount));
+        } else if (type === "olm_create_outbound") {
+            this._sendReply(message, this._olmCreateOutbound(message.accountPickle, message.theirIdentityKey, message.theirOneTimeKey));
         }
     }
 }

src/utils/Lock.js

@@ -20,7 +20,7 @@ export class Lock {
         this._resolve = null;
     }
 
-    take() {
+    tryTake() {
         if (!this._promise) {
             this._promise = new Promise(resolve => {
                 this._resolve = resolve;
@@ -30,6 +30,12 @@ export class Lock {
         return false;
     }
 
+    async take() {
+        while(!this.tryTake()) {
+            await this.released();
+        }
+    }
+
     get isTaken() {
         return !!this._promise;
     }
@@ -52,25 +58,25 @@ export function tests() {
     return {
         "taking a lock twice returns false": assert => {
             const lock = new Lock();
-            assert.equal(lock.take(), true);
+            assert.equal(lock.tryTake(), true);
             assert.equal(lock.isTaken, true);
-            assert.equal(lock.take(), false);
+            assert.equal(lock.tryTake(), false);
         },
         "can take a released lock again": assert => {
             const lock = new Lock();
-            lock.take();
+            lock.tryTake();
             lock.release();
             assert.equal(lock.isTaken, false);
-            assert.equal(lock.take(), true);
+            assert.equal(lock.tryTake(), true);
         },
         "2 waiting for lock, only first one gets it": async assert => {
             const lock = new Lock();
-            lock.take();
+            lock.tryTake();
 
             let first;
-            lock.released().then(() => first = lock.take());
+            lock.released().then(() => first = lock.tryTake());
             let second;
-            lock.released().then(() => second = lock.take());
+            lock.released().then(() => second = lock.tryTake());
             const promise = lock.released();
             lock.release();
             await promise;

src/utils/LockMap.js

@@ -24,12 +24,10 @@ export class LockMap {
     async takeLock(key) {
         let lock = this._map.get(key);
         if (lock) {
-            while (!lock.take()) {
+            await lock.take();
-                await lock.released();
-            }
         } else {
             lock = new Lock();
-            lock.take();
+            lock.tryTake();
             this._map.set(key, lock);
         }
         // don't leave old locks lying around