fix: Fix ratchet null issue

Fixes #26.

lib/src/omemo/events.dart

@@ -5,13 +5,16 @@ abstract class OmemoEvent {}
 
 /// Triggered when a ratchet has been modified
 class RatchetModifiedEvent extends OmemoEvent {
-  RatchetModifiedEvent(this.jid, this.deviceId, this.ratchet, this.added);
+  RatchetModifiedEvent(this.jid, this.deviceId, this.ratchet, this.added, this.replaced);
   final String jid;
   final int deviceId;
   final OmemoDoubleRatchet ratchet;
 
   /// Indicates whether the ratchet has just been created (true) or just modified (false).
   final bool added;
+
+  /// Indicates whether the ratchet has been replaced (true) or not.
+  final bool replaced;
 }
 
 /// Triggered when a ratchet has been removed and should be removed from storage.

lib/src/omemo/omemomanager.dart

@@ -29,8 +29,13 @@ import 'package:omemo_dart/src/x3dh/x3dh.dart';
 import 'package:synchronized/synchronized.dart';
 
 class _InternalDecryptionResult {
-  const _InternalDecryptionResult(this.ratchetCreated, this.payload);
+  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;
 }
 
@@ -162,7 +167,7 @@ class OmemoManager {
     _ratchetMap[key] = ratchet;
 
     // Commit the ratchet
-    _eventStreamController.add(RatchetModifiedEvent(jid, deviceId, ratchet, true));
+    _eventStreamController.add(RatchetModifiedEvent(jid, deviceId, ratchet, true, false));
   }
 
   /// Build a new session with the user at [jid] with the device [deviceId] using data
@@ -243,6 +248,7 @@ class OmemoManager {
         mapKey.deviceId,
         oldRatchet,
         false,
+        false,
       ),
     );
   }
@@ -267,17 +273,17 @@ class OmemoManager {
       throw NotEncryptedForDeviceException();
     }
  
-    final ratchetKey = RatchetMapKey(senderJid, senderDeviceId);
     final decodedRawKey = base64.decode(rawKey.value);
     List<int>? keyAndHmac;
     OmemoAuthenticatedMessage authMessage;
-    OmemoDoubleRatchet? oldRatchet;
     OmemoMessage? message;
     var ratchetCreated = false;
