feat: Decrypt events on megolm key receiving better

lib/encryption/encryption.dart

@@ -232,30 +232,40 @@ class Encryption {
     if (event.type != EventTypes.Encrypted) {
       return event;
     }
-    if (client.database != null &&
+    try {
-        keyManager.getInboundGroupSession(roomId, event.content['session_id'],
+      if (client.database != null &&
-                event.content['sender_key']) ==
+          keyManager.getInboundGroupSession(roomId, event.content['session_id'],
-            null) {
+                  event.content['sender_key']) ==
-      await keyManager.loadInboundGroupSession(
+              null) {
-          roomId, event.content['session_id'], event.content['sender_key']);
+        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);
       }
-      await client.database?.storeEventUpdate(
+      event = decryptRoomEventSync(roomId, event);
-        client.id,
+      if (event.type == EventTypes.Encrypted &&
-        EventUpdate(
+          event.content['can_request_session'] == true) {
-          eventType: event.type,
+        keyManager.maybeAutoRequest(
-          content: event.toJson(),
+            roomId, event.content['session_id'], event.content['sender_key']);
-          roomID: event.roomId,
+      }
-          type: updateType,
+      if (event.type != EventTypes.Encrypted && store) {
-          sortOrder: event.sortOrder,
+        if (updateType != EventUpdateType.history) {
-        ),
+          event.room?.setState(event);
-      );
+        }
+        await client.database?.storeEventUpdate(
+          client.id,
+          EventUpdate(
+            eventType: event.type,
+            content: event.toJson(),
+            roomID: event.roomId,
+            type: updateType,
+            sortOrder: event.sortOrder,
+          ),
+        );
+      }
+      return event;
+    } catch (e, s) {
+      Logs.error('[Decrypt] Could not decrpyt event: ' + e.toString(), s);
+      return event;
     }
-    return event;
   }
 
   /// Encrypts the given json payload and creates a send-ready m.room.encrypted

lib/encryption/key_manager.dart

@@ -19,7 +19,6 @@
 import 'dart:convert';
 
 import 'package:olm/olm.dart' as olm;
-import 'package:pedantic/pedantic.dart';
 
 import './encryption.dart';
 import './utils/outbound_group_session.dart';
@@ -51,18 +50,31 @@ class KeyManager {
         if (info.algorithm != RoomKeysAlgorithmType.v1Curve25519AesSha2) {
           return false;
         }
-        if (keyObj.init_with_private_key(base64.decode(secret)) ==
+        return keyObj.init_with_private_key(base64.decode(secret)) ==
-            info.authData['public_key']) {
+            info.authData['public_key'];
-          _requestedSessionIds.clear();
-          return true;
-        }
-        return false;
       } catch (_) {
         return false;
       } finally {
         keyObj.free();
       }
     });
+    encryption.ssss.setCacheCallback(MEGOLM_KEY, (String secret) {
+      // we got a megolm key cached, clear our requested keys and try to re-decrypt
+      // last events
+      _requestedSessionIds.clear();
+      for (final room in client.rooms) {
+        final lastEvent = room.lastEvent;
+        if (lastEvent.type == EventTypes.Encrypted &&
+            lastEvent.content['can_request_session'] == true) {
+          try {
+            maybeAutoRequest(room.id, lastEvent.content['session_id'],
+                lastEvent.content['sener_key']);
+          } catch (_) {
+            // dispose
+          }
+        }
+      }
+    });
   }
 
   bool get enabled => client.accountData[MEGOLM_KEY] != null;
@@ -189,6 +201,21 @@ class KeyManager {
     return null;
   }
 
+  /// Attempt auto-request for a key
+  void maybeAutoRequest(String roomId, String sessionId, String senderKey) {
+    final room = client.getRoomById(roomId);
+    final requestIdent = '$roomId|$sessionId|$senderKey';
+    if (client.enableE2eeRecovery &&
+        room != null &&
+        !_requestedSessionIds.contains(requestIdent) &&
+        !client.isUnknownSession) {
+      // do e2ee recovery
+      _requestedSessionIds.add(requestIdent);
+      runInRoot(
+          () => request(room, sessionId, senderKey, onlineKeyBackupOnly: true));
+    }
+  }
+
   /// Loads an inbound group session
   Future<SessionKey> loadInboundGroupSession(
       String roomId, String sessionId, String senderKey) async {
@@ -206,17 +233,6 @@ class KeyManager {
     final session = await client.database
         ?.getDbInboundGroupSession(client.id, roomId, sessionId);
     if (session == null) {
-      final room = client.getRoomById(roomId);
-      final requestIdent = '$roomId|$sessionId|$senderKey';
-      if (client.enableE2eeRecovery &&
-          room != null &&
-          !_requestedSessionIds.contains(requestIdent) &&
-          !client.isUnknownSession) {
-        // do e2ee recovery
-        _requestedSessionIds.add(requestIdent);
-        unawaited(runInRoot(() =>
-            request(room, sessionId, senderKey, onlineKeyBackupOnly: true)));
-      }
       return null;
     }
     if (!_inboundGroupSessions.containsKey(roomId)) {

lib/encryption/ssss.dart

@@ -19,6 +19,7 @@
 import 'dart:core';
 import 'dart:convert';
 import 'dart:typed_data';
+import 'dart:async';
 
 import 'package:base58check/base58.dart';
 import 'package:crypto/crypto.dart';
@@ -48,7 +49,8 @@ class SSSS {
   final Encryption encryption;
   Client get client => encryption.client;
   final pendingShareRequests = <String, _ShareRequest>{};
-  final _validators = <String, Future<bool> Function(String)>{};
+  final _validators = <String, FutureOr<bool> Function(String)>{};
+  final _cacheCallbacks = <String, FutureOr<void> Function(String)>{};
   final Map<String, DbSSSSCache> _cache = <String, DbSSSSCache>{};
   SSSS(this.encryption);
 
@@ -145,10 +147,14 @@ class SSSS {
         info.iterations, info.bits != null ? (info.bits / 8).ceil() : 32));
   }
 
-  void setValidator(String type, Future<bool> Function(String) validator) {
+  void setValidator(String type, FutureOr<bool> Function(String) validator) {
     _validators[type] = validator;
   }
 
+  void setCacheCallback(String type, FutureOr<void> Function(String) callback) {
+    _cacheCallbacks[type] = callback;
+  }
+
   String get defaultKeyId {
     final keyData = client.accountData['m.secret_storage.default_key'];
     if (keyData == null || !(keyData.content['key'] is String)) {
@@ -221,12 +227,17 @@ class SSSS {
       // cache the thing
       await client.database
           .storeSSSSCache(client.id, type, keyId, enc['ciphertext'], decrypted);
+      if (_cacheCallbacks.containsKey(type) && await getCached(type) == null) {
+        _cacheCallbacks[type](decrypted);
+      }
     }
     return decrypted;
   }
 
   Future<void> store(
       String type, String secret, String keyId, Uint8List key) async {
+    final triggerCacheCallback =
+        _cacheCallbacks.containsKey(type) && await getCached(type) == null;
     final encrypted = encryptAes(secret, key, type);
     final content = <String, dynamic>{
       'encrypted': <String, dynamic>{},
@@ -242,6 +253,9 @@ class SSSS {
       // cache the thing
       await client.database
           .storeSSSSCache(client.id, type, keyId, encrypted.ciphertext, secret);
+      if (triggerCacheCallback) {
+        _cacheCallbacks[type](secret);
+      }
     }
   }
 
@@ -404,6 +418,9 @@ class SSSS {
               .content['encrypted'][keyId]['ciphertext'];
           await client.database.storeSSSSCache(
               client.id, request.type, keyId, ciphertext, secret);
+          if (_cacheCallbacks.containsKey(request.type)) {
+            _cacheCallbacks[request.type](secret);
+          }
         }
       }
     }

lib/src/timeline.dart

@@ -135,6 +135,22 @@ class Timeline {
     if (decryptAtLeastOneEvent) onUpdate();
   }
 
+  /// Request the keys for undecryptable events of this timeline
+  void requestKeys() {
+    for (final event in events) {
+      if (event.type == EventTypes.Encrypted &&
+          event.messageType == MessageTypes.BadEncrypted &&
+          event.content['can_request_session'] == true) {
+        try {
+          room.client.encryption.keyManager.maybeAutoRequest(room.id,
+              event.content['session_id'], event.content['sender_key']);
+        } catch (_) {
+          // dispose
+        }
+      }
+    }
+  }
+
   int _findEvent({String event_id, String unsigned_txid}) {
     // we want to find any existing event where either the passed event_id or the passed unsigned_txid
     // matches either the event_id or transaction_id of the existing event.