feat: Rework the double ratchet

lib/omemo_dart.dart

@@ -10,7 +10,7 @@ export 'src/omemo/encrypted_key.dart';
 export 'src/omemo/encryption_result.dart';
 export 'src/omemo/events.dart';
 export 'src/omemo/fingerprint.dart';
-export 'src/omemo/omemomanager.dart';
+export 'src/omemo/omemo.dart';
 export 'src/omemo/ratchet_map_key.dart';
 export 'src/omemo/stanza.dart';
 export 'src/trust/base.dart';

lib/src/common/result.dart

@@ -0,0 +1,19 @@
+// TODO: Pull into moxlib
+class Result<T, V> {
+  const Result(this._data)
+      : assert(
+          _data is T || _data is V,
+          'Invalid data type: Must be either $T or $V',
+        );
+  final dynamic _data;
+
+  bool isType<S>() => _data is S;
+
+  S get<S>() {
+    assert(_data is S, 'Data is not $S');
+
+    return _data as S;
+  }
+
+  Object get dataRuntimeType => _data.runtimeType;
+}

lib/src/double_ratchet/crypto.dart

@@ -1,59 +0,0 @@
-import 'package:omemo_dart/protobuf/schema.pb.dart';
-import 'package:omemo_dart/src/crypto.dart';
-import 'package:omemo_dart/src/errors.dart';
-import 'package:omemo_dart/src/helpers.dart';
-
-/// Info string for ENCRYPT
-const encryptHkdfInfoString = 'OMEMO Message Key Material';
-
-/// Signals ENCRYPT function as specified by OMEMO 0.8.3.
-/// Encrypt [plaintext] using the message key [mk], given associated_data [associatedData]
-/// and the AD output from the X3DH [sessionAd].
-Future<List<int>> encrypt(
-  List<int> mk,
-  List<int> plaintext,
-  List<int> associatedData,
-  List<int> sessionAd,
-) async {
-  // Generate encryption, authentication key and IV
-  final keys = await deriveEncryptionKeys(mk, encryptHkdfInfoString);
-  final ciphertext =
-      await aes256CbcEncrypt(plaintext, keys.encryptionKey, keys.iv);
-
-  final header =
-      OMEMOMessage.fromBuffer(associatedData.sublist(sessionAd.length))
-        ..ciphertext = ciphertext;
-  final headerBytes = header.writeToBuffer();
-  final hmacInput = concat([sessionAd, headerBytes]);
-  final hmacResult = await truncatedHmac(hmacInput, keys.authenticationKey);
-  final message = OMEMOAuthenticatedMessage()
-    ..mac = hmacResult
-    ..message = headerBytes;
-  return message.writeToBuffer();
-}
-
-/// Signals DECRYPT function as specified by OMEMO 0.8.3.
-/// Decrypt [ciphertext] with the message key [mk], given the associated_data [associatedData]
-/// and the AD output from the X3DH.
-Future<List<int>> decrypt(
-  List<int> mk,
-  List<int> ciphertext,
-  List<int> associatedData,
-  List<int> sessionAd,
-) async {
-  // Generate encryption, authentication key and IV
-  final keys = await deriveEncryptionKeys(mk, encryptHkdfInfoString);
-
-  // Assumption ciphertext is a OMEMOAuthenticatedMessage
-  final message = OMEMOAuthenticatedMessage.fromBuffer(ciphertext);
-  final header = OMEMOMessage.fromBuffer(message.message);
-
-  final hmacInput = concat([sessionAd, header.writeToBuffer()]);
-  final hmacResult = await truncatedHmac(hmacInput, keys.authenticationKey);
-
-  if (!listsEqual(hmacResult, message.mac)) {
-    throw InvalidMessageHMACException();
-  }
-
-  return aes256CbcDecrypt(header.ciphertext, keys.encryptionKey, keys.iv);
-}

lib/src/double_ratchet/double_ratchet.dart

@@ -2,46 +2,29 @@ import 'dart:convert';
 import 'package:cryptography/cryptography.dart';
 import 'package:hex/hex.dart';
 import 'package:meta/meta.dart';
-import 'package:omemo_dart/protobuf/schema.pb.dart';
+import 'package:omemo_dart/src/common/result.dart';
 import 'package:omemo_dart/src/crypto.dart';
-import 'package:omemo_dart/src/double_ratchet/crypto.dart';
 import 'package:omemo_dart/src/double_ratchet/kdf.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/protobuf/schema.pb.dart';
 
 /// Amount of messages we may skip per session
 const maxSkip = 1000;
 
