fix: Track device Ids as integers

lib/src/omemo/bundle.dart

@@ -12,16 +12,16 @@ class OmemoBundle {
     this.ikEncoded,
     this.opksEncoded,
   );
-  final String id;
+  final int id;
   /// The SPK but base64 encoded
   final String spkEncoded;
-  final String spkId;
+  final int spkId;
   /// The SPK signature but base64 encoded
   final String spkSignatureEncoded;
   /// The IK but base64 encoded
   final String ikEncoded;
   /// The mapping of a OPK's id to the base64 encoded data
-  final Map<String, String> opksEncoded;
+  final Map<int, String> opksEncoded;
 
   OmemoPublicKey get spk {
     final data = base64Decode(spkEncoded);
@@ -33,7 +33,7 @@ class OmemoBundle {
     return OmemoPublicKey.fromBytes(data, KeyPairType.ed25519);
   }
 
-  OmemoPublicKey getOpk(String id) {
+  OmemoPublicKey getOpk(int id) {
     final data = base64Decode(opksEncoded[id]!);
     return OmemoPublicKey.fromBytes(data, KeyPairType.x25519);
   }

lib/src/omemo/device.dart

@@ -20,16 +20,16 @@ class Device {
     final spkId = generateRandom32BitNumber();
     final signature = await sig(ik, await spk.pk.getBytes());
 
-    final opks = <String, OmemoKeyPair>{};
+    final opks = <int, OmemoKeyPair>{};
     for (var i = 0; i < opkAmount; i++) {
-      opks[i.toString()] = await OmemoKeyPair.generateNewPair(KeyPairType.x25519);
+      opks[i] = await OmemoKeyPair.generateNewPair(KeyPairType.x25519);
     }
 
-    return Device(id.toString(), ik, spk, spkId.toString(), signature, opks);
+    return Device(id, ik, spk, spkId, signature, opks);
   }
   
   /// The device Id
-  final String id;
+  final int id;
 
   /// The identity key
   final OmemoKeyPair ik;
@@ -37,16 +37,16 @@ class Device {
   /// The signed prekey...
   final OmemoKeyPair spk;
   /// ...its Id, ...
-  final String spkId;
+  final int spkId;
   /// ...and its signature
   final List<int> spkSignature;
 
   /// Map of an id to the associated Onetime-Prekey
-  final Map<String, OmemoKeyPair> opks;
+  final Map<int, OmemoKeyPair> opks;
 
   /// This replaces the Onetime-Prekey with id [id] with a completely new one. Returns
   /// a new Device object that copies over everything but replaces said key.
-  Future<Device> replaceOnetimePrekey(String id) async {
+  Future<Device> replaceOnetimePrekey(int id) async {
     final newOpk = await OmemoKeyPair.generateNewPair(KeyPairType.x25519);
 
     return Device(
@@ -76,14 +76,14 @@ class Device {
       id,
       ik,
       newSpk,
-      newSpkId.toString(),
+      newSpkId,
       newSignature,
       opks,
     );
   }
 
   Future<OmemoBundle> toBundle() async {
-    final encodedOpks = <String, String>{};
+    final encodedOpks = <int, String>{};
 
     for (final opkKey in opks.keys) {
       encodedOpks[opkKey] = base64.encode(await opks[opkKey]!.pk.getBytes());

lib/src/omemo/sessionmanager.dart

@@ -1,10 +1,12 @@
 import 'dart:convert';
 import 'package:collection/collection.dart';
+import 'package:cryptography/cryptography.dart';
 import 'package:omemo_dart/protobuf/schema.pb.dart';
 import 'package:omemo_dart/src/crypto.dart';
 import 'package:omemo_dart/src/double_ratchet/double_ratchet.dart';
 import 'package:omemo_dart/src/errors.dart';
 import 'package:omemo_dart/src/helpers.dart';
+import 'package:omemo_dart/src/keys.dart';
 import 'package:omemo_dart/src/omemo/bundle.dart';
 import 'package:omemo_dart/src/omemo/device.dart';
 import 'package:omemo_dart/src/x3dh/x3dh.dart';
@@ -22,13 +24,13 @@ class EncryptionResult {
 
   /// Mapping of the device Id to the key for decrypting ciphertext, encrypted
   /// for the ratchet with said device Id
-  final Map<String, List<int>> encryptedKeys;
+  final Map<int, List<int>> encryptedKeys;
 }
 
 class EncryptedKey {
 
   const EncryptedKey(this.rid, this.value);
-  final String rid;
+  final int rid;
   final String value;
 }
 
@@ -48,16 +50,16 @@ class OmemoSessionManager {
   final Lock _lock;
   
   /// Mapping of the Device Id to its OMEMO session
-  final Map<String, OmemoDoubleRatchet> _ratchetMap;
+  final Map<int, OmemoDoubleRatchet> _ratchetMap;
 
   /// Mapping of a bare Jid to its Device Ids
-  final Map<String, List<String>> _deviceMap;
+  final Map<String, List<int>> _deviceMap;
 
   /// Our own keys
   Device device;
 
   /// Add a session [ratchet] with the [deviceId] to the internal tracking state.
-  Future<void> addSession(String jid, String deviceId, OmemoDoubleRatchet ratchet) async {
+  Future<void> addSession(String jid, int deviceId, OmemoDoubleRatchet ratchet) async {
     await _lock.synchronized(() async {
       // Add the bundle Id
       if (!_deviceMap.containsKey(jid)) {
@@ -78,7 +80,7 @@ class OmemoSessionManager {
 
   /// Create a ratchet session initiated by Alice to the user with Jid [jid] and the device
   /// [deviceId] from the bundle [bundle].
-  Future<X3DHAliceResult> addSessionFromBundle(String jid, String deviceId, OmemoBundle bundle) async {
+  Future<OMEMOKeyExchange> addSessionFromBundle(String jid, int deviceId, OmemoBundle bundle) async {
     final kexResult = await x3dhFromBundle(
       bundle,
       device.ik,
@@ -91,19 +93,26 @@ class OmemoSessionManager {
 
     await addSession(jid, deviceId, ratchet);
 
-    return kexResult;
+    return OMEMOKeyExchange()
+      ..pkId = kexResult.opkId
+      // TODO(PapaTutuWawa): Fix
+      ..spkId = 0
+      ..ik = await device.ik.pk.getBytes()
+      ..ek = await kexResult.ek.pk.getBytes();
   }
 
   /// Build a new session with the user at [jid] with the device [deviceId] using data
   /// from the key exchange [kex].
-  // TODO(PapaTutuWawa): Use OMEMOKeyExchange
   // TODO(PapaTutuWawa): Replace the OPK
-  Future<void> addSessionFromKeyExchange(String jid, String deviceId, X3DHMessage kex) async {
+  Future<void> addSessionFromKeyExchange(String jid, int deviceId, OMEMOKeyExchange kex) async {
     final kexResult = await x3dhFromInitialMessage(
-      kex,
+      X3DHMessage(
+        OmemoPublicKey.fromBytes(kex.ik, KeyPairType.ed25519),
+        OmemoPublicKey.fromBytes(kex.ek, KeyPairType.x25519),
+        kex.pkId,
+      ),
       device.spk,
-      // TODO(PapaTutuWawa): Fix
+      device.opks.values.elementAt(kex.pkId),
-      device.opks.values.elementAt(0),
       device.ik,
     );
     final ratchet = await OmemoDoubleRatchet.acceptNewSession(
@@ -118,7 +127,7 @@ class OmemoSessionManager {
   /// Encrypt the key [plaintext] for all known bundles of [jid]. Returns a map that
   /// maps the Bundle Id to the ciphertext of [plaintext].
   Future<EncryptionResult> encryptToJid(String jid, String plaintext) async {
-    final encryptedKeys = <String, List<int>>{};
+    final encryptedKeys = <int, List<int>>{};
 
     // Generate the key and encrypt the plaintext
     final key = generateRandomBytes(32);
@@ -149,7 +158,7 @@ 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.
-  Future<String> decryptMessage(List<int> ciphertext, String senderJid, String senderDeviceId, List<EncryptedKey> keys) async {
+  Future<String> decryptMessage(List<int> ciphertext, String senderJid, int senderDeviceId, List<EncryptedKey> keys) async {
     // Try to find a session we can decrypt with.
     final rawKey = keys.firstWhereOrNull((key) => key.rid == device.id);
     if (rawKey == null) {

lib/src/x3dh/x3dh.dart

@@ -16,7 +16,7 @@ class X3DHAliceResult {
   const X3DHAliceResult(this.ek, this.sk, this.opkId, this.ad);
   final OmemoKeyPair ek;
   final List<int> sk;
-  final String opkId;
+  final int opkId;
   final List<int> ad;
 }
 
@@ -26,7 +26,7 @@ class X3DHMessage {
   const X3DHMessage(this.ik, this.ek, this.opkId);
   final OmemoPublicKey ik;
   final OmemoPublicKey ek;
-  final String opkId;
+  final int opkId;
 }
 
 class X3DHBobResult {

test/double_ratchet_test.dart

@@ -41,16 +41,16 @@ void main() {
     final spkBob = await OmemoKeyPair.generateNewPair(KeyPairType.x25519);
     final opkBob = await OmemoKeyPair.generateNewPair(KeyPairType.x25519);
     final bundleBob = OmemoBundle(
-      '1',
+      1,
       await spkBob.pk.asBase64(),
-      '3',
+      3,
       base64Encode(
         await sig(ikBob, await spkBob.pk.getBytes()),
       ),
       //'Q5in+/L4kJixEX692h6mJkPMyp4I3SlQ84L0E7ipPzqfPHOMiraUlqG2vG/O8wvFjLsKYZpPBraga9IvwhqVDA==',
       await ikBob.pk.asBase64(),
       {
-        '2': await opkBob.pk.asBase64(),
+        2: await opkBob.pk.asBase64(),
       },
     );
     
@@ -65,7 +65,7 @@ void main() {
       X3DHMessage(
         ikAlice.pk,
         resultAlice.ek.pk,
-        '2',
+        2,
       ),
       spkBob,
       opkBob,

test/omemo_test.dart

@@ -12,7 +12,7 @@ void main() {
     final bobSession = await OmemoSessionManager.generateNewIdentity(opkAmount: 1);
 
     // Perform the X3DH
-    final x3dhAliceResult = await aliceSession.addSessionFromBundle(
+    final kex = await aliceSession.addSessionFromBundle(
       bobJid,
       bobSession.device.id,
       await bobSession.device.toBundle(),
@@ -20,11 +20,7 @@ void main() {
     await bobSession.addSessionFromKeyExchange(
       aliceJid,
       aliceSession.device.id,
-      X3DHMessage(
+      kex,
-        aliceSession.device.ik.pk,
-        x3dhAliceResult.ek.pk,
-        '2',
-      ),
     );
 
     // Alice encrypts a message for Bob

test/x3dh_test.dart

@@ -11,16 +11,16 @@ void main() {
     final spkBob = await OmemoKeyPair.generateNewPair(KeyPairType.x25519);
     final opkBob = await OmemoKeyPair.generateNewPair(KeyPairType.x25519);
     final bundleBob = OmemoBundle(
-      '1',
+      1,
       await spkBob.pk.asBase64(),
-      '3',
+      3,
       base64Encode(
         await sig(ikBob, await spkBob.pk.getBytes()),
       ),
       //'Q5in+/L4kJixEX692h6mJkPMyp4I3SlQ84L0E7ipPzqfPHOMiraUlqG2vG/O8wvFjLsKYZpPBraga9IvwhqVDA==',
       await ikBob.pk.asBase64(),
       {
-        '2': await opkBob.pk.asBase64(),
+        2: await opkBob.pk.asBase64(),
       },
     );
     
@@ -35,7 +35,7 @@ void main() {
       X3DHMessage(
         ikAlice.pk,
         resultAlice.ek.pk,
-        '2',
+        2,
       ),
       spkBob,
       opkBob,
@@ -53,15 +53,15 @@ void main() {
     final spkBob = await OmemoKeyPair.generateNewPair(KeyPairType.x25519);
     final opkBob = await OmemoKeyPair.generateNewPair(KeyPairType.x25519);
     final bundleBob = OmemoBundle(
-      '1',
+      1,
       await spkBob.pk.asBase64(),
-      '3',
+      3,
       // NOTE: A bit flakey, but it is highly unlikely that the same keypair as this one
       //       gets generated.
       'Q5in+/L4kJixEX692h6mJkPMyp4I3SlQ84L0E7ipPzqfPHOMiraUlqG2vG/O8wvFjLsKYZpPBraga9IvwhqVDA==',
       await ikBob.pk.asBase64(),
       {
-        '2': await opkBob.pk.asBase64(),
+        2: await opkBob.pk.asBase64(),
       },
     );