fix: Properly rotate megolm session on device changes

lib/encryption/encryption.dart

@@ -159,17 +159,7 @@ class Encryption {
       var haveIndex = inboundGroupSession.indexes.containsKey(messageIndexKey);
       if (haveIndex &&
           inboundGroupSession.indexes[messageIndexKey] != messageIndexValue) {
-        // TODO: maybe clear outbound session, if it is ours
-        // TODO: Make it so that we can't re-request the session keys, this is just for debugging
         Logs.error('[Decrypt] Could not decrypt due to a corrupted session.');
-        Logs.error('[Decrypt] Want session: $roomId $sessionId $senderKey');
-        Logs.error(
-            '[Decrypt] Have sessoin: ${inboundGroupSession.roomId} ${inboundGroupSession.sessionId} ${inboundGroupSession.senderKey}');
-        Logs.error(
-            '[Decrypt] Want indexes: $messageIndexKey $messageIndexValue');
-        Logs.error(
-            '[Decrypt] Have indexes: $messageIndexKey ${inboundGroupSession.indexes[messageIndexKey]}');
-        canRequestSession = true;
         throw (DecryptError.CHANNEL_CORRUPTED);
       }
       inboundGroupSession.indexes[messageIndexKey] = messageIndexValue;
@@ -188,13 +178,15 @@ class Encryption {
     } catch (exception) {
       // alright, if this was actually by our own outbound group session, we might as well clear it
       if (client.enableE2eeRecovery &&
+          exception != DecryptError.UNKNOWN_SESSION &&
           (keyManager
                       .getOutboundGroupSession(roomId)
                       ?.outboundGroupSession
                       ?.session_id() ??
                   '') ==
               event.content['session_id']) {
-        keyManager.clearOutboundGroupSession(roomId, wipe: true);
+        runInRoot(() =>
+            keyManager.clearOrUseOutboundGroupSession(roomId, wipe: true));
       }
       if (canRequestSession) {
         decryptedPayload = {
@@ -237,7 +229,18 @@ class Encryption {
   Future<Event> decryptRoomEvent(String roomId, Event event,
       {bool store = false,
       EventUpdateType updateType = EventUpdateType.timeline}) async {
-    final doStore = () async {
+    if (event.type != EventTypes.Encrypted) {
+      return event;
+    }
+    if (client.database != null &&
+        keyManager.getInboundGroupSession(roomId, event.content['session_id'],
+                event.content['sender_key']) ==
+            null) {
+      await keyManager.loadInboundGroupSession(
+          roomId, event.content['session_id'], event.content['sender_key']);
+    }
+    event = decryptRoomEventSync(roomId, event);
+    if (event.type != EventTypes.Encrypted && store) {
       if (updateType != EventUpdateType.history) {
         event.room?.setState(event);
       }
@@ -251,25 +254,6 @@ class Encryption {
           sortOrder: event.sortOrder,
         ),
       );
-    };
-    if (event.type != EventTypes.Encrypted) {
-      return event;
-    }
-    event = decryptRoomEventSync(roomId, event);
-    if (event.type != EventTypes.Encrypted) {
-      if (store) {
-        await doStore();
-      }
-      return event;
-    }
-    if (client.database == null) {
-      return event;
-    }
-    await keyManager.loadInboundGroupSession(
-        roomId, event.content['session_id'], event.content['sender_key']);
-    event = decryptRoomEventSync(roomId, event);
-    if (event.type != EventTypes.Encrypted && store) {
-      await doStore();
     }
     return event;
   }
@@ -289,7 +273,7 @@ class Encryption {
     if (keyManager.getOutboundGroupSession(roomId) == null) {
       await keyManager.loadOutboundGroupSession(roomId);
     }
-    await keyManager.clearOutboundGroupSession(roomId);
+    await keyManager.clearOrUseOutboundGroupSession(roomId);
     if (keyManager.getOutboundGroupSession(roomId) == null) {
       await keyManager.createOutboundGroupSession(roomId);
     }
@@ -315,7 +299,6 @@ class Encryption {
       'session_id': sess.outboundGroupSession.session_id(),
       if (mRelatesTo != null) 'm.relates_to': mRelatesTo,
     };
-    sess.sentMessages++;
     await keyManager.storeOutboundGroupSession(roomId, sess);
     return encryptedPayload;
   }

lib/encryption/key_manager.dart

@@ -231,6 +231,18 @@ class KeyManager {
     return sess;
   }
 
+  Map<String, Map<String, bool>> _getDeviceKeyIdMap(
+      List<DeviceKeys> deviceKeys) {
+    final deviceKeyIds = <String, Map<String, bool>>{};
+    for (final device in deviceKeys) {
+      if (!deviceKeyIds.containsKey(device.userId)) {
+        deviceKeyIds[device.userId] = <String, bool>{};
+      }
+      deviceKeyIds[device.userId][device.deviceId] = device.blocked;
+    }
+    return deviceKeyIds;
+  }
+
   /// clear all cached inbound group sessions. useful for testing
   void clearOutboundGroupSessions() {
     _outboundGroupSessions.clear();
@@ -238,8 +250,8 @@ class KeyManager {
 
   /// Clears the existing outboundGroupSession but first checks if the participating
   /// devices have been changed. Returns false if the session has not been cleared because
-  /// it wasn't necessary.
+  /// it wasn't necessary. Otherwise returns true.
-  Future<bool> clearOutboundGroupSession(String roomId,
+  Future<bool> clearOrUseOutboundGroupSession(String roomId,
       {bool wipe = false}) async {
     final room = client.getRoomById(roomId);
     final sess = getOutboundGroupSession(roomId);
@@ -247,15 +259,7 @@ class KeyManager {
       return true;
     }
     if (!wipe) {
-      // first check if the devices in the room changed
+      // first check if it needs to be rotated
-      final deviceKeys = await room.getUserDeviceKeys();
-      deviceKeys.removeWhere((k) => k.blocked);
-      final deviceKeyIds = deviceKeys.map((k) => k.deviceId).toList();
-      deviceKeyIds.sort();
-      if (deviceKeyIds.toString() != sess.devices.toString()) {
-        wipe = true;
-      }
-      // next check if it needs to be rotated
       final encryptionContent = room.getState(EventTypes.Encryption)?.content;
       final maxMessages = encryptionContent != null &&
               encryptionContent['rotation_period_msgs'] is int
@@ -271,7 +275,76 @@ class KeyManager {
               .isBefore(DateTime.now())) {
         wipe = true;
       }
+    }
     if (!wipe) {
+      // next check if the devices in the room changed
+      final devicesToReceive = <DeviceKeys>[];
+      final newDeviceKeys = await room.getUserDeviceKeys();
+      final newDeviceKeyIds = _getDeviceKeyIdMap(newDeviceKeys);
+      // first check for user differences
+      final oldUserIds = Set.from(sess.devices.keys);
+      final newUserIds = Set.from(newDeviceKeyIds.keys);
+      if (oldUserIds.difference(newUserIds).isNotEmpty) {
+        // a user left the room, we must wipe the session
+        wipe = true;
+      } else {
+        final newUsers = newUserIds.difference(oldUserIds);
+        if (newUsers.isNotEmpty) {
+          // new user! Gotta send the megolm session to them
+          devicesToReceive
+              .addAll(newDeviceKeys.where((d) => newUsers.contains(d.userId)));
+        }
+        // okay, now we must test all the individual user devices, if anything new got blocked
+        // or if we need to send to any new devices.
+        // for this it is enough if we iterate over the old user Ids, as the new ones already have the needed keys in the list.
+        // we also know that all the old user IDs appear in the old one, else we have already wiped the session
+        for (final userId in oldUserIds) {
+          final oldBlockedDevices = Set.from(sess.devices[userId].entries
+              .where((e) => e.value)
+              .map((e) => e.key));
+          final newBlockedDevices = Set.from(newDeviceKeyIds[userId]
+              .entries
+              .where((e) => e.value)
+              .map((e) => e.key));
+          // we don't really care about old devices that got dropped (deleted), we only care if new ones got added and if new ones got blocked
+          // check if new devices got blocked
+          if (newBlockedDevices.difference(oldBlockedDevices).isNotEmpty) {
+            wipe = true;
+            break;
+          }
+          // and now add all the new devices!
+          final oldDeviceIds = Set.from(sess.devices[userId].keys);
+          final newDeviceIds = Set.from(newDeviceKeyIds[userId].keys);
+          final newDevices = newDeviceIds.difference(oldDeviceIds);
+          if (newDeviceIds.isNotEmpty) {
+            devicesToReceive.addAll(newDeviceKeys.where(
+                (d) => d.userId == userId && newDevices.contains(d.deviceId)));
+          }
+        }
+      }
+
+      if (!wipe) {
+        // okay, we use the outbound group session!
+        sess.sentMessages++;
+        sess.devices = newDeviceKeyIds;
+        final rawSession = <String, dynamic>{
+          'algorithm': 'm.megolm.v1.aes-sha2',
+          'room_id': room.id,
+          'session_id': sess.outboundGroupSession.session_id(),
+          'session_key': sess.outboundGroupSession.session_key(),
+        };
+        try {
+          devicesToReceive.removeWhere((k) => k.blocked);
+          if (devicesToReceive.isNotEmpty) {
+            await client.sendToDeviceEncrypted(
+                devicesToReceive, 'm.room_key', rawSession);
+          }
+        } catch (e, s) {
+          Logs.error(
+              '[LibOlm] Unable to re-send the session key at later index to new devices: ' +
+                  e.toString(),
+              s);
+        }
         return false;
       }
     }
@@ -296,15 +369,13 @@ class KeyManager {
   }
 
   Future<OutboundGroupSession> createOutboundGroupSession(String roomId) async {
-    await clearOutboundGroupSession(roomId, wipe: true);
+    await clearOrUseOutboundGroupSession(roomId, wipe: true);
     final room = client.getRoomById(roomId);
     if (room == null) {
       return null;
     }
     final deviceKeys = await room.getUserDeviceKeys();
-    deviceKeys.removeWhere((k) => k.blocked);
+    final deviceKeyIds = _getDeviceKeyIdMap(deviceKeys);
-    final deviceKeyIds = deviceKeys.map((k) => k.deviceId).toList();
-    deviceKeyIds.sort();
     final outboundGroupSession = olm.OutboundGroupSession();
     try {
       outboundGroupSession.create();
@@ -331,6 +402,7 @@ class KeyManager {
       key: client.userID,
     );
     try {
+      deviceKeys.removeWhere((k) => k.blocked);
       await client.sendToDeviceEncrypted(deviceKeys, 'm.room_key', rawSession);
       await storeOutboundGroupSession(roomId, sess);
       _outboundGroupSessions[roomId] = sess;

lib/encryption/utils/outbound_group_session.dart

@@ -24,7 +24,11 @@ import '../../src/database/database.dart' show DbOutboundGroupSession;
 import '../../src/utils/logs.dart';
 
 class OutboundGroupSession {
-  List<String> devices;
+  /// The devices is a map from user id to device id to if the device is blocked.
+  /// This way we can easily know if a new user is added, leaves, a new devices is added, and,
+  /// very importantly, if we block a device. These are all important for determining if/when
+  /// an outbound session needs to be rotated.
+  Map<String, Map<String, bool>> devices;
   DateTime creationTime;
   olm.OutboundGroupSession outboundGroupSession;
   int sentMessages;
@@ -40,10 +44,21 @@ class OutboundGroupSession {
 
   OutboundGroupSession.fromDb(DbOutboundGroupSession dbEntry, String key)
       : key = key {
+    try {
+      devices = {};
+      for (final entry in json.decode(dbEntry.deviceIds).entries) {
+        devices[entry.key] = Map<String, bool>.from(entry.value);
+      }
+    } catch (e) {
+      // devices is bad (old data), so just not use this session
+      Logs.info(
+          '[OutboundGroupSession] Session in database is old, not using it. ' +
+              e.toString());
+      return;
+    }
     outboundGroupSession = olm.OutboundGroupSession();
     try {
       outboundGroupSession.unpickle(key, dbEntry.pickle);
-      devices = List<String>.from(json.decode(dbEntry.deviceIds));
       creationTime = DateTime.fromMillisecondsSinceEpoch(dbEntry.creationTime);
       sentMessages = dbEntry.sentMessages;
     } catch (e, s) {

test/encryption/key_manager_test.dart

@@ -102,7 +102,7 @@ void main() {
       expect(
           client.encryption.keyManager.getOutboundGroupSession(roomId) != null,
           true);
-      await client.encryption.keyManager.clearOutboundGroupSession(roomId);
+      await client.encryption.keyManager.clearOrUseOutboundGroupSession(roomId);
       expect(
           client.encryption.keyManager.getOutboundGroupSession(roomId) != null,
           true);
@@ -114,17 +114,17 @@ void main() {
 
       // rotate after too many messages
       sess.sentMessages = 300;
-      await client.encryption.keyManager.clearOutboundGroupSession(roomId);
+      await client.encryption.keyManager.clearOrUseOutboundGroupSession(roomId);
       expect(
           client.encryption.keyManager.getOutboundGroupSession(roomId) != null,
           false);
 
-      // rotate if devices in room change
+      // rotate if device is blocked
       sess =
           await client.encryption.keyManager.createOutboundGroupSession(roomId);
       client.userDeviceKeys['@alice:example.com'].deviceKeys['JLAFKJWSCS']
           .blocked = true;
-      await client.encryption.keyManager.clearOutboundGroupSession(roomId);
+      await client.encryption.keyManager.clearOrUseOutboundGroupSession(roomId);
       expect(
           client.encryption.keyManager.getOutboundGroupSession(roomId) != null,
           false);
@@ -135,16 +135,61 @@ void main() {
       sess =
           await client.encryption.keyManager.createOutboundGroupSession(roomId);
       sess.creationTime = DateTime.now().subtract(Duration(days: 30));
-      await client.encryption.keyManager.clearOutboundGroupSession(roomId);
+      await client.encryption.keyManager.clearOrUseOutboundGroupSession(roomId);
       expect(
           client.encryption.keyManager.getOutboundGroupSession(roomId) != null,
           false);
 
+      // rotate if user leaves
+      sess =
+          await client.encryption.keyManager.createOutboundGroupSession(roomId);
+      final room = client.getRoomById(roomId);
+      final member = room.getState('m.room.member', '@alice:example.com');
+      member.content['membership'] = 'leave';
+      room.mJoinedMemberCount--;
+      await client.encryption.keyManager.clearOrUseOutboundGroupSession(roomId);
+      expect(
+          client.encryption.keyManager.getOutboundGroupSession(roomId) != null,
+          false);
+      member.content['membership'] = 'join';
+      room.mJoinedMemberCount++;
+
+      // do not rotate if new device is added
+      sess =
+          await client.encryption.keyManager.createOutboundGroupSession(roomId);
+      client.userDeviceKeys['@alice:example.com'].deviceKeys['NEWDEVICE'] =
+          DeviceKeys.fromJson({
+        'user_id': '@alice:example.com',
+        'device_id': 'NEWDEVICE',
+        'algorithms': ['m.olm.v1.curve25519-aes-sha2', 'm.megolm.v1.aes-sha2'],
+        'keys': {
+          'curve25519:JLAFKJWSCS':
+              '3C5BFWi2Y8MaVvjM8M22DBmh24PmgR0nPvJOIArzgyI',
+          'ed25519:JLAFKJWSCS': 'lEuiRJBit0IG6nUf5pUzWTUEsRVVe/HJkoKuEww9ULI'
+        },
+      }, client);
+      await client.encryption.keyManager.clearOrUseOutboundGroupSession(roomId);
+      expect(
+          client.encryption.keyManager.getOutboundGroupSession(roomId) != null,
+          true);
+
+      // do not rotate if new user is added
+      member.content['membership'] = 'leave';
+      room.mJoinedMemberCount--;
+      sess =
+          await client.encryption.keyManager.createOutboundGroupSession(roomId);
+      member.content['membership'] = 'join';
+      room.mJoinedMemberCount++;
+      await client.encryption.keyManager.clearOrUseOutboundGroupSession(roomId);
+      expect(
+          client.encryption.keyManager.getOutboundGroupSession(roomId) != null,
+          true);
+
       // force wipe
       sess =
           await client.encryption.keyManager.createOutboundGroupSession(roomId);
       await client.encryption.keyManager
-          .clearOutboundGroupSession(roomId, wipe: true);
+          .clearOrUseOutboundGroupSession(roomId, wipe: true);
       expect(
           client.encryption.keyManager.getOutboundGroupSession(roomId) != null,
           false);