fix: Hopefully fix all tests being flaky

It seems that the varint encoding function would not work for
some integers as input. This should in theory fix this issue. Since
the SPK IDs are randomly between 0 and 2**32 - 1, it makes sense that
the tests fail only sometimes.

lib/src/omemo/ratchet_map_key.dart

@@ -4,9 +4,24 @@ import 'package:meta/meta.dart';
 class RatchetMapKey {
 
   const RatchetMapKey(this.jid, this.deviceId);
+
+  factory RatchetMapKey.fromJsonKey(String key) {
+    final parts = key.split(':');
+    final deviceId = int.parse(parts.first);
+
+    return RatchetMapKey(
+      parts.sublist(1).join(':'),
+      deviceId,
+    );
+  }
+
   final String jid;
   final int deviceId;
 
+  String toJsonKey() {
+    return '$deviceId:$jid';
+  }
+  
   @override
   bool operator ==(Object other) {
     return other is RatchetMapKey && jid == other.jid && deviceId == other.deviceId;

lib/src/omemo/sessionmanager.dart

@@ -177,16 +177,18 @@ class OmemoSessionManager {
     // Pick the correct SPK
     final device = await getDevice();
     OmemoKeyPair? spk;
-    if (kex.spkId == device.spkId) {
+
-      spk = device.spk;
+    await _lock.synchronized(() async {
-    } else if (kex.spkId == device.oldSpkId) {
+      if (kex.spkId == _device.spkId) {
-      spk = device.oldSpk;
+        spk = _device.spk;
-    } else {
+      } else if (kex.spkId == _device.oldSpkId) {
+        spk = _device.oldSpk;
+      }
+    });
+    if (spk == null) {
       throw UnknownSignedPrekeyException();
     }
     
-    assert(spk != null, 'The used SPK must be found');
-    
     final kexResult = await x3dhFromInitialMessage(
       X3DHMessage(
         OmemoPublicKey.fromBytes(kex.ik!, KeyPairType.ed25519),
@@ -198,7 +200,7 @@ class OmemoSessionManager {
       device.ik,
     );
     final ratchet = await OmemoDoubleRatchet.acceptNewSession(
-      spk,
+      spk!,
       OmemoPublicKey.fromBytes(kex.ik!, KeyPairType.ed25519),
       kexResult.sk,
       kexResult.ad,
@@ -297,6 +299,25 @@ class OmemoSessionManager {
     );
   }
 
+  /// In case a decryption error occurs, the Double Ratchet spec says to just restore
+  /// the ratchet to its old state. As such, this function restores the ratchet at
+  /// [mapKey] with [oldRatchet].
+  Future<void> _restoreRatchet(RatchetMapKey mapKey, OmemoDoubleRatchet oldRatchet) async {
+    await _lock.synchronized(() {
+      print('RESTORING RATCHETS');
+      _ratchetMap[mapKey] = oldRatchet;
+
+      // Commit the ratchet
+      _eventStreamController.add(
+        RatchetModifiedEvent(
+          mapKey.jid,
+          mapKey.deviceId,
+          oldRatchet,
+        ),
+      );
+    });
+  }
+  
   /// Attempt to decrypt [ciphertext]. [keys] refers to the <key /> elements inside the
   /// <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
@@ -314,9 +335,15 @@ class OmemoSessionManager {
       throw NotEncryptedForDeviceException();
     }
 
+    final ratchetKey = RatchetMapKey(senderJid, senderDeviceId);
     final decodedRawKey = base64.decode(rawKey.value);
     OmemoAuthenticatedMessage authMessage;
+    OmemoDoubleRatchet? oldRatchet;
     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 kex = OmemoKeyExchange.fromBuffer(decodedRawKey);
       await _addSessionFromKeyExchange(
@@ -347,31 +374,38 @@ class OmemoSessionManager {
     }
 
     final message = OmemoMessage.fromBuffer(authMessage.message!);
-    final ratchetKey = RatchetMapKey(senderJid, senderDeviceId);
     List<int>? keyAndHmac;
-    await _lock.synchronized(() async {
+    // We can guarantee that the ratchet exists at this point in time
-      final ratchet = _ratchetMap[ratchetKey]!;
+    final ratchet = (await _getRatchet(ratchetKey))!;
+    oldRatchet ??= ratchet ;
+
+    try {
       if (rawKey.kex) {
         keyAndHmac = await ratchet.ratchetDecrypt(message, authMessage.writeToBuffer());
       } else {
         keyAndHmac = await ratchet.ratchetDecrypt(message, decodedRawKey);
       }
+    } on InvalidMessageHMACException {
+      await _restoreRatchet(ratchetKey, oldRatchet);
+      rethrow;
+    }
 
-      // Commit the ratchet
+    // Commit the ratchet
-      _eventStreamController.add(RatchetModifiedEvent(senderJid, senderDeviceId, ratchet));
+    _eventStreamController.add(RatchetModifiedEvent(senderJid, senderDeviceId, ratchet));
-    });
 
     // Empty OMEMO messages should just have the key decrypted and/or session set up.
     if (ciphertext == null) {
       return null;
     }
     
-    final key = keyAndHmac!.sublist(0, 32);
+    final key = keyAndHmac.sublist(0, 32);
-    final hmac = keyAndHmac!.sublist(32, 48);
+    final hmac = keyAndHmac.sublist(32, 48);
     final derivedKeys = await deriveEncryptionKeys(key, omemoPayloadInfoString);
 
     final computedHmac = await truncatedHmac(ciphertext, derivedKeys.authenticationKey);
     if (!listsEqual(hmac, computedHmac)) {
+      // TODO(PapaTutuWawa): I am unsure if we should restore the ratchet here
+      await _restoreRatchet(ratchetKey, oldRatchet);
       throw InvalidMessageHMACException();
     }
     
@@ -441,6 +475,24 @@ class OmemoSessionManager {
     });
   }
 
+  /// Removes all ratchets for Jid [jid]. Triggers a DeviceMapModified event at the end and an
+  /// RatchetRemovedEvent for each ratchet.
+  Future<void> removeAllRatchets(String jid) async {
+    await _lock.synchronized(() async {
+      for (final deviceId in _deviceMap[jid]!) {
+        // Remove the ratchet
+        _ratchetMap.remove(RatchetMapKey(jid, deviceId));
+        // Commit it
+        _eventStreamController.add(RatchetRemovedEvent(jid, deviceId));
+      }
+
+      // Remove the device from jid
+      _deviceMap.remove(jid);
+      // Commit it
+      _eventStreamController.add(DeviceMapModifiedEvent(_deviceMap));
+    });
+  }
+  
   /// Returns the list of device identifiers belonging to [jid] that are yet unacked, i.e.
   /// we have not yet received an empty OMEMO message from.
   Future<List<int>?> getUnacknowledgedRatchets(String jid) async {
@@ -497,6 +549,12 @@ class OmemoSessionManager {
     });
   }
 
+  Future<OmemoDoubleRatchet?> _getRatchet(RatchetMapKey key) async {
+    return _lock.synchronized(() async {
+      return _ratchetMap[key];
+    });
+  }
+  
   @visibleForTesting
   OmemoDoubleRatchet getRatchet(String jid, int deviceId) => _ratchetMap[RatchetMapKey(jid, deviceId)]!;
 

lib/src/protobuf/protobuf.dart

@@ -46,21 +46,18 @@ List<int> encodeVarint(int i) {
   assert(i >= 0, "Two's complement is not implemented");
   final ret = List<int>.empty(growable: true);
 
-  var j = 0;
+  // Thanks to https://github.com/hathibelagal-dev/LEB128 for the trick with toRadixString!
-  while (true) {
+  final numSevenBlocks = (i.toRadixString(2).length / 7).ceil();
+  for (var j = 0; j < numSevenBlocks; j++) {
     // The 7 LSB of the byte we're creating
     final x = (i & (lsb7Mask << j * 7)) >> j * 7;
-    // The next bits
-    final next = i & (lsb7Mask << (j + 1) * 7);
 
-    if (next == 0) {
+    if (j == numSevenBlocks - 1) {
       // If we were to shift further, we only get zero, so we're at the end
       ret.add(x);
-      break;
     } else {
       // We still have at least one bit more to go, so set the MSB to 1
       ret.add(x + msb);
-      j++;
     }
   }
 

lib/src/trust/always.dart

@@ -17,4 +17,7 @@ class AlwaysTrustingTrustManager extends TrustManager {
 
   @override
   Future<void> setEnabled(String jid, int deviceId, bool enabled) async {}
+
+  @override
+  Future<Map<String, dynamic>> toJson() async => <String, dynamic>{};
 }

lib/src/trust/base.dart

@@ -16,4 +16,7 @@ abstract class TrustManager {
   /// Mark the device with id [deviceId] of Jid [jid] as enabled if [enabled] is true or as disabled
   /// if [enabled] is false.
   Future<void> setEnabled(String jid, int deviceId, bool enabled);
+
+  /// Serialize the trust manager to JSON.
+  Future<Map<String, dynamic>> toJson();
 }

lib/src/trust/btbv.dart

@@ -8,29 +8,52 @@ import 'package:synchronized/synchronized.dart';
 /// - blindTrust: The fingerprint is not verified using OOB means
 /// - verified: The fingerprint has been verified using OOB means
 enum BTBVTrustState {
-  notTrusted,
+  notTrusted, // = 1
-  blindTrust,
+  blindTrust, // = 2
-  verified,
+  verified,   // = 3
+}
+
+int _trustToInt(BTBVTrustState state) {
+  switch (state) {
+    case BTBVTrustState.notTrusted: return 1;
+    case BTBVTrustState.blindTrust: return 2;
+    case BTBVTrustState.verified:   return 3;
+  }
+}
+
+BTBVTrustState _trustFromInt(int i) {
+  switch (i) {
+    case 1: return BTBVTrustState.notTrusted;
+    case 2: return BTBVTrustState.blindTrust;
+    case 3: return BTBVTrustState.verified;
+    default: return BTBVTrustState.notTrusted;
+  }
 }
 
 /// A TrustManager that implements the idea of Blind Trust Before Verification.
 /// See https://gultsch.de/trust.html for more details.
 abstract class BlindTrustBeforeVerificationTrustManager extends TrustManager {
-  BlindTrustBeforeVerificationTrustManager()
+  BlindTrustBeforeVerificationTrustManager({
-    : trustCache = {},
+    Map<RatchetMapKey, BTBVTrustState>? trustCache,
-      enablementCache = {},
+    Map<RatchetMapKey, bool>? enablementCache,
-      devices = {},
+    Map<String, List<int>>? devices,
-      _lock = Lock();
+  }) : trustCache = trustCache ?? {},
+       enablementCache = enablementCache ?? {},
+       devices = devices ?? {},
+       _lock = Lock();
 
   /// The cache for mapping a RatchetMapKey to its trust state
+  @visibleForTesting
   @protected
   final Map<RatchetMapKey, BTBVTrustState> trustCache;
 
   /// The cache for mapping a RatchetMapKey to whether it is enabled or not
+  @visibleForTesting
   @protected
   final Map<RatchetMapKey, bool> enablementCache;
   
   /// Mapping of Jids to their device identifiers
+  @visibleForTesting
   @protected
   final Map<String, List<int>> devices;
 
@@ -142,6 +165,50 @@ abstract class BlindTrustBeforeVerificationTrustManager extends TrustManager {
     await commitState();
   }
 
+  @override
+  Future<Map<String, dynamic>> toJson() async {
+    return {
+      'devices': devices,
+      'trust': trustCache.map((key, value) => MapEntry(
+        key.toJsonKey(), _trustToInt(value),
+      ),),
+      'enable': enablementCache.map((key, value) => MapEntry(key.toJsonKey(), value)),
+    };
+  }
+
+  /// From a serialized version of a BTBV trust manager, extract the device list.
+  /// NOTE: This is needed as Dart cannot just cast a List<dynamic> to List<int> and so on.
+  static Map<String, List<int>> deviceListFromJson(Map<String, dynamic> json) {
+    return (json['devices']! as Map<String, dynamic>).map<String, List<int>>(
+      (key, value) => MapEntry(
+        key,
+        (value as List<dynamic>).map<int>((i) => i as int).toList(),
+      ),
+    );
+  }
+
+  /// From a serialized version of a BTBV trust manager, extract the trust cache.
+  /// NOTE: This is needed as Dart cannot just cast a List<dynamic> to List<int> and so on.
+  static Map<RatchetMapKey, BTBVTrustState> trustCacheFromJson(Map<String, dynamic> json) {
+    return (json['trust']! as Map<String, dynamic>).map<RatchetMapKey, BTBVTrustState>(
+      (key, value) => MapEntry(
+        RatchetMapKey.fromJsonKey(key),
+        _trustFromInt(value as int),
+      ),
+    );
+  }
+
+  /// From a serialized version of a BTBV trust manager, extract the enable cache.
+  /// NOTE: This is needed as Dart cannot just cast a List<dynamic> to List<int> and so on.
+  static Map<RatchetMapKey, bool> enableCacheFromJson(Map<String, dynamic> json) {
+    return (json['enable']! as Map<String, dynamic>).map<RatchetMapKey, bool>(
+      (key, value) => MapEntry(
+        RatchetMapKey.fromJsonKey(key),
+        value as bool,
+      ),
+    );
+  }
+  
   /// Called when the state of the trust manager has been changed. Allows the user to
   /// commit the trust state to persistent storage.
   @visibleForOverriding

lib/src/trust/never.dart

@@ -17,4 +17,7 @@ class NeverTrustingTrustManager extends TrustManager {
 
   @override
   Future<void> setEnabled(String jid, int deviceId, bool enabled) async {}
+
+  @override
+  Future<Map<String, dynamic>> toJson() async => <String, dynamic>{};
 }

test/omemo_test.dart

@@ -220,7 +220,7 @@ void main() {
     expect(messagePlaintext, aliceMessage2);
   });
 
-  test('Test using sending empty OMEMO messages', () async {
+  test('Test sending empty OMEMO messages', () async {
     const aliceJid = 'alice@server.example';
     const bobJid = 'bob@other.server.example';
       
@@ -614,4 +614,145 @@ void main() {
     expect(await aliceRatchet1.equals(aliceRatchet2), false);
     expect(await bobRatchet1.equals(bobRatchet2), false);
   });
+
+  test('Test receiving an old message that contains 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 Bob received is
+    // received again
+    try {
+      await bobSession.decryptMessage(
+        msg1.ciphertext,
+        aliceJid,
+        await aliceSession.getDeviceId(),
+        msg1.encryptedKeys,
+      );
+      expect(true, false);
+    } on InvalidMessageHMACException {
+      // NOOP
+    }
+
+    final msg3 = await aliceSession.encryptToJid(
+      bobJid,
+      'Are you okay?',
+    );
+    final result = await bobSession.decryptMessage(
+      msg3.ciphertext,
+      aliceJid,
+      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(
+        msg2.ciphertext,
+        bobJid,
+        await bobSession.getDeviceId(),
+        msg2.encryptedKeys,
+      );
+      expect(true, false);
+    } catch (_) {
+      // NOOP
+    }
+
+    final msg3 = await aliceSession.encryptToJid(
+      bobJid,
+      'Are you okay?',
+    );
+    final result = await bobSession.decryptMessage(
+      msg3.ciphertext,
+      aliceJid,
+      await aliceSession.getDeviceId(),
+      msg3.encryptedKeys,
+    );
+
+    expect(result, 'Are you okay?');
+  });
 }

test/protobuf_test.dart

@@ -48,6 +48,10 @@ void main() {
         <int>[172, 2],
       );
     });
+
+    test('Test some special cases', () {
+      expect(decodeVarint(encodeVarint(1042464893), 0).n, 1042464893);
+    });
   });
 
   group('OMEMOMessage', () {
@@ -170,5 +174,13 @@ void main() {
       expect(decoded.message!.mac, <int>[5, 6, 8, 0]);
       expect(decoded.message!.message, <int>[4, 5, 7, 3, 2]);
     });
+    test('Test decoding an issue', () {
+      /*
+      final data = 'CAAQfRogc2GwslU219dUkrMHNM4KdZRmuFnBTae+bQaJ+55IsAMiII7aZKj2sUpb6xR/3Ari7WZUmKFV0G6czUc4NMvjKDBaKnwKEM2ZpI8X3TgcxhxwENANnlsSaAgAEAAaICy8T9WPgLb7RdYd8/4JkrLF0RahEkC3ZaEfk5jw3dsLIkBMILzLyByweLgF4lCn0oNea+kbdrFr6rY7r/7WyI8hXEQz38QpnN+jyGGwC7Ga0dq70WuyqE7VpiFArQwqZh2G';
+      final kex = OmemoKeyExchange.fromBuffer(base64Decode(data));
+
+      expect(kex.spkId!, 1042464893);
+      */
+    });
   });
 }

test/serialisation_test.dart

@@ -1,3 +1,4 @@
+import 'dart:convert';
 import 'package:omemo_dart/omemo_dart.dart';
 import 'package:omemo_dart/src/trust/always.dart';
 import 'package:test/test.dart';
@@ -104,4 +105,84 @@ void main() {
       expect(await oldRatchet.equals(newRatchet), true);
     }
   });
