fix: Fix receiving an old key exchange breaking decryption

This was mostly caused by Dart not copying values but referencing
them. AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA.

We know make some assumptions about received key exchanges, so this
needs some field testing.

lib/src/double_ratchet/double_ratchet.dart

@@ -167,7 +167,7 @@ class OmemoDoubleRatchet {
   /// Create an OMEMO session using the Signed Pre Key [spk], the shared secret [sk] that
   /// was obtained using a X3DH and the associated data [ad] that was also obtained through
   /// a X3DH. [ik] refers to Bob's (the receiver's) IK public key.
-  static Future<OmemoDoubleRatchet> initiateNewSession(OmemoPublicKey spk, OmemoPublicKey ik, List<int> sk, List<int> ad) async {
+  static Future<OmemoDoubleRatchet> initiateNewSession(OmemoPublicKey spk, OmemoPublicKey ik, List<int> sk, List<int> ad, int pn) async {
     final dhs = await OmemoKeyPair.generateNewPair(KeyPairType.x25519);
     final dhr = spk;
     final rk  = await kdfRk(sk, await omemoDH(dhs, dhr, 0));
@@ -181,7 +181,7 @@ class OmemoDoubleRatchet {
       null,
       0,
       0,
-      0,
+      pn,
       ik,
       ad,
       {},
@@ -330,12 +330,39 @@ class OmemoDoubleRatchet {
     return decrypt(mk, ciphertext, concat([sessionAd, header.writeToBuffer()]), sessionAd);
   }
 
+  OmemoDoubleRatchet clone() {
+    return OmemoDoubleRatchet(
+      dhs,
+      dhr,
+      rk,
+      cks != null ?
+        List<int>.from(cks!) :
+        null,
+      ckr != null ?
+        List<int>.from(ckr!) :
+        null,
+      ns,
+      nr,
+      pn,
+      ik,
+      sessionAd,
+      Map<SkippedKey, List<int>>.from(mkSkipped),
+      acknowledged,
+    );
+  }
+  
   @visibleForTesting
   Future<bool> equals(OmemoDoubleRatchet other) async {
+    final dhrMatch = dhr == null ?
+      other.dhr == null :
       // ignore: invalid_use_of_visible_for_testing_member
-    final dhrMatch = dhr == null ? other.dhr == null : await dhr!.equals(other.dhr!);
+      other.dhr != null && await dhr!.equals(other.dhr!);
-    final ckrMatch = ckr == null ? other.ckr == null : listsEqual(ckr!, other.ckr!);
+    final ckrMatch = ckr == null ?
-    final cksMatch = cks == null ? other.cks == null : listsEqual(cks!, other.cks!);
+      other.ckr == null :
+      other.ckr != null && listsEqual(ckr!, other.ckr!);
+    final cksMatch = cks == null ?
+      other.cks == null :
+      other.cks != null && listsEqual(cks!, other.cks!);
  
     // ignore: invalid_use_of_visible_for_testing_member
     final dhsMatch = await dhs.equals(other.dhs);

lib/src/errors.dart

@@ -30,3 +30,13 @@ class NoDecryptionKeyException implements Exception {
 class UnknownSignedPrekeyException implements Exception {
   String errMsg() => 'Unknown Signed Prekey used.';
 }
+
+/// Triggered by the Session Manager when the received Key Exchange message does not
+/// meet our expectations. This happens when the PN attribute of the message is not equal
+/// to our receive number.
+class InvalidKeyExchangeException implements Exception {
+  const InvalidKeyExchangeException(this.expectedPn, this.actualPn);
+  final int expectedPn;
+  final int actualPn;
+  String errMsg() => 'The pn attribute of the key exchange is invalid. Expected $expectedPn, got $actualPn';
+}

lib/src/omemo/sessionmanager.dart

@@ -143,7 +143,7 @@ class OmemoSessionManager {
   /// Create a ratchet session initiated by Alice to the user with Jid [jid] and the device
   /// [deviceId] from the bundle [bundle].
   @visibleForTesting
-  Future<OmemoKeyExchange> addSessionFromBundle(String jid, int deviceId, OmemoBundle bundle) async {
+  Future<OmemoKeyExchange> addSessionFromBundle(String jid, int deviceId, OmemoBundle bundle, int pn) async {
     final device = await getDevice();
     final kexResult = await x3dhFromBundle(
       bundle,
@@ -154,6 +154,7 @@ class OmemoSessionManager {
       bundle.ik,
       kexResult.sk,
       kexResult.ad,
+      pn,
     );
 
     await _trustManager.onNewSession(jid, deviceId);
@@ -240,7 +241,22 @@ class OmemoSessionManager {
     final kex = <int, OmemoKeyExchange>{};
     if (newSessions != null) {
       for (final newSession in newSessions) {
-        kex[newSession.id] = await addSessionFromBundle(newSession.jid, newSession.id, newSession);
+        final session = await _getRatchet(
+          RatchetMapKey(
+            newSession.jid,
+            newSession.id,
+          ),
+        );
+
+        final pn = session != null ?
+          session.ns :
+          0;
+        kex[newSession.id] = await addSessionFromBundle(
+          newSession.jid,
+          newSession.id,
+          newSession,
+          pn,
+        );
       }
     }
     
@@ -300,7 +316,7 @@ class OmemoSessionManager {
   /// [mapKey] with [oldRatchet].
   Future<void> _restoreRatchet(RatchetMapKey mapKey, OmemoDoubleRatchet oldRatchet) async {
     await _lock.synchronized(() {
-      _log.finest('Restoring ratchet ${mapKey.jid}:${mapKey.deviceId}');
+      _log.finest('Restoring ratchet ${mapKey.jid}:${mapKey.deviceId} to ${oldRatchet.nr}');
       _ratchetMap[mapKey] = oldRatchet;
 
       // Commit the ratchet
@@ -335,22 +351,32 @@ class OmemoSessionManager {
     final decodedRawKey = base64.decode(rawKey.value);
     OmemoAuthenticatedMessage authMessage;
     OmemoDoubleRatchet? oldRatchet;
+    OmemoMessage? message;
     if (rawKey.kex) {
       // If the ratchet already existed, we store it. If it didn't, oldRatchet will stay
       // null.
-      oldRatchet = await _getRatchet(ratchetKey);
+      final oldRatchet = (await _getRatchet(ratchetKey))?.clone();
+      final kex = OmemoKeyExchange.fromBuffer(decodedRawKey);
+      authMessage = kex.message!;
+      message = OmemoMessage.fromBuffer(authMessage.message!);
+
+      // Guard against old key exchanges
+      if (oldRatchet != null) {
+        _log.finest('KEX for existent ratchet. ${oldRatchet.pn}');
+        if (message.pn != oldRatchet.nr) {
+          throw InvalidKeyExchangeException(oldRatchet.nr, message.pn!);
+        }
+      }
       
       // TODO(PapaTutuWawa): Only do this when we should
-      final kex = OmemoKeyExchange.fromBuffer(decodedRawKey);
       await _addSessionFromKeyExchange(
         senderJid,
         senderDeviceId,
         kex,
       );
 
-      authMessage = kex.message!;
-
       // 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!);
 
@@ -359,6 +385,7 @@ class OmemoSessionManager {
       });
     } else {
       authMessage = OmemoAuthenticatedMessage.fromBuffer(decodedRawKey);
+      message = OmemoMessage.fromBuffer(authMessage.message!);
     }
     
     final devices = _deviceMap[senderJid];
@@ -369,11 +396,10 @@ class OmemoSessionManager {
       throw NoDecryptionKeyException();
     }
 
-    final message = OmemoMessage.fromBuffer(authMessage.message!);
     List<int>? keyAndHmac;
     // We can guarantee that the ratchet exists at this point in time
     final ratchet = (await _getRatchet(ratchetKey))!;
-    oldRatchet ??= ratchet ;
+    oldRatchet ??= ratchet.clone();
 
     try {
       if (rawKey.kex) {

test/double_ratchet_test.dart

@@ -84,6 +84,7 @@ void main() {
       ikBob.pk,
       resultAlice.sk,
       resultAlice.ad,
+      0,
     );
     final bobsRatchet = await OmemoDoubleRatchet.acceptNewSession(
       spkBob,

test/omemo_test.dart

@@ -653,7 +653,7 @@ void main() {
     expect(await bobRatchet1.equals(bobRatchet2), false);
   });
 
-  test('Test receiving an old message that contains a KEX', () async {
+  test('Test receiving old messages including a KEX', () async {
     const aliceJid = 'alice@server.example';
     const bobJid = 'bob@other.server.example';
     // Alice and Bob generate their sessions
@@ -668,6 +668,8 @@ void main() {
       opkAmount: 2,
     );
 
+    final bobsReceivedMessages = List<EncryptionResult>.empty(growable: true);
+    
     // Alice sends Bob a message
     final msg1 = await aliceSession.encryptToJid(
       bobJid,
@@ -676,7 +678,7 @@ void main() {
         await bobSession.getDeviceBundle(),
       ],
     );
-
+    bobsReceivedMessages.add(msg1);
     await bobSession.decryptMessage(
       msg1.ciphertext,
       aliceJid,
@@ -696,89 +698,50 @@ void main() {
       msg2.encryptedKeys,
     );
 
-    // Due to some issue with the transport protocol, the first message Bob received is
+    // Send some messages between the two
-    // received again
+    for (var i = 0; i < 100; i++) {
-    try {
+      final msg = await aliceSession.encryptToJid(
-      await bobSession.decryptMessage(
+        bobJid,
-        msg1.ciphertext,
+        'Hello $i',
+      );
+      bobsReceivedMessages.add(msg);
+      final result = await bobSession.decryptMessage(
+        msg.ciphertext,
         aliceJid,
         await aliceSession.getDeviceId(),
-        msg1.encryptedKeys,
+        msg.encryptedKeys,
       );
-      expect(true, false);
+
-    } on InvalidMessageHMACException {
+      expect(result, 'Hello $i');
-      // NOOP
     }
 
-    final msg3 = await aliceSession.encryptToJid(
+    // Due to some issue with the transport protocol, the messages to Bob are received
-      bobJid,
+    // again.
-      'Are you okay?',
+    final ratchetPreError = bobSession
-    );
+      .getRatchet(aliceJid, await aliceSession.getDeviceId())
-    final result = await bobSession.decryptMessage(
+      .clone();
-      msg3.ciphertext,
+    var errorCounter = 0;
-      aliceJid,
+    for (final msg in bobsReceivedMessages) {
-      await aliceSession.getDeviceId(),
-      msg3.encryptedKeys,
-    );
-
-    expect(result, 'Are you okay?');
-  });
-
-  test('Test receiving an old message that does not contain a KEX', () 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 bobSession.getDeviceBundle(),
-      ],
-    );
-    await bobSession.decryptMessage(
-      msg1.ciphertext,
-      aliceJid,
-      await aliceSession.getDeviceId(),
-      msg1.encryptedKeys,
-    );
-
-    // Bob responds
-    final msg2 = await bobSession.encryptToJid(
-      aliceJid,
-      'Hello!',
-    );
-    await aliceSession.decryptMessage(
-      msg2.ciphertext,
-      bobJid,
-      await bobSession.getDeviceId(),
-      msg2.encryptedKeys,
-    );
-
-    // Due to some issue with the transport protocol, the first message Alice received is
-    // received again.
       try {
-      await aliceSession.decryptMessage(
+        await bobSession.decryptMessage(
-        msg2.ciphertext,
+          msg.ciphertext,
-        bobJid,
+          aliceJid,
-        await bobSession.getDeviceId(),
+          await aliceSession.getDeviceId(),
-        msg2.encryptedKeys,
+          msg.encryptedKeys,
         );
         expect(true, false);
       } catch (_) {
-      // NOOP
+        errorCounter++;
       }
+    }
+    final ratchetPostError = bobSession
+      .getRatchet(aliceJid, await aliceSession.getDeviceId())
+      .clone();
+
+    // The 100 messages including the initial KEX message
+    expect(errorCounter, 101);
+    expect(await ratchetPreError.equals(ratchetPostError), true);
+
     
     final msg3 = await aliceSession.encryptToJid(
       bobJid,

test/serialisation_test.dart

@@ -86,6 +86,7 @@ void main() {
       'bob@localhost',
       await bobSession.getDeviceId(),
       await bobSession.getDeviceBundle(),
+      0,
     );
 
     // Serialise and deserialise