From 4211775c9044e24c9d5ecedd1d665af62ac92ba9 Mon Sep 17 00:00:00 2001
From: "Alexander \"PapaTutuWawa" <papatutuwawa@polynom.me>
Date: Tue, 2 Aug 2022 15:40:26 +0200
Subject: [PATCH] feat: Add signature validation

---
 lib/omemo_dart.dart |  1 +
 lib/src/errors.dart |  4 ++++
 lib/src/x3dh.dart   | 14 ++++++++++++++
 test/x3dh_test.dart | 42 ++++++++++++++++++++++++++++++++++++++----
 4 files changed, 57 insertions(+), 4 deletions(-)
 create mode 100644 lib/src/errors.dart

diff --git a/lib/omemo_dart.dart b/lib/omemo_dart.dart
index 383f2bf..0b63122 100644
--- a/lib/omemo_dart.dart
+++ b/lib/omemo_dart.dart
@@ -1,5 +1,6 @@
 library omemo_dart;
 
 export 'src/bundle.dart';
+export 'src/errors.dart';
 export 'src/key.dart';
 export 'src/x3dh.dart';
diff --git a/lib/src/errors.dart b/lib/src/errors.dart
new file mode 100644
index 0000000..f4377a3
--- /dev/null
+++ b/lib/src/errors.dart
@@ -0,0 +1,4 @@
+class InvalidSignatureException implements Exception {
+  @override
+  String errMsg() => 'The signature of the SPK does not match the provided signature';
+}
diff --git a/lib/src/x3dh.dart b/lib/src/x3dh.dart
index 4c020ec..72feadd 100644
--- a/lib/src/x3dh.dart
+++ b/lib/src/x3dh.dart
@@ -2,6 +2,7 @@ import 'dart:convert';
 import 'dart:math';
 import 'package:cryptography/cryptography.dart';
 import 'package:omemo_dart/src/bundle.dart';
+import 'package:omemo_dart/src/errors.dart';
 import 'package:omemo_dart/src/key.dart';
 
 /// The overarching assumption is that we use Ed25519 keys for the identity keys
@@ -99,6 +100,19 @@ List<int> concat(List<List<int>> inputs) {
 /// Alice builds a session with Bob using his bundle [bundle] and Alice's identity key
 /// pair [ik].
 Future<X3DHAliceResult> x3dhFromBundle(OmemoBundle bundle, OmemoKeyPair ik) async {
+  // Check the signature first
+  final signatureValue = await Ed25519().verify(
+    await bundle.spk.getBytes(),
+    signature: Signature(
+      bundle.spkSignature,
+      publicKey: bundle.ik.asPublicKey(),
+    ),
+  );
+
+  if (!signatureValue) {
+    throw InvalidSignatureException();
+  }
+
   // Generate EK
   final ek = await OmemoKeyPair.generateNewPair(KeyPairType.x25519);
 
diff --git a/test/x3dh_test.dart b/test/x3dh_test.dart
index e7173a5..53c2de5 100644
--- a/test/x3dh_test.dart
+++ b/test/x3dh_test.dart
@@ -1,9 +1,10 @@
+import 'dart:convert';
 import 'package:cryptography/cryptography.dart';
 import 'package:omemo_dart/omemo_dart.dart';
 import 'package:test/test.dart';
 
 void main() {
-  test('X3DH', () async {
+  test('X3DH with correct signature', () async {
     // Generate keys
     final ikAlice = await OmemoKeyPair.generateNewPair(KeyPairType.ed25519);
     final ikBob = await OmemoKeyPair.generateNewPair(KeyPairType.ed25519);
@@ -13,14 +14,16 @@ void main() {
       '1',
       await spkBob.pk.asBase64(),
       '3',
-      // TODO(PapaTutuWawa): Do
-      'n/a',
+      base64Encode(
+        await sig(ikBob, await spkBob.pk.getBytes()),
+      ),
+      //'Q5in+/L4kJixEX692h6mJkPMyp4I3SlQ84L0E7ipPzqfPHOMiraUlqG2vG/O8wvFjLsKYZpPBraga9IvwhqVDA==',
       await ikBob.pk.asBase64(),
       {
         '2': await opkBob.pk.asBase64(),
       },
     );
-
+    
     // Alice does X3DH
     final resultAlice = await x3dhFromBundle(bundleBob, ikAlice);
 
@@ -42,4 +45,35 @@ void main() {
     expect(resultAlice.sk, resultBob.sk);
     expect(resultAlice.ad, resultBob.ad);
   });
+
+  test('X3DH with incorrect signature', () async {
+    // Generate keys
+    final ikAlice = await OmemoKeyPair.generateNewPair(KeyPairType.ed25519);
+    final ikBob = await OmemoKeyPair.generateNewPair(KeyPairType.ed25519);
+    final spkBob = await OmemoKeyPair.generateNewPair(KeyPairType.x25519);
+    final opkBob = await OmemoKeyPair.generateNewPair(KeyPairType.x25519);
+    final bundleBob = OmemoBundle(
+      '1',
+      await spkBob.pk.asBase64(),
+      '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(),
+      },
+    );
+    
+    // Alice does X3DH
+    var exception = false;
+    try {
+      await x3dhFromBundle(bundleBob, ikAlice);
+    } catch(e) {
+      exception = true;
+      expect(e is InvalidSignatureException, true, reason: 'Expected InvalidSignatureException, but got $e');
+    }
+
+    expect(exception, true, reason: 'Expected test failure');
+  });
 }