-    if (rawKey.kex) {
+
     // If the ratchet already existed, we store it. If it didn't, oldRatchet will stay
     // null.
+    final ratchetKey = RatchetMapKey(senderJid, senderDeviceId);
     final oldRatchet = getRatchet(ratchetKey)?.clone();
+    if (rawKey.kex) {
       final kex = OmemoKeyExchange.fromBuffer(decodedRawKey);
       authMessage = kex.message!;
       message = OmemoMessage.fromBuffer(authMessage.message!);
@@ -288,48 +294,45 @@ class OmemoManager {
         if (oldRatchet.kexTimestamp > timestamp) {
           throw InvalidKeyExchangeException();
         }
-        
-        // Try to decrypt it
-        try {
-          final decrypted = await oldRatchet.ratchetDecrypt(message, authMessage.writeToBuffer());
-
-          // Commit the ratchet
-          _eventStreamController.add(
-            RatchetModifiedEvent(
-              senderJid,
-              senderDeviceId,
-              oldRatchet,
-              false,
-            ),
-          );
-          
-          final plaintext = await _decryptAndVerifyHmac(
-            ciphertext,
-            decrypted,
-          );
-          _addSession(senderJid, senderDeviceId, oldRatchet);
-          return _InternalDecryptionResult(
-            true,
-            plaintext,
-          );
-        } catch (_) {
-          _log.finest('Failed to use old ratchet with KEX for existing ratchet');
-        }
       }
 
       final r = await _addSessionFromKeyExchange(senderJid, senderDeviceId, kex);
-      await _trustManager.onNewSession(senderJid, senderDeviceId);
+
+      // Try to decrypt with the new ratchet r
+      try {
+        keyAndHmac = await r.ratchetDecrypt(message, authMessage.writeToBuffer());
+        final result = await _decryptAndVerifyHmac(ciphertext, keyAndHmac);
+
+        // Add the new ratchet
         _addSession(senderJid, senderDeviceId, r);
-      ratchetCreated = true;
 
         // Replace the OPK
-      // TODO(PapaTutuWawa): Replace the OPK when we know that the KEX worked
         await _deviceLock.synchronized(() async {
           device = await device.replaceOnetimePrekey(kex.pkId!);
 
           // Commit the device
           _eventStreamController.add(DeviceModifiedEvent(device));
         });
+
+        // Commit the ratchet
+        _eventStreamController.add(
+          RatchetModifiedEvent(
+            senderJid,
+            senderDeviceId,
+            r,
+            oldRatchet == null,
+            oldRatchet != null,
+          ),
+        );
+
+        return _InternalDecryptionResult(
+          oldRatchet == null,
+          oldRatchet != null,
+          result,
+        );
+      } catch (ex) {
+        _log.finest('Kex failed due to $ex. Not proceeding with kex.');
+      }
     } else {
       authMessage = OmemoAuthenticatedMessage.fromBuffer(decodedRawKey);
       message = OmemoMessage.fromBuffer(authMessage.message!);
@@ -340,9 +343,10 @@ class OmemoManager {
       throw NoDecryptionKeyException();
     }
 
+    // TODO(PapaTutuWawa): When receiving a message that is not an OMEMOKeyExchange from a device there is no session with, clients SHOULD create a session with that device and notify it about the new session by responding with an empty OMEMO message as per Sending a message.
+    
     // We can guarantee that the ratchet exists at this point in time
     final ratchet = getRatchet(ratchetKey)!;
-    oldRatchet ??= ratchet.clone();
 
     try {
       if (rawKey.kex) {
@@ -351,7 +355,7 @@ class OmemoManager {
         keyAndHmac = await ratchet.ratchetDecrypt(message, decodedRawKey);
       }
     } catch (_) {
-      _restoreRatchet(ratchetKey, oldRatchet);
+      _restoreRatchet(ratchetKey, oldRatchet!);
       rethrow;
     }
 
@@ -362,16 +366,18 @@ class OmemoManager {
         senderDeviceId,
         ratchet,
         false,
+        false,
       ),
     );
 
     try {
       return _InternalDecryptionResult(
         ratchetCreated,
+        false,
         await _decryptAndVerifyHmac(ciphertext, keyAndHmac),
       );
     } catch (_) { 
-      _restoreRatchet(ratchetKey, oldRatchet);
+      _restoreRatchet(ratchetKey, oldRatchet!);
       rethrow;
     }
   }
@@ -551,7 +557,7 @@ class OmemoManager {
         }
 
         // Commit the ratchet
-        _eventStreamController.add(RatchetModifiedEvent(jid, deviceId, ratchet, false));
+        _eventStreamController.add(RatchetModifiedEvent(jid, deviceId, ratchet, false, false));
       }
     }
 
@@ -600,7 +606,7 @@ class OmemoManager {
 
     if (ratchet!.acknowledged) {
       // Ratchet is acknowledged
-      if (ratchet.nr > 53 || result.ratchetCreated) {
+      if (ratchet.nr > 53 || result.ratchetCreated || result.ratchetReplaced) {
         await sendEmptyOmemoMessageImpl(
           await _encryptToJids(
             [stanza.bareSenderJid],
@@ -677,7 +683,7 @@ class OmemoManager {
         ..acknowledged = true;
 
       // Commit it
-      _eventStreamController.add(RatchetModifiedEvent(jid, deviceId, ratchet, false));
+      _eventStreamController.add(RatchetModifiedEvent(jid, deviceId, ratchet, false, false));
     } else {
       _log.severe('Attempted to acknowledge ratchet ${key.toJsonKey()}, even though it does not exist');
     }

lib/src/omemo/sessionmanager.dart

@@ -136,7 +136,7 @@ class OmemoSessionManager {
       _ratchetMap[key] = ratchet;
 
       // Commit the ratchet
-      _eventStreamController.add(RatchetModifiedEvent(jid, deviceId, ratchet, true));
+      _eventStreamController.add(RatchetModifiedEvent(jid, deviceId, ratchet, true, false));
     });
   }
 
@@ -321,7 +321,7 @@ class OmemoSessionManager {
           }
 
           // Commit the ratchet
-          _eventStreamController.add(RatchetModifiedEvent(jid, deviceId, ratchet, false));
+          _eventStreamController.add(RatchetModifiedEvent(jid, deviceId, ratchet, false, false));
         }
       }
     });
