From cc314756718654b1eaf6a206fbf88bbfa653b299 Mon Sep 17 00:00:00 2001
From: "Alexander \"PapaTutuWawa" <papatutuwawa@polynom.me>
Date: Sat, 7 Jan 2023 22:14:10 +0100
Subject: [PATCH] fix: Ratchets are not removed when using removeAllRatchets

---
 lib/src/omemo/omemomanager.dart | 11 +++--
 test/omemomanager_test.dart     | 83 +++++++++++++++++++++++++++++++++
 2 files changed, 89 insertions(+), 5 deletions(-)

diff --git a/lib/src/omemo/omemomanager.dart b/lib/src/omemo/omemomanager.dart
index 1f1816c..2912bef 100644
--- a/lib/src/omemo/omemomanager.dart
+++ b/lib/src/omemo/omemomanager.dart
@@ -277,7 +277,7 @@ class OmemoManager {
     if (rawKey.kex) {
       // If the ratchet already existed, we store it. If it didn't, oldRatchet will stay
       // null.
-      final oldRatchet = _getRatchet(ratchetKey)?.clone();
+      final oldRatchet = getRatchet(ratchetKey)?.clone();
       final kex = OmemoKeyExchange.fromBuffer(decodedRawKey);
       authMessage = kex.message!;
       message = OmemoMessage.fromBuffer(authMessage.message!);
@@ -341,7 +341,7 @@ class OmemoManager {
     }
 
     // We can guarantee that the ratchet exists at this point in time
-    final ratchet = _getRatchet(ratchetKey)!;
+    final ratchet = getRatchet(ratchetKey)!;
     oldRatchet ??= ratchet.clone();
 
     try {
@@ -378,7 +378,8 @@ class OmemoManager {
 
   /// Returns, if it exists, the ratchet associated with [key].
   /// NOTE: Must be called from within the ratchet critical section.
-  OmemoDoubleRatchet? _getRatchet(RatchetMapKey key) => _ratchetMap[key];
+  @visibleForTesting
+  OmemoDoubleRatchet? getRatchet(RatchetMapKey key) => _ratchetMap[key];
 
   /// Figure out what bundles we have to still build a session with.
   Future<List<OmemoBundle>> _fetchNewBundles(String jid) async {
@@ -594,7 +595,7 @@ class OmemoManager {
     }
 
     // Check if the ratchet is acked
-    final ratchet = _getRatchet(ratchetKey);
+    final ratchet = getRatchet(ratchetKey);
     assert(ratchet != null, 'We decrypted the message, so the ratchet must exist');
 
     if (ratchet!.acknowledged) {
@@ -762,7 +763,7 @@ class OmemoManager {
 
     for (final deviceId in _deviceList[jid]!) {
       // Remove the ratchet and commit it
-      _ratchetMap.remove(jid);
+      _ratchetMap.remove(RatchetMapKey(jid, deviceId));
       _eventStreamController.add(RatchetRemovedEvent(jid, deviceId));
     }
 
diff --git a/test/omemomanager_test.dart b/test/omemomanager_test.dart
index 448d6fa..61dccd1 100644
--- a/test/omemomanager_test.dart
+++ b/test/omemomanager_test.dart
@@ -942,4 +942,87 @@ void main() {
       expect(aliceReceivedMessage.payload, messageText);
     }
   });
+
+  test('Test removing all ratchets and sending a message', () 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 => 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);
+
+    final aliceResult2 = await aliceManager.onOutgoingStanza(
+      const OmemoOutgoingStanza(
+        [bobJid],
+        'I did not trust your last device, Bob!',
+      ),
+    );
+
+    // Bob decrypts it
+    final bobResult2 = await bobManager.onIncomingStanza(
+      OmemoIncomingStanza(
+        aliceJid,
+        aliceDevice.id,
+        DateTime.now().millisecondsSinceEpoch,
+        aliceResult2.encryptedKeys,
+        base64.encode(aliceResult2.ciphertext!),
+      ),
+    );
+
+    expect(bobResult2.payload, 'I did not trust your last device, Bob!');
+  });
 }