+
+  test('Test serialising and deserialising the BlindTrustBeforeVerificationTrustManager', () async {
+    // Caroline's BTBV manager
+    final btbv = MemoryBTBVTrustManager();
+    // Example data
+    const aliceJid = 'alice@some.server';
+    const bobJid = 'bob@other.server';
+    
+    // Caroline starts a chat a device from Alice
+    await btbv.onNewSession(aliceJid, 1);
+    expect(await btbv.isTrusted(aliceJid, 1), true);
+    expect(await btbv.isEnabled(aliceJid, 1), true);
+
+    // Caroline meets with Alice and verifies her fingerprint
+    await btbv.setDeviceTrust(aliceJid, 1, BTBVTrustState.verified);
+    expect(await btbv.isTrusted(aliceJid, 1), true);
+
+    // Alice adds a new device
+    await btbv.onNewSession(aliceJid, 2);
+    expect(await btbv.isTrusted(aliceJid, 2), false);
+    expect(btbv.getDeviceTrust(aliceJid, 2), BTBVTrustState.notTrusted);
+    expect(await btbv.isEnabled(aliceJid, 2), false);
+
+    // Caronline starts a chat with Bob but since they live far apart, Caroline cannot
+    // verify his fingerprint.
+    await btbv.onNewSession(bobJid, 3);
+
+    // Bob adds a new device
+    await btbv.onNewSession(bobJid, 4);
+    expect(await btbv.isTrusted(bobJid, 3), true);
+    expect(await btbv.isTrusted(bobJid, 4), true);
+    expect(btbv.getDeviceTrust(bobJid, 3), BTBVTrustState.blindTrust);
+    expect(btbv.getDeviceTrust(bobJid, 4), BTBVTrustState.blindTrust);
+    expect(await btbv.isEnabled(bobJid, 3), true);
+    expect(await btbv.isEnabled(bobJid, 4), true);
+  });
+
+  test('Test serializing and deserializing RatchetMapKey', () {
+    const test1 = RatchetMapKey('user@example.org', 1234);
+    final result1 = RatchetMapKey.fromJsonKey(test1.toJsonKey());
+    expect(result1.jid, test1.jid);
+    expect(result1.deviceId, test1.deviceId);
+
+    const test2 = RatchetMapKey('user@example.org/hallo:welt', 3333);
+    final result2 = RatchetMapKey.fromJsonKey(test2.toJsonKey());
+    expect(result2.jid, test2.jid);
+    expect(result2.deviceId, test2.deviceId);
+  });
+
+  test('Test serializing and deserializing the components of the BTBV manager', () async {
+    // Caroline's BTBV manager
+    final btbv = MemoryBTBVTrustManager();
+    // Example data
+    const aliceJid = 'alice@some.server';
+    const bobJid = 'bob@other.server';
+    
+    await btbv.onNewSession(aliceJid, 1);
+    await btbv.setDeviceTrust(aliceJid, 1, BTBVTrustState.verified);
+    await btbv.onNewSession(aliceJid, 2);
+    await btbv.onNewSession(bobJid, 3);
+    await btbv.onNewSession(bobJid, 4);
+
+    final managerJson = await btbv.toJson();
+    final managerString = jsonEncode(managerJson);
+    final managerPostJson = jsonDecode(managerString) as Map<String, dynamic>;
+    final deviceList = BlindTrustBeforeVerificationTrustManager.deviceListFromJson(
+      managerPostJson,
+    );
+    expect(btbv.devices, deviceList);
+
+    final trustCache = BlindTrustBeforeVerificationTrustManager.trustCacheFromJson(
+      managerPostJson,
+    );
+    expect(btbv.trustCache, trustCache);
+
+    final enableCache = BlindTrustBeforeVerificationTrustManager.enableCacheFromJson(
+      managerPostJson,
+    );
+    expect(btbv.enablementCache, enableCache);
+  });
 }