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();