From 6a9e98665dd647cb045d4c1749f5adee16274d9a Mon Sep 17 00:00:00 2001
From: "Alexander \"PapaTutuWawa" <papatutuwawa@polynom.me>
Date: Tue, 27 Dec 2022 12:47:45 +0100
Subject: [PATCH] fix: Track new KEX entries using a ratchet key

---
 lib/src/omemo/omemomanager.dart | 12 ++++++------
 1 file changed, 6 insertions(+), 6 deletions(-)

diff --git a/lib/src/omemo/omemomanager.dart b/lib/src/omemo/omemomanager.dart
index 788977a..204b019 100644
--- a/lib/src/omemo/omemomanager.dart
+++ b/lib/src/omemo/omemomanager.dart
@@ -435,10 +435,10 @@ class OmemoManager {
       keyPayload = List<int>.filled(32, 0x0);
     }
 
-    final kex = <int, OmemoKeyExchange>{};
+    final kex = <RatchetMapKey, OmemoKeyExchange>{};
     for (final jid in jids) {
       for (final newSession in await _fetchNewBundles(jid)) {
-        kex[newSession.id] = await addSessionFromBundle(
+        kex[RatchetMapKey(jid, newSession.id)] = await addSessionFromBundle(
           newSession.jid,
           newSession.id,
           newSession,
@@ -477,9 +477,9 @@ class OmemoManager {
 
         final ciphertext = (await ratchet.ratchetEncrypt(keyPayload)).ciphertext;
  
-        if (kex.isNotEmpty && kex.containsKey(deviceId)) {
+        if (kex.containsKey(ratchetKey)) {
           // The ratchet did not exist
-          final k = kex[deviceId]!
+          final k = kex[ratchetKey]!
             ..message = OmemoAuthenticatedMessage.fromBuffer(ciphertext);
           final buffer = base64.encode(k.writeToBuffer());
           encryptedKeys.add(
@@ -497,7 +497,7 @@ class OmemoManager {
           // The ratchet exists but is not acked
           if (ratchet.kex != null) {
             final oldKex = OmemoKeyExchange.fromBuffer(base64.decode(ratchet.kex!))
-            ..message = OmemoAuthenticatedMessage.fromBuffer(ciphertext);
+              ..message = OmemoAuthenticatedMessage.fromBuffer(ciphertext);
             
             encryptedKeys.add(
               EncryptedKey(
@@ -656,7 +656,7 @@ class OmemoManager {
   Future<Device> getDevice() => _deviceLock.synchronized(() => _device);
 
   /// Returns the id of the device used for encryption and decryption.
-  Future<int> getDeviceId() => (await getDevice()).id;
+  Future<int> getDeviceId() async => (await getDevice()).id;
 
   /// Returns the fingerprints for all devices of [jid] that we have a session with.
   /// If there are not sessions with [jid], then returns null.