fix: Pass all tests

lib/src/double_ratchet/double_ratchet.dart

@@ -40,13 +40,20 @@ class KeyExchangeData {
   const KeyExchangeData(
     this.pkId,
     this.spkId,
-    this.ek,
     this.ik,
+    this.ek,
   );
 
+  /// The id of the used OPK.
   final int pkId;
+
+  /// The id of the used SPK.
   final int spkId;
+
+  /// The ephemeral key used while the key exchange.
   final OmemoPublicKey ek;
+
+  /// The identity key used in the key exchange.
   final OmemoPublicKey ik;
 }
 
@@ -61,7 +68,6 @@ class OmemoDoubleRatchet {
     this.nr, // Nr
     this.pn, // Pn
     this.ik,
-    this.ek,
     this.sessionAd,
     this.mkSkipped, // MKSKIPPED
     this.acknowledged,
@@ -93,13 +99,10 @@ class OmemoDoubleRatchet {
   /// for verification purposes
   final OmemoPublicKey ik;
 
-  /// The ephemeral public key of the chat partner. Not used for encryption but for possible
+  /// Associated data for this ratchet.
-  /// checks when replacing the ratchet. As such, this is only non-null for the initiating
-  /// side.
-  final OmemoPublicKey? ek;
-
   final List<int> sessionAd;
 
+  /// List of skipped message keys.
   final Map<SkippedKey, List<int>> mkSkipped;
 
   /// The point in time at which we performed the kex exchange to create this ratchet.
@@ -107,7 +110,7 @@ class OmemoDoubleRatchet {
   int kexTimestamp;
 
   /// The key exchange that was used for initiating the session.
-  final KeyExchangeData? kex;
+  final KeyExchangeData kex;
 
   /// Indicates whether we received an empty OMEMO message after building a session with
   /// the device.
@@ -118,13 +121,14 @@ class OmemoDoubleRatchet {
   /// a X3DH. [ik] refers to Bob's (the receiver's) IK public key.
   static Future<OmemoDoubleRatchet> initiateNewSession(
     OmemoPublicKey spk,
+    int spkId,
     OmemoPublicKey ik,
+    OmemoPublicKey ownIk,
     OmemoPublicKey ek,
     List<int> sk,
     List<int> ad,
     int timestamp,
     int pkId,
-    int spkId,
   ) async {
     final dhs = await OmemoKeyPair.generateNewPair(KeyPairType.x25519);
     final rk = await kdfRk(sk, await omemoDH(dhs, spk, 0));
@@ -139,7 +143,6 @@ class OmemoDoubleRatchet {
       0,
       0,
       ik,
-      ek,
       ad,
       {},
       false,
@@ -147,7 +150,7 @@ class OmemoDoubleRatchet {
       KeyExchangeData(
         pkId,
         spkId,
-        ik,
+        ownIk,
         ek,
       ),
     );
@@ -159,7 +162,10 @@ class OmemoDoubleRatchet {
   /// Alice's (the initiator's) IK public key.
   static Future<OmemoDoubleRatchet> acceptNewSession(
     OmemoKeyPair spk,
+    int spkId,
     OmemoPublicKey ik,
+    int pkId,
+    OmemoPublicKey ek,
     List<int> sk,
     List<int> ad,
     int kexTimestamp,
@@ -174,12 +180,16 @@ class OmemoDoubleRatchet {
       0,
       0,
       ik,
-      null,
       ad,
       {},
       true,
       kexTimestamp,
-      null,
+      KeyExchangeData(
+        pkId,
+        spkId,
+        ik,
+        ek,
+      ),
     );
   }
 
@@ -357,7 +367,6 @@ class OmemoDoubleRatchet {
       nr,
       pn,
       ik,
-      ek,
       sessionAd,
       Map<SkippedKey, List<int>>.from(mkSkipped),
       acknowledged,

lib/src/omemo/encryption_result.dart

@@ -22,8 +22,9 @@ class EncryptionResult {
   /// Mapping of a JID to 
   final Map<String, List<EncryptToJidError>> deviceEncryptionErrors;
 
+  // TODO: Turn this into a property that is computed in [onOutgoingStanza].
   /// True if the encryption was a success. This means that we could encrypt for
-  /// at least one ratchet.
+  /// at least one ratchet per recipient. [recipients] is the number of recipients
-  /// TODO:
+  /// that the message should've been encrypted for.
-  bool isSuccess(int numberOfRecipients) => true;
+  bool isSuccess(int recipients) => encryptedKeys.length == recipients;
 }

lib/src/omemo/events.dart

@@ -39,8 +39,13 @@ class RatchetRemovedEvent extends OmemoEvent {
 
 /// Triggered when the device map has been modified
 class DeviceListModifiedEvent extends OmemoEvent {
-  DeviceListModifiedEvent(this.list);
+  DeviceListModifiedEvent(this.jid, this.devices);
-  final Map<String, List<int>> list;
+
+  /// The JID of the user.
+  final String jid;
+
+  /// The list of devices for [jid].
+  final List<int> devices;
 }
 
 /// Triggered by the OmemoSessionManager when our own device bundle was modified

lib/src/omemo/omemo.dart

@@ -3,7 +3,6 @@ import 'dart:collection';
 import 'dart:convert';
 import 'package:collection/collection.dart';
 import 'package:cryptography/cryptography.dart';
-import 'package:hex/hex.dart';
 import 'package:logging/logging.dart';
 import 'package:meta/meta.dart';
 import 'package:omemo_dart/src/common/result.dart';
@@ -28,20 +27,6 @@ import 'package:omemo_dart/src/trust/base.dart';
 import 'package:omemo_dart/src/x3dh/x3dh.dart';
 import 'package:synchronized/synchronized.dart';
 
-class _InternalDecryptionResult {
-  const _InternalDecryptionResult(
-    this.ratchetCreated,
-    this.ratchetReplaced,
-    this.payload,
-  ) : assert(
-          !ratchetCreated || !ratchetReplaced,
-          'Ratchet must be either replaced or created',
-        );
-  final bool ratchetCreated;
-  final bool ratchetReplaced;
-  final String? payload;
-}
-
 extension AppendToListOrCreateExtension<K, V> on Map<K, List<V>> {
   void appendOrCreate(K key, V value) {
     if (containsKey(key)) {
@@ -185,26 +170,27 @@ class OmemoManager {
   /// Returns a list of new bundles, that may be empty.
   Future<List<OmemoBundle>> _fetchNewOmemoBundles(String jid) async {
     // Do we have to request the device list or are we already up-to-date?
-    if (_deviceListRequested.containsKey(jid) && _deviceList.containsKey(jid)) {
+    if (!_deviceListRequested.containsKey(jid) || !_deviceList.containsKey(jid)) {
-      return [];
-    }
-
       final newDeviceList = await fetchDeviceListImpl(jid);
-    if (newDeviceList == null) {
+      if (newDeviceList != null) {
-      return [];
-    }
-
         // Figure out what bundles we must fetch
         _deviceList[jid] = newDeviceList;
         _deviceListRequested[jid] = true;
 
-    // TODO: Maybe do this per JID?
         _eventStreamController.add(
-      DeviceListModifiedEvent(_deviceList),
+          DeviceListModifiedEvent(jid, newDeviceList),
         );
+      }
+    }
+
+    // Check that we have the device list
+    if (!_deviceList.containsKey(jid)) {
+      _log.warning('$jid not tracked in device list.');
+      return [];
+    }
 
     final ownDevice = await getDevice();
-    final bundlesToFetch = newDeviceList.where((device) {
+    final bundlesToFetch = _deviceList[jid]!.where((device) {
       // Do not include our current device, if we request bundles for our own JID.
       if (ownDevice.jid == jid && device == ownDevice.id) {
         return false;
@@ -285,14 +271,28 @@ class OmemoManager {
       );
     }
 
+    // Check how we should process the message
     final ratchetKey = RatchetMapKey(stanza.bareSenderJid, stanza.senderDeviceId);
-    if (key.kex) {
+    var processAsKex = key.kex;
+    if (key.kex && _ratchetMap.containsKey(ratchetKey)) {
+      final ratchet = _ratchetMap[ratchetKey]!;
+      final kexMessage = OMEMOKeyExchange.fromBuffer(base64Decode(key.value));
+      final ratchetEk = await ratchet.kex.ek.getBytes();
+      final sameEk = listsEqual(kexMessage.ek, ratchetEk);
+
+      if (sameEk) {
+        processAsKex = false;
+      } else {
+        processAsKex = true;
+      }
+      _log.finest('kexMessage.ek == ratchetEk: $sameEk');
+    }
+
+    // Process the message
+    if (processAsKex) {
       _log.finest('Decoding message as OMEMOKeyExchange');
       final kexMessage = OMEMOKeyExchange.fromBuffer(base64Decode(key.value));
 
-      // TODO: Check if we already have such a session and if we can build it
-      // See XEP-0384 4.3
-
       // Find the correct SPK
       final device = await getDevice();
       OmemoKeyPair spk;
@@ -312,13 +312,16 @@ class OmemoManager {
         kexMessage.ik,
         KeyPairType.ed25519,
       );
+      final kexEk = OmemoPublicKey.fromBytes(
+        kexMessage.ek,
+        KeyPairType.x25519,
+      );
+
+      // TODO: Guard against invalid signatures
       final kex = await x3dhFromInitialMessage(
         X3DHMessage(
           kexIk, 
-          OmemoPublicKey.fromBytes(
+          kexEk,
-            kexMessage.ek,
-            KeyPairType.x25519,
-          ),
           kexMessage.pkId,
         ),
         spk,
@@ -327,7 +330,10 @@ class OmemoManager {
       );
       final ratchet = await OmemoDoubleRatchet.acceptNewSession(
         spk,
+        kexMessage.spkId,
         kexIk,
+        kexMessage.pkId,
+        kexEk,
         kex.sk,
         kex.ad,
         getTimestamp(),
@@ -395,8 +401,7 @@ class OmemoManager {
       }
 
       // Send the hearbeat, if we have to
-      // TODO: Handle replace
+      await _maybeSendEmptyMessage(ratchetKey, true, _ratchetMap.containsKey(ratchetKey));
-      await _maybeSendEmptyMessage(ratchetKey, true, false);
 
       return DecryptionResult(
         result.get<String?>(),
@@ -415,7 +420,16 @@ class OmemoManager {
 
       _log.finest('Decoding message as OMEMOAuthenticatedMessage');
       final ratchet = _ratchetMap[ratchetKey]!.clone();
-      final authMessage = OMEMOAuthenticatedMessage.fromBuffer(base64Decode(key.value));
+
+      // Correctly decode the message
+      OMEMOAuthenticatedMessage authMessage;
+      if (key.kex) {
+        _log.finest('Extracting OMEMOAuthenticatedMessage from OMEMOKeyExchange');
+        authMessage = OMEMOKeyExchange.fromBuffer(base64Decode(key.value)).message;
+      } else {
+        authMessage = OMEMOAuthenticatedMessage.fromBuffer(base64Decode(key.value));
+      }
+
       final keyAndHmac = await ratchet.ratchetDecrypt(authMessage);
       if (keyAndHmac.isType<OmemoError>()) {
         final error = keyAndHmac.get<OmemoError>();
@@ -504,6 +518,7 @@ class OmemoManager {
       }
 
       for (final bundle in newBundles) {
+        _log.finest('Building new ratchet $jid:${bundle.id}');
         final ratchetKey = RatchetMapKey(jid, bundle.id);
         final ownDevice = await getDevice();
         final kexResult = await x3dhFromBundle(
@@ -512,13 +527,14 @@ class OmemoManager {
         );
         final newRatchet = await OmemoDoubleRatchet.initiateNewSession(
           bundle.spk,
+          bundle.spkId,
           bundle.ik,
+          ownDevice.ik.pk,
           kexResult.ek.pk,
           kexResult.sk,
           kexResult.ad,
           getTimestamp(),
           kexResult.opkId,
-          bundle.spkId,
         );
 
         // Track the ratchet
@@ -529,11 +545,13 @@ class OmemoManager {
         await trustManager.onNewSession(jid, bundle.id);
 
         // Track the KEX for later
+        final ik = await ownDevice.ik.pk.getBytes();
+        final ek = await kexResult.ek.pk.getBytes();
         kex[ratchetKey] = OMEMOKeyExchange()
-          ..pkId = kexResult.opkId
+          ..pkId = newRatchet.kex.pkId
-          ..spkId = bundle.spkId
+          ..spkId = newRatchet.kex.spkId
-          ..ik = await ownDevice.ik.pk.getBytes()
+          ..ik = ik
-          ..ek = await kexResult.ek.pk.getBytes();
+          ..ek = ek;
       }
     }
 
@@ -615,30 +633,16 @@ class OmemoManager {
             ),
           );
         } else if (!ratchet.acknowledged) {
-          // The ratchet as not yet been acked
+          // The ratchet as not yet been acked.
-          if (ratchet.kex == null) {
-            // The ratchet is not acked but we also don't have an old KEX to send with it
-            _log.warning('Ratchet $jid:$device is not acked but has no previous KEX.');
-
-            encryptedKeys.appendOrCreate(
-              jid,
-              EncryptedKey(
-                jid,
-                device,
-                base64Encode(authMessage.writeToBuffer()),
-                false,
-              ),
-            );
-            continue;
-          }
-
           // Keep sending the old KEX
+          _log.finest('Using old KEX data for OMEMOKeyExchange');
           final kexMessage = OMEMOKeyExchange()
-            ..pkId = ratchet.kex!.pkId
+            ..pkId = ratchet.kex.pkId
-            ..spkId = ratchet.kex!.spkId
+            ..spkId = ratchet.kex.spkId
-            ..ik = await ratchet.kex!.ik.getBytes()
+            ..ik = await ratchet.kex.ik.getBytes()
-            ..ek = await ratchet.kex!.ek.getBytes()
+            ..ek = await ratchet.kex.ek.getBytes()
             ..message = authMessage;
+
           encryptedKeys.appendOrCreate(
             jid,
             EncryptedKey(
@@ -681,14 +685,43 @@ class OmemoManager {
     await sendEmptyOmemoMessageImpl(result, jid);
   }
 
-  // TODO
+  /// Removes all ratchets associated with [jid].
-  Future<void> removeAllRatchets(String jid) async {}
+  Future<void> removeAllRatchets(String jid) async {
+    await _enterRatchetCriticalSection(jid);
 
-  // TODO
+    for (final device in _deviceList[jid] ?? <int>[]) {
-  Future<void> onDeviceListUpdate(String jid, List<int> devices) async {}
+      // Remove the ratchet and commit
+      _ratchetMap.remove(RatchetMapKey(jid, device));
+      _eventStreamController.add(RatchetRemovedEvent(jid, device));
+    }
 
-  // TODO
+    // Clear the device list
-  Future<void> onNewConnection() async {}
+    _deviceList.remove(jid);
+    _deviceListRequested.remove(jid);
+    _eventStreamController.add(DeviceListModifiedEvent(jid, []));
+
+    await _leaveRatchetCriticalSection(jid);
+  }
+
+  /// To be called when a update to the device list of [jid] is returned.
+  /// [devices] is the list of device identifiers contained in the update.
+  Future<void> onDeviceListUpdate(String jid, List<int> devices) async {
+    // Update our state
+    _deviceList[jid] = devices;
+    _deviceListRequested[jid] = true;
+
+    // Commit the device list
+    _eventStreamController.add(
+      DeviceListModifiedEvent(jid, devices),
+    );
+  }
+
+  /// To be called when a new connection is made, i.e. when the previous stream could
+  /// previous stream could not be resumed using XEP-0198.
+  Future<void> onNewConnection() async {
+    _deviceListRequested.clear();
+    _subscriptionMap.clear();
+  }
 
   // Mark the ratchet [jid]:[device] as acknowledged.
   Future<void> ratchetAcknowledged(String jid, int device) async {
@@ -709,7 +742,7 @@ class OmemoManager {
   }
 
   // TODO
-  Future<List<DeviceFingerprint>> getFingerprintsForJid(String jid) async => [];
+  Future<List<DeviceFingerprint>?> getFingerprintsForJid(String jid) async => null;
 
   /// Returns the device used for encryption and decryption.
   Future<OmemoDevice> getDevice() => _deviceLock.synchronized(() => _device);
@@ -718,5 +751,5 @@ class OmemoManager {
   Future<int> getDeviceId() async => (await getDevice()).id;
 
   @visibleForTesting
-  OmemoDoubleRatchet getRatchet(RatchetMapKey key) => _ratchetMap[key]!;
+  OmemoDoubleRatchet? getRatchet(RatchetMapKey key) => _ratchetMap[key];
 }

test/double_ratchet_test.dart

@@ -51,17 +51,21 @@ void main() {
     // Build a session
     final alicesRatchet = await OmemoDoubleRatchet.initiateNewSession(
       spkBob.pk,
+      bundleBob.spkId,
       ikBob.pk,
+      ikAlice.pk,
       resultAlice.ek.pk,
       resultAlice.sk,
       resultAlice.ad,
       0,
       resultAlice.opkId,
-      bundleBob.spkId,
     );
     final bobsRatchet = await OmemoDoubleRatchet.acceptNewSession(
       spkBob,
+      bundleBob.spkId,
       ikAlice.pk,
+      2,
+      resultAlice.ek.pk,
       resultBob.sk,
       resultBob.ad,
       0,

test/omemo_test.dart

@@ -1,6 +1,7 @@
 import 'dart:convert';
 import 'package:logging/logging.dart';
 import 'package:omemo_dart/omemo_dart.dart';
+import 'package:omemo_dart/src/protobuf/schema.pb.dart';
 import 'package:omemo_dart/src/trust/always.dart';
 import 'package:test/test.dart';
 
@@ -210,7 +211,7 @@ void main() {
         ),
       );
 
-      expect(aliceResultLoop.encryptedKeys[bobJid]!.first.kex, false);
+      expect(aliceResultLoop.encryptedKeys[bobJid]!.first.kex, isFalse);
 
       final bobResultLoop = await bobManager.onIncomingStanza(
         OmemoIncomingStanza(
@@ -376,7 +377,7 @@ void main() {
     );
 
     expect(bobResult1.payload, null);
-    expect(bobResult1.error is NotEncryptedForDeviceError, true);
+    expect(bobResult1.error, const TypeMatcher<NotEncryptedForDeviceError>());
 
     // Now Alice's client loses and regains the connection
     await aliceManager.onNewConnection();
@@ -401,6 +402,7 @@ void main() {
     );
 
     expect(aliceResult2.encryptedKeys.length, 1);
+    expect(bobResult2.error, null);
     expect(bobResult2.payload, 'Hello Bob x2');
   });
 
@@ -506,7 +508,7 @@ void main() {
         bobJid,
         bobDevice2.id,
         DateTime.now().millisecondsSinceEpoch,
-        bobResult2.encryptedKeys[bobJid]!,
+        bobResult2.encryptedKeys[aliceJid]!,
         base64.encode(bobResult2.ciphertext!),
         false,
       ),
@@ -617,7 +619,7 @@ void main() {
       ),
     );
 
-    expect(aliceResult2.encryptedKeys.length, 2);
+    expect(aliceResult2.encryptedKeys[bobJid]!.length, 2);
 
     // And Bob decrypts it
     final bobResult21 = await bobManager1.onIncomingStanza(
@@ -658,7 +660,7 @@ void main() {
         bobJid,
         bobDevice2.id,
         DateTime.now().millisecondsSinceEpoch,
-        bobResult32.encryptedKeys[bobJid]!,
+        bobResult32.encryptedKeys[aliceJid]!,
         base64.encode(bobResult32.ciphertext!),
         false,
       ),
@@ -878,13 +880,10 @@ void main() {
       ),
     );
 
-    expect(aliceResult.isSuccess(1), false);
+    expect(aliceResult.isSuccess(1), isFalse);
-    // TODO
+    expect(aliceResult.deviceEncryptionErrors[bobJid]!.length, 1);
-    /*expect(
+    final error = aliceResult.deviceEncryptionErrors[bobJid]!.first;
-      aliceResult.jidEncryptionErrors[bobJid]
+    expect(error.error, const TypeMatcher<NoKeyMaterialAvailableError>());
-          is NoKeyMaterialAvailableException,
-      true,
-    );*/
   });
 
   test('Test sending a message two two JIDs with failed lookups', () async {
@@ -933,14 +932,9 @@ void main() {
       ),
     );
 
-    expect(aliceResult.isSuccess(2), true);
+    expect(aliceResult.isSuccess(2), isFalse);
-    // TODO
+    expect(aliceResult.deviceEncryptionErrors[cocoJid]!.length, 1);
-    /*
+    expect(aliceResult.deviceEncryptionErrors[cocoJid]!.first.error, const TypeMatcher<NoKeyMaterialAvailableError>(),);
-    expect(
-      aliceResult.jidEncryptionErrors[cocoJid]
-          is NoKeyMaterialAvailableException,
-      true,
-    );*/
 
     // Bob decrypts it
     final bobResult = await bobManager.onIncomingStanza(
@@ -1025,7 +1019,7 @@ void main() {
           messageText,
         ),
       );
-      expect(bobResponseMessage.isSuccess(1), true);
+      expect(bobResponseMessage.isSuccess(1), isTrue);
 
       final aliceReceivedMessage = await aliceManager.onIncomingStanza(
         OmemoIncomingStanza(
@@ -1181,7 +1175,7 @@ void main() {
       'Test removing all ratchets and sending a message without post-heartbeat ack',
       () async {
     // This test is the same as "Test removing all ratchets and sending a message" except
-    // that bob does not ack the ratchet after Alice's heartbeat after she recreated
+    // that Bob does not ack the ratchet after Alice's heartbeat after she recreated
     // all ratchets.
     const aliceJid = 'alice@server1';
     const bobJid = 'bob@server2';
@@ -1227,7 +1221,7 @@ void main() {
     );
 
     // And Bob decrypts it
-    await bobManager.onIncomingStanza(
+    final bobResult1 = await bobManager.onIncomingStanza(
       OmemoIncomingStanza(
         aliceJid,
         aliceDevice.id,
@@ -1237,6 +1231,7 @@ void main() {
         false,
       ),
     );
+    expect(bobResult1.error, isNull);
 
     // Ratchets are acked
     await aliceManager.ratchetAcknowledged(
@@ -1245,9 +1240,10 @@ void main() {
     );
 
     // Alice now removes all ratchets for Bob and sends another new message
+    Logger.root.info('Removing all ratchets for $bobJid');
     await aliceManager.removeAllRatchets(bobJid);
 
-    expect(aliceManager.getRatchet(RatchetMapKey(bobJid, bobDevice.id)), null);
+    expect(aliceManager.getRatchet(RatchetMapKey(bobJid, bobDevice.id)), isNull);
 
     // Alice prepares an empty OMEMO message
     await aliceManager.sendOmemoHeartbeat(bobJid);
@@ -1266,12 +1262,14 @@ void main() {
     expect(bobResult2.error, null);
 
     // Alice sends another message
+    Logger.root.info('Sending final message');
     final aliceResult3 = await aliceManager.onOutgoingStanza(
       const OmemoOutgoingStanza(
         [bobJid],
         'I did not trust your last device, Bob!',
       ),
     );
+    expect(aliceResult3.encryptedKeys[bobJid]!.first.kex, isTrue);
 
     // Bob decrypts it
     final bobResult3 = await bobManager.onIncomingStanza(
@@ -1358,8 +1356,7 @@ void main() {
     );
 
     // The first message must be a KEX message
-    // TODO
+    expect(aliceResult1.encryptedKeys[bobJid]!.first.kex, isTrue);
-    //expect(aliceResult1.encryptedKeys.first.kex, true);
 
     // Bob decrypts Alice's message
     final bobResult1 = await bobManager.onIncomingStanza(
@@ -1384,24 +1381,21 @@ void main() {
     );
 
     // The response should contain a KEX
-    // TODO
+    expect(aliceResult2.encryptedKeys[bobJid]!.first.kex, isTrue);
-    /*
-    expect(aliceResult2.encryptedKeys.first.kex, true);
 
     // The basic data should be the same
     final parsedFirstKex = OMEMOKeyExchange.fromBuffer(
-      base64.decode(aliceResult1.encryptedKeys.first.value),
+      base64.decode(aliceResult1.encryptedKeys[bobJid]!.first.value),
     );
     final parsedSecondKex = OMEMOKeyExchange.fromBuffer(
-      base64.decode(aliceResult2.encryptedKeys.first.value),
+      base64.decode(aliceResult2.encryptedKeys[bobJid]!.first.value),
     );
     expect(parsedSecondKex.pkId, parsedFirstKex.pkId);
     expect(parsedSecondKex.spkId, parsedFirstKex.spkId);
     expect(parsedSecondKex.ik, parsedFirstKex.ik);
     expect(parsedSecondKex.ek, parsedFirstKex.ek);
-    */
 
-    // Alice decrypts it
+    // Bob decrypts it
     final bobResult2 = await bobManager.onIncomingStanza(
       OmemoIncomingStanza(
         aliceJid,
@@ -1429,7 +1423,7 @@ void main() {
         bobJid,
         bobDevice.id,
         DateTime.now().millisecondsSinceEpoch,
-        bobResult3.encryptedKeys[bobJid]!,
+        bobResult3.encryptedKeys[aliceJid]!,
         base64.encode(bobResult3.ciphertext!),
         false,
       ),
@@ -1452,8 +1446,7 @@ void main() {
     );
 
     // The response should contain no KEX
-    // TODO
+    expect(aliceResult4.encryptedKeys[bobJid]!.first.kex, isFalse);
-    //expect(aliceResult4.encryptedKeys.first.kex, false);
 
     // Bob decrypts it
     final bobResult4 = await bobManager.onIncomingStanza(

test/serialisation_test.dart

@@ -1,6 +1,6 @@
 import 'dart:convert';
 import 'package:omemo_dart/omemo_dart.dart';
-import 'package:omemo_dart/protobuf/schema.pb.dart';
+import 'package:omemo_dart/src/protobuf/schema.pb.dart';
 import 'package:omemo_dart/src/trust/always.dart';
 import 'package:test/test.dart';