fix: Fix ratchets going out of sync

analysis_options.yaml

@@ -9,8 +9,6 @@ linter:
 
 analyzer:
   exclude:
-    - "lib/protobuf/*.dart"
+    - "lib/src/protobuf/*.dart"
-    # TODO: Remove once OmemoSessionManager is gone
-    - "test/omemo_test.dart"
     - "example/omemo_dart_example.dart"
     - "test/serialisation_test.dart"

lib/src/double_ratchet/double_ratchet.dart

@@ -200,6 +200,7 @@ class OmemoDoubleRatchet {
     );
     rk = List.from(newRk1);
     ckr = List.from(newRk1);
+
     dhs = await OmemoKeyPair.generateNewPair(KeyPairType.x25519);
     final newRk2 = await kdfRk(
       rk,
@@ -226,6 +227,7 @@ class OmemoDoubleRatchet {
         final newCkr = await kdfCk(ckr!, kdfCkNextChainKey);
         final mk = await kdfCk(ckr!, kdfCkNextMessageKey);
         ckr = newCkr;
+
         mkSkipped[SkippedKey(dhr!, nr)] = mk;
         nr++;
       }
@@ -309,6 +311,7 @@ class OmemoDoubleRatchet {
     final ck = await kdfCk(ckr!, kdfCkNextChainKey);
     final mk = await kdfCk(ckr!, kdfCkNextMessageKey);
     ckr = ck;
+    nr++;
 
     return _decrypt(message, header.ciphertext, mk);
   }

lib/src/omemo/omemo.dart

@@ -118,6 +118,7 @@ class OmemoManager {
 
   /// Enter the critical section for performing cryptographic operations on the ratchets
   Future<void> _enterRatchetCriticalSection(String jid) async {
+    return;
     final completer = await _ratchetCriticalSectionLock.synchronized(() {
       if (_ratchetCriticalSectionQueue.containsKey(jid)) {
         final c = Completer<void>();
@@ -136,6 +137,7 @@ class OmemoManager {
 
   /// Leave the critical section for the ratchets.
   Future<void> _leaveRatchetCriticalSection(String jid) async {
+    return;
     await _ratchetCriticalSectionLock.synchronized(() {
       if (_ratchetCriticalSectionQueue.containsKey(jid)) {
         if (_ratchetCriticalSectionQueue[jid]!.isEmpty) {
@@ -229,6 +231,38 @@ class OmemoManager {
     return bundles;
   }
 
+  Future<void> _maybeSendEmptyMessage(RatchetMapKey key, bool created, bool replaced) async {
+    final ratchet = _ratchetMap[key]!;
+    if (ratchet.acknowledged) {
+      // The ratchet is acknowledged
+      _log.finest('Checking whether to heartbeat to ${key.jid}, ratchet.nr (${ratchet.nr}) >= 53: ${ratchet.nr >= 53}, created: $created, replaced: $replaced');
+      if (ratchet.nr >= 53 || created || replaced) {
+        await sendEmptyOmemoMessageImpl(
+          await _onOutgoingStanzaImpl(
+            OmemoOutgoingStanza(
+              [key.jid],
+              null,
+            ),
+          ),
+          key.jid,
+        );
+      }
+    } else {
+      // Ratchet is not acknowledged
+      _log.finest('Sending acknowledgement heartbeat to ${key.jid}');
+      await ratchetAcknowledged(key.jid, key.deviceId);
+      await sendEmptyOmemoMessageImpl(
+        await _onOutgoingStanzaImpl(
+          OmemoOutgoingStanza(
+            [key.jid],
+            null,
+          ),
+        ),
+        key.jid,
+      );
+    }
+  }
+
   /// 
   Future<DecryptionResult> onIncomingStanza(OmemoIncomingStanza stanza) async {
     // NOTE: We do this so that we cannot forget to acquire and free the critical
@@ -253,6 +287,7 @@ class OmemoManager {
 
     final ratchetKey = RatchetMapKey(stanza.bareSenderJid, stanza.senderDeviceId);
     if (key.kex) {
+      _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
@@ -328,6 +363,13 @@ class OmemoManager {
         stanza.senderDeviceId,
       );
 
+      // If we received an empty OMEMO message, mark the ratchet as acknowledged
+      if (result.get<String?>() == null) {
+        if (!ratchet.acknowledged) {
+          ratchet.acknowledged = true;
+        }
+      }
+
       // Commit the ratchet
       _ratchetMap[ratchetKey] = ratchet;
       _deviceList.appendOrCreate(stanza.bareSenderJid, stanza.senderDeviceId);
@@ -352,6 +394,10 @@ class OmemoManager {
         });
       }
 
+      // Send the hearbeat, if we have to
+      // TODO: Handle replace
+      await _maybeSendEmptyMessage(ratchetKey, true, false);
+
       return DecryptionResult(
         result.get<String?>(),
         null,
@@ -367,7 +413,8 @@ class OmemoManager {
         );
       }
 
-      final ratchet = _ratchetMap[key]!.clone();
+      _log.finest('Decoding message as OMEMOAuthenticatedMessage');
+      final ratchet = _ratchetMap[ratchetKey]!.clone();
       final authMessage = OMEMOAuthenticatedMessage.fromBuffer(base64Decode(key.value));
       final keyAndHmac = await ratchet.ratchetDecrypt(authMessage);
       if (keyAndHmac.isType<OmemoError>()) {
@@ -389,7 +436,15 @@ class OmemoManager {
         );
       }
 
+      // If we received an empty OMEMO message, mark the ratchet as acknowledged
+      if (result.get<String?>() == null) {
+        if (!ratchet.acknowledged) {
+          ratchet.acknowledged = true;
+        }
+      }
+
       // Message was successfully decrypted, so commit the ratchet
+      _ratchetMap[ratchetKey] = ratchet;
       _eventStreamController.add(
         RatchetModifiedEvent(
           stanza.bareSenderJid,
@@ -400,6 +455,9 @@ class OmemoManager {
         ),
       );
 
+      // Send a heartbeat, if required.
+      await _maybeSendEmptyMessage(ratchetKey, false, false);
+
       return DecryptionResult(
         result.get<String?>(),
         null,
@@ -541,21 +599,21 @@ class OmemoManager {
         }
 
         // Encrypt
-        final ratchet = _ratchetMap[ratchetKey]!.clone();
+        final ratchet = _ratchetMap[ratchetKey]!;
         final authMessage = await ratchet.ratchetEncrypt(payloadKey);
 
         // Package
         if (kex.containsKey(ratchetKey)) {
           final kexMessage = kex[ratchetKey]!..message = authMessage;
-        encryptedKeys.appendOrCreate(
+          encryptedKeys.appendOrCreate(
-          jid,
-          EncryptedKey(
             jid,
-            device,
+            EncryptedKey(
-            base64Encode(kexMessage.writeToBuffer()),
+              jid,
-            true,
+              device,
-          ),
+              base64Encode(kexMessage.writeToBuffer()),
-        );
+              true,
+            ),
+          );
         } else if (!ratchet.acknowledged) {
           // The ratchet as not yet been acked
           if (ratchet.kex == null) {
@@ -612,8 +670,16 @@ class OmemoManager {
     );
   }
 
-  // TODO
+  // Sends an empty OMEMO message (heartbeat) to [jid].
-  Future<void> sendOmemoHeartbeat(String jid) async {}
+  Future<void> sendOmemoHeartbeat(String jid) async {
+    final result = await onOutgoingStanza(
+      OmemoOutgoingStanza(
+        [jid],
+        null,
+      ),
+    );
+    await sendEmptyOmemoMessageImpl(result, jid);
+  }
 
   // TODO
   Future<void> removeAllRatchets(String jid) async {}
@@ -624,8 +690,23 @@ class OmemoManager {
   // TODO
   Future<void> onNewConnection() async {}
 
-  // TODO
+  // Mark the ratchet [jid]:[device] as acknowledged.
-  Future<void> ratchetAcknowledged(String jid, int device) async {}
+  Future<void> ratchetAcknowledged(String jid, int device) async {
+    await _enterRatchetCriticalSection(jid);
+
+    final ratchetKey = RatchetMapKey(jid, device);
+    if (!_ratchetMap.containsKey(ratchetKey)) {
+      _log.warning('Cannot mark $jid:$device as acknowledged as the ratchet does not exist');
+    } else {
+      // Commit
+      final ratchet = _ratchetMap[ratchetKey]!..acknowledged = true;
+      _eventStreamController.add(
+        RatchetModifiedEvent(jid, device, ratchet, false, false),
+      );
+    }
+
+    await _leaveRatchetCriticalSection(jid);
+  }
 
   // TODO
   Future<List<DeviceFingerprint>> getFingerprintsForJid(String jid) async => [];

test/omemo_test.dart

@@ -1,7 +1,6 @@
 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';
 
@@ -108,7 +107,7 @@ void main() {
         bobJid,
         bobDevice.id,
         DateTime.now().millisecondsSinceEpoch,
-        bobResult2.encryptedKeys[bobJid]!,
+        bobResult2.encryptedKeys[aliceJid]!,
         base64.encode(bobResult2.ciphertext!),
         false,
       ),
@@ -130,6 +129,7 @@ void main() {
         await OmemoDevice.generateNewDevice(aliceJid, opkAmount: 1);
     final bobDevice = await OmemoDevice.generateNewDevice(bobJid, opkAmount: 1);
 
+    EncryptionResult? bobEmptyMessage;
     final aliceManager = OmemoManager(
       aliceDevice,
       AlwaysTrustingTrustManager(),
@@ -151,6 +151,7 @@ void main() {
       AlwaysTrustingTrustManager(),
       (result, recipientJid) async {
         bobEmptyMessageSent++;
+        bobEmptyMessage = result;
       },
       (jid) async {
         expect(jid, aliceJid);
@@ -188,10 +189,20 @@ void main() {
     expect(bobResult.payload, 'Hello world');
 
     // Bob acknowledges the message
-    await aliceManager.ratchetAcknowledged(bobJid, bobDevice.id);
+    await aliceManager.onIncomingStanza(
+      OmemoIncomingStanza(
+        bobJid,
+        bobDevice.id,
+        getTimestamp(),
+        bobEmptyMessage!.encryptedKeys[aliceJid]!,
+        null,
+        false,
+      ),
+    );
 
     // Alice now sends 52 messages that Bob decrypts
-    for (var i = 0; i <= 51; i++) {
+    for (var i = 0; i < 52; i++) {
+      Logger.root.finest('${i+1}/52');
       final aliceResultLoop = await aliceManager.onOutgoingStanza(
         OmemoOutgoingStanza(
           [bobJid],
@@ -199,6 +210,8 @@ void main() {
         ),
       );
 
+      expect(aliceResultLoop.encryptedKeys[bobJid]!.first.kex, false);
+
       final bobResultLoop = await bobManager.onIncomingStanza(
         OmemoIncomingStanza(
           aliceJid,
@@ -210,6 +223,7 @@ void main() {
         ),
       );
 
+      expect(bobResultLoop.error, null);
       expect(aliceEmptyMessageSent, 0);
       expect(bobEmptyMessageSent, 1);
       expect(bobResultLoop.payload, 'Test message $i');
@@ -237,6 +251,42 @@ void main() {
     expect(aliceEmptyMessageSent, 0);
     expect(bobEmptyMessageSent, 2);
     expect(bobResultFinal.payload, 'Test message last');
+
+    // Alice receives it and sends another message
+    final aliceResultPostFinal = await aliceManager.onIncomingStanza(
+      OmemoIncomingStanza(
+        bobJid,
+        bobDevice.id,
+        getTimestamp(),
+        bobEmptyMessage!.encryptedKeys[aliceJid]!,
+        null,
+        false,
+      ),
+    );
+    expect(aliceResultPostFinal.error, null);
+    final aliceMessagePostFinal = await aliceManager.onOutgoingStanza(
+      const OmemoOutgoingStanza(
+        [bobJid],
+        "I'm not done yet!",
+      ),
+    );
+
+    // And Bob decrypts it
+    final bobResultPostFinal = await bobManager.onIncomingStanza(
+      OmemoIncomingStanza(
+        aliceJid,
+        aliceDevice.id,
+        getTimestamp(),
+        aliceMessagePostFinal.encryptedKeys[bobJid]!,
+        base64Encode(aliceMessagePostFinal.ciphertext!),
+        false,
+      ),
+    );
+
+    expect(bobResultPostFinal.error, null);
+    expect(bobResultPostFinal.payload, "I'm not done yet!");
+    expect(aliceEmptyMessageSent, 0);
+    expect(bobEmptyMessageSent, 2);
   });
 
   test('Test accessing data without it existing', () async {
@@ -770,7 +820,7 @@ void main() {
 
     // Alice has to reconnect but has no connection yet
     failure = true;
-    aliceManager.onNewConnection();
+    await aliceManager.onNewConnection();
 
     // Alice sends another message to Bob
     final aliceResult2 = await aliceManager.onOutgoingStanza(