@@ -349,6 +349,7 @@ class OmemoSessionManager {
           mapKey.deviceId,
           oldRatchet,
           false,
+          false,
         ),
       );
     });
@@ -425,6 +426,7 @@ class OmemoSessionManager {
               senderDeviceId,
               oldRatchet,
               false,
+              false,
             ),
           );
           
@@ -486,6 +488,7 @@ class OmemoSessionManager {
         senderDeviceId,
         ratchet,
         false,
+        false,
       ),
     );
 
@@ -617,7 +620,7 @@ class OmemoSessionManager {
         ..acknowledged = true;
 
       // Commit it
-      _eventStreamController.add(RatchetModifiedEvent(jid, deviceId, ratchet, false));
+      _eventStreamController.add(RatchetModifiedEvent(jid, deviceId, ratchet, false, false));
     });
   }
 

test/omemomanager_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/omemo_key_exchange.dart';
 import 'package:omemo_dart/src/trust/always.dart';
 import 'package:test/test.dart';
 
@@ -950,10 +951,13 @@ void main() {
     final aliceDevice = await OmemoDevice.generateNewDevice(aliceJid, opkAmount: 1);
     final bobDevice = await OmemoDevice.generateNewDevice(bobJid, opkAmount: 1);
 
+    EncryptionResult? aliceEmptyMessage;
     final aliceManager = OmemoManager(
       aliceDevice,
       AlwaysTrustingTrustManager(),
-      (result, recipientJid) async {},
+      (result, recipientJid) async {
+        aliceEmptyMessage = result;
+      },
       (jid) async {
         expect(jid, bobJid);
 
@@ -1005,7 +1009,29 @@ void main() {
 
     expect(aliceManager.getRatchet(RatchetMapKey(bobJid, bobDevice.id)), null);
 
-    final aliceResult2 = await aliceManager.onOutgoingStanza(
+    // Alice prepares an empty OMEMO message
+    await aliceManager.sendOmemoHeartbeat(bobJid);
+
+    // And Bob receives it
+    final bobResult2 = await bobManager.onIncomingStanza(
+      OmemoIncomingStanza(
+        aliceJid,
+        await aliceManager.getDeviceId(),
+        DateTime.now().millisecondsSinceEpoch,
+        aliceEmptyMessage!.encryptedKeys,
+        null,
+      ),
+    );
+    expect(bobResult2.error, null);
+
+    // Bob acks the new ratchet
+    await aliceManager.ratchetAcknowledged(
+      bobJid,
+      await bobManager.getDeviceId(),
+    );
+    
+    // Alice sends another message
+    final aliceResult3 = await aliceManager.onOutgoingStanza(
       const OmemoOutgoingStanza(
         [bobJid],
         'I did not trust your last device, Bob!',
@@ -1013,6 +1039,254 @@ void main() {
     );
 
     // Bob decrypts it
+    final bobResult3 = await bobManager.onIncomingStanza(
+      OmemoIncomingStanza(
+        aliceJid,
+        aliceDevice.id,
+        DateTime.now().millisecondsSinceEpoch,
+        aliceResult3.encryptedKeys,
+        base64.encode(aliceResult3.ciphertext!),
+      ),
+    );
+
+    expect(bobResult3.error, null);
+    expect(bobResult3.payload, 'I did not trust your last device, Bob!');
+
+    // Bob responds
+    final bobResult4 = await bobManager.onOutgoingStanza(
+      OmemoOutgoingStanza(
+        [aliceJid],
+        "That's okay.",
+      ),
+    );
+
+    // Alice decrypts
+    final aliceResult4 = await aliceManager.onIncomingStanza(
+      OmemoIncomingStanza(
+        bobJid,
+        bobDevice.id,
+        DateTime.now().millisecondsSinceEpoch,
+        bobResult4.encryptedKeys,
+        base64.encode(bobResult4.ciphertext!),
+      ),
+    );
+
+    expect(aliceResult4.error, null);
+    expect(aliceResult4.payload, "That's okay.");
+  });
+
+  test('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
+    // all ratchets.
+    const aliceJid = 'alice@server1';
+    const bobJid = 'bob@server2';
+
+    final aliceDevice = await OmemoDevice.generateNewDevice(aliceJid, opkAmount: 1);
+    final bobDevice = await OmemoDevice.generateNewDevice(bobJid, opkAmount: 1);
+
+    EncryptionResult? aliceEmptyMessage;
+    final aliceManager = OmemoManager(
+      aliceDevice,
+      AlwaysTrustingTrustManager(),
+      (result, recipientJid) async {
+        aliceEmptyMessage = result;
+      },
+      (jid) async {
+        expect(jid, bobJid);
+
+        return [bobDevice.id];
+      },
+      (jid, id) async {
+        expect(jid, bobJid);
+
+        return bobDevice.toBundle();
+      },
+      (jid) async {},
+    );
+    final bobManager = OmemoManager(
+      bobDevice,
+      AlwaysTrustingTrustManager(),
+      (result, recipientJid) async {},
+      (jid) async => null,
+      (jid, id) async => null,
+      (jid) async {},
+    );
+
+    // Alice encrypts a message for Bob
+    final aliceResult1 = await aliceManager.onOutgoingStanza(
+      const OmemoOutgoingStanza(
+        [bobJid],
+        'Hello Bob!',
+      ),
+    );
+
+    // And Bob decrypts it
+    await bobManager.onIncomingStanza(
+      OmemoIncomingStanza(
+        aliceJid,
+        aliceDevice.id,
+        DateTime.now().millisecondsSinceEpoch,
+        aliceResult1.encryptedKeys,
+        base64.encode(aliceResult1.ciphertext!),
+      ),
+    );
+
+    // Ratchets are acked
+    await aliceManager.ratchetAcknowledged(
+      bobJid,
+      bobDevice.id,
+    );
+
+    // Alice now removes all ratchets for Bob and sends another new message
+    await aliceManager.removeAllRatchets(bobJid);
+
+    expect(aliceManager.getRatchet(RatchetMapKey(bobJid, bobDevice.id)), null);
+
+    // Alice prepares an empty OMEMO message
+    await aliceManager.sendOmemoHeartbeat(bobJid);
+
+    // And Bob receives it
+    final bobResult2 = await bobManager.onIncomingStanza(
+      OmemoIncomingStanza(
+        aliceJid,
+        await aliceManager.getDeviceId(),
+        DateTime.now().millisecondsSinceEpoch,
+        aliceEmptyMessage!.encryptedKeys,
+        null,
+      ),
+    );
+    expect(bobResult2.error, null);
+
+    // Alice sends another message
+    final aliceResult3 = await aliceManager.onOutgoingStanza(
+      const OmemoOutgoingStanza(
+        [bobJid],
+        'I did not trust your last device, Bob!',
+      ),
+    );
+
+    // Bob decrypts it
+    final bobResult3 = await bobManager.onIncomingStanza(
+      OmemoIncomingStanza(
+        aliceJid,
+        aliceDevice.id,
+        DateTime.now().millisecondsSinceEpoch,
+        aliceResult3.encryptedKeys,
+        base64.encode(aliceResult3.ciphertext!),
+      ),
+    );
+
+    expect(bobResult3.error, null);
+    expect(bobResult3.payload, 'I did not trust your last device, Bob!');
+
+    // Bob responds
+    final bobResult4 = await bobManager.onOutgoingStanza(
+      OmemoOutgoingStanza(
+        [aliceJid],
+        "That's okay.",
+      ),
+    );
+
+    // Alice decrypts
+    final aliceResult4 = await aliceManager.onIncomingStanza(
+      OmemoIncomingStanza(
+        bobJid,
+        bobDevice.id,
+        DateTime.now().millisecondsSinceEpoch,
+        bobResult4.encryptedKeys,
+        base64.encode(bobResult4.ciphertext!),
+      ),
+    );
+
+    expect(aliceResult4.error, null);
+    expect(aliceResult4.payload, "That's okay.");
+  });
+
+  test('Test resending key exchanges', () async {
+    const aliceJid = 'alice@server1';
+    const bobJid = 'bob@server2';
+
+    final aliceDevice = await OmemoDevice.generateNewDevice(aliceJid, opkAmount: 1);
+    final bobDevice = await OmemoDevice.generateNewDevice(bobJid, opkAmount: 1);
+
+    final aliceManager = OmemoManager(
+      aliceDevice,
+      AlwaysTrustingTrustManager(),
+      (result, recipientJid) async {},
+      (jid) async {
+        expect(jid, bobJid);
+        return [ bobDevice.id ];
+      },
+      (jid, id) async {
+        expect(jid, bobJid);
+        return bobDevice.toBundle();
+      },
+      (jid) async {},
+    );
+    final bobManager = OmemoManager(
+      bobDevice,
+      AlwaysTrustingTrustManager(),
+      (result, recipientJid) async {},
+      (jid) async {
+        expect(jid, aliceJid);
+        return [aliceDevice.id];
+      },
+      (jid, id) async {
+        expect(jid, aliceJid);
+        return aliceDevice.toBundle();
+      },
+      (jid) async {},
+    );
+    
+    // Alice sends Bob a message
+    final aliceResult1 = await aliceManager.onOutgoingStanza(
+      OmemoOutgoingStanza(
+        [bobJid],
+        'Hello World!',
+      ),
+    );
+
+    // The first message must be a KEX message
+    expect(aliceResult1.encryptedKeys.first.kex, true);
+
+    // Bob decrypts Alice's message
+    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 World!');
+
+    // Alice immediately sends another message 
+    final aliceResult2 = await aliceManager.onOutgoingStanza(
+      OmemoOutgoingStanza(
+        [bobJid],
+        'Hello Bob',
+      ),
+    );
+
+    // The response should contain a KEX
+    expect(aliceResult2.encryptedKeys.first.kex, true);
+
+    // The basic data should be the same
+    final parsedFirstKex = OmemoKeyExchange.fromBuffer(
+      base64.decode(aliceResult1.encryptedKeys.first.value),
+    );
+    final parsedSecondKex = OmemoKeyExchange.fromBuffer(
+      base64.decode(aliceResult2.encryptedKeys.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
     final bobResult2 = await bobManager.onIncomingStanza(
       OmemoIncomingStanza(
         aliceJid,
@@ -1022,7 +1296,58 @@ void main() {
         base64.encode(aliceResult2.ciphertext!),
       ),
     );
+    expect(bobResult2.error, null);
+    expect(bobResult2.payload, 'Hello Bob');
 
-    expect(bobResult2.payload, 'I did not trust your last device, Bob!');
+    // Bob also sends a message
+    final bobResult3 = await bobManager.onOutgoingStanza(
+      OmemoOutgoingStanza(
+        [aliceJid],
+        'Hello Alice!',
+      ),
+    );
+
+    // Alice decrypts it
+    final aliceResult3 = await aliceManager.onIncomingStanza(
+      OmemoIncomingStanza(
+        bobJid,
+        bobDevice.id,
+        DateTime.now().millisecondsSinceEpoch,
+        bobResult3.encryptedKeys,
+        base64.encode(bobResult3.ciphertext!),
+      ),
+    );
+    expect(aliceResult3.error, null);
+    expect(aliceResult3.payload, 'Hello Alice!'); 
+    
+    // Bob now acks the ratchet
+    await aliceManager.ratchetAcknowledged(
+      bobJid,
+      bobDevice.id,
+    );
+
+    // Alice replies
+    final aliceResult4 = await aliceManager.onOutgoingStanza(
+      OmemoOutgoingStanza(
+        [bobJid],
+        'Hi Bob',
+      ),
+    );
+
+    // The response should contain no KEX
+    expect(aliceResult4.encryptedKeys.first.kex, false);
+
+    // Bob decrypts it
+    final bobResult4 = await bobManager.onIncomingStanza(
+      OmemoIncomingStanza(
+        aliceJid,
+        aliceDevice.id,
+        DateTime.now().millisecondsSinceEpoch,
+        aliceResult4.encryptedKeys,
+        base64.encode(aliceResult4.ciphertext!),
+      ),
+    );
+    expect(bobResult4.error, null);
+    expect(bobResult4.payload, 'Hi Bob');
   });
 }