-class RatchetStep {
+/// Info string for ENCRYPT
-  const RatchetStep(this.header, this.ciphertext);
+const encryptHkdfInfoString = 'OMEMO Message Key Material';
-  final OMEMOMessage header;
-  final List<int> ciphertext;
-}
 
 @immutable
 class SkippedKey {
   const SkippedKey(this.dh, this.n);
 
-  factory SkippedKey.fromJson(Map<String, dynamic> data) {
+  /// The DH public key for which we skipped a message key.
-    return SkippedKey(
-      OmemoPublicKey.fromBytes(
-        base64.decode(data['public']! as String),
-        KeyPairType.x25519,
-      ),
-      data['n']! as int,
-    );
-  }
-
   final OmemoPublicKey dh;
-  final int n;
 
-  Future<Map<String, dynamic>> toJson() async {
+  /// The associated number of the message key we skipped.
-    return {
+  final int n;
-      'public': base64.encode(await dh.getBytes()),
-      'n': n,
-    };
-  }
 
   @override
   bool operator ==(Object other) {
@@ -63,6 +46,7 @@ class OmemoDoubleRatchet {
     this.nr, // Nr
     this.pn, // Pn
     this.ik,
+    this.ek,
     this.sessionAd,
     this.mkSkipped, // MKSKIPPED
     this.acknowledged,
@@ -70,73 +54,6 @@ class OmemoDoubleRatchet {
     this.kex,
   );
 
-  factory OmemoDoubleRatchet.fromJson(Map<String, dynamic> data) {
-    /*
-    {
-      'dhs': 'base/64/encoded',
-      'dhs_pub': 'base/64/encoded',
-      'dhr': null | 'base/64/encoded',
-      'rk': 'base/64/encoded',
-      'cks': null | 'base/64/encoded',
-      'ckr': null | 'base/64/encoded',
-      'ns': 0,
-      'nr': 0,
-      'pn': 0,
-      'ik_pub': null | 'base/64/encoded',
-      'session_ad': 'base/64/encoded',
-      'acknowledged': true | false,
-      'kex_timestamp': int,
-      'kex': 'base/64/encoded',
-      'mkskipped': [
-        {
-          'key': 'base/64/encoded',
-          'public': 'base/64/encoded',
-          'n': 0
-        }, ...
-      ]
-    }
-    */
-    // NOTE: Dart has some issues with just casting a List<dynamic> to List<Map<...>>, as
-    //       such we need to convert the items by hand.
-    final mkSkipped = Map<SkippedKey, List<int>>.fromEntries(
-      (data['mkskipped']! as List<dynamic>)
-          .map<MapEntry<SkippedKey, List<int>>>(
-        (entry) {
-          final map = entry as Map<String, dynamic>;
-          final key = SkippedKey.fromJson(map);
-          return MapEntry(
-            key,
-            base64.decode(map['key']! as String),
-          );
-        },
-      ),
-    );
-
-    return OmemoDoubleRatchet(
-      OmemoKeyPair.fromBytes(
-        base64.decode(data['dhs_pub']! as String),
-        base64.decode(data['dhs']! as String),
-        KeyPairType.x25519,
-      ),
-      decodeKeyIfNotNull(data, 'dhr', KeyPairType.x25519),
-      base64.decode(data['rk']! as String),
-      base64DecodeIfNotNull(data, 'cks'),
-      base64DecodeIfNotNull(data, 'ckr'),
-      data['ns']! as int,
-      data['nr']! as int,
-      data['pn']! as int,
-      OmemoPublicKey.fromBytes(
-        base64.decode(data['ik_pub']! as String),
-        KeyPairType.ed25519,
-      ),
-      base64.decode(data['session_ad']! as String),
-      mkSkipped,
-      data['acknowledged']! as bool,
-      data['kex_timestamp']! as int,
-      data['kex'] as String?,
-    );
-  }
-
   /// Sending DH keypair
   OmemoKeyPair dhs;
 
@@ -161,6 +78,11 @@ class OmemoDoubleRatchet {
   /// for verification purposes
   final OmemoPublicKey ik;
 
+  /// The ephemeral public key of the chat partner. Not used for encryption but for possible
+  /// checks when replacing the ratchet. As such, this is only non-null for the initiating
+  /// side.
+  final OmemoPublicKey? ek;
+
   final List<int> sessionAd;
 
   final Map<SkippedKey, List<int>> mkSkipped;
@@ -182,25 +104,25 @@ class OmemoDoubleRatchet {
   static Future<OmemoDoubleRatchet> initiateNewSession(
     OmemoPublicKey spk,
     OmemoPublicKey ik,
+    OmemoPublicKey ek,
     List<int> sk,
     List<int> ad,
     int timestamp,
   ) async {
     final dhs = await OmemoKeyPair.generateNewPair(KeyPairType.x25519);
-    final dhr = spk;
+    final rk = await kdfRk(sk, await omemoDH(dhs, spk, 0));
-    final rk = await kdfRk(sk, await omemoDH(dhs, dhr, 0));
-    final cks = rk;
 
     return OmemoDoubleRatchet(
       dhs,
-      dhr,
+      spk,
-      rk,
+      List.from(rk),
-      cks,
+      List.from(rk),
       null,
       0,
       0,
       0,
       ik,
+      ek,
       ad,
       {},
       false,
@@ -230,6 +152,7 @@ class OmemoDoubleRatchet {
       0,
       0,
       ik,
+      null,
       ad,
       {},
       true,
@@ -238,67 +161,42 @@ class OmemoDoubleRatchet {
     );
   }
 
-  Future<Map<String, dynamic>> toJson() async {
+  /// Performs a single ratchet step in case we received a new
-    final mkSkippedSerialised =
+  /// public key in [header].
-        List<Map<String, dynamic>>.empty(growable: true);
+  Future<void> _dhRatchet(OMEMOMessage header) async {
-    for (final entry in mkSkipped.entries) {
+    pn = ns;
-      final result = await entry.key.toJson();
+    ns = 0;
-      result['key'] = base64.encode(entry.value);
+    nr = 0;
-
+    dhr = OmemoPublicKey.fromBytes(header.dhPub, KeyPairType.x25519);
-      mkSkippedSerialised.add(result);
+    final newRk1 = await kdfRk(
-    }
+      rk,
-
+      await omemoDH(
-    return {
+        dhs,
-      'dhs': base64.encode(await dhs.sk.getBytes()),
+        dhr!,
-      'dhs_pub': base64.encode(await dhs.pk.getBytes()),
+        0,
-      'dhr': dhr != null ? base64.encode(await dhr!.getBytes()) : null,
+      ),
-      'rk': base64.encode(rk),
-      'cks': cks != null ? base64.encode(cks!) : null,
-      'ckr': ckr != null ? base64.encode(ckr!) : null,
-      'ns': ns,
-      'nr': nr,
-      'pn': pn,
-      'ik_pub': base64.encode(await ik.getBytes()),
-      'session_ad': base64.encode(sessionAd),
-      'mkskipped': mkSkippedSerialised,
-      'acknowledged': acknowledged,
-      'kex_timestamp': kexTimestamp,
-      'kex': kex,
-    };
-  }
-
-  /// Returns the OMEMO compatible fingerprint of the ratchet session.
-  Future<String> getOmemoFingerprint() async {
-    final curveKey = await ik.toCurve25519();
-    return HEX.encode(await curveKey.getBytes());
-  }
-
-  Future<List<int>?> _trySkippedMessageKeys(
-    OMEMOMessage header,
-    List<int> ciphertext,
-  ) async {
-    final key = SkippedKey(
-      OmemoPublicKey.fromBytes(header.dhPub, KeyPairType.x25519),
-      header.n,
     );
-    if (mkSkipped.containsKey(key)) {
+    rk = List.from(newRk1);
-      final mk = mkSkipped[key]!;
+    ckr = List.from(newRk1);
-      mkSkipped.remove(key);
+    dhs = await OmemoKeyPair.generateNewPair(KeyPairType.x25519);
-
+    final newRk2 = await kdfRk(
-      return decrypt(
+      rk,
-        mk,
+      await omemoDH(
-        ciphertext,
+        dhs,
-        concat([sessionAd, header.writeToBuffer()]),
+        dhr!,
-        sessionAd,
+        0,
-      );
+      ),
-    }
+    );
-
+    rk = List.from(newRk2);
-    return null;
+    cks = List.from(newRk2);
   }
 
-  Future<void> _skipMessageKeys(int until) async {
+  /// Skip (and keep track of) message keys until our receive counter is
+  /// equal to [until]. If we would skip too many messages, returns
+  /// a [SkippingTooManyKeysError]. If not, returns null.
+  Future<OmemoError?> _skipMessageKeys(int until) async {
     if (nr + maxSkip < until) {
-      throw SkippingTooManyMessagesException();
+      return SkippingTooManyKeysError();
     }
 
     if (ckr != null) {
@@ -310,121 +208,119 @@ class OmemoDoubleRatchet {
         nr++;
       }
     }
+
+    return null;
   }
 
-  Future<void> _dhRatchet(OMEMOMessage header) async {
+  /// Decrypt [ciphertext] using keys derived from the message key [mk]. Also computes the
-    pn = ns;
+  /// HMAC from the [OMEMOMessage] embedded in [message].
-    ns = 0;
+  /// 
-    nr = 0;
+  /// If the computed HMAC does not match the HMAC in [message], returns
-    dhr = OmemoPublicKey.fromBytes(header.dhPub, KeyPairType.x25519);
+  /// [InvalidMessageHMACError]. If it matches, returns the decrypted
+  /// payload.
+  Future<Result<OmemoError, List<int>>> _decrypt(OMEMOAuthenticatedMessage message, List<int> ciphertext, List<int> mk) async {
+    final keys = await deriveEncryptionKeys(mk, encryptHkdfInfoString);
 
-    final newRk = await kdfRk(rk, await omemoDH(dhs, dhr!, 0));
+    final hmacInput = concat([sessionAd, message.message]);
-    rk = List.from(newRk);
+    final hmacResult = await truncatedHmac(hmacInput, keys.authenticationKey);
-    ckr = List.from(newRk);
+    if (!listsEqual(hmacResult, message.mac)) {
-    dhs = await OmemoKeyPair.generateNewPair(KeyPairType.x25519);
+      return Result(InvalidMessageHMACError());
-    final newNewRk = await kdfRk(rk, await omemoDH(dhs, dhr!, 0));
-    rk = List.from(newNewRk);
-    cks = List.from(newNewRk);
-  }
-
-  /// Encrypt [plaintext] using the Double Ratchet.
-  Future<RatchetStep> ratchetEncrypt(List<int> plaintext) async {
-    final newCks = await kdfCk(cks!, kdfCkNextChainKey);
-    final mk = await kdfCk(cks!, kdfCkNextMessageKey);
-
-    cks = newCks;
-    final header = OMEMOMessage()
-      ..dhPub = await dhs.pk.getBytes()
-      ..pn = pn
-      ..n = ns;
-
-    ns++;
-
-    return RatchetStep(
-      header,
-      await encrypt(
-        mk,
-        plaintext,
-        concat([sessionAd, header.writeToBuffer()]),
-        sessionAd,
-      ),
-    );
-  }
-
-  /// Decrypt a [ciphertext] that was sent with the header [header] using the Double
-  /// Ratchet. Returns the decrypted (raw) plaintext.
-  ///
-  /// Throws an SkippingTooManyMessagesException if too many messages were to be skipped.
-  Future<List<int>> ratchetDecrypt(
-    OMEMOMessage header,
-    List<int> ciphertext,
-  ) async {
-    // Check if we skipped too many messages
-    final plaintext = await _trySkippedMessageKeys(header, ciphertext);
-    if (plaintext != null) {
-      return plaintext;
     }
 
-    final dhPubMatches = listsEqual(
+    final plaintext = await aes256CbcDecrypt(ciphertext, keys.encryptionKey, keys.iv);
-      header.dhPub,
+    return Result(plaintext);
-      (await dhr?.getBytes()) ?? <int>[],
+  }
+
+  /// Checks whether we could decrypt the payload in [header] with a skipped key. If yes,
+  /// attempts to decrypt it. If not, returns null.
+  /// 
+  /// If the decryption is successful, returns the plaintext payload. If an error occurs, like
+  /// an [InvalidMessageHMACError], that is returned instead.
+  Future<Result<OmemoError, List<int>?>> _trySkippedMessageKeys(OMEMOAuthenticatedMessage message, OMEMOMessage header) async {
+    final key = SkippedKey(
+      OmemoPublicKey.fromBytes(header.dhPub, KeyPairType.x25519),
+      header.n,
     );
-    if (!dhPubMatches) {
+    if (mkSkipped.containsKey(key)) {
-      await _skipMessageKeys(header.pn);
+      final mk = mkSkipped[key]!;
+      mkSkipped.remove(key);
+
+      return _decrypt(message, header.ciphertext, mk);
+    }
+
+    return const Result(null);
+  }
+
+  /// Decrypt the payload (deeply) embedded in [message].
+  /// 
+  /// If everything goes well, returns the plaintext payload. If an error occurs, that
+  /// is returned instead.
+  Future<Result<OmemoError, List<int>>> ratchetDecrypt(OMEMOAuthenticatedMessage message) async {
+    final header = OMEMOMessage.fromBuffer(message.message);
+
+    // Try skipped keys
+    final plaintextRaw = await _trySkippedMessageKeys(message, header);
+    if (plaintextRaw.isType<OmemoError>()) {
+      // Propagate the error
+      return Result(plaintextRaw.get<OmemoError>());
+    }
+
+    final plaintext = plaintextRaw.get<List<int>?>();
+    if (plaintext != null) {
+      return Result(plaintext);
+    }
+
+    if (dhr == null || !listsEqual(header.dhPub, await dhr!.getBytes())) {
+      final skipResult1 = await _skipMessageKeys(header.pn);
+      if (skipResult1 != null) {
+        return Result(skipResult1);
+      }
+
       await _dhRatchet(header);
     }
 
-    await _skipMessageKeys(header.n);
+    final skipResult2 = await _skipMessageKeys(header.n);
-    final newCkr = await kdfCk(ckr!, kdfCkNextChainKey);
+    if (skipResult2 != null) {
+      return Result(skipResult2);
+    }
+
+    final ck = await kdfCk(ckr!, kdfCkNextChainKey);
     final mk = await kdfCk(ckr!, kdfCkNextMessageKey);
-    ckr = newCkr;
+    ckr = ck;
-    nr++;
 
-    return decrypt(
+    return _decrypt(message, header.ciphertext, mk);
-      mk,
-      ciphertext,
-      concat([sessionAd, header.writeToBuffer()]),
-      sessionAd,
-    );
   }
 
-  OmemoDoubleRatchet clone() {
+  /// Encrypt the payload [plaintext] using the double ratchet session.
-    return OmemoDoubleRatchet(
+  Future<OMEMOAuthenticatedMessage> ratchetEncrypt(List<int> plaintext) async {
-      dhs,
+    // Advance the ratchet
-      dhr,
+    final ck = await kdfCk(cks!, kdfCkNextChainKey);
-      rk,
+    final mk = await kdfCk(cks!, kdfCkNextMessageKey);
-      cks != null ? List<int>.from(cks!) : null,
+    cks = ck;
-      ckr != null ? List<int>.from(ckr!) : null,
+
-      ns,
+    // Generate encryption, authentication key and IV
-      nr,
+    final keys = await deriveEncryptionKeys(mk, encryptHkdfInfoString);
-      pn,
+    final ciphertext =
-      ik,
+        await aes256CbcEncrypt(plaintext, keys.encryptionKey, keys.iv);
-      sessionAd,
+
-      Map<SkippedKey, List<int>>.from(mkSkipped),
+    // Fill-in the header and serialize it here so we do it only once
-      acknowledged,
+    final header = OMEMOMessage()
-      kexTimestamp,
+        ..dhPub = await dhs.pk.getBytes()
-      kex,
+        ..pn = pn
-    );
+        ..n = ns
+        ..ciphertext = ciphertext;
+    final headerBytes = header.writeToBuffer();
+
+    // Increment the send counter
+    ns++;
+
+    final newAd = concat([sessionAd, headerBytes]);
+    final hmac = await truncatedHmac(newAd, keys.authenticationKey);
+    return OMEMOAuthenticatedMessage()
+      ..mac = hmac
+      ..message = headerBytes;
   }
 
-  OmemoDoubleRatchet cloneWithKex(String kex) {
-    return OmemoDoubleRatchet(
-      dhs,
-      dhr,
-      rk,
-      cks != null ? List<int>.from(cks!) : null,
-      ckr != null ? List<int>.from(ckr!) : null,
-      ns,
-      nr,
-      pn,
-      ik,
-      sessionAd,
-      Map<SkippedKey, List<int>>.from(mkSkipped),
-      acknowledged,
-      kexTimestamp,
-      kex,
-    );
-  }
 
   @visibleForTesting
   Future<bool> equals(OmemoDoubleRatchet other) async {

lib/src/double_ratchet/kdf.dart

@@ -8,7 +8,7 @@ const kdfRkInfoString = 'OMEMO Root Chain';
 const kdfCkNextMessageKey = 0x01;
 const kdfCkNextChainKey = 0x02;
 
-/// Signals KDF_CK function as specified by OMEMO 0.8.0.
+/// Signals KDF_CK function as specified by OMEMO 0.8.3.
 Future<List<int>> kdfCk(List<int> ck, int constant) async {
   final hkdf = Hkdf(hmac: Hmac(Sha256()), outputLength: 32);
   final result = await hkdf.deriveKey(
@@ -19,7 +19,7 @@ Future<List<int>> kdfCk(List<int> ck, int constant) async {
   return result.extractBytes();
 }
 
-/// Signals KDF_RK function as specified by OMEMO 0.8.0.
+/// Signals KDF_RK function as specified by OMEMO 0.8.3.
 Future<List<int>> kdfRk(List<int> rk, List<int> dhOut) async {
   final algorithm = Hkdf(
     hmac: Hmac(Sha256()),

lib/src/errors.dart

@@ -1,61 +1,47 @@
-abstract class OmemoException {}
+abstract class OmemoError {}
 
 /// Triggered during X3DH if the signature if the SPK does verify to the actual SPK.
-class InvalidSignatureException extends OmemoException implements Exception {
+class InvalidSignatureException extends OmemoError implements Exception {
   String errMsg() =>
       'The signature of the SPK does not match the provided signature';
 }
 
 /// Triggered by the Double Ratchet if the computed HMAC does not match the attached HMAC.
-/// Triggered by the Session Manager if the computed HMAC does not match the attached HMAC.
+class InvalidMessageHMACError extends OmemoError {}
-class InvalidMessageHMACException extends OmemoException implements Exception {
-  String errMsg() => 'The computed HMAC does not match the provided HMAC';
-}
 
 /// Triggered by the Double Ratchet if skipping messages would cause skipping more than
 /// MAXSKIP messages
-class SkippingTooManyMessagesException extends OmemoException
+class SkippingTooManyKeysError extends OmemoError {}
-    implements Exception {
-  String errMsg() => 'Skipping messages would cause a skip bigger than MAXSKIP';
-}
 
 /// Triggered by the Session Manager if the message key is not encrypted for the device.
-class NotEncryptedForDeviceException extends OmemoException
+class NotEncryptedForDeviceError extends OmemoError {}
-    implements Exception {
-  String errMsg() => 'Not encrypted for this device';
-}
 
 /// Triggered by the Session Manager when there is no key for decrypting the message.
-class NoDecryptionKeyException extends OmemoException implements Exception {
+class NoDecryptionKeyException extends OmemoError implements Exception {
   String errMsg() => 'No key available for decrypting the message';
 }
 
 /// Triggered by the Session Manager when the identifier of the used Signed Prekey
 /// is neither the current SPK's identifier nor the old one's.
-class UnknownSignedPrekeyException extends OmemoException implements Exception {
+class UnknownSignedPrekeyError extends OmemoError {}
-  String errMsg() => 'Unknown Signed Prekey used.';
-}
 
 /// Triggered by the Session Manager when the received Key Exchange message does not meet
 /// the requirement that a key exchange, given that the ratchet already exists, must be
 /// sent after its creation.
-class InvalidKeyExchangeException extends OmemoException implements Exception {
+class InvalidKeyExchangeException extends OmemoError implements Exception {
   String errMsg() => 'The key exchange was sent before the last kex finished';
 }
 
-/// Triggered by the Session Manager when a message's sequence number is smaller than we
-/// expect it to be.
-class MessageAlreadyDecryptedException extends OmemoException
-    implements Exception {
-  String errMsg() => 'The message has already been decrypted';
-}
-
 /// Triggered by the OmemoManager when we could not encrypt a message as we have
 /// no key material available. That happens, for example, when we want to create a
 /// ratchet session with a JID we had no session with but fetching the device bundle
 /// failed.
-class NoKeyMaterialAvailableException extends OmemoException
+class NoKeyMaterialAvailableException extends OmemoError
     implements Exception {
   String errMsg() =>
       'No key material available to create a ratchet session with';
 }
+
+/// A non-key-exchange message was received that was encrypted for our device, but we have no ratchet with
+/// the device that sent the message.
+class NoSessionWithDeviceError extends OmemoError {}

lib/src/omemo/decryption_result.dart

@@ -5,5 +5,5 @@ import 'package:omemo_dart/src/errors.dart';
 class DecryptionResult {
   const DecryptionResult(this.payload, this.error);
   final String? payload;
-  final OmemoException? error;
+  final OmemoError? error;
 }

lib/src/omemo/encryption_result.dart

@@ -20,10 +20,10 @@ class EncryptionResult {
   final List<EncryptedKey> encryptedKeys;
 
   /// Mapping of a ratchet map keys to a possible exception.
-  final Map<RatchetMapKey, OmemoException> deviceEncryptionErrors;
+  final Map<RatchetMapKey, OmemoError> deviceEncryptionErrors;
 
   /// Mapping of a JID to a possible exception.
-  final Map<String, OmemoException> jidEncryptionErrors;
+  final Map<String, OmemoError> jidEncryptionErrors;
 
   /// True if the encryption was a success. This means that we could encrypt for
   /// at least one ratchet.

lib/src/omemo/omemo.dart

@@ -0,0 +1,362 @@
+import 'dart:async';
+import 'dart:collection';
+import 'dart:convert';
+import 'package:collection/collection.dart';
+import 'package:cryptography/cryptography.dart';
+import 'package:hex/hex.dart';
+import 'package:logging/logging.dart';
+import 'package:meta/meta.dart';
+import 'package:omemo_dart/src/common/result.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/constants.dart';
+import 'package:omemo_dart/src/omemo/decryption_result.dart';
+import 'package:omemo_dart/src/omemo/device.dart';
+import 'package:omemo_dart/src/omemo/encrypted_key.dart';
+import 'package:omemo_dart/src/omemo/encryption_result.dart';
+import 'package:omemo_dart/src/omemo/events.dart';
+import 'package:omemo_dart/src/omemo/fingerprint.dart';
+import 'package:omemo_dart/src/omemo/ratchet_map_key.dart';
+import 'package:omemo_dart/src/omemo/stanza.dart';
+import 'package:omemo_dart/src/protobuf/schema.pb.dart';
+import 'package:omemo_dart/src/trust/base.dart';
+import 'package:omemo_dart/src/x3dh/x3dh.dart';
+import 'package:synchronized/synchronized.dart';
+
+class _InternalDecryptionResult {
+  const _InternalDecryptionResult(
+    this.ratchetCreated,
+    this.ratchetReplaced,
+    this.payload,
+  ) : assert(
+          !ratchetCreated || !ratchetReplaced,
+          'Ratchet must be either replaced or created',
+        );
+  final bool ratchetCreated;
+  final bool ratchetReplaced;
+  final String? payload;
+}
+
+extension AppendToListOrCreateExtension<K, V> on Map<K, List<V>> {
+  void appendOrCreate(K key, V value) {
+    if (containsKey(key)) {
+      this[key]!.add(value);
+    } else {
+      this[key] = [value];
+    }
+  }
+}
+
+extension StringFromBase64Extension on String {
+  List<int> fromBase64() => base64Decode(this);
+}
+
+class OmemoManager {
+  OmemoManager(
+    this._device,
+    this._trustManager,
+    this.sendEmptyOmemoMessageImpl,
+    this.fetchDeviceListImpl,
+    this.fetchDeviceBundleImpl,
+    this.subscribeToDeviceListNodeImpl,
+  );
+
+  final Logger _log = Logger('OmemoManager');
+
+  /// Functions for connecting with the OMEMO library
+
+  /// Send an empty OMEMO:2 message using the encrypted payload @result to
+  /// @recipientJid.
+  final Future<void> Function(EncryptionResult result, String recipientJid)
+      sendEmptyOmemoMessageImpl;
+
+  /// Fetch the list of device ids associated with @jid. If the device list cannot be
+  /// fetched, return null.
+  final Future<List<int>?> Function(String jid) fetchDeviceListImpl;
+
+  /// Fetch the device bundle for the device with id @id of jid. If it cannot be fetched, return null.
+  final Future<OmemoBundle?> Function(String jid, int id) fetchDeviceBundleImpl;
+
+  /// Subscribe to the device list PEP node of @jid.
+  final Future<void> Function(String jid) subscribeToDeviceListNodeImpl;
+
+  /// Map bare JID to its known devices
+  Map<String, List<int>> _deviceList = {};
+
+  /// Map bare JIDs to whether we already requested the device list once
+  final Map<String, bool> _deviceListRequested = {};
+
+  /// Map bare a ratchet key to its ratchet. Note that this is also locked by
+  /// _ratchetCriticalSectionLock.
+  Map<RatchetMapKey, OmemoDoubleRatchet> _ratchetMap = {};
+
+  /// Map bare JID to whether we already tried to subscribe to the device list node.
+  final Map<String, bool> _subscriptionMap = {};
+
+  /// For preventing a race condition in encryption/decryption
+  final Map<String, Queue<Completer<void>>> _ratchetCriticalSectionQueue = {};
+  final Lock _ratchetCriticalSectionLock = Lock();
+
+  /// The OmemoManager's trust management
+  final TrustManager _trustManager;
+  TrustManager get trustManager => _trustManager;
+
+  /// Our own keys...
+  final Lock _deviceLock = Lock();
+  // ignore: prefer_final_fields
+  OmemoDevice _device;
+
+  /// The event bus of the session manager
+  final StreamController<OmemoEvent> _eventStreamController =
+      StreamController<OmemoEvent>.broadcast();
+  Stream<OmemoEvent> get eventStream => _eventStreamController.stream;
+
+  /// Enter the critical section for performing cryptographic operations on the ratchets
+  Future<void> _enterRatchetCriticalSection(String jid) async {
+    final completer = await _ratchetCriticalSectionLock.synchronized(() {
+      if (_ratchetCriticalSectionQueue.containsKey(jid)) {
+        final c = Completer<void>();
+        _ratchetCriticalSectionQueue[jid]!.addLast(c);
+        return c;
+      }
+
+      _ratchetCriticalSectionQueue[jid] = Queue();
+      return null;
+    });
+
+    if (completer != null) {
+      await completer.future;
+    }
+  }
+
+  /// Leave the critical section for the ratchets.
+  Future<void> _leaveRatchetCriticalSection(String jid) async {
+    await _ratchetCriticalSectionLock.synchronized(() {
+      if (_ratchetCriticalSectionQueue.containsKey(jid)) {
+        if (_ratchetCriticalSectionQueue[jid]!.isEmpty) {
+          _ratchetCriticalSectionQueue.remove(jid);
+        } else {
+          _ratchetCriticalSectionQueue[jid]!.removeFirst().complete();
+        }
+      }
+    });
+  }
+
+  Future<Result<OmemoError, String?>> _decryptAndVerifyHmac(
+    List<int>? ciphertext,
+    List<int> keyAndHmac,
+  ) async {
+    // Empty OMEMO messages should just have the key decrypted and/or session set up.
+    if (ciphertext == null) {
+      return const Result(null);
+    }
+
+    final key = keyAndHmac.sublist(0, 32);
+    final hmac = keyAndHmac.sublist(32, 48);
+    final derivedKeys = await deriveEncryptionKeys(key, omemoPayloadInfoString);
+    final computedHmac =
+        await truncatedHmac(ciphertext, derivedKeys.authenticationKey);
+    if (!listsEqual(hmac, computedHmac)) {
+      return Result(InvalidMessageHMACError());
+    }
+
+    // TODO: Handle an exception from the crypto implementation
+    return Result(
+      utf8.decode(
+        await aes256CbcDecrypt(
+          ciphertext,
+          derivedKeys.encryptionKey,
+          derivedKeys.iv,
+        ),
+      ),
+    );
+  }
+
+  /// 
+  Future<DecryptionResult> onIncomingStanza(OmemoIncomingStanza stanza) async {
+    // NOTE: We do this so that we cannot forget to acquire and free the critical
+    //       section.
+    await _enterRatchetCriticalSection(stanza.bareSenderJid);
+    final result = await _onIncomingStanzaImpl(stanza);
+    await _leaveRatchetCriticalSection(stanza.bareSenderJid);
+
+    return result;
+  }
+
+  Future<DecryptionResult> _onIncomingStanzaImpl(OmemoIncomingStanza stanza) async {
+    // Find the correct key for our device
+    final deviceId = await getDeviceId();
+    final key = stanza.keys.firstWhereOrNull((key) => key.rid == deviceId);
+    if (key == null) {
+      return DecryptionResult(
+        null,
+        NotEncryptedForDeviceError(),
+      );
+    }
+
+    final ratchetKey = RatchetMapKey(stanza.bareSenderJid, stanza.senderDeviceId);
+    if (key.kex) {
+      final kexMessage = OMEMOKeyExchange.fromBuffer(base64Decode(key.value));
+
+      // TODO: Check if we already have such a session and if we can build it
+      // See XEP-0384 4.3
+
+      // Find the correct SPK
+      final device = await getDevice();
+      OmemoKeyPair spk;
+      if (kexMessage.spkId == device.spkId) {
+        spk = device.spk;
+      } else if (kexMessage.spkId == device.oldSpkId) {
+        spk = device.oldSpk!;
+      } else {
+        return DecryptionResult(
+          null,
+          UnknownSignedPrekeyError(),
+        );
+      }
+
+      // Build the new ratchet session
+      final kexIk = OmemoPublicKey.fromBytes(
+        kexMessage.ik,
+        KeyPairType.ed25519,
+      );
+      final kex = await x3dhFromInitialMessage(
+        X3DHMessage(
+          kexIk, 
+          OmemoPublicKey.fromBytes(
+            kexMessage.ek,
+            KeyPairType.ed25519,
+          ),
+          kexMessage.pkId,
+        ),
+        spk,
+        device.opks[kexMessage.pkId]!,
+        device.ik,
+      );
+      final ratchet = await OmemoDoubleRatchet.acceptNewSession(
+        spk,
+        kexIk,
+        kex.sk,
+        kex.ad,
+        getTimestamp(),
+      );
+
+      final keyAndHmac = await ratchet.ratchetDecrypt(
+        kexMessage.message,
+      );
+      if (keyAndHmac.isType<OmemoError>()) {
+        final error = keyAndHmac.get<OmemoError>();
+        _log.warning('Failed to decrypt symmetric key: $error');
+
+        return DecryptionResult(null, error);
+      }
+
+      final result = await _decryptAndVerifyHmac(
+        stanza.payload != null ? base64Decode(stanza.payload!) : null,
+        keyAndHmac.get<List<int>>(),
+      );
+      if (result.isType<OmemoError>()) {
+        final error = result.get<OmemoError>();
+        _log.warning('Decrypting payload failed: $error');
+
+        return DecryptionResult(
+          null,
+          error,
+        );
+      }
+
+      // Notify the trust manager
+      await trustManager.onNewSession(
+        stanza.bareSenderJid,
+        stanza.senderDeviceId,
+      );
+
+      // Commit the ratchet
+      _ratchetMap[ratchetKey] = ratchet;
+      _deviceList.appendOrCreate(stanza.bareSenderJid, stanza.senderDeviceId);
+      _eventStreamController.add(
+        RatchetModifiedEvent(
+          stanza.bareSenderJid,
+          stanza.senderDeviceId,
+          ratchet,
+          true,
+          false,
+        ),
+      );
+
+      // Replace the OPK if we're not doing a catchup.
+      if (!stanza.isCatchup) {
+        await _deviceLock.synchronized(() async {
+          await _device.replaceOnetimePrekey(kexMessage.pkId);
+
+          _eventStreamController.add(
+            DeviceModifiedEvent(_device),
+          );
+        });
+      }
+
+      return DecryptionResult(
+        result.get<String?>(),
+        null,
+      );
+    } else {
+      // Check if we even have a ratchet
+      final ratchet = _ratchetMap[ratchetKey];
+      if (ratchet == null) {
+        // TODO: Build a session with the device
+
+        return DecryptionResult(
+          null,
+          NoSessionWithDeviceError(),
+        );
+      }
+
+      final authMessage = OMEMOAuthenticatedMessage.fromBuffer(base64Decode(key.value));
+      final keyAndHmac = await ratchet.ratchetDecrypt(authMessage);
+      if (keyAndHmac.isType<OmemoError>()) {
+        final error = keyAndHmac.get<OmemoError>();
+        _log.warning('Failed to decrypt symmetric key: $error');
+        return DecryptionResult(null, error);
+      }
+
+      final result = await _decryptAndVerifyHmac(
+        stanza.payload?.fromBase64(),
+        keyAndHmac.get<List<int>>(),
+      );
+      if (result.isType<OmemoError>()) {
+        final error = result.get<OmemoError>();
+        _log.warning('Failed to decrypt message: $error');
+        return DecryptionResult(
+          null,
+          error,
+        );
+      }
+
+      // Message was successfully decrypted, so commit the ratchet
+      _eventStreamController.add(
+        RatchetModifiedEvent(
+          stanza.bareSenderJid,
+          stanza.senderDeviceId,
+          ratchet,
+          false,
+          false,
+        ),
+      );
+
+      return DecryptionResult(
+        result.get<String?>(),
+        null,
+      );
+    }
+  }
+
+  /// Returns the device used for encryption and decryption.
+  Future<OmemoDevice> getDevice() => _deviceLock.synchronized(() => _device);
+
+  /// Returns the id of the device used for encryption and decryption.
+  Future<int> getDeviceId() async => (await getDevice()).id;
+}

lib/src/omemo/omemomanager.dart

@@ -1,852 +0,0 @@
-import 'dart:async';
-import 'dart:collection';
-import 'dart:convert';
-import 'package:collection/collection.dart';
-import 'package:cryptography/cryptography.dart';
-import 'package:hex/hex.dart';
-import 'package:logging/logging.dart';
-import 'package:meta/meta.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/constants.dart';
-import 'package:omemo_dart/src/omemo/decryption_result.dart';
-import 'package:omemo_dart/src/omemo/device.dart';
-import 'package:omemo_dart/src/omemo/encrypted_key.dart';
-import 'package:omemo_dart/src/omemo/encryption_result.dart';
-import 'package:omemo_dart/src/omemo/events.dart';
-import 'package:omemo_dart/src/omemo/fingerprint.dart';
-import 'package:omemo_dart/src/omemo/ratchet_map_key.dart';
-import 'package:omemo_dart/src/omemo/stanza.dart';
-import 'package:omemo_dart/src/trust/base.dart';
-import 'package:omemo_dart/src/x3dh/x3dh.dart';
-import 'package:synchronized/synchronized.dart';
-
-class _InternalDecryptionResult {
-  const _InternalDecryptionResult(
-    this.ratchetCreated,
-    this.ratchetReplaced,
-    this.payload,
-  ) : assert(
-          !ratchetCreated || !ratchetReplaced,
-          'Ratchet must be either replaced or created',
-        );
-  final bool ratchetCreated;
-  final bool ratchetReplaced;
-  final String? payload;
-}
-
-class OmemoManager {
-  OmemoManager(
-    this._device,
-    this._trustManager,
-    this.sendEmptyOmemoMessageImpl,
-    this.fetchDeviceListImpl,
-    this.fetchDeviceBundleImpl,
-    this.subscribeToDeviceListNodeImpl,
-  );
-
-  final Logger _log = Logger('OmemoManager');
-
-  /// Functions for connecting with the OMEMO library
-
-  /// Send an empty OMEMO:2 message using the encrypted payload @result to
-  /// @recipientJid.
-  final Future<void> Function(EncryptionResult result, String recipientJid)
-      sendEmptyOmemoMessageImpl;
-
-  /// Fetch the list of device ids associated with @jid. If the device list cannot be
-  /// fetched, return null.
-  final Future<List<int>?> Function(String jid) fetchDeviceListImpl;
-
-  /// Fetch the device bundle for the device with id @id of jid. If it cannot be fetched, return null.
-  final Future<OmemoBundle?> Function(String jid, int id) fetchDeviceBundleImpl;
-
-  /// Subscribe to the device list PEP node of @jid.
-  final Future<void> Function(String jid) subscribeToDeviceListNodeImpl;
-
-  /// Map bare JID to its known devices
-  Map<String, List<int>> _deviceList = {};
-
-  /// Map bare JIDs to whether we already requested the device list once
-  final Map<String, bool> _deviceListRequested = {};
-
-  /// Map bare a ratchet key to its ratchet. Note that this is also locked by
-  /// _ratchetCriticalSectionLock.
-  Map<RatchetMapKey, OmemoDoubleRatchet> _ratchetMap = {};
-
-  /// Map bare JID to whether we already tried to subscribe to the device list node.
-  final Map<String, bool> _subscriptionMap = {};
-
-  /// For preventing a race condition in encryption/decryption
-  final Map<String, Queue<Completer<void>>> _ratchetCriticalSectionQueue = {};
-  final Lock _ratchetCriticalSectionLock = Lock();
-
-  /// The OmemoManager's trust management
-  final TrustManager _trustManager;
-  TrustManager get trustManager => _trustManager;
-
-  /// Our own keys...
-  final Lock _deviceLock = Lock();
-  // ignore: prefer_final_fields
-  OmemoDevice _device;
-
-  /// The event bus of the session manager
-  final StreamController<OmemoEvent> _eventStreamController =
-      StreamController<OmemoEvent>.broadcast();
-  Stream<OmemoEvent> get eventStream => _eventStreamController.stream;
-
-  /// Enter the critical section for performing cryptographic operations on the ratchets
-  Future<void> _enterRatchetCriticalSection(String jid) async {
-    final completer = await _ratchetCriticalSectionLock.synchronized(() {
-      if (_ratchetCriticalSectionQueue.containsKey(jid)) {
-        final c = Completer<void>();
-        _ratchetCriticalSectionQueue[jid]!.addLast(c);
-        return c;
-      }
-
-      _ratchetCriticalSectionQueue[jid] = Queue();
-      return null;
-    });
-
-    if (completer != null) {
-      await completer.future;
-    }
-  }
-
-  /// Leave the critical section for the ratchets.
-  Future<void> _leaveRatchetCriticalSection(String jid) async {
-    await _ratchetCriticalSectionLock.synchronized(() {
-      if (_ratchetCriticalSectionQueue.containsKey(jid)) {
-        if (_ratchetCriticalSectionQueue[jid]!.isEmpty) {
-          _ratchetCriticalSectionQueue.remove(jid);
-        } else {
-          _ratchetCriticalSectionQueue[jid]!.removeFirst().complete();
-        }
-      }
-    });
-  }
-
-  Future<String?> _decryptAndVerifyHmac(
-    List<int>? ciphertext,
-    List<int> keyAndHmac,
-  ) async {
-    // 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 hmac = keyAndHmac.sublist(32, 48);
-    final derivedKeys = await deriveEncryptionKeys(key, omemoPayloadInfoString);
-    final computedHmac =
-        await truncatedHmac(ciphertext, derivedKeys.authenticationKey);
-    if (!listsEqual(hmac, computedHmac)) {
-      throw InvalidMessageHMACException();
-    }
-
-    return utf8.decode(
-      await aes256CbcDecrypt(
-        ciphertext,
-        derivedKeys.encryptionKey,
-        derivedKeys.iv,
-      ),
-    );
-  }
-
-  /// Add a session [ratchet] with the [deviceId] to the internal tracking state.
-  /// NOTE: Must be called from within the ratchet critical section.
-  void _addSession(String jid, int deviceId, OmemoDoubleRatchet ratchet) {
-    // Add the bundle Id
-    if (!_deviceList.containsKey(jid)) {
-      _deviceList[jid] = [deviceId];
-
-      // Commit the device map
-      _eventStreamController.add(DeviceListModifiedEvent(_deviceList));
-    } else {
-      // Prevent having the same device multiple times in the list
-      if (!_deviceList[jid]!.contains(deviceId)) {
-        _deviceList[jid]!.add(deviceId);
-
-        // Commit the device map
-        _eventStreamController.add(DeviceListModifiedEvent(_deviceList));
-      }
-    }
-
-    // Add the ratchet session
-    final key = RatchetMapKey(jid, deviceId);
-    _ratchetMap[key] = ratchet;
-
-    // Commit the ratchet
-    _eventStreamController
-        .add(RatchetModifiedEvent(jid, deviceId, ratchet, true, false));
-  }
-
-  /// Build a new session with the user at [jid] with the device [deviceId] using data
-  /// from the key exchange [kex]. In case [kex] contains an unknown Signed Prekey
-  /// identifier an UnknownSignedPrekeyException will be thrown.
-  Future<OmemoDoubleRatchet> _addSessionFromKeyExchange(
-    String jid,
-    int deviceId,
-    OMEMOKeyExchange kex,
-  ) async {
-    // Pick the correct SPK
-    final device = await getDevice();
-    OmemoKeyPair spk;
-    if (kex.spkId == _device.spkId) {
-      spk = _device.spk;
-    } else if (kex.spkId == _device.oldSpkId) {
-      spk = _device.oldSpk!;
-    } else {
-      throw UnknownSignedPrekeyException();
-    }
-
-    final kexResult = await x3dhFromInitialMessage(
-      X3DHMessage(
-        OmemoPublicKey.fromBytes(kex.ik, KeyPairType.ed25519),
-        OmemoPublicKey.fromBytes(kex.ek, KeyPairType.x25519),
-        kex.pkId,
-      ),
-      spk,
-      device.opks.values.elementAt(kex.pkId),
-      device.ik,
-    );
-    final ratchet = await OmemoDoubleRatchet.acceptNewSession(
-      spk,
-      OmemoPublicKey.fromBytes(kex.ik, KeyPairType.ed25519),
-      kexResult.sk,
-      kexResult.ad,
-      getTimestamp(),
-    );
-
-    // Notify the trust manager
-    await trustManager.onNewSession(jid, deviceId);
-
-    return ratchet;
-  }
-
-  /// Create a ratchet session initiated by Alice to the user with Jid [jid] and the device
-  /// [deviceId] from the bundle [bundle].
-  @visibleForTesting
-  Future<OMEMOKeyExchange> addSessionFromBundle(
-    String jid,
-    int deviceId,
-    OmemoBundle bundle,
-  ) async {
-    final device = await getDevice();
-    final kexResult = await x3dhFromBundle(
-      bundle,
-      device.ik,
-    );
-    final ratchet = await OmemoDoubleRatchet.initiateNewSession(
-      bundle.spk,
-      bundle.ik,
-      kexResult.sk,
-      kexResult.ad,
-      getTimestamp(),
-    );
-
-    await _trustManager.onNewSession(jid, deviceId);
-    _addSession(jid, deviceId, ratchet);
-
-    return OMEMOKeyExchange()
-      ..pkId = kexResult.opkId
-      ..spkId = bundle.spkId
-      ..ik = await device.ik.pk.getBytes()
-      ..ek = await kexResult.ek.pk.getBytes();
-  }
-
-  /// 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].
-  /// NOTE: Must be called from within the ratchet critical section
-  void _restoreRatchet(RatchetMapKey mapKey, OmemoDoubleRatchet oldRatchet) {
-    _log.finest(
-      'Restoring ratchet ${mapKey.jid}:${mapKey.deviceId} to ${oldRatchet.nr}',
-    );
-    _ratchetMap[mapKey] = oldRatchet;
-
-    // Commit the ratchet
-    _eventStreamController.add(
-      RatchetModifiedEvent(
-        mapKey.jid,
-        mapKey.deviceId,
-        oldRatchet,
-        false,
-        false,
-      ),
-    );
-  }
-
-  /// 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
-  /// <encrypted /> element.
-  /// [timestamp] refers to the time the message was sent. This might be either what the
-  /// server tells you via "XEP-0203: Delayed Delivery" or the point in time at which
-  /// you received the stanza, if no Delayed Delivery element was found.
-  ///
-  /// If the received message is an empty OMEMO message, i.e. there is no <payload />
-  /// element, then [ciphertext] must be set to null. In this case, this function
-  /// will return null as there is no message to be decrypted. This, however, is used
-  /// to set up sessions or advance the ratchets.
-  Future<_InternalDecryptionResult> _decryptMessage(
-    List<int>? ciphertext,
-    String senderJid,
-    int senderDeviceId,
-    List<EncryptedKey> keys,
-    int timestamp,
-  ) async {
-    // Try to find a session we can decrypt with.
-    var device = await getDevice();
-    final rawKey = keys.firstWhereOrNull((key) => key.rid == device.id);
-    if (rawKey == null) {
-      throw NotEncryptedForDeviceException();
-    }
-
-    final decodedRawKey = base64.decode(rawKey.value);
-    List<int>? keyAndHmac;
-    OMEMOAuthenticatedMessage authMessage;
-    OMEMOMessage? message;
-
-    // If the ratchet already existed, we store it. If it didn't, oldRatchet will stay
-    // null.
-    final ratchetKey = RatchetMapKey(senderJid, senderDeviceId);
-    final oldRatchet = getRatchet(ratchetKey)?.clone();
-    if (rawKey.kex) {
-      final kex = OMEMOKeyExchange.fromBuffer(decodedRawKey);
-      authMessage = kex.message;
-      message = OMEMOMessage.fromBuffer(authMessage.message);
-
-      // Guard against old key exchanges
-      if (oldRatchet != null) {
-        _log.finest(
-          'KEX for existent ratchet ${ratchetKey.toJsonKey()}. ${oldRatchet.kexTimestamp} > $timestamp: ${oldRatchet.kexTimestamp > timestamp}',
-        );
-        if (oldRatchet.kexTimestamp > timestamp) {
-          throw InvalidKeyExchangeException();
-        }
-      }
-
-      final r =
-          await _addSessionFromKeyExchange(senderJid, senderDeviceId, kex);
-
-      // Try to decrypt with the new ratchet r
-      try {
-        keyAndHmac =
-            await r.ratchetDecrypt(message, authMessage.writeToBuffer());
-        final result = await _decryptAndVerifyHmac(ciphertext, keyAndHmac);
-
-        // Add the new ratchet
-        _addSession(senderJid, senderDeviceId, r);
-
-        // Replace the OPK
-        await _deviceLock.synchronized(() async {
-          device = await device.replaceOnetimePrekey(kex.pkId);
-
-          // Commit the device
-          _eventStreamController.add(DeviceModifiedEvent(device));
-        });
-
-        // Commit the ratchet
-        _eventStreamController.add(
-          RatchetModifiedEvent(
-            senderJid,
-            senderDeviceId,
-            r,
-            oldRatchet == null,
-            oldRatchet != null,
-          ),
-        );
-
-        return _InternalDecryptionResult(
-          oldRatchet == null,
-          oldRatchet != null,
-          result,
-        );
-      } catch (ex) {
-        _log.finest('Kex failed due to $ex. Not proceeding with kex.');
-      }
-    } else {
-      authMessage = OMEMOAuthenticatedMessage.fromBuffer(decodedRawKey);
-      message = OMEMOMessage.fromBuffer(authMessage.message);
-    }
-
-    final devices = _deviceList[senderJid];
-    if (devices?.contains(senderDeviceId) != true) {
-      throw NoDecryptionKeyException();
-    }
-
-    // TODO(PapaTutuWawa): When receiving a message that is not an OMEMOKeyExchange from a device there is no session with, clients SHOULD create a session with that device and notify it about the new session by responding with an empty OMEMO message as per Sending a message.
-
-    // We can guarantee that the ratchet exists at this point in time
-    final ratchet = getRatchet(ratchetKey)!;
-
-    try {
-      if (rawKey.kex) {
-        keyAndHmac =
-            await ratchet.ratchetDecrypt(message, authMessage.writeToBuffer());
-      } else {
-        keyAndHmac = await ratchet.ratchetDecrypt(message, decodedRawKey);
-      }
-    } catch (_) {
-      _restoreRatchet(ratchetKey, oldRatchet!);
-      rethrow;
-    }
-
-    // Commit the ratchet
-    _eventStreamController.add(
-      RatchetModifiedEvent(
-        senderJid,
-        senderDeviceId,
-        ratchet,
-        false,
-        false,
-      ),
-    );
-
-    try {
-      return _InternalDecryptionResult(
-        false,
-        false,
-        await _decryptAndVerifyHmac(ciphertext, keyAndHmac),
-      );
-    } catch (_) {
-      _restoreRatchet(ratchetKey, oldRatchet!);
-      rethrow;
-    }
-  }
-
-  /// Returns, if it exists, the ratchet associated with [key].
-  /// NOTE: Must be called from within the ratchet critical section.
-  @visibleForTesting
-  OmemoDoubleRatchet? getRatchet(RatchetMapKey key) => _ratchetMap[key];
-
-  /// Figure out what bundles we have to still build a session with.
-  Future<List<OmemoBundle>> _fetchNewBundles(String jid) async {
-    // Check if we already requested the device list for [jid]
-    List<int> bundlesToFetch;
-    if (!_deviceListRequested.containsKey(jid) ||
-        !_deviceList.containsKey(jid)) {
-      // We don't have an up-to-date version of the device list
-      final newDeviceList = await fetchDeviceListImpl(jid);
-      if (newDeviceList == null) return [];
-
-      _deviceList[jid] = newDeviceList;
-      bundlesToFetch = newDeviceList.where((id) {
-        return !_ratchetMap.containsKey(RatchetMapKey(jid, id)) ||
-            _deviceList[jid]?.contains(id) == false;
-      }).toList();
-
-      // Trigger an event with the new device list
-      _eventStreamController.add(DeviceListModifiedEvent(_deviceList));
-    } else {
-      // We already have an up-to-date version of the device list
-      bundlesToFetch = _deviceList[jid]!
-          .where((id) => !_ratchetMap.containsKey(RatchetMapKey(jid, id)))
-          .toList();
-    }
-
-    if (bundlesToFetch.isNotEmpty) {
-      _log.finest('Fetching bundles $bundlesToFetch for $jid');
-    }
-
-    final device = await getDevice();
-    final newBundles = List<OmemoBundle>.empty(growable: true);
-    for (final id in bundlesToFetch) {
-      if (jid == device.jid && id == device.id) continue;
-
-      final bundle = await fetchDeviceBundleImpl(jid, id);
-      if (bundle != null) newBundles.add(bundle);
-    }
-
-    return newBundles;
-  }
-
-  /// Encrypt the key [plaintext] for all known bundles of the Jids in [jids]. Returns a
-  /// map that maps the device Id to the ciphertext of [plaintext].
-  ///
-  /// If [plaintext] is null, then the result will be an empty OMEMO message, i.e. one that
-  /// does not contain a <payload /> element. This means that the ciphertext attribute of
-  /// the result will be null as well.
-  /// NOTE: Must be called within the ratchet critical section
-  Future<EncryptionResult> _encryptToJids(
-    List<String> jids,
-    String? plaintext,
-  ) async {
-    final encryptedKeys = List<EncryptedKey>.empty(growable: true);
-
-    var ciphertext = const <int>[];
-    var keyPayload = const <int>[];
-    if (plaintext != null) {
-      // Generate the key and encrypt the plaintext
-      final key = generateRandomBytes(32);
-      final keys = await deriveEncryptionKeys(key, omemoPayloadInfoString);
-      ciphertext = await aes256CbcEncrypt(
-        utf8.encode(plaintext),
-        keys.encryptionKey,
-        keys.iv,
-      );
-      final hmac = await truncatedHmac(ciphertext, keys.authenticationKey);
-      keyPayload = concat([key, hmac]);
-    } else {
-      keyPayload = List<int>.filled(32, 0x0);
-    }
-
-    final kex = <RatchetMapKey, OMEMOKeyExchange>{};
-    for (final jid in jids) {
-      for (final newSession in await _fetchNewBundles(jid)) {
-        kex[RatchetMapKey(jid, newSession.id)] = await addSessionFromBundle(
-          newSession.jid,
-          newSession.id,
-          newSession,
-        );
-      }
-    }
-
-    // We assume that the user already checked if the session exists
-    final deviceEncryptionErrors = <RatchetMapKey, OmemoException>{};
-    final jidEncryptionErrors = <String, OmemoException>{};
-    for (final jid in jids) {
-      final devices = _deviceList[jid];
-      if (devices == null) {
-        _log.severe('Device list does not exist for $jid.');
-        jidEncryptionErrors[jid] = NoKeyMaterialAvailableException();
-        continue;
-      }
-
-      if (!_subscriptionMap.containsKey(jid)) {
-        unawaited(subscribeToDeviceListNodeImpl(jid));
-        _subscriptionMap[jid] = true;
-      }
-
-      for (final deviceId in devices) {
-        // Empty OMEMO messages are allowed to bypass trust
-        if (plaintext != null) {
-          // Only encrypt to devices that are trusted
-          if (!(await _trustManager.isTrusted(jid, deviceId))) continue;
-
-          // Only encrypt to devices that are enabled
-          if (!(await _trustManager.isEnabled(jid, deviceId))) continue;
-        }
-
-        final ratchetKey = RatchetMapKey(jid, deviceId);
-        var ratchet = _ratchetMap[ratchetKey];
-        if (ratchet == null) {
-          _log.severe('Ratchet ${ratchetKey.toJsonKey()} does not exist.');
-          deviceEncryptionErrors[ratchetKey] =
-              NoKeyMaterialAvailableException();
-          continue;
-        }
-
-        final ciphertext =
-            (await ratchet.ratchetEncrypt(keyPayload)).ciphertext;
-
-        if (kex.containsKey(ratchetKey)) {
-          // The ratchet did not exist
-          final k = kex[ratchetKey]!
-            ..message = OMEMOAuthenticatedMessage.fromBuffer(ciphertext);
-          final buffer = base64.encode(k.writeToBuffer());
-          encryptedKeys.add(
-            EncryptedKey(
-              jid,
-              deviceId,
-              buffer,
-              true,
-            ),
-          );
-
-          ratchet = ratchet.cloneWithKex(buffer);
-          _ratchetMap[ratchetKey] = ratchet;
-        } else if (!ratchet.acknowledged) {
-          // The ratchet exists but is not acked
-          if (ratchet.kex != null) {
-            final oldKex =
-                OMEMOKeyExchange.fromBuffer(base64.decode(ratchet.kex!))
-                  ..message = OMEMOAuthenticatedMessage.fromBuffer(ciphertext);
-
-            encryptedKeys.add(
-              EncryptedKey(
-                jid,
-                deviceId,
-                base64.encode(oldKex.writeToBuffer()),
-                true,
-              ),
-            );
-          } else {
-            // The ratchet is not acked but we don't have the old key exchange
-            _log.warning(
-              'Ratchet for $jid:$deviceId is not acked but the kex attribute is null',
-            );
-            encryptedKeys.add(
-              EncryptedKey(
-                jid,
-                deviceId,
-                base64.encode(ciphertext),
-                false,
-              ),
-            );
-          }
-        } else {
-          // The ratchet exists and is acked
-          encryptedKeys.add(
-            EncryptedKey(
-              jid,
-              deviceId,
-              base64.encode(ciphertext),
-              false,
-            ),
-          );
-        }
-
-        // Commit the ratchet
-        _eventStreamController
-            .add(RatchetModifiedEvent(jid, deviceId, ratchet, false, false));
-      }
-    }
-
-    return EncryptionResult(
-      plaintext != null ? ciphertext : null,
-      encryptedKeys,
-      deviceEncryptionErrors,
-      jidEncryptionErrors,
-    );
-  }
-
-  /// Call when receiving an OMEMO:2 encrypted stanza. Will handle everything and
-  /// decrypt it.
-  Future<DecryptionResult> onIncomingStanza(OmemoIncomingStanza stanza) async {
-    await _enterRatchetCriticalSection(stanza.bareSenderJid);
-
-    if (!_subscriptionMap.containsKey(stanza.bareSenderJid)) {
-      unawaited(subscribeToDeviceListNodeImpl(stanza.bareSenderJid));
-      _subscriptionMap[stanza.bareSenderJid] = true;
-    }
-
-    final ratchetKey =
-        RatchetMapKey(stanza.bareSenderJid, stanza.senderDeviceId);
-    final _InternalDecryptionResult result;
-    try {
-      result = await _decryptMessage(
-        stanza.payload != null ? base64.decode(stanza.payload!) : null,
-        stanza.bareSenderJid,
-        stanza.senderDeviceId,
-        stanza.keys,
-        stanza.timestamp,
-      );
-    } on OmemoException catch (ex) {
-      await _leaveRatchetCriticalSection(stanza.bareSenderJid);
-      return DecryptionResult(
-        null,
-        ex,
-      );
-    }
-
-    // Check if the ratchet is acked
-    final ratchet = getRatchet(ratchetKey);
-    assert(
-      ratchet != null,
-      'We decrypted the message, so the ratchet must exist',
-    );
-
-    if (ratchet!.acknowledged) {
-      // Ratchet is acknowledged
-      if (ratchet.nr > 53 || result.ratchetCreated || result.ratchetReplaced) {
-        await sendEmptyOmemoMessageImpl(
-          await _encryptToJids(
-            [stanza.bareSenderJid],
-            null,
-          ),
-          stanza.bareSenderJid,
-        );
-      }
-
-      // Ratchet is acked
-      await _leaveRatchetCriticalSection(stanza.bareSenderJid);
-      return DecryptionResult(
-        result.payload,
-        null,
-      );
-    } else {
-      // Ratchet is not acked.
-      // Mark as acked and send an empty OMEMO message.
-      await ratchetAcknowledged(
-        stanza.bareSenderJid,
-        stanza.senderDeviceId,
-        enterCriticalSection: false,
-      );
-      await sendEmptyOmemoMessageImpl(
-        await _encryptToJids(
-          [stanza.bareSenderJid],
-          null,
-        ),
-        stanza.bareSenderJid,
-      );
-
-      await _leaveRatchetCriticalSection(stanza.bareSenderJid);
-      return DecryptionResult(
-        result.payload,
-        null,
-      );
-    }
-  }
-
-  /// Call when sending out an encrypted stanza. Will handle everything and
-  /// encrypt it.
-  Future<EncryptionResult> onOutgoingStanza(OmemoOutgoingStanza stanza) async {
-    _log.finest('Waiting to enter critical section');
-    await _enterRatchetCriticalSection(stanza.recipientJids.first);
-    _log.finest('Entered critical section');
-
-    final result = _encryptToJids(
-      stanza.recipientJids,
-      stanza.payload,
-    );
-
-    await _leaveRatchetCriticalSection(stanza.recipientJids.first);
-
-    return result;
-  }
-
-  // Sends a hearbeat message as specified by XEP-0384 to [jid].
-  Future<void> sendOmemoHeartbeat(String jid) async {
-    // TODO(Unknown): Include some error handling
-    final result = await _encryptToJids(
-      [jid],
-      null,
-    );
-    await sendEmptyOmemoMessageImpl(result, jid);
-  }
-
-  /// Mark the ratchet for device [deviceId] from [jid] as acked.
-  Future<void> ratchetAcknowledged(
-    String jid,
-    int deviceId, {
-    bool enterCriticalSection = true,
-  }) async {
-    if (enterCriticalSection) await _enterRatchetCriticalSection(jid);
-
-    final key = RatchetMapKey(jid, deviceId);
-    if (_ratchetMap.containsKey(key)) {
-      final ratchet = _ratchetMap[key]!..acknowledged = true;
-
-      // Commit it
-      _eventStreamController
-          .add(RatchetModifiedEvent(jid, deviceId, ratchet, false, false));
-    } else {
-      _log.severe(
-        'Attempted to acknowledge ratchet ${key.toJsonKey()}, even though it does not exist',
-      );
-    }
-
-    if (enterCriticalSection) await _leaveRatchetCriticalSection(jid);
-  }
-
-  /// Generates an entirely new device. May be useful when the user wants to reset their cryptographic
-  /// identity. Triggers an event to commit it to storage.
-  Future<void> regenerateDevice() async {
-    await _deviceLock.synchronized(() async {
-      _device = await OmemoDevice.generateNewDevice(_device.jid);
-
-      // Commit it
-      _eventStreamController.add(DeviceModifiedEvent(_device));
-    });
-  }
-
-  /// Returns the device used for encryption and decryption.
-  Future<OmemoDevice> getDevice() => _deviceLock.synchronized(() => _device);
-
-  /// Returns the id of the device used for encryption and decryption.
-  Future<int> getDeviceId() async => (await getDevice()).id;
-
-  /// Directly aquire the current device as a OMEMO device bundle.
-  Future<OmemoBundle> getDeviceBundle() async => (await getDevice()).toBundle();
-
-  /// Directly aquire the current device's fingerprint.
-  Future<String> getDeviceFingerprint() async =>
-      (await getDevice()).getFingerprint();
-
-  /// Returns the fingerprints for all devices of [jid] that we have a session with.
-  /// If there are not sessions with [jid], then returns null.
-  Future<List<DeviceFingerprint>?> getFingerprintsForJid(String jid) async {
-    if (!_deviceList.containsKey(jid)) return null;
-
-    await _enterRatchetCriticalSection(jid);
-
-    final fingerprintKeys = _deviceList[jid]!
-        .map((id) => RatchetMapKey(jid, id))
-        .where((key) => _ratchetMap.containsKey(key));
-
-    final fingerprints = List<DeviceFingerprint>.empty(growable: true);
-    for (final key in fingerprintKeys) {
-      final curveKey = await _ratchetMap[key]!.ik.toCurve25519();
-      fingerprints.add(
-        DeviceFingerprint(
-          key.deviceId,
-          HEX.encode(await curveKey.getBytes()),
-        ),
-      );
-    }
-
-    await _leaveRatchetCriticalSection(jid);
-    return fingerprints;
-  }
-
-  /// Ensures that the device list is fetched again on the next message sending.
-  void onNewConnection() {
-    _deviceListRequested.clear();
-    _subscriptionMap.clear();
-  }
-
-  /// Sets the device list for [jid] to [devices]. Triggers a DeviceListModifiedEvent.
-  void onDeviceListUpdate(String jid, List<int> devices) {
-    _deviceList[jid] = devices;
-    _deviceListRequested[jid] = true;
-
-    // Trigger an event
-    _eventStreamController.add(DeviceListModifiedEvent(_deviceList));
-  }
-
-  void initialize(
-    Map<RatchetMapKey, OmemoDoubleRatchet> ratchetMap,
-    Map<String, List<int>> deviceList,
-  ) {
-    _deviceList = deviceList;
-    _ratchetMap = ratchetMap;
-  }
-
-  /// Removes all ratchets for JID [jid]. This also removes all trust decisions for
-  /// [jid] from the trust manager. This function triggers a RatchetRemovedEvent for
-  /// every removed ratchet and a DeviceListModifiedEvent afterwards. Behaviour for
-  /// the trust manager is dependent on its implementation.
-  Future<void> removeAllRatchets(String jid) async {
-    await _enterRatchetCriticalSection(jid);
-
-    for (final deviceId in _deviceList[jid]!) {
-      // Remove the ratchet and commit it
-      _ratchetMap.remove(RatchetMapKey(jid, deviceId));
-      _eventStreamController.add(RatchetRemovedEvent(jid, deviceId));
-    }
-
-    // Remove the devices from the device list cache and commit it
-    _deviceList.remove(jid);
-    _deviceListRequested.remove(jid);
-    _eventStreamController.add(DeviceListModifiedEvent(_deviceList));
-
-    // Remove trust decisions
-    await _trustManager.removeTrustDecisionsForJid(jid);
-
-    await _leaveRatchetCriticalSection(jid);
-  }
-
-  /// Replaces the internal device with [newDevice]. Does not trigger an event.
-  Future<void> replaceDevice(OmemoDevice newDevice) async {
-    await _deviceLock.synchronized(() {
-      _device = newDevice;
-    });
-  }
-}

lib/src/omemo/stanza.dart

@@ -8,6 +8,7 @@ class OmemoIncomingStanza {
     this.timestamp,
     this.keys,
     this.payload,
+    this.isCatchup,
   );
 
   /// The bare JID of the sender of the stanza.
@@ -19,11 +20,14 @@ class OmemoIncomingStanza {
   /// The timestamp when the stanza was received.
   final int timestamp;
 
-  /// The included encrypted keys
+  /// The included encrypted keys for our own JID
   final List<EncryptedKey> keys;
 
   /// The string payload included in the <encrypted /> element.
   final String? payload;
+
+  /// Flag indicating whether the message was received due to a catchup.
+  final bool isCatchup;
 }
 
 /// Describes a stanza that is to be sent out

test/double_ratchet_test.dart

@@ -2,38 +2,10 @@
 import 'dart:convert';
 import 'package:cryptography/cryptography.dart';
 import 'package:omemo_dart/omemo_dart.dart';
-import 'package:omemo_dart/protobuf/schema.pb.dart';
+import 'package:omemo_dart/src/protobuf/schema.pb.dart';
-import 'package:omemo_dart/src/double_ratchet/crypto.dart';
 import 'package:test/test.dart';
 
 void main() {
-  test('Test encrypting and decrypting', () async {
-    final sessionAd = List<int>.filled(32, 0x0);
-    final mk = List<int>.filled(32, 0x1);
-    final plaintext = utf8.encode('Hallo');
-    final header = OMEMOMessage()
-      ..n = 0
-      ..pn = 0
-      ..dhPub = List<int>.empty();
-    final asd = concat([sessionAd, header.writeToBuffer()]);
-
-    final ciphertext = await encrypt(
-      mk,
-      plaintext,
-      asd,
-      sessionAd,
-    );
-
-    final decrypted = await decrypt(
-      mk,
-      ciphertext,
-      asd,
-      sessionAd,
-    );
-
-    expect(decrypted, plaintext);
-  });
-
   test('Test the Double Ratchet', () async {
     // Generate keys
     const bobJid = 'bob@other.example.server';
@@ -81,6 +53,7 @@ void main() {
     final alicesRatchet = await OmemoDoubleRatchet.initiateNewSession(
       spkBob.pk,
       ikBob.pk,
+      resultAlice.ek.pk,
       resultAlice.sk,
       resultAlice.ad,
       0,
@@ -98,6 +71,7 @@ void main() {
     for (var i = 0; i < 100; i++) {
       final messageText = 'Hello, dear $i';
 
+      print('${i + 1}/100');
       if (i.isEven) {
         // Alice encrypts a message
         final aliceRatchetResult =
@@ -109,12 +83,12 @@ void main() {
 
         // Bob tries to decrypt it
         final bobRatchetResult = await bobsRatchet.ratchetDecrypt(
-          aliceRatchetResult.header,
+          aliceRatchetResult,
-          aliceRatchetResult.ciphertext,
         );
         print('Bob decrypted the message');
 
-        expect(utf8.encode(messageText), bobRatchetResult);
+        expect(bobRatchetResult.isType<List<int>>(), true);
+        expect(bobRatchetResult.get<List<int>>(), utf8.encode(messageText));
       } else {
         // Bob sends a message to Alice
         final bobRatchetResult =
@@ -126,12 +100,13 @@ void main() {
 
         // Alice tries to decrypt it
         final aliceRatchetResult = await alicesRatchet.ratchetDecrypt(
-          bobRatchetResult.header,
+          bobRatchetResult,
-          bobRatchetResult.ciphertext,
         );
         print('Alice decrypted the message');
 
-        expect(utf8.encode(messageText), aliceRatchetResult);
+        expect(aliceRatchetResult.isType<List<int>>(), true);
+        expect(aliceRatchetResult.get<List<int>>(), utf8.encode(messageText));
+        expect(utf8.encode(messageText), aliceRatchetResult.get<List<int>>());
       }
     }
   });