From af33ed51d14e0316c69278043b55f28fbec71f9c Mon Sep 17 00:00:00 2001
From: "Alexander \"PapaTutuWawa" <papatutuwawa@polynom.me>
Date: Thu, 15 Jun 2023 16:07:23 +0200
Subject: [PATCH] feat: Guard against invalid X3DH signatures

---
 lib/src/errors.dart           |  5 +----
 lib/src/omemo/omemo.dart      | 18 ++++++++++++++++--
 lib/src/x3dh/x3dh.dart        |  7 ++++---
 test/double_ratchet_test.dart |  3 ++-
 test/x3dh_test.dart           | 18 ++++--------------
 5 files changed, 27 insertions(+), 24 deletions(-)

diff --git a/lib/src/errors.dart b/lib/src/errors.dart
index 1aafd61..858a618 100644
--- a/lib/src/errors.dart
+++ b/lib/src/errors.dart
@@ -1,10 +1,7 @@
 abstract class OmemoError {}
 
 /// Triggered during X3DH if the signature if the SPK does verify to the actual SPK.
-class InvalidSignatureException extends OmemoError implements Exception {
-  String errMsg() =>
-      'The signature of the SPK does not match the provided signature';
-}
+class InvalidKeyExchangeSignatureError extends OmemoError {}
 
 /// Triggered by the Double Ratchet if the computed HMAC does not match the attached HMAC.
 class InvalidMessageHMACError extends OmemoError {}
diff --git a/lib/src/omemo/omemo.dart b/lib/src/omemo/omemo.dart
index ed9cf4f..1c5fd33 100644
--- a/lib/src/omemo/omemo.dart
+++ b/lib/src/omemo/omemo.dart
@@ -509,6 +509,7 @@ class OmemoManager {
       ciphertext = [];
     }
 
+    final encryptionErrors = <String, List<EncryptToJidError>>{};
     final addedRatchetKeys = List<RatchetMapKey>.empty(growable: true);
     final kex = <RatchetMapKey, OMEMOKeyExchange>{};
     for (final jid in stanza.recipientJids) {
@@ -521,10 +522,24 @@ class OmemoManager {
         _log.finest('Building new ratchet $jid:${bundle.id}');
         final ratchetKey = RatchetMapKey(jid, bundle.id);
         final ownDevice = await getDevice();
-        final kexResult = await x3dhFromBundle(
+        final kexResultRaw = await x3dhFromBundle(
           bundle,
           ownDevice.ik,
         );
+        // TODO: Track the failure and do not attempt to encrypt to this device
+        //       on every send.
+        if (kexResultRaw.isType<InvalidKeyExchangeSignatureError>()) {
+          encryptionErrors.appendOrCreate(
+            jid,
+            EncryptToJidError(
+              bundle.id,
+              kexResultRaw.get<InvalidKeyExchangeSignatureError>(),
+            ),
+          );
+          continue;
+        }
+
+        final kexResult = kexResultRaw.get<X3DHAliceResult>();
         final newRatchet = await OmemoDoubleRatchet.initiateNewSession(
           bundle.spk,
           bundle.spkId,
@@ -567,7 +582,6 @@ class OmemoManager {
     }
 
     // Encrypt the symmetric key for all devices.
-    final encryptionErrors = <String, List<EncryptToJidError>>{};
     final encryptedKeys = <String, List<EncryptedKey>>{};
     for (final jid in stanza.recipientJids) {
       // Check if we know about any devices to use
diff --git a/lib/src/x3dh/x3dh.dart b/lib/src/x3dh/x3dh.dart
index 670a805..8f261a1 100644
--- a/lib/src/x3dh/x3dh.dart
+++ b/lib/src/x3dh/x3dh.dart
@@ -1,6 +1,7 @@
 import 'dart:convert';
 import 'dart:math';
 import 'package:cryptography/cryptography.dart';
+import 'package:omemo_dart/src/common/result.dart';
 import 'package:omemo_dart/src/crypto.dart';
 import 'package:omemo_dart/src/errors.dart';
 import 'package:omemo_dart/src/helpers.dart';
@@ -70,7 +71,7 @@ Future<List<int>> kdf(List<int> km) async {
 
 /// Alice builds a session with Bob using his bundle [bundle] and Alice's identity key
 /// pair [ik].
-Future<X3DHAliceResult> x3dhFromBundle(
+Future<Result<InvalidKeyExchangeSignatureError, X3DHAliceResult>> x3dhFromBundle(
   OmemoBundle bundle,
   OmemoKeyPair ik,
 ) async {
@@ -84,7 +85,7 @@ Future<X3DHAliceResult> x3dhFromBundle(
   );
 
   if (!signatureValue) {
-    throw InvalidSignatureException();
+    return Result(InvalidKeyExchangeSignatureError());
   }
 
   // Generate EK
@@ -106,7 +107,7 @@ Future<X3DHAliceResult> x3dhFromBundle(
     await bundle.ik.getBytes(),
   ]);
 
-  return X3DHAliceResult(ek, sk, opkId, ad);
+  return Result(X3DHAliceResult(ek, sk, opkId, ad));
 }
 
 /// Bob builds the X3DH shared secret from the inital message [msg], the SPK [spk], the
diff --git a/test/double_ratchet_test.dart b/test/double_ratchet_test.dart
index 1347e57..f630bae 100644
--- a/test/double_ratchet_test.dart
+++ b/test/double_ratchet_test.dart
@@ -28,7 +28,8 @@ void main() {
     );
 
     // Alice does X3DH
-    final resultAlice = await x3dhFromBundle(bundleBob, ikAlice);
+    final resultAliceRaw = await x3dhFromBundle(bundleBob, ikAlice);
+    final resultAlice = resultAliceRaw.get<X3DHAliceResult>();
 
     // Alice sends the inital message to Bob
     // ...
diff --git a/test/x3dh_test.dart b/test/x3dh_test.dart
index 68b249d..a8b91c5 100644
--- a/test/x3dh_test.dart
+++ b/test/x3dh_test.dart
@@ -26,7 +26,8 @@ void main() {
     );
 
     // Alice does X3DH
-    final resultAlice = await x3dhFromBundle(bundleBob, ikAlice);
+    final resultAliceRaw = await x3dhFromBundle(bundleBob, ikAlice);
+    final resultAlice = resultAliceRaw.get<X3DHAliceResult>();
 
     // Alice sends the inital message to Bob
     // ...
@@ -68,18 +69,7 @@ void main() {
     );
 
     // 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');
+    final result = await x3dhFromBundle(bundleBob, ikAlice);
+    expect(result.isType<InvalidKeyExchangeSignatureError>(), isTrue);
   });
 }