fix: load keys for decryption separately for [Room].searchEvents

Author: coder-with-a-bushido

lib/encryption/key_manager.dart

@@ -28,7 +28,8 @@ class KeyManager {
   final _inboundGroupSessions = <String, Map<String, SessionKey>>{};
   final _outboundGroupSessions = <String, OutboundGroupSession>{};
   final Set<String> _loadedOutboundGroupSessions = <String>{};
-  final Set<String> _requestedSessionIds = <String>{};
+  final Map<String, Completer<void>> _requestedSessionIdCompleters =
+      <String, Completer<void>>{};
 
   KeyManager(this.encryption) {
     encryption.ssss.setValidator(megolmKey, (String secret) async {
@@ -49,7 +50,7 @@ class KeyManager {
     encryption.ssss.setCacheCallback(megolmKey, (String secret) {
       // we got a megolm key cached, clear our requested keys and try to re-decrypt
       // last events
-      _requestedSessionIds.clear();
+      _requestedSessionIdCompleters.clear();
       for (final room in client.rooms) {
         final lastEvent = room.lastEvent;
         if (lastEvent != null &&
@@ -219,30 +220,44 @@ class KeyManager {
   }
 
   /// Attempt auto-request for a key
-  void maybeAutoRequest(
+  FutureOr<void> maybeAutoRequest(
     String roomId,
     String sessionId,
     String? senderKey, {
     bool tryOnlineBackup = true,
     bool onlineKeyBackupOnly = true,
-  }) {
+    bool awaitRequest = false,
+  }) async {
     final room = client.getRoomById(roomId);
     final requestIdent = '$roomId|$sessionId';
-    if (room != null &&
-        !_requestedSessionIds.contains(requestIdent) &&
-        !client.isUnknownSession) {
-      // do e2ee recovery
-      _requestedSessionIds.add(requestIdent);
-
-      runInRoot(
-        () async => request(
+    if (room != null && !client.isUnknownSession) {
+      final existingReqCompleter = _requestedSessionIdCompleters[requestIdent];
+      if (existingReqCompleter != null) {
+        if (awaitRequest) {
+          await existingReqCompleter.future;
+        }
+        // return if already exists (after awaiting completion if required)
+        return;
+      }
+
+      Future<void> requestFn() async {
+        final completer = Completer<void>();
+        _requestedSessionIdCompleters[requestIdent] = completer;
+        await request(
           room,
           sessionId,
           senderKey,
           tryOnlineBackup: tryOnlineBackup,
           onlineKeyBackupOnly: onlineKeyBackupOnly,
-        ),
-      );
+        );
+        completer.complete();
+      }
+
+      if (awaitRequest) {
+        await requestFn();
+      } else {
+        runInRoot(requestFn);
+      }
     }
   }
 
@@ -744,9 +759,13 @@ class KeyManager {
     bool tryOnlineBackup = true,
     bool onlineKeyBackupOnly = false,
   }) async {
+    Logs().i(
+      '[KeyManager] Requesting session key for $sessionId for room ${room.id}',
+    );
     if (tryOnlineBackup && await isCached()) {
       // let's first check our online key backup store thingy...
-      final hadPreviously = getInboundGroupSession(room.id, sessionId) != null;
+      final hadSessionBefore =
+          getInboundGroupSession(room.id, sessionId) != null;
       try {
         await loadSingleKey(room.id, sessionId);
       } catch (err, stacktrace) {
@@ -762,10 +781,17 @@ class KeyManager {
           );
         }
       }
-      // TODO: also don't request from others if we have an index of 0 now
-      if (!hadPreviously &&
-          getInboundGroupSession(room.id, sessionId) != null) {
-        return; // we managed to load the session from online backup, no need to care about it now
+
+      final storedSession = getInboundGroupSession(room.id, sessionId);
+      // We loaded *something* from backup we didn't have before — peers are
+      // unlikely to do better, so don't spam to-device requests.
+      final loadedFreshFromBackup = !hadSessionBefore && storedSession != null;
+      // We have the session from its very first index; peers literally cannot
+      // give us anything better.
+      final hasFullSession =
+          storedSession?.inboundGroupSession?.firstKnownIndex == 0;
+      if (loadedFreshFromBackup || hasFullSession) {
+        return;
       }
     }
     if (onlineKeyBackupOnly) {

lib/src/room.dart

@@ -2837,15 +2837,41 @@ class Room {
         .map((matrixEvent) => Event.fromMatrixEvent(matrixEvent, this))
         .toList();
 
+    /// Attempt decrypting the event, if it fails, load the key and try again once.
+    /// Returns null if event cannot be decrypted.
+    Future<Event?> loadKeysAndDecryptEvent(
+      Event event, {
+      bool keyAlreadyLoaded = false,
+    }) async {
+      final decrypted = await client.encryption!.decryptRoomEvent(
+        event,
+        store: false,
+        updateType: EventUpdateType.history,
+      );
+      if (decrypted.type != EventTypes.Encrypted) {
+        return decrypted;
+      } else if (!keyAlreadyLoaded) {
+        final content = event.parsedRoomEncryptedContent;
+        if (content.sessionId != null) {
+          await client.encryption!.keyManager.maybeAutoRequest(
+            id,
+            content.sessionId!,
+            content.senderKey,
+            tryOnlineBackup: true,
+            onlineKeyBackupOnly: true,
+            awaitRequest: true,
+          );
+          return loadKeysAndDecryptEvent(event, keyAlreadyLoaded: true);
+        }
+      }
+      return null;
+    }
+
     // Decrypt all events one after another:
     for (final (index, event) in events.indexed) {
-      if (event.type == EventTypes.Encrypted) {
-        final decrypted = await client.encryption?.decryptRoomEvent(
-          event,
-          store: false,
-          updateType: EventUpdateType.history,
-        );
-        if (decrypted != null && decrypted.type != EventTypes.Encrypted) {
+      if (event.type == EventTypes.Encrypted && client.encryption != null) {
+        final decrypted = await loadKeysAndDecryptEvent(event);
+        if (decrypted != null) {
           events[index] = decrypted;
         }
       }

test/encryption/online_key_backup_test.dart

@@ -121,6 +121,102 @@ void main() {
       expect(ret != null, true);
     });
 
+    test(
+        'Room.searchEvents decrypts an event whose session is only in online key backup',
+        () async {
+      // 1. Fresh outbound + inbound megolm session the client has never seen.
+      final outbound = vod.GroupSession();
+      final inbound = vod.InboundGroupSession(outbound.sessionKey);
+      final newSessionId = inbound.sessionId;
+      final newSenderKey = client.identityKey;
+
+      // 2. Real ciphertext that the matching inbound session can decrypt.
+      final ciphertext = outbound.encrypt(
+        json.encode({
+          'type': 'm.room.message',
+          'content': {
+            'msgtype': 'm.text',
+            'body': 'a needle in the encrypted haystack',
+          },
+        }),
+      );
+
+      // 3. Encrypt the inbound session for the FakeMatrixApi backup using
+      //    the same backup public key the FakeMatrixApi advertises.
+      const backupPubKey = 'GXYaxqhNhUK28zUdxOmEsFRguz+PzBsDlTLlF0O0RkM';
+      final encryptor = vod.PkEncryption.fromPublicKey(
+        vod.Curve25519PublicKey.fromBase64(backupPubKey),
+      );
+      final encryptedSession = encryptor.encrypt(
+        json.encode({
+          'algorithm': AlgorithmTypes.megolmV1AesSha2,
+          'forwarding_curve25519_key_chain': <String>[],
+          'sender_key': newSenderKey,
+          'sender_claimed_keys': {'ed25519': client.fingerprintKey},
+          'session_key': inbound.exportAt(0),
+        }),
+      );
+      final (ct, mac, ephemeral) = encryptedSession.toBase64();
+
+      // 4. A fresh room that exists in client.rooms (required for
+      //    keyManager.maybeAutoRequest's lookup) and has no DB history.
+      const newRoomId = '!searchBackupOnly:example.com';
+      final newRoom = Room(client: client, id: newRoomId, prev_batch: '');
+      client.rooms.add(newRoom);
+
+      // 5. Mock the online key backup GET for this specific session.
+      FakeMatrixApi.currentApi!.api['GET']![
+              '/client/v3/room_keys/keys/${Uri.encodeComponent(newRoomId)}/${Uri.encodeComponent(newSessionId)}?version=5'] =
+          (_) => {
+                'first_message_index': 0,
+                'forwarded_count': 0,
+                'is_verified': true,
+                'session_data': {
+                  'ephemeral': ephemeral,
+                  'ciphertext': ct,
+                  'mac': mac,
+                },
+              };
+
+      // 6. Mock /messages with our encrypted event.
+      FakeMatrixApi.currentApi!.api['GET']![
+              '/client/v3/rooms/${Uri.encodeComponent(newRoomId)}/messages?from&dir=b&limit=1000&filter=%7B%22types%22%3A%5B%22m.room.message%22%2C%22m.room.encrypted%22%5D%7D'] =
+          (_) => {
+                'chunk': [
+                  {
+                    'content': {
+                      'algorithm': AlgorithmTypes.megolmV1AesSha2,
+                      'sender_key': newSenderKey,
+                      'session_id': newSessionId,
+                      'ciphertext': ciphertext,
+                      'device_id': client.deviceID,
+                    },
+                    'type': EventTypes.Encrypted,
+                    'event_id': '\$searchEnc',
+                    'origin_server_ts': 1432735824653,
+                    'sender': '@alice:example.com',
+                  },
+                ],
+                'end': 't_search_end',
+                'start': 't_search_start',
+              };
+
+      // Pre-condition: the session is in neither memory nor DB.
+      expect(
+        client.encryption!.keyManager
+            .getInboundGroupSession(newRoomId, newSessionId),
+        isNull,
+      );
+
+      final result = await newRoom.searchEvents(searchFunc: (_) => true);
+
+      // The event was decrypted via the backup-load path that searchEvents
+      // performs internally.
+      final decrypted = result.events.where((e) => e.eventId == '\$searchEnc');
+      expect(decrypted, isNotEmpty);
+      expect(decrypted.first.body, 'a needle in the encrypted haystack');
+    });
+
     test('dispose client', () async {
       await client.dispose(closeDatabase: false);
     });