fix: Use stanza receival timestamps to guard against stale kex messages

example/omemo_dart_example.dart

@@ -119,6 +119,8 @@ void main() async {
     aliceDevice.id,
     // The deserialised keys
     keys,
+    // Since the message was not delayed, we use the current time
+    DateTime.now().millisecondsSinceEpoch,
   );
 
   // All Bob has to do now is replace the OMEMO wrapper element 

lib/src/double_ratchet/double_ratchet.dart

@@ -68,6 +68,7 @@ class OmemoDoubleRatchet {
     this.sessionAd,
     this.mkSkipped, // MKSKIPPED
     this.acknowledged,
+    this.kexTimestamp,
   );
   
   factory OmemoDoubleRatchet.fromJson(Map<String, dynamic> data) {
@@ -85,6 +86,7 @@ class OmemoDoubleRatchet {
       'ik_pub': 'base/64/encoded',
       'session_ad': 'base/64/encoded',
       'acknowledged': true | false,
+      'kex_timestamp': int,
       'mkskipped': [
         {
           'key': 'base/64/encoded',
@@ -129,6 +131,7 @@ class OmemoDoubleRatchet {
       base64.decode(data['session_ad']! as String),
       mkSkipped,
       data['acknowledged']! as bool,
+      data['kex_timestamp']! as int,
     );
   }
   
@@ -160,14 +163,18 @@ class OmemoDoubleRatchet {
 
   final Map<SkippedKey, List<int>> mkSkipped;
 
+  /// The point in time at which we performed the kex exchange to create this ratchet.
+  /// Precision is milliseconds since epoch.
+  int kexTimestamp;
+
   /// Indicates whether we received an empty OMEMO message after building a session with
-  /// the device.
+  /// the device. 
   bool acknowledged;
 
   /// 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, int pn) async {
+  static Future<OmemoDoubleRatchet> initiateNewSession(OmemoPublicKey spk, OmemoPublicKey ik, List<int> sk, List<int> ad, int timestamp) async {
     final dhs = await OmemoKeyPair.generateNewPair(KeyPairType.x25519);
     final dhr = spk;
     final rk  = await kdfRk(sk, await omemoDH(dhs, dhr, 0));
@@ -181,11 +188,12 @@ class OmemoDoubleRatchet {
       null,
       0,
       0,
-      pn,
+      0,
       ik,
       ad,
       {},
       false,
+      timestamp,
     );
   }
 
@@ -193,7 +201,7 @@ class OmemoDoubleRatchet {
   /// Pre Key keypair [spk], the shared secret [sk] that was obtained through a X3DH and
   /// the associated data [ad] that was also obtained through a X3DH. [ik] refers to
   /// Alice's (the initiator's) IK public key.
-  static Future<OmemoDoubleRatchet> acceptNewSession(OmemoKeyPair spk, OmemoPublicKey ik, List<int> sk, List<int> ad) async {
+  static Future<OmemoDoubleRatchet> acceptNewSession(OmemoKeyPair spk, OmemoPublicKey ik, List<int> sk, List<int> ad, int kexTimestamp) async {
     return OmemoDoubleRatchet(
       spk,
       null,
@@ -207,6 +215,7 @@ class OmemoDoubleRatchet {
       ad,
       {},
       false,
+      kexTimestamp,
     );
   }
 
@@ -233,6 +242,7 @@ class OmemoDoubleRatchet {
       'session_ad': base64.encode(sessionAd),
       'mkskipped': mkSkippedSerialised,
       'acknowledged': acknowledged,
+      'kex_timestamp': kexTimestamp,
     };
   }
   
@@ -268,18 +278,18 @@ class OmemoDoubleRatchet {
   }
 
   Future<void> _dhRatchet(OmemoMessage header) async {
-    pn = header.n!;
+    pn = ns;
     ns = 0;
     nr = 0;
     dhr = OmemoPublicKey.fromBytes(header.dhPub!, KeyPairType.x25519);
 
     final newRk = await kdfRk(rk, await omemoDH(dhs, dhr!, 0));
-    rk = newRk;
-    ckr = newRk;
+    rk = List.from(newRk);
+    ckr = List.from(newRk);
     dhs = await OmemoKeyPair.generateNewPair(KeyPairType.x25519);
     final newNewRk = await kdfRk(rk, await omemoDH(dhs, dhr!, 0));
-    rk = newNewRk;
-    cks = newNewRk;
+    rk = List.from(newNewRk);
+    cks = List.from(newNewRk);
   }
 
   /// Encrypt [plaintext] using the Double Ratchet.
@@ -313,8 +323,8 @@ class OmemoDoubleRatchet {
     }
 
     final dhPubMatches = listsEqual(
-      header.dhPub ?? <int>[],
-      await dhr?.getBytes() ?? <int>[],
+      header.dhPub!,
+      (await dhr?.getBytes()) ?? <int>[],
     );
     if (!dhPubMatches) {
       await _skipMessageKeys(header.pn!);
@@ -348,6 +358,7 @@ class OmemoDoubleRatchet {
       sessionAd,
       Map<SkippedKey, List<int>>.from(mkSkipped),
       acknowledged,
+      kexTimestamp,
     );
   }
   
@@ -378,6 +389,7 @@ class OmemoDoubleRatchet {
       ns == other.ns &&
       nr == other.nr &&
       pn == other.pn &&
-      listsEqual(sessionAd, other.sessionAd);
+      listsEqual(sessionAd, other.sessionAd) &&
+      kexTimestamp == other.kexTimestamp;
   }
 }

lib/src/errors.dart

@@ -31,14 +31,11 @@ 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.
+/// Triggered by the Session Manager when the received Key Exchange message does not meet
+/// the requirement that a key exchange, given that the ratchet already exists, must be
+/// sent after its creation.
 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';
+  String errMsg() => 'The key exchange was sent before the last kex finished';
 }
 
 /// Triggered by the Session Manager when a message's sequence number is smaller than we

lib/src/helpers.dart

@@ -73,3 +73,7 @@ OmemoKeyPair? decodeKeyPairIfNotNull(String? pk, String? sk, KeyPairType type) {
     type,
   );
 }
+
+int getTimestamp() {
+  return DateTime.now().millisecondsSinceEpoch;
+}

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, int pn) async {
+  Future<OmemoKeyExchange> addSessionFromBundle(String jid, int deviceId, OmemoBundle bundle) async {
     final device = await getDevice();
     final kexResult = await x3dhFromBundle(
       bundle,
@@ -154,7 +154,7 @@ class OmemoSessionManager {
       bundle.ik,
       kexResult.sk,
       kexResult.ad,
-      pn,
+      getTimestamp(),
     );
 
     await _trustManager.onNewSession(jid, deviceId);
@@ -201,6 +201,7 @@ class OmemoSessionManager {
       OmemoPublicKey.fromBytes(kex.ik!, KeyPairType.ed25519),
       kexResult.sk,
       kexResult.ad,
+      getTimestamp(),
     );
 
     await _trustManager.onNewSession(jid, deviceId);
@@ -241,21 +242,10 @@ class OmemoSessionManager {
     final kex = <int, OmemoKeyExchange>{};
     if (newSessions != null) {
       for (final newSession in newSessions) {
-        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,
         );
       }
     }
@@ -334,12 +324,15 @@ class OmemoSessionManager {
   /// <keys /> element with a "jid" attribute matching our own. [senderJid] refers to the
   /// bare Jid of the sender. [senderDeviceId] refers to the "sid" attribute of the
   /// <encrypted /> element.
+  /// [timestamp] refers to the time the message was sent. This might be either what the
+  /// server tells you via "XEP-0203: Delayed Delivery" or the point in time at which
+  /// you received the stanza, if no Delayed Delivery element was found.
   ///
   /// If the received message is an empty OMEMO message, i.e. there is no <payload />
   /// element, then [ciphertext] must be set to null. In this case, this function
   /// will return null as there is no message to be decrypted. This, however, is used
   /// to set up sessions or advance the ratchets.
-  Future<String?> decryptMessage(List<int>? ciphertext, String senderJid, int senderDeviceId, List<EncryptedKey> keys) async {
+  Future<String?> decryptMessage(List<int>? ciphertext, String senderJid, int senderDeviceId, List<EncryptedKey> keys, int timestamp) async {
     // Try to find a session we can decrypt with.
     var device = await getDevice();
     final rawKey = keys.firstWhereOrNull((key) => key.rid == device.id);
@@ -363,8 +356,8 @@ class OmemoSessionManager {
       // 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!);
+        if (oldRatchet.kexTimestamp > timestamp) {
+          throw InvalidKeyExchangeException();
         }
       }
       
@@ -401,12 +394,6 @@ class OmemoSessionManager {
     final ratchet = (await _getRatchet(ratchetKey))!;
     oldRatchet ??= ratchet.clone();
 
-    if (!rawKey.kex) {
-      if (message.n! < ratchet.nr - 1) {
-        throw MessageAlreadyDecryptedException();
-      }
-    }
-    
     try {
       if (rawKey.kex) {
         keyAndHmac = await ratchet.ratchetDecrypt(message, authMessage.writeToBuffer());

test/double_ratchet_test.dart

@@ -91,6 +91,7 @@ void main() {
       ikAlice.pk,
       resultBob.sk,
       resultBob.ad,
+      0,
     );
 
     expect(alicesRatchet.sessionAd, bobsRatchet.sessionAd);

test/omemo_test.dart

@@ -90,6 +90,7 @@ void main() {
       aliceJid,
       await aliceSession.getDeviceId(),
       aliceMessage.encryptedKeys,
+      0,
     );
     expect(messagePlaintext, bobMessage);
     // The ratchet should be modified two times: Once for when the ratchet is created and
@@ -121,6 +122,7 @@ void main() {
       bobJid,
       await bobSession.getDeviceId(),
       bobResponseMessage.encryptedKeys,
+      0,
     );
     expect(bobResponseText, aliceReceivedMessage);
   });
@@ -170,6 +172,7 @@ void main() {
       aliceJid,
       await aliceSession.getDeviceId(),
       aliceMessage.encryptedKeys,
+      0,
     );
     expect(messagePlaintext, bobMessage);
 
@@ -189,6 +192,7 @@ void main() {
       bobJid,
       await bobSession.getDeviceId(),
       bobResponseMessage.encryptedKeys,
+      0,
     );
     expect(bobResponseText, aliceReceivedMessage);
 
@@ -245,6 +249,7 @@ void main() {
       aliceJid,
       await aliceSession1.getDeviceId(),
       aliceMessage.encryptedKeys,
+      0,
     );
     expect(messagePlaintext, bobMessage);
 
@@ -254,6 +259,7 @@ void main() {
       aliceJid,
       await aliceSession1.getDeviceId(),
       aliceMessage.encryptedKeys,
+      0,
     );
     expect(messagePlaintext, aliceMessage2);
   });
@@ -294,6 +300,7 @@ void main() {
       aliceJid,
       await aliceSession.getDeviceId(),
       aliceMessage.encryptedKeys,
+      0,
     );
     expect(bobMessage, null);
 
@@ -371,6 +378,7 @@ void main() {
       aliceJid,
       await aliceSession.getDeviceId(),
       aliceMessage.encryptedKeys,
+      0,
     );
     expect(messagePlaintext, bobMessage);
   });
@@ -437,6 +445,7 @@ void main() {
       aliceJid,
       await aliceSession.getDeviceId(),
       aliceMessage.encryptedKeys,
+      0,
     );
 
     for (var i = 0; i < 100; i++) {
@@ -456,6 +465,7 @@ void main() {
         bobJid,
         await bobSession.getDeviceId(),
         bobResponseMessage.encryptedKeys,
+        0,
       );
       expect(messageText, aliceReceivedMessage);
     }
@@ -610,6 +620,7 @@ void main() {
       aliceJid,
       await aliceSession.getDeviceId(),
       msg1.encryptedKeys,
+      0,
     );
     final aliceRatchet1 = aliceSession.getRatchet(
       bobJid,
@@ -634,6 +645,7 @@ void main() {
       aliceJid,
       await aliceSession.getDeviceId(),
       msg2.encryptedKeys,
+      getTimestamp(),
     );
     final aliceRatchet2 = aliceSession.getRatchet(
       bobJid,
@@ -669,6 +681,7 @@ void main() {
     );
 
     final bobsReceivedMessages = List<EncryptionResult>.empty(growable: true);
+    final bobsReceivedMessagesTimestamps = List<int>.empty(growable: true);
     
     // Alice sends Bob a message
     final msg1 = await aliceSession.encryptToJid(
@@ -679,11 +692,15 @@ void main() {
       ],
     );
     bobsReceivedMessages.add(msg1);
+    final t1 = getTimestamp();
+    bobsReceivedMessagesTimestamps.add(t1);
+
     await bobSession.decryptMessage(
       msg1.ciphertext,
       aliceJid,
       await aliceSession.getDeviceId(),
       msg1.encryptedKeys,
+      t1,
     );
 
     // Bob responds
@@ -691,13 +708,15 @@ void main() {
       aliceJid,
       'Hello!',
     );
+
     await aliceSession.decryptMessage(
       msg2.ciphertext,
       bobJid,
       await bobSession.getDeviceId(),
       msg2.encryptedKeys,
+      getTimestamp(),
     );
-
+    
     // Send some messages between the two
     for (var i = 0; i < 100; i++) {
       final msg = await aliceSession.encryptToJid(
@@ -705,11 +724,14 @@ void main() {
         'Hello $i',
       );
       bobsReceivedMessages.add(msg);
+      final t = getTimestamp();
+      bobsReceivedMessagesTimestamps.add(t);
       final result = await bobSession.decryptMessage(
         msg.ciphertext,
         aliceJid,
         await aliceSession.getDeviceId(),
         msg.encryptedKeys,
+        t,
       );
 
       expect(result, 'Hello $i');
@@ -720,20 +742,23 @@ void main() {
     final ratchetPreError = bobSession
       .getRatchet(aliceJid, await aliceSession.getDeviceId())
       .clone();
+    var invalidKex = 0;
     var errorCounter = 0;
-    for (final msg in bobsReceivedMessages) {
+    for (var i = 0; i < bobsReceivedMessages.length; i++) {
+      final msg = bobsReceivedMessages[i];
       try {
         await bobSession.decryptMessage(
           msg.ciphertext,
           aliceJid,
           await aliceSession.getDeviceId(),
           msg.encryptedKeys,
+          bobsReceivedMessagesTimestamps[i],
         );
         expect(true, false);
-      } on MessageAlreadyDecryptedException catch (_) {
+      } on InvalidMessageHMACException catch (_) {
         errorCounter++;
       } on InvalidKeyExchangeException catch (_) {
-        errorCounter++;
+        invalidKex++;
       }
     }
     final ratchetPostError = bobSession
@@ -741,7 +766,8 @@ void main() {
       .clone();
 
     // The 100 messages including the initial KEX message
-    expect(errorCounter, 101);
+    expect(invalidKex, 1);
+    expect(errorCounter, 100);
     expect(await ratchetPreError.equals(ratchetPostError), true);
 
     
@@ -754,6 +780,7 @@ void main() {
       aliceJid,
       await aliceSession.getDeviceId(),
       msg3.encryptedKeys,
+      104,
     );
 
     expect(result, 'Are you okay?');

test/serialisation_test.dart

@@ -62,6 +62,7 @@ void main() {
       aliceJid,
       await aliceSession.getDeviceId(),
       aliceMessage.encryptedKeys,
+      getTimestamp(),
     );
     final aliceOld = aliceSession.getRatchet(bobJid, await bobSession.getDeviceId());
     final aliceSerialised = jsonify(await aliceOld.toJson());
@@ -86,7 +87,6 @@ void main() {
       'bob@localhost',
       await bobSession.getDeviceId(),
       await bobSession.getDeviceBundle(),
-      0,
     );
 
     // Serialise and deserialise