feat: Better guard against failed lookups

lib/src/errors.dart

@@ -44,3 +44,11 @@ class InvalidKeyExchangeException extends OmemoException implements Exception {
 class MessageAlreadyDecryptedException extends OmemoException implements Exception {
   String errMsg() => 'The message has already been decrypted';
 }
+
+/// Triggered by the OmemoManager when we could not encrypt a message as we have
+/// no key material available. That happens, for example, when we want to create a
+/// ratchet session with a JID we had no session with but fetching the device bundle
+/// failed.
+class NoKeyMaterialAvailableException extends OmemoException implements Exception {
+  String errMsg() => 'No key material available to create a ratchet session with';
+}

lib/src/omemo/encryption_result.dart

@@ -1,14 +1,26 @@
 import 'package:meta/meta.dart';
+import 'package:omemo_dart/src/errors.dart';
 import 'package:omemo_dart/src/omemo/encrypted_key.dart';
+import 'package:omemo_dart/src/omemo/ratchet_map_key.dart';
 
 @immutable
 class EncryptionResult {
-  const EncryptionResult(this.ciphertext, this.encryptedKeys);
+  const EncryptionResult(this.ciphertext, this.encryptedKeys, this.deviceEncryptionErrors, this.jidEncryptionErrors);
   
-  /// The actual message that was encrypted
+  /// The actual message that was encrypted.
   final List<int>? ciphertext;
-
+  
   /// Mapping of the device Id to the key for decrypting ciphertext, encrypted
-  /// for the ratchet with said device Id
+  /// for the ratchet with said device Id.
   final List<EncryptedKey> encryptedKeys;
+
+  /// Mapping of a ratchet map keys to a possible exception.
+  final Map<RatchetMapKey, OmemoException> deviceEncryptionErrors;
+
+  /// Mapping of a JID to a possible exception.
+  final Map<String, OmemoException> jidEncryptionErrors;
+
+  /// True if the encryption was a success. This means that we could encrypt for
+  /// at least one ratchet.
+  bool isSuccess(int numberOfRecipients) => encryptedKeys.isNotEmpty && jidEncryptionErrors.length < numberOfRecipients;
 }

lib/src/omemo/omemomanager.dart

@@ -447,8 +447,17 @@ class OmemoManager {
     }
 
     // We assume that the user already checked if the session exists
+    final deviceEncryptionErrors = <RatchetMapKey, OmemoException>{};
+    final jidEncryptionErrors = <String, OmemoException>{};
     for (final jid in jids) {
-      for (final deviceId in _deviceList[jid]!) {
+      final devices = _deviceList[jid];
+      if (devices == null) {
+        _log.severe('Device list does not exist for $jid.');
+        jidEncryptionErrors[jid] = NoKeyMaterialAvailableException();
+        continue;
+      }
+
+      for (final deviceId in devices) {
         // Empty OMEMO messages are allowed to bypass trust
         if (plaintext != null) {
           // Only encrypt to devices that are trusted
@@ -459,7 +468,13 @@ class OmemoManager {
         }
 
         final ratchetKey = RatchetMapKey(jid, deviceId);
-        var ratchet = _ratchetMap[ratchetKey]!;
+        var ratchet = _ratchetMap[ratchetKey];
+        if (ratchet == null) {
+          _log.severe('Ratchet ${ratchetKey.toJsonKey()} does not exist.');
+          deviceEncryptionErrors[ratchetKey] = NoKeyMaterialAvailableException();
+          continue;
+        }
+
         final ciphertext = (await ratchet.ratchetEncrypt(keyPayload)).ciphertext;
  
         if (kex.isNotEmpty && kex.containsKey(deviceId)) {
@@ -522,8 +537,11 @@ class OmemoManager {
     }
 
     return EncryptionResult(
-      plaintext != null ? ciphertext : null,
+      plaintext != null ?
+        ciphertext : null,
       encryptedKeys,
+      deviceEncryptionErrors,
+      jidEncryptionErrors,
     );
   }
 

lib/src/omemo/sessionmanager.dart

@@ -329,6 +329,8 @@ class OmemoSessionManager {
     return EncryptionResult(
       plaintext != null ? ciphertext : null,
       encryptedKeys,
+      const <RatchetMapKey, OmemoException>{},
+      const <String, OmemoException>{},
     );
   }
 

test/omemomanager_test.dart

@@ -574,4 +574,269 @@ void main() {
 
     expect(aliceResult3.payload, 'Hello Alice!');
   });
+
+  test('Test sending a message to two different JIDs', () async {
+    const aliceJid = 'alice@server1';
+    const bobJid = 'bob@server2';
+    const cocoJid = 'coco@server3';
+
+    final aliceDevice = await Device.generateNewDevice(aliceJid, opkAmount: 1);
+    final bobDevice = await Device.generateNewDevice(bobJid, opkAmount: 1);
+    final cocoDevice = await Device.generateNewDevice(cocoJid, opkAmount: 1);
+
+    final aliceManager = OmemoManager(
+      aliceDevice,
+      AlwaysTrustingTrustManager(),
+      (result, recipientJid) async {},
+      (jid) async {
+        if (jid == bobJid) {
+          return [bobDevice.id];
+        } else if (jid == cocoJid) {
+          return [cocoDevice.id];
+        }
+
+        return null;
+      },
+      (jid, id) async {
+        if (jid == bobJid) {
+          return bobDevice.toBundle();
+        } else if (jid == cocoJid) {
+          return cocoDevice.toBundle();
+        }
+
+        return null;
+      },
+    );
+    final bobManager = OmemoManager(
+      bobDevice,
+      AlwaysTrustingTrustManager(),
+      (result, recipientJid) async {},
+      (jid) async => null,
+      (jid, id) async => null,
+    );
+    final cocoManager = OmemoManager(
+      cocoDevice,
+      AlwaysTrustingTrustManager(),
+      (result, recipientJid) async {},
+      (jid) async => null,
+      (jid, id) async => null,
+    );
+
+    // Alice sends a message to Bob and Coco
+    final aliceResult = await aliceManager.onOutgoingStanza(
+      const OmemoOutgoingStanza(
+        [bobJid, cocoJid],
+        'Hello Bob and Coco!',
+      ),
+    );
+
+    // Bob and Coco decrypt them
+    final bobResult = await bobManager.onIncomingStanza(
+      OmemoIncomingStanza(
+        aliceJid,
+        aliceDevice.id,
+        DateTime.now().millisecondsSinceEpoch,
+        aliceResult!.encryptedKeys,
+        base64.encode(aliceResult.ciphertext!),
+      ),
+    );
+    final cocoResult = await cocoManager.onIncomingStanza(
+      OmemoIncomingStanza(
+        aliceJid,
+        aliceDevice.id,
+        DateTime.now().millisecondsSinceEpoch,
+        aliceResult.encryptedKeys,
+        base64.encode(aliceResult.ciphertext!),
+      ),
+    );
+
+    expect(bobResult.error, null);
+    expect(cocoResult.error, null);
+    expect(bobResult.payload, 'Hello Bob and Coco!');
+    expect(cocoResult.payload, 'Hello Bob and Coco!');
+  });
+
+  test('Test a fetch failure', () async {
+    const aliceJid = 'alice@server1';
+    const bobJid = 'bob@server2';
+    var failure = false;
+
+    final aliceDevice = await Device.generateNewDevice(aliceJid, opkAmount: 1);
+    final bobDevice = await Device.generateNewDevice(bobJid, opkAmount: 1);
+
+    final aliceManager = OmemoManager(
+      aliceDevice,
+      AlwaysTrustingTrustManager(),
+      (result, recipientJid) async {},
+      (jid) async {
+        expect(jid, bobJid);
+
+        return failure ?
+          null :
+          [bobDevice.id];
+      },
+      (jid, id) async {
+        expect(jid, bobJid);
+
+        return failure ?
+          null :
+          bobDevice.toBundle();
+      },
+    );
+    final bobManager = OmemoManager(
+      bobDevice,
+      AlwaysTrustingTrustManager(),
+      (result, recipientJid) async {},
+      (jid) async => null,
+      (jid, id) async => null,
+    );
+
+    // Alice sends a message to Bob and Coco
+    final aliceResult1 = await aliceManager.onOutgoingStanza(
+      const OmemoOutgoingStanza(
+        [bobJid],
+        'Hello Bob!',
+      ),
+    );
+
+    // Bob decrypts it
+    final bobResult1 = await bobManager.onIncomingStanza(
+      OmemoIncomingStanza(
+        aliceJid,
+        aliceDevice.id,
+        DateTime.now().millisecondsSinceEpoch,
+        aliceResult1!.encryptedKeys,
+        base64.encode(aliceResult1.ciphertext!),
+      ),
+    );
+
+    expect(bobResult1.error, null);
+    expect(bobResult1.payload, 'Hello Bob!');
+
+    // Bob acks the message
+    await aliceManager.ratchetAcknowledged(
+      bobJid,
+      bobDevice.id,
+    );
+
+    // Alice has to reconnect but has no connection yet
+    failure = true;
+    aliceManager.onNewConnection();
+
+    // Alice sends another message to Bob
+    final aliceResult2 = await aliceManager.onOutgoingStanza(
+      const OmemoOutgoingStanza(
+        [bobJid],
+        'Hello Bob! x2',
+      ),
+    );
+
+    // And Bob decrypts it
+    final bobResult2 = await bobManager.onIncomingStanza(
+      OmemoIncomingStanza(
+        aliceJid,
+        aliceDevice.id,
+        DateTime.now().millisecondsSinceEpoch,
+        aliceResult2!.encryptedKeys,
+        base64.encode(aliceResult2.ciphertext!),
+      ),
+    );
+
+    expect(bobResult2.error, null);
+    expect(bobResult2.payload, 'Hello Bob! x2');
+  });
+
+  test('Test sending a message with failed lookups', () async {
+    const aliceJid = 'alice@server1';
+    const bobJid = 'bob@server2';
+
+    final aliceDevice = await Device.generateNewDevice(aliceJid, opkAmount: 1);
+
+    final aliceManager = OmemoManager(
+      aliceDevice,
+      AlwaysTrustingTrustManager(),
+      (result, recipientJid) async {},
+      (jid) async {
+        expect(jid, bobJid);
+
+        return null;
+      },
+      (jid, id) async {
+        expect(jid, bobJid);
+
+        return null;
+      },
+    );
+
+    // Alice sends a message to Bob
+    final aliceResult = await aliceManager.onOutgoingStanza(
+      const OmemoOutgoingStanza(
+        [bobJid],
+        'Hello Bob!',
+      ),
+    );
+
+    expect(aliceResult!.isSuccess(1), false);
+    expect(aliceResult.jidEncryptionErrors[bobJid] is NoKeyMaterialAvailableException, true);
+  });
+
+  test('Test sending a message two two JIDs with failed lookups', () async {
+    const aliceJid = 'alice@server1';
+    const bobJid = 'bob@server2';
+    const cocoJid = 'coco@server3';
+
+    final aliceDevice = await Device.generateNewDevice(aliceJid, opkAmount: 1);
+    final bobDevice = await Device.generateNewDevice(bobJid, opkAmount: 1);
+
+    final aliceManager = OmemoManager(
+      aliceDevice,
+      AlwaysTrustingTrustManager(),
+      (result, recipientJid) async {},
+      (jid) async {
+        if (jid == bobJid) {
+          return [bobDevice.id];
+        }
+
+        return null;
+      },
+      (jid, id) async {
+        if (jid == bobJid) {
+          return bobDevice.toBundle();
+        }
+
+        return null;
+      },
+    );
+    final bobManager = OmemoManager(
+      bobDevice,
+      AlwaysTrustingTrustManager(),
+      (result, recipientJid) async {},
+      (jid) async => null,
+      (jid, id) async => null,
+    );
+
+    // Alice sends a message to Bob and Coco
+    final aliceResult = await aliceManager.onOutgoingStanza(
+      const OmemoOutgoingStanza(
+        [bobJid, cocoJid],
+        'Hello Bob and Coco!',
+      ),
+    );
+
+    expect(aliceResult!.isSuccess(2), true);
+    expect(aliceResult.jidEncryptionErrors[cocoJid] is NoKeyMaterialAvailableException, true);
+
+    // Bob decrypts it
+    final bobResult = await bobManager.onIncomingStanza(
+      OmemoIncomingStanza(
+        aliceJid,
+        aliceDevice.id,
+        DateTime.now().millisecondsSinceEpoch,
+        aliceResult.encryptedKeys,
+        base64.encode(aliceResult.ciphertext!),
+      ),
+    );
+
+    expect(bobResult.payload, 'Hello Bob and Coco!');
+  });
 }