support sending out room key in room encryption for newly joined members

src/matrix/e2ee/DeviceTracker.js

@@ -65,7 +65,7 @@ export class DeviceTracker {
     }
 
     async trackRoom(room) {
-        if (room.isTrackingMembers) {
+        if (room.isTrackingMembers || !room.isEncrypted) {
             return;
         }
         const memberList = await room.loadMemberList();
@@ -230,8 +230,7 @@ export class DeviceTracker {
      * @param  {String} roomId [description]
      * @return {[type]}        [description]
      */
-    async deviceIdentitiesForTrackedRoom(roomId, hsApi) {
+    async devicesForTrackedRoom(roomId, hsApi) {
-        let identities;
         const txn = await this._storage.readTxn([
             this._storage.storeNames.roomMembers,
             this._storage.storeNames.userIdentities,
@@ -243,8 +242,27 @@ export class DeviceTracker {
         
         // So, this will also contain non-joined memberships
         const userIds = await txn.roomMembers.getAllUserIds(roomId);
-        const allMemberIdentities = await Promise.all(userIds.map(userId => txn.userIdentities.get(userId)));
+
-        identities = allMemberIdentities.filter(identity => {
+        return await this._devicesForUserIds(roomId, userIds, txn, hsApi);
+    }
+
+    async devicesForRoomMembers(roomId, userIds, hsApi) {
+        const txn = await this._storage.readTxn([
+            this._storage.storeNames.userIdentities,
+        ]);
+        return await this._devicesForUserIds(roomId, userIds, txn, hsApi);
+    }
+
+    /**
+     * @param  {string} roomId  [description]
+     * @param  {Array<string>} userIds a set of user ids to try and find the identity for. Will be check to belong to roomId.
+     * @param  {Transaction} userIdentityTxn to read the user identities
+     * @param  {HomeServerApi} hsApi
+     * @return {Array<DeviceIdentity>}
+     */
+    async _devicesForUserIds(roomId, userIds, userIdentityTxn, hsApi) {
+        const allMemberIdentities = await Promise.all(userIds.map(userId => userIdentityTxn.userIdentities.get(userId)));
+        const identities = allMemberIdentities.filter(identity => {
             // identity will be missing for any userIds that don't have 
             // membership join in any of your encrypted rooms
             return identity && identity.roomIds.includes(roomId);

src/matrix/e2ee/RoomEncryption.js

@@ -21,7 +21,7 @@ import {makeTxnId} from "../common.js";
 const ENCRYPTED_TYPE = "m.room.encrypted";
 
 export class RoomEncryption {
-    constructor({room, deviceTracker, olmEncryption, megolmEncryption, megolmDecryption, encryptionParams}) {
+    constructor({room, deviceTracker, olmEncryption, megolmEncryption, megolmDecryption, encryptionParams, storage}) {
         this._room = room;
         this._deviceTracker = deviceTracker;
         this._olmEncryption = olmEncryption;
@@ -35,6 +35,7 @@ export class RoomEncryption {
         // not `event_id`, but an internal event id passed in to the decrypt methods
         this._eventIdsByMissingSession = new Map();
         this._senderDeviceCache = new Map();
+        this._storage = storage;
     }
 
     notifyTimelineClosed() {
@@ -114,10 +115,12 @@ export class RoomEncryption {
         // share the new megolm session if needed
         if (megolmResult.roomKeyMessage) {
             await this._deviceTracker.trackRoom(this._room);
-            const devices = await this._deviceTracker.deviceIdentitiesForTrackedRoom(this._room.id, hsApi);
+            const devices = await this._deviceTracker.devicesForTrackedRoom(this._room.id, hsApi);
-            const messages = await this._olmEncryption.encrypt(
+            await this._sendRoomKey(megolmResult.roomKeyMessage, devices, hsApi);
-                "m.room_key", megolmResult.roomKeyMessage, devices, hsApi);
+            // if we happen to rotate the session before we have sent newly joined members the room key
-            await this._sendMessagesToDevices(ENCRYPTED_TYPE, messages, hsApi);
+            // then mark those members as not needing the key anymore
+            const userIds = Array.from(devices.reduce((set, device) => set.add(device.userId), new Set()));
+            await this._clearNeedsRoomKeyFlag(userIds);
         }
         return {
             type: ENCRYPTED_TYPE,
@@ -125,6 +128,58 @@ export class RoomEncryption {
         };
     }
 
+    async shareRoomKeyForMemberChanges(memberChanges, hsApi) {
+        const pendingUserIds = [];
+        for (const m of memberChanges.values()) {
+            if (m.member.needsRoomKey) {
+                pendingUserIds.push(m.userId);
+            }
+        }
+        return await this._shareRoomKey(pendingUserIds, hsApi);
+    }
+
+    async _shareRoomKey(userIds, hsApi) {
+        if (userIds.length === 0) {
+            return;
+        }
+        const readRoomKeyTxn = await this._storage.readTxn([this._storage.storeNames.outboundGroupSessions]);
+        const roomKeyMessage = await this._megolmEncryption.createRoomKeyMessage(this._room.id, readRoomKeyTxn);
+        // no room key if we haven't created a session yet
+        // (or we removed it and will create a new one on the next send)
+        if (roomKeyMessage) {
+            const devices = await this._deviceTracker.devicesForRoomMembers(this._room.id, userIds, hsApi);
+            await this._sendRoomKey(roomKeyMessage, devices, hsApi);
+            const actuallySentUserIds = Array.from(devices.reduce((set, device) => set.add(device.userId), new Set()));
+            await this._clearNeedsRoomKeyFlag(actuallySentUserIds);
+        } else {
+            // we don't have a session yet, clear them all
+            await this._clearNeedsRoomKeyFlag(userIds);
+        }
+    }
+
+    async _clearNeedsRoomKeyFlag(userIds) {
+        const txn = await this._storage.readWriteTxn([this._storage.storeNames.roomMembers]);
+        try {
+            await Promise.all(userIds.map(async userId => {
+                const memberData = await txn.roomMembers.get(this._room.id, userId);
+                if (memberData.needsRoomKey) {
+                    memberData.needsRoomKey = false;
+                    txn.roomMembers.set(memberData);
+                }
+            }));
+        } catch (err) {
+            txn.abort();
+            throw err;
+        }
+        await txn.complete();
+    }
+
+    async _sendRoomKey(roomKeyMessage, devices, hsApi) {
+        const messages = await this._olmEncryption.encrypt(
+            "m.room_key", roomKeyMessage, devices, hsApi);
+        await this._sendMessagesToDevices(ENCRYPTED_TYPE, messages, hsApi);
+    }
+
     async _sendMessagesToDevices(type, messages, hsApi) {
         const messagesByUser = groupBy(messages, message => message.device.userId);
         const payload = {

src/matrix/e2ee/megolm/Encryption.js

@@ -30,6 +30,19 @@ export class Encryption {
         txn.outboundGroupSessions.remove(roomId);
     }
 
+    async createRoomKeyMessage(roomId, txn) {
+        let sessionEntry = await txn.outboundGroupSessions.get(roomId);
+        if (sessionEntry) {
+            const session = new this._olm.OutboundGroupSession();
+            try {
+                session.unpickle(this._pickleKey, sessionEntry.session);
+                return this._createRoomKeyMessage(session, roomId);
+            } finally {
+                session.free();
+            }
+        }
+    }
+
     /**
      * Encrypts a message with megolm
      * @param  {string} roomId           
@@ -127,12 +140,9 @@ export class Encryption {
             session_id: session.session_id(),
             session_key: session.session_key(),
             algorithm: MEGOLM_ALGORITHM,
-            // if we need to do this, do we need to create
+            // chain_index is ignored by element-web if not all clients
-            // the room key message after or before having encrypted
+            // but let's send it anyway, as element-web does so
-            // with the new session? I guess before as we do now
+            chain_index: session.message_index()
-            // because the chain_index is where you should start decrypting?
-            // 
-            // chain_index: session.message_index()
         }
     }