refactor: Clean up the X3DH implementation

example/omemo_dart_example.dart

@@ -1,6 +1,5 @@
-import 'package:omemo_dart/omemo_dart.dart';
+//import 'package:omemo_dart/omemo_dart.dart';
 
 void main() {
-  var awesome = Awesome();
+  // TODO(PapaTutuWawa): Currently NOOP
-  print('awesome: ${awesome.isAwesome}');
 }

flake.lock

@@ -17,16 +17,16 @@
     },
     "nixpkgs": {
       "locked": {
-        "lastModified": 1650015244,
+        "lastModified": 1657540956,
-        "narHash": "sha256-f6sgDj9A8FXTVyA2zkxA66YX+j6BftxE9VHTeIMhEKE=",
+        "narHash": "sha256-ihGbOFWtAkENwxBE5kV/yWt2MncvW+BObLDsmxCLo/Q=",
-        "owner": "PapaTutuWawa",
+        "owner": "NANASHI0X74",
         "repo": "nixpkgs",
-        "rev": "766f4f20760651ab523e716abe164317445b24ab",
+        "rev": "043de04db8a6b0391b3fefaaade160514d866946",
         "type": "github"
       },
       "original": {
-        "owner": "PapaTutuWawa",
+        "owner": "NANASHI0X74",
-        "ref": "nixos-unstable",
+        "ref": "flutter-3-0-0",
         "repo": "nixpkgs",
         "type": "github"
       }

flake.nix

@@ -1,13 +1,16 @@
 {
   description = "omemo_dart";
   inputs = {
-    nixpkgs.url = "github:PapaTutuWawa/nixpkgs/nixos-unstable";
+    nixpkgs.url = "github:NANASHI0X74/nixpkgs/flutter-3-0-0";
     flake-utils.url = "github:numtide/flake-utils";
   };
   outputs = { self, nixpkgs, flake-utils }: flake-utils.lib.eachDefaultSystem (system: let
     pkgs = import nixpkgs {
       inherit system;
-      config.android_sdk.accept_license = true;
+      config = {
+        android_sdk.accept_license = true;
+        allowUnfree = true;
+      };
     };
     android = pkgs.androidenv.composeAndroidPackages {
       # TODO: Find a way to pin these

lib/omemo_dart.dart

@@ -1,8 +1,5 @@
-/// Support for doing something awesome.
-///
-/// More dartdocs go here.
 library omemo_dart;
 
-export 'src/omemo_dart_base.dart';
+export 'src/bundle.dart';
-
+export 'src/key.dart';
-// TODO: Export any libraries intended for clients of this package.
+export 'src/x3dh.dart';

lib/src/bundle.dart

@@ -0,0 +1,42 @@
+import 'dart:convert';
+import 'package:cryptography/cryptography.dart';
+import 'key.dart';
+
+class OmemoBundle {
+
+  const OmemoBundle(
+    this.id,
+    this.spkEncoded,
+    this.spkId,
+    this.spkSignatureEncoded,
+    this.ikEncoded,
+    this.opksEncoded,
+  );
+  final String id;
+  /// The SPK but base64 encoded
+  final String spkEncoded;
+  final String 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;
+
+  OmemoPublicKey get spk {
+    final data = base64Decode(spkEncoded);
+    return OmemoPublicKey.fromBytes(data, KeyPairType.x25519);
+  }
+
+  OmemoPublicKey get ik {
+    final data = base64Decode(ikEncoded);
+    return OmemoPublicKey.fromBytes(data, KeyPairType.ed25519);
+  }
+
+  OmemoPublicKey getOpk(String id) {
+    final data = base64Decode(opksEncoded[id]!);
+    return OmemoPublicKey.fromBytes(data, KeyPairType.x25519);
+  }
+
+  List<int> get spkSignature => base64Decode(spkSignatureEncoded);
+}

lib/src/key.dart

@@ -0,0 +1,123 @@
+import 'dart:convert';
+import 'package:cryptography/cryptography.dart';
+import 'package:pinenacl/api.dart';
+import 'package:pinenacl/tweetnacl.dart';
+
+const privateKeyLength = 32;
+const publicKeyLength = 32;
+
+class OmemoPublicKey {
+
+  const OmemoPublicKey(this._pubkey);
+  final SimplePublicKey _pubkey;
+
+  factory OmemoPublicKey.fromBytes(List<int> bytes, KeyPairType type) {
+    return OmemoPublicKey(
+      SimplePublicKey(
+        bytes,
+        type: type,
+      ),
+    );
+  }
+  
+  KeyPairType get type => _pubkey.type;
+
+  /// Return the bytes that comprise the public key.
+  Future<List<int>> getBytes() async => _pubkey.bytes;
+
+  /// Returns the public key encoded as base64.
+  Future<String> asBase64() async => base64Encode(_pubkey.bytes);
+
+  Future<OmemoPublicKey> toCurve25519() async {
+    assert(type == KeyPairType.ed25519, 'Cannot convert non-Ed25519 public key to X25519');
+
+    final pkc = Uint8List(publicKeyLength);
+    TweetNaClExt.crypto_sign_ed25519_pk_to_x25519_pk(
+      pkc,
+      Uint8List.fromList(await getBytes()),
+    );
+
+    return OmemoPublicKey(SimplePublicKey(List<int>.from(pkc), type: KeyPairType.x25519));
+  }
+
+  SimplePublicKey asPublicKey() => _pubkey;
+  
+  /// Convert the public key into a [SimpleKeyPairData] with a stub private key. Useful
+  /// for when cryptography calls for a KeyPair, but only uses the public key.
+  //SimpleKeyPairData asPseudoKeypair() {
+  //
+  //}
+}
+
+class OmemoPrivateKey {
+
+  const OmemoPrivateKey(this._privkey, this.type);
+  final List<int> _privkey;
+  final KeyPairType type;
+
+  Future<List<int>> getBytes() async => _privkey;
+  
+  Future<OmemoPrivateKey> toCurve25519() async {
+    assert(type == KeyPairType.ed25519, 'Cannot convert non-Ed25519 private key to X25519');
+
+    final skc = Uint8List(privateKeyLength);
+    TweetNaClExt.crypto_sign_ed25519_sk_to_x25519_sk(
+      skc,
+      Uint8List.fromList(await getBytes()),
+    );
+
+    return OmemoPrivateKey(List<int>.from(skc), KeyPairType.x25519);
+  }
+}
+
+/// A generic wrapper class for both Ed25519 and X25519 keypairs
+class OmemoKeyPair {
+
+  const OmemoKeyPair(this.pk, this.sk, this.type);
+  final KeyPairType type;
+  final OmemoPublicKey pk;
+  final OmemoPrivateKey sk;
+
+  static Future<OmemoKeyPair> generateNewPair(KeyPairType type) async {
+    assert(type == KeyPairType.ed25519 || type == KeyPairType.x25519);
+
+    SimpleKeyPair kp;
+    if (type == KeyPairType.ed25519) {
+      final ed = Ed25519();
+      kp = await ed.newKeyPair();
+    } else if (type == KeyPairType.x25519) {
+      final x = Cryptography.instance.x25519();
+      kp = await x.newKeyPair();
+    } else {
+      // Should never happen
+      throw Exception();
+    }
+
+    final kpd = await kp.extract();
+    
+    return OmemoKeyPair(
+      OmemoPublicKey(await kp.extractPublicKey()),
+      OmemoPrivateKey(await kpd.extractPrivateKeyBytes(), type),
+      type,
+    );
+  }
+  
+  /// Return the bytes that comprise the public key.
+  Future<OmemoKeyPair> toCurve25519() async {
+    assert(type == KeyPairType.ed25519, 'Cannot convert non-Ed25519 keypair to X25519');
+
+    return OmemoKeyPair(
+      await pk.toCurve25519(),
+      await sk.toCurve25519(),
+      KeyPairType.x25519,
+    );
+  }
+
+  Future<SimpleKeyPairData> asKeyPair() async {
+    return SimpleKeyPairData(
+      await sk.getBytes(),
+      publicKey: pk.asPublicKey(),
+      type: type,
+    );
+  }
+}

lib/src/x3dh.dart

@@ -0,0 +1,122 @@
+import 'dart:convert';
+import 'dart:math';
+import 'package:cryptography/cryptography.dart';
+import 'bundle.dart';
+import 'key.dart';
+
+/// The overarching assumption is that we use Ed25519 keys for the identity keys
+
+/// Performed by Alice
+class X3DHResult {
+
+  const X3DHResult(this.ek, this.sk, this.opkId);
+  final OmemoKeyPair ek;
+  final List<int> sk;
+  final String opkId;
+}
+
+/// Received by Bob
+class X3DHMessage {
+
+  const X3DHMessage(this.ik, this.ek, this.opkId);
+  final OmemoPublicKey ik;
+  final OmemoPublicKey ek;
+  final String opkId;
+}
+
+/// Sign [message] using the keypair [keyPair]. Note that [keyPair] must be
+/// a Ed25519 keypair.
+Future<List<int>> sig(OmemoKeyPair keyPair, List<int> message) async {
+  assert(keyPair.type == KeyPairType.ed25519);
+  final signature = await Ed25519().sign(
+    message,
+    keyPair: await keyPair.asKeyPair(),
+  );
+
+  return signature.bytes;
+}
+
+/// Performs X25519 with [pk1] and [pk2]. If [identityKey] is set, then
+/// it indicates which of [pk1] ([identityKey] == 1) or [pk2] ([identityKey] == 2)
+/// is the identity key.
+Future<List<int>> dh(OmemoKeyPair kp, OmemoPublicKey pk, int identityKey) async {
+  var ckp = kp;
+  var cpk = pk;
+
+  if (identityKey == 1) {
+    ckp = await kp.toCurve25519();
+  } else if (identityKey == 2) {
+    cpk = await pk.toCurve25519();
+  }
+
+  final shared = await Cryptography.instance.x25519().sharedSecretKey(
+    keyPair: await ckp.asKeyPair(),
+    remotePublicKey: cpk.asPublicKey(),
+  );
+
+  return shared.extractBytes();
+}
+
+/// Derive a secret from the key material [km].
+Future<List<int>> kdf(List<int> km) async {
+  final f = List<int>.filled(32, 0xFF);
+  final input = List<int>.empty(growable: true);
+  input
+    ..addAll(f)
+    ..addAll(km);
+
+  final algorithm = Hkdf(
+    hmac: Hmac(Sha256()),
+    outputLength: 32,
+  );
+  final output = await algorithm.deriveKey(
+    secretKey: SecretKey(input),
+    // TODO: Fix
+    nonce: List<int>.filled(32, 0x00),
+    info: utf8.encode('OMEMO X3DH'),
+  );
+
+  return output.extractBytes();
+}
+
+/// Flattens [inputs] and concatenates the elements.
+List<int> concat(List<List<int>> inputs) {
+  final tmp = List<int>.empty(growable: true);
+  for (final input in inputs) {
+    tmp.addAll(input);
+  }
+
+  return tmp;
+}
+
+/// Alice builds a session with Bob using his bundle [bundle] and Alice's identity key
+/// pair [ika].
+Future<X3DHResult> x3dhFromBundle(OmemoBundle bundle, OmemoKeyPair ik) async {
+  // Generate EK
+  final ek = await OmemoKeyPair.generateNewPair(KeyPairType.x25519);
+
+  final random = Random.secure();
+  final opkIndex = random.nextInt(bundle.opksEncoded.length);
+  final opkId = bundle.opksEncoded.keys.elementAt(opkIndex);
+  final opk = bundle.getOpk(opkId);
+  
+  final dh1 = await dh(ik, bundle.spk, 1);
+  final dh2 = await dh(ek, bundle.ik,  2);
+  final dh3 = await dh(ek, bundle.spk, 0);
+  final dh4 = await dh(ek, opk, 0);
+
+  final sk = await kdf(concat([dh1, dh2, dh3, dh4]));
+
+  return X3DHResult(ek, sk, opkId);
+}
+
+/// Bob builds the X3DH shared secret from the inital message [msg], the SPK [spk], the
+/// OPK [opk] that was selected by Alice and our IK [ik]. Returns the shared secret.
+Future<List<int>> x3dhFromInitialMessage(X3DHMessage msg, OmemoKeyPair spk, OmemoKeyPair opk, OmemoKeyPair ik) async {
+  final dh1 = await dh(spk, msg.ik, 2);
+  final dh2 = await dh(ik,  msg.ek, 1);
+  final dh3 = await dh(spk, msg.ek, 0);
+  final dh4 = await dh(opk, msg.ek, 0);
+
+  return kdf(concat([dh1, dh2, dh3, dh4]));
+}

lib/src/x3dh/x3dh.dart

@@ -1,125 +0,0 @@
-import 'dart:convert';
-import 'dart:typed_data';
-import 'package:cryptography/cryptography.dart';
-import 'package:pinenacl/api.dart';
-import 'package:pinenacl/tweetnacl.dart';
-
-/// The overarching assumption is that we use Ed25519 keys for the identity keys
-
-class X3DHRun {
-
-  const X3DHRun(this.epk, this.sharedSecret);
-  final SimpleKeyPair epk;
-  final List<int> sharedSecret;
-}
-
-SimpleKeyPairData fromPublicKey(SimplePublicKey pk) {
-  return SimpleKeyPairData([], publicKey: pk, type: KeyPairType.x25519);
-}
-
-Future<List<int>> sig(SimpleKeyPair keyPair, List<int> message) async {
-  final signature = await Ed25519().sign(
-    message,
-    keyPair: keyPair,
-  );
-
-  return signature.bytes;
-}
-
-/// Performs X25519 with [pk1] and [pk2]. If [identityKey] is set, then
-/// it indicates which of [pk1] ([identityKey] == 1) or [pk2] ([identityKey] == 2)
-/// is the identity key.
-Future<List<int>> dh(SimpleKeyPair kp, SimplePublicKey pk, int identityKey) async {
-  var ckp = kp;
-  var cpk = pk;
-
-  if (identityKey == 1) {
-    final pubkeyBytes = (await kp.extractPublicKey()).bytes;
-    final pkc = Uint8List(32);
-    TweetNaClExt.crypto_sign_ed25519_pk_to_x25519_pk(pkc, Uint8List.fromList(pubkeyBytes));
-
-    final keyPairData = await kp.extract();
-    final privateKeyBytes = await keyPairData.extractPrivateKeyBytes();
-    final skc = Uint8List(32);
-    TweetNaClExt.crypto_sign_ed25519_sk_to_x25519_sk(skc, Uint8List.fromList(privateKeyBytes));
-
-    ckp = SimpleKeyPairData(
-      List<int>.from(skc),
-      publicKey: SimplePublicKey(List<int>.from(pkc), type: KeyPairType.x25519),
-      type: KeyPairType.x25519,
-    );
-  } else if (identityKey == 2) {
-    final pubkeyBytes = pk.bytes;
-    final pkc = Uint8List(32);
-    TweetNaClExt.crypto_sign_ed25519_pk_to_x25519_pk(pkc, Uint8List.fromList(pubkeyBytes));
-
-    cpk = SimplePublicKey(
-      List<int>.from(pkc),
-      type: KeyPairType.x25519,
-    );
-  }
-
-  final shared = await Cryptography.instance.x25519().sharedSecretKey(
-    keyPair: ckp,
-    remotePublicKey: cpk,
-  );
-
-  return shared.extractBytes();
-}
-
-Future<List<int>> kdf(List<int> km) async {
-  final f = List<int>.filled(32, 0xFF);
-  final input = List<int>.empty(growable: true);
-  input
-    ..addAll(f)
-    ..addAll(km);
-
-  final algorithm = Hkdf(
-    hmac: Hmac(Sha256()),
-    outputLength: 32,
-  );
-  final output = await algorithm.deriveKey(
-    secretKey: SecretKey(input),
-    // TODO: Fix
-    nonce: List<int>.filled(32, 0x00),
-    info: utf8.encode('OMEMO X3DH'),
-  );
-
-  return output.extractBytes();
-}
-
-List<int> concat(List<List<int>> inputs) {
-  final tmp = List<int>.empty(growable: true);
-  for (final input in inputs) {
-    tmp.addAll(input);
-  }
-
-  return tmp;
-}
-
-// Alice -> Bob
-Future<X3DHRun> x3dhFromPrekeyBundle(SimplePublicKey ikb, SimplePublicKey spkb, SimplePublicKey opkb, SimpleKeyPair ika) async {
-  // Generate EPK
-  final epk = await Cryptography.instance.x25519().newKeyPair();
-
-  final dh1 = await dh(ika, spkb, 1);
-  final dh2 = await dh(epk, ikb,  2);
-  final dh3 = await dh(epk, spkb, 0);
-  final dh4 = await dh(epk, opkb, 0);
-
-  final sk = await kdf(concat([dh1, dh2, dh3, dh4]));
-
-  return X3DHRun(
-    epk,
-    sk,
-  );
-}
-
-Future<List<int>> x3dhFromInitialMessage(SimplePublicKey ika, SimplePublicKey epk, SimpleKeyPair opkb, SimpleKeyPair spk, SimpleKeyPair ikb) async {
-  final dh1 = await dh(spk, ika, 2);
-  final dh2 = await dh(ikb, epk, 1);
-  final dh3 = await dh(spk, epk, 0);
-  final dh4 = await dh(opkb, epk, 0);
-
-  return kdf(concat([dh1, dh2, dh3, dh4]));
-}

pubspec.yaml

@@ -3,13 +3,13 @@ description: An XMPP library independent OMEMO library
 version: 0.1.0
 
 environment:
-  sdk: '>=2.15.1 <3.0.0'
+  sdk: '>=2.17.0 <3.0.0'
 
 
 dependencies:
-  cryptography:
+  cryptography: ^2.0.5
-  pinenacl:
+  pinenacl: ^0.5.1
 
 dev_dependencies:
-  lints: ^1.0.0
+  lints: ^2.0.0
-  test: ^1.16.0
+  test: ^1.21.0

test/omemo_dart_test.dart

@@ -1,48 +1,44 @@
 import 'package:cryptography/cryptography.dart';
 import 'package:test/test.dart';
-import 'package:omemo_dart/src/x3dh/x3dh.dart';
+import 'package:omemo_dart/omemo_dart.dart';
-
-Future<List<int>> publicKeyBytes(SimpleKeyPair kp) async {
-  final pk = await kp.extractPublicKey();
-  return pk.bytes;
-}
 
 void main() {
   test("X3DH", () async {
-    final ed = Ed25519();
+    // Generate keys
-    final x = Cryptography.instance.x25519();
+    final ikAlice = await OmemoKeyPair.generateNewPair(KeyPairType.ed25519);
-
+    final ikBob = await OmemoKeyPair.generateNewPair(KeyPairType.ed25519);
-    // Generate IKs for Alice and Bob
+    final spkBob = await OmemoKeyPair.generateNewPair(KeyPairType.x25519);
-    final ikAlice = await ed.newKeyPair();
+    final opkBob = await OmemoKeyPair.generateNewPair(KeyPairType.x25519);
-    final ikBob = await ed.newKeyPair();
+    final bundleBob = OmemoBundle(
-    
+      '1',
-    // Generate SPKs for Alice and Bob
+      await spkBob.pk.asBase64(),
-    final spkAlice = await x.newKeyPair();
+      '3',
-    final spkSigAlice = await sig(ikAlice, await publicKeyBytes(spkAlice));
+      // TODO(PapaTutuWawa):
-    final spkBob = await x.newKeyPair();
+      'n/a',
-    final spkSigBob = await sig(ikBob, await publicKeyBytes(spkBob));
+      await ikBob.pk.asBase64(),
-
+      {
-    // Generate an OPK for Alice and Bob
+        '2': await opkBob.pk.asBase64(),
-    final opkAlice = await x.newKeyPair();
+      },
-    final opkBob = await x.newKeyPair();
-
-    
-    // Perform X3DH
-    final aliceMessage = await x3dhFromPrekeyBundle(
-      await ikBob.extractPublicKey(),
-      await spkBob.extractPublicKey(),
-      await opkBob.extractPublicKey(),
-      ikAlice,
     );
 
-    final bobDh = await x3dhFromInitialMessage(
+    // Alice does X3DH
-      await ikAlice.extractPublicKey(),
+    final resultAlice = await x3dhFromBundle(bundleBob, ikAlice);
-      await aliceMessage.epk.extractPublicKey(),
+
-      opkBob,
+    // Alice sends the inital message to Bob
+    // ...
+    
+    // Bob does X3DH
+    final skBob = await x3dhFromInitialMessage(
+      X3DHMessage(
+        ikAlice.pk,
+        resultAlice.ek.pk,
+        '2',
+      ),
       spkBob,
+      opkBob,
       ikBob,
     );
     
-    expect(aliceMessage.sharedSecret, bobDh);
+    expect(resultAlice.sk, skBob);
   });
 }