From ee7b09bdb0510a6d62254315765180a9a942e835 Mon Sep 17 00:00:00 2001
From: "Alexander \"PapaTutuWawa" <papatutuwawa@polynom.me>
Date: Thu, 18 Aug 2022 15:20:32 +0200
Subject: [PATCH] feat: Ratchets should overwrite each other

---
 lib/src/omemo/sessionmanager.dart | 23 +++++-----
 test/omemo_test.dart              | 71 +++++++++++++++++++++++++++++++
 2 files changed, 83 insertions(+), 11 deletions(-)

diff --git a/lib/src/omemo/sessionmanager.dart b/lib/src/omemo/sessionmanager.dart
index 8190caf..ca7825d 100644
--- a/lib/src/omemo/sessionmanager.dart
+++ b/lib/src/omemo/sessionmanager.dart
@@ -110,21 +110,22 @@ class OmemoSessionManager {
       // Add the bundle Id
       if (!_deviceMap.containsKey(jid)) {
         _deviceMap[jid] = [deviceId];
-      } else {
-        _deviceMap[jid]!.add(deviceId);
-      }
 
-      // Commit the device map
-      _eventStreamController.add(DeviceMapModifiedEvent(_deviceMap));
+        // Commit the device map
+        _eventStreamController.add(DeviceMapModifiedEvent(_deviceMap));
+      } else {
+        // Prevent having the same device multiple times in the list
+        if (!_deviceMap[jid]!.contains(deviceId)) {
+          _deviceMap[jid]!.add(deviceId);
+
+          // Commit the device map
+          _eventStreamController.add(DeviceMapModifiedEvent(_deviceMap));
+        }
+      }
 
       // Add the ratchet session
       final key = RatchetMapKey(jid, deviceId);
-      if (!_ratchetMap.containsKey(key)) {
-        _ratchetMap[key] = ratchet;
-      } else {
-        // TODO(PapaTutuWawa): What do we do now?
-        throw Exception();
-      }
+      _ratchetMap[key] = ratchet;
 
       // Commit the ratchet
       _eventStreamController.add(RatchetModifiedEvent(jid, deviceId, ratchet));
diff --git a/test/omemo_test.dart b/test/omemo_test.dart
index a9cd102..7d8e319 100644
--- a/test/omemo_test.dart
+++ b/test/omemo_test.dart
@@ -543,4 +543,75 @@ void main() {
       true,
     );
   });
+
+  test('Test overwriting sessions', () async {
+    const aliceJid = 'alice@server.example';
+    const bobJid = 'bob@other.server.example';
+    // Alice and Bob generate their sessions
+    final aliceSession = await OmemoSessionManager.generateNewIdentity(
+      aliceJid,
+      AlwaysTrustingTrustManager(),
+      opkAmount: 1,
+    );
+    final bobSession = await OmemoSessionManager.generateNewIdentity(
+      bobJid,
+      AlwaysTrustingTrustManager(),
+      opkAmount: 2,
+    );
+
+    // Alice sends Bob a message
+    final msg1 = await aliceSession.encryptToJid(
+      bobJid,
+      'Hallo Welt',
+      newSessions: [
+        await (await bobSession.getDevice()).toBundle(),
+      ],
+    );
+    await bobSession.decryptMessage(
+      msg1.ciphertext,
+      aliceJid,
+      (await aliceSession.getDevice()).id,
+      msg1.encryptedKeys,
+    );
+    final aliceRatchet1 = aliceSession.getRatchet(
+      bobJid,
+      (await bobSession.getDevice()).id,
+    );
+    final bobRatchet1 = bobSession.getRatchet(
+      aliceJid,
+      (await aliceSession.getDevice()).id,
+    );
+
+    // Alice is impatient and immediately sends another message before the original one
+    // can be acknowledged by Bob
+    final msg2 = await aliceSession.encryptToJid(
+      bobJid,
+      "Why don't you answer?",
+      newSessions: [
+        await (await bobSession.getDevice()).toBundle(),
+      ],
+    );
+    await bobSession.decryptMessage(
+      msg2.ciphertext,
+      aliceJid,
+      (await aliceSession.getDevice()).id,
+      msg2.encryptedKeys,
+    );
+    final aliceRatchet2 = aliceSession.getRatchet(
+      bobJid,
+      (await bobSession.getDevice()).id,
+    );
+    final bobRatchet2 = bobSession.getRatchet(
+      aliceJid,
+      (await aliceSession.getDevice()).id,
+    );
+
+    // Both should only have one ratchet
+    expect(aliceSession.getRatchetMap().length, 1);
+    expect(bobSession.getRatchetMap().length, 1);
+    
+    // The ratchets should both be different
+    expect(await aliceRatchet1.equals(aliceRatchet2), false);
+    expect(await bobRatchet1.equals(bobRatchet2), false);
+  });
 }