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

@@ -69,7 +69,7 @@ class OmemoDoubleRatchet {
     this.mkSkipped, // MKSKIPPED
     this.acknowledged,
   );
-
+  
   factory OmemoDoubleRatchet.fromJson(Map<String, dynamic> data) {
     /*
     {
@@ -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,18 +330,45 @@ 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 {
-    // ignore: invalid_use_of_visible_for_testing_member
-    final dhrMatch = dhr == null ? other.dhr == null : await dhr!.equals(other.dhr!);
-    final ckrMatch = ckr == null ? other.ckr == null : listsEqual(ckr!, other.ckr!);
-    final cksMatch = cks == null ? other.cks == null : listsEqual(cks!, other.cks!);
+    final dhrMatch = dhr == null ?
+      other.dhr == null :
+      // ignore: invalid_use_of_visible_for_testing_member
+      other.dhr != null && await dhr!.equals(other.dhr!);
+    final ckrMatch = ckr == null ?
+      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);
     // ignore: invalid_use_of_visible_for_testing_member
     final ikMatch = await ik.equals(other.ik);
-    
+
     return dhsMatch &&
       ikMatch &&
       dhrMatch &&

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);
-
-      // TODO(PapaTutuWawa): Only do this when we should
+      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
       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) {