docs: Migrate the example to the "new" OmemoManager

Fixes #30.
style: Formattiing issues
2023-06-12 19:37:51 +02:00 · 2023-06-12 19:37:45 +02:00 · 2023-06-12 19:21:07 +02:00 · 2023-06-12 19:20:43 +02:00 · 2023-06-12 19:13:38 +02:00 · 2023-01-22 19:13:56 +01:00
49 changed files with 4347 additions and 752 deletions
--- a/.github/FUNDING.yml
+++ b/.github/FUNDING.yml
@@ -0,0 +1 @@
+ko_fi: papatutuwawa
--- a/.gitignore
+++ b/.gitignore
@@ -12,6 +12,3 @@ pubspec.lock
 # NixOS
 .direnv
 .envrc
-
-# Protobuf build artifacts
-lib/protobuf/*.dart
--- a/.gitlint
+++ b/.gitlint
@@ -5,6 +5,6 @@ line-length=72
 [title-trailing-punctuation]
 [title-hard-tab]
 [title-match-regex]
-regex=^(feat|fix|test|release|chore|security|docs|refactor):.*$
+regex=^(feat|fix|test|release|chore|security|docs|refactor|style):.*$
 [body-trailing-whitespace]
 [body-first-line-empty]
--- a/.pubignore
+++ b/.pubignore
@@ -0,0 +1 @@
+lib/protobuf
--- a/.woodpecker.yml
+++ b/.woodpecker.yml
@@ -0,0 +1,11 @@
+pipeline:
+  lint:
+    image: dart:2.18.1
+    commands:
+      - dart pub get
+      - dart analyze --fatal-infos --fatal-warnings
+  test:
+    image: dart:2.18.1
+    commands:
+      - dart pub get
+      - dart test
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -4,3 +4,53 @@
 - Implement the Double Ratchet, X3DH and OMEMO specific bits
 - Add a Blind-Trust-Before-Verification TrustManager
 - Supported OMEMO version: 0.8.3
+
+## 0.1.3
+
+- Fix bug with the Double Ratchet causing only the initial message to be decryptable
+- Expose `getDeviceMap` as a developer usable function
+## 0.2.0
+
+- Add convenience functions `getDeviceId` and `getDeviceBundle`
+- Creating a new ratchet with an id for which we already have a ratchet will now overwrite the old ratchet
+- Ratchet now carry an "acknowledged" attribute
+
+## 0.2.1
+
+- Add `isRatchetAcknowledged`
+- Ratchets that are created due to accepting a kex are now unacknowledged
+
+## 0.3.0
+
+- Implement enabling and disabling ratchets via the TrustManager interface
+- Fix deserialization of the various objects
+- Remove the BTBV TrustManager's loadState method. Just use the constructor
+- Allow removing all ratchets for a given Jid
+- If an error occurs while decrypting the message, the ratchet will now be reset to its prior state
+- Fix a bug within the Varint encoding function. This should fix some occasional UnknownSignedPrekeyExceptions
+- Remove OmemoSessionManager's toJson and fromJson. Use toJsonWithoutSessions and fromJsonWithoutSessions. Restoring sessions is not out-of-scope for that function
+
+## 0.3.1
+
+- Fix a bug that caused the device's id to change when replacing a OPK
+- Every decryption failure now causes the ratchet to be restored to a pre-decryption state
+- Add method to get the device's fingerprint
+
+## 0.4.0
+
+- Deprecate `OmemoSessionManager`. Use `OmemoManager` instead.
+- Implement queued access to the ratchets inside the `OmemoManager`.
+- Implement heartbeat messages.
+- [BREAKING] Rename `Device` to `OmemoDevice`.
+
+## 0.4.1
+
+- Fix fetching the current device and building a ratchet session with it when encrypting for our own JID
+
+## 0.4.2
+
+- Fix removeAllRatchets not removing, well, all ratchets. In fact, it did not remove any ratchet.
+
+## 0.4.3
+
+- Fix bug that causes ratchets to be unable to decrypt anything after receiving a heartbeat with a completely new session
--- a/README.md
+++ b/README.md
@@ -26,10 +26,10 @@ Include `omemo_dart` in your `pubspec.yaml` like this:
 # [...]

 dependencies:
-	omemo_dart:
-		hosted: https://git.polynom.me/api/packages/PapaTutuWawa/pub
-		version: ^0.1.0
-	# [...]
+  omemo_dart:
+    hosted: https://git.polynom.me/api/packages/PapaTutuWawa/pub
+    version: ^0.4.3
+  # [...]

 # [...]
 ```
@@ -56,3 +56,9 @@ messages' formatting.
 Licensed under the MIT license.

 See `LICENSE`.
+
+## Support
+
+If you like what I do and you want to support me, feel free to donate to me on Ko-Fi.
+
+[<img src="https://codeberg.org/moxxy/moxxyv2/raw/branch/master/assets/repo/kofi.png" height="36" style="height: 36px; border: 0px;"></img>](https://ko-fi.com/papatutuwawa)
--- a/analysis_options.yaml
+++ b/analysis_options.yaml
@@ -10,3 +10,7 @@ linter:
 analyzer:
  exclude:
    - "lib/protobuf/*.dart"
+    # TODO: Remove once OmemoSessionManager is gone
+    - "test/omemo_test.dart"
+    - "example/omemo_dart_example.dart"
+    - "test/serialisation_test.dart"
--- a/example/omemo_dart_example.dart
+++ b/example/omemo_dart_example.dart
@@ -8,28 +8,45 @@ void main() async {
  const aliceJid = 'alice@some.server';
  const bobJid = 'bob@other.serve';

-  // You are Alice and want to begin using OMEMO, so you first create a SessionManager
-  final aliceSession = await OmemoSessionManager.generateNewIdentity(
-    // The bare Jid of Alice as a String
-    aliceJid,
+  // You are Alice and want to begin using OMEMO, so you first create an OmemoManager.
+  final aliceManager = OmemoManager(
+    // Generate Alice's OMEMO device bundle. We can specify how many One-time Prekeys we want, but
+    // per default, omemo_dart generates 100 (recommended by XEP-0384).
+    await OmemoDevice.generateNewDevice(aliceJid),
    // The trust manager we want to use. In this case, we use the provided one that
    // implements "Blind Trust Before Verification". To make things simpler, we keep
    // no persistent data and can thus use the MemoryBTBVTrustManager. If we wanted to keep
    // the state, we would have to override BlindTrustBeforeVerificationTrustManager.
    MemoryBTBVTrustManager(),
-    // Here we specify how many Onetime Prekeys we want to have. XEP-0384 recommends around
-    // 100 OPKs, so let's generate 100. The parameter defaults to 100.
-    //opkAmount: 100,
+    // This function is called whenever we need to send an OMEMO heartbeat to [recipient].
+    // [result] is the encryted data to include. This needs to be wired into your XMPP library's
+    // OMEMO implementation.
+    // For simplicity, we use an empty function and imagine it works.
+    (result, recipient) async => {},
+    // This function is called whenever we need to fetch the device list for [jid].
+    // This needs to be wired into your XMPP library's OMEMO implementation.
+    // For simplicity, we use an empty function and imagine it works.
+    (jid) async => [],
+    // This function is called whenever we need to fetch the device bundle with id [id] from [jid].
+    // This needs to be wired into your XMPP library's OMEMO implementation.
+    // For simplicity, we use an empty function and imagine it works.
+    (jid, id) async => null,
+    // This function is called whenever we need to subscribe to [jid]'s device list PubSub node.
+    // This needs to be wired into your XMPP library's OMEMO implementation.
+    // For simplicity, we use an empty function and imagine it works.
+    (jid) async {},
  );

  // Alice now wants to chat with Bob at his bare Jid "bob@other.server". To make things
  // simple, we just generate the identity bundle ourselves. In the real world, we would
  // request it using PEP and then convert the device bundle into a OmemoBundle object.
-  final bobSession = await OmemoSessionManager.generateNewIdentity(
-    bobJid,
+  final bobManager = OmemoManager(
+    await OmemoDevice.generateNewDevice(bobJid),
    MemoryBTBVTrustManager(),
-    // Just for illustrative purposes
-    opkAmount: 1,
+    (result, recipient) async => {},
+    (jid) async => [],
+    (jid, id) async => null,
+    (jid) async {},
  );

  // Alice prepares to send the message to Bob, so she builds the message stanza and
@@ -56,19 +73,18 @@ void main() async {

  // Since Alice has no open session with Bob, we need to tell the session manager to build
  // it when sending the message.
-  final message = await aliceSession.encryptToJid(
-    // The bare receiver Jid
-    bobJid,
-    // The envelope we want to encrypt
-    envelope,
-    // Since this is the first time Alice contacts Bob from this device, we need to create
-    // a new session. Let's also assume that Bob only has one device. We may, however,
-    // add more bundles to newSessions, if we know of more.
-    newSessions: [
-      await (await bobSession.getDevice()).toBundle(),
-    ],
+  final message = await aliceSession.onOutgoingStanza(
+    OmemoOutgoingStanza(
+      // The bare receiver Jid
+      [bobJid],
+
+      // The payload we want to encrypt, i.e. the envelope.
+      envelope,
+    ),
  );

+  // In a proper implementation, we should also do some error checking here.
+
  // Alice now builds the actual message stanza for Bob
  final payload = base64.encode(message.ciphertext!);
  final aliceDevice = await aliceSession.getDevice();
@@ -110,15 +126,26 @@ void main() async {

  // Bob extracts the payload and attempts to decrypt it.
  // ignore: unused_local_variable
-  final bobMessage = await bobSession.decryptMessage(
-    // base64 decode the payload
-    base64.decode(payload),
-    // Specify the Jid of the sender
-    aliceJid,
-    // Specify the device identifier of the sender (the "sid" attribute of <header />)
-    aliceDevice.id,
-    // The deserialised keys
-    keys,
+  final bobMessage = await bobManager.onIncomingStanza(
+    OmemoIncomingStanza(
+      // The bare sender JID of the message. In this case, it's Alice's.
+      aliceJid,
+
+      // The 'sid' attribute of the <header /> element. Here, we know that Alice only has one device.
+      aliceDevice.id,
+
+      // Time the message was sent. Since the message was not delayed, we use the
+      // current time.
+      DateTime.now().millisecondsSinceEpoch,
+
+      /// The decoded <key /> elements. from the header. Note that we only include the ones
+      /// relevant for Bob, so all children of <keys jid='$bobJid' />.
+      keys,
+
+      /// The text of the <payload /> element, if it exists. If not, then the message might be
+      /// a hearbeat, where no payload is sent. In that case, use null.
+      payload,
+    ),
  );

  // All Bob has to do now is replace the OMEMO wrapper element
--- a/lib/omemo_dart.dart
+++ b/lib/omemo_dart.dart
@@ -10,8 +10,9 @@ 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/ratchet_map_key.dart';
-export 'src/omemo/sessionmanager.dart';
+export 'src/omemo/stanza.dart';
 export 'src/trust/base.dart';
 export 'src/trust/btbv.dart';
 export 'src/x3dh/x3dh.dart';
--- a/lib/protobuf/schema.pb.dart
+++ b/lib/protobuf/schema.pb.dart
@@ -0,0 +1,379 @@
+///
+//  Generated code. Do not modify.
+//  source: schema.proto
+//
+// @dart = 2.12
+// ignore_for_file: annotate_overrides,camel_case_types,constant_identifier_names,directives_ordering,library_prefixes,non_constant_identifier_names,prefer_final_fields,return_of_invalid_type,unnecessary_const,unnecessary_import,unnecessary_this,unused_import,unused_shown_name
+
+import 'dart:core' as $core;
+
+import 'package:protobuf/protobuf.dart' as $pb;
+
+class OMEMOMessage extends $pb.GeneratedMessage {
+  static final $pb.BuilderInfo _i = $pb.BuilderInfo(
+      const $core.bool.fromEnvironment('protobuf.omit_message_names')
+          ? ''
+          : 'OMEMOMessage',
+      createEmptyInstance: create)
+    ..a<$core.int>(
+        1,
+        const $core.bool.fromEnvironment('protobuf.omit_field_names')
+            ? ''
+            : 'n',
+        $pb.PbFieldType.QU3)
+    ..a<$core.int>(
+        2,
+        const $core.bool.fromEnvironment('protobuf.omit_field_names')
+            ? ''
+            : 'pn',
+        $pb.PbFieldType.QU3)
+    ..a<$core.List<$core.int>>(
+        3,
+        const $core.bool.fromEnvironment('protobuf.omit_field_names')
+            ? ''
+            : 'dhPub',
+        $pb.PbFieldType.QY)
+    ..a<$core.List<$core.int>>(
+        4,
+        const $core.bool.fromEnvironment('protobuf.omit_field_names')
+            ? ''
+            : 'ciphertext',
+        $pb.PbFieldType.OY);
+
+  OMEMOMessage._() : super();
+  factory OMEMOMessage({
+    $core.int? n,
+    $core.int? pn,
+    $core.List<$core.int>? dhPub,
+    $core.List<$core.int>? ciphertext,
+  }) {
+    final _result = create();
+    if (n != null) {
+      _result.n = n;
+    }
+    if (pn != null) {
+      _result.pn = pn;
+    }
+    if (dhPub != null) {
+      _result.dhPub = dhPub;
+    }
+    if (ciphertext != null) {
+      _result.ciphertext = ciphertext;
+    }
+    return _result;
+  }
+  factory OMEMOMessage.fromBuffer($core.List<$core.int> i,
+          [$pb.ExtensionRegistry r = $pb.ExtensionRegistry.EMPTY]) =>
+      create()..mergeFromBuffer(i, r);
+  factory OMEMOMessage.fromJson($core.String i,
+          [$pb.ExtensionRegistry r = $pb.ExtensionRegistry.EMPTY]) =>
+      create()..mergeFromJson(i, r);
+  @$core.Deprecated('Using this can add significant overhead to your binary. '
+      'Use [GeneratedMessageGenericExtensions.deepCopy] instead. '
+      'Will be removed in next major version')
+  OMEMOMessage clone() => OMEMOMessage()..mergeFromMessage(this);
+  @$core.Deprecated('Using this can add significant overhead to your binary. '
+      'Use [GeneratedMessageGenericExtensions.rebuild] instead. '
+      'Will be removed in next major version')
+  OMEMOMessage copyWith(void Function(OMEMOMessage) updates) =>
+      super.copyWith((message) => updates(message as OMEMOMessage))
+          as OMEMOMessage; // ignore: deprecated_member_use
+  $pb.BuilderInfo get info_ => _i;
+  @$core.pragma('dart2js:noInline')
+  static OMEMOMessage create() => OMEMOMessage._();
+  OMEMOMessage createEmptyInstance() => create();
+  static $pb.PbList<OMEMOMessage> createRepeated() =>
+      $pb.PbList<OMEMOMessage>();
+  @$core.pragma('dart2js:noInline')
+  static OMEMOMessage getDefault() => _defaultInstance ??=
+      $pb.GeneratedMessage.$_defaultFor<OMEMOMessage>(create);
+  static OMEMOMessage? _defaultInstance;
+
+  @$pb.TagNumber(1)
+  $core.int get n => $_getIZ(0);
+  @$pb.TagNumber(1)
+  set n($core.int v) {
+    $_setUnsignedInt32(0, v);
+  }
+
+  @$pb.TagNumber(1)
+  $core.bool hasN() => $_has(0);
+  @$pb.TagNumber(1)
+  void clearN() => clearField(1);
+
+  @$pb.TagNumber(2)
+  $core.int get pn => $_getIZ(1);
+  @$pb.TagNumber(2)
+  set pn($core.int v) {
+    $_setUnsignedInt32(1, v);
+  }
+
+  @$pb.TagNumber(2)
+  $core.bool hasPn() => $_has(1);
+  @$pb.TagNumber(2)
+  void clearPn() => clearField(2);
+
+  @$pb.TagNumber(3)
+  $core.List<$core.int> get dhPub => $_getN(2);
+  @$pb.TagNumber(3)
+  set dhPub($core.List<$core.int> v) {
+    $_setBytes(2, v);
+  }
+
+  @$pb.TagNumber(3)
+  $core.bool hasDhPub() => $_has(2);
+  @$pb.TagNumber(3)
+  void clearDhPub() => clearField(3);
+
+  @$pb.TagNumber(4)
+  $core.List<$core.int> get ciphertext => $_getN(3);
+  @$pb.TagNumber(4)
+  set ciphertext($core.List<$core.int> v) {
+    $_setBytes(3, v);
+  }
+
+  @$pb.TagNumber(4)
+  $core.bool hasCiphertext() => $_has(3);
+  @$pb.TagNumber(4)
+  void clearCiphertext() => clearField(4);
+}
+
+class OMEMOAuthenticatedMessage extends $pb.GeneratedMessage {
+  static final $pb.BuilderInfo _i = $pb.BuilderInfo(
+      const $core.bool.fromEnvironment('protobuf.omit_message_names')
+          ? ''
+          : 'OMEMOAuthenticatedMessage',
+      createEmptyInstance: create)
+    ..a<$core.List<$core.int>>(
+        1,
+        const $core.bool.fromEnvironment('protobuf.omit_field_names')
+            ? ''
+            : 'mac',
+        $pb.PbFieldType.QY)
+    ..a<$core.List<$core.int>>(
+        2,
+        const $core.bool.fromEnvironment('protobuf.omit_field_names')
+            ? ''
+            : 'message',
+        $pb.PbFieldType.QY);
+
+  OMEMOAuthenticatedMessage._() : super();
+  factory OMEMOAuthenticatedMessage({
+    $core.List<$core.int>? mac,
+    $core.List<$core.int>? message,
+  }) {
+    final _result = create();
+    if (mac != null) {
+      _result.mac = mac;
+    }
+    if (message != null) {
+      _result.message = message;
+    }
+    return _result;
+  }
+  factory OMEMOAuthenticatedMessage.fromBuffer($core.List<$core.int> i,
+          [$pb.ExtensionRegistry r = $pb.ExtensionRegistry.EMPTY]) =>
+      create()..mergeFromBuffer(i, r);
+  factory OMEMOAuthenticatedMessage.fromJson($core.String i,
+          [$pb.ExtensionRegistry r = $pb.ExtensionRegistry.EMPTY]) =>
+      create()..mergeFromJson(i, r);
+  @$core.Deprecated('Using this can add significant overhead to your binary. '
+      'Use [GeneratedMessageGenericExtensions.deepCopy] instead. '
+      'Will be removed in next major version')
+  OMEMOAuthenticatedMessage clone() =>
+      OMEMOAuthenticatedMessage()..mergeFromMessage(this);
+  @$core.Deprecated('Using this can add significant overhead to your binary. '
+      'Use [GeneratedMessageGenericExtensions.rebuild] instead. '
+      'Will be removed in next major version')
+  OMEMOAuthenticatedMessage copyWith(
+          void Function(OMEMOAuthenticatedMessage) updates) =>
+      super.copyWith((message) => updates(message as OMEMOAuthenticatedMessage))
+          as OMEMOAuthenticatedMessage; // ignore: deprecated_member_use
+  $pb.BuilderInfo get info_ => _i;
+  @$core.pragma('dart2js:noInline')
+  static OMEMOAuthenticatedMessage create() => OMEMOAuthenticatedMessage._();
+  OMEMOAuthenticatedMessage createEmptyInstance() => create();
+  static $pb.PbList<OMEMOAuthenticatedMessage> createRepeated() =>
+      $pb.PbList<OMEMOAuthenticatedMessage>();
+  @$core.pragma('dart2js:noInline')
+  static OMEMOAuthenticatedMessage getDefault() => _defaultInstance ??=
+      $pb.GeneratedMessage.$_defaultFor<OMEMOAuthenticatedMessage>(create);
+  static OMEMOAuthenticatedMessage? _defaultInstance;
+
+  @$pb.TagNumber(1)
+  $core.List<$core.int> get mac => $_getN(0);
+  @$pb.TagNumber(1)
+  set mac($core.List<$core.int> v) {
+    $_setBytes(0, v);
+  }
+
+  @$pb.TagNumber(1)
+  $core.bool hasMac() => $_has(0);
+  @$pb.TagNumber(1)
+  void clearMac() => clearField(1);
+
+  @$pb.TagNumber(2)
+  $core.List<$core.int> get message => $_getN(1);
+  @$pb.TagNumber(2)
+  set message($core.List<$core.int> v) {
+    $_setBytes(1, v);
+  }
+
+  @$pb.TagNumber(2)
+  $core.bool hasMessage() => $_has(1);
+  @$pb.TagNumber(2)
+  void clearMessage() => clearField(2);
+}
+
+class OMEMOKeyExchange extends $pb.GeneratedMessage {
+  static final $pb.BuilderInfo _i = $pb.BuilderInfo(
+      const $core.bool.fromEnvironment('protobuf.omit_message_names')
+          ? ''
+          : 'OMEMOKeyExchange',
+      createEmptyInstance: create)
+    ..a<$core.int>(
+        1,
+        const $core.bool.fromEnvironment('protobuf.omit_field_names')
+            ? ''
+            : 'pkId',
+        $pb.PbFieldType.QU3)
+    ..a<$core.int>(
+        2,
+        const $core.bool.fromEnvironment('protobuf.omit_field_names')
+            ? ''
+            : 'spkId',
+        $pb.PbFieldType.QU3)
+    ..a<$core.List<$core.int>>(
+        3,
+        const $core.bool.fromEnvironment('protobuf.omit_field_names')
+            ? ''
+            : 'ik',
+        $pb.PbFieldType.QY)
+    ..a<$core.List<$core.int>>(
+        4,
+        const $core.bool.fromEnvironment('protobuf.omit_field_names')
+            ? ''
+            : 'ek',
+        $pb.PbFieldType.QY)
+    ..aQM<OMEMOAuthenticatedMessage>(
+        5,
+        const $core.bool.fromEnvironment('protobuf.omit_field_names')
+            ? ''
+            : 'message',
+        subBuilder: OMEMOAuthenticatedMessage.create);
+
+  OMEMOKeyExchange._() : super();
+  factory OMEMOKeyExchange({
+    $core.int? pkId,
+    $core.int? spkId,
+    $core.List<$core.int>? ik,
+    $core.List<$core.int>? ek,
+    OMEMOAuthenticatedMessage? message,
+  }) {
+    final _result = create();
+    if (pkId != null) {
+      _result.pkId = pkId;
+    }
+    if (spkId != null) {
+      _result.spkId = spkId;
+    }
+    if (ik != null) {
+      _result.ik = ik;
+    }
+    if (ek != null) {
+      _result.ek = ek;
+    }
+    if (message != null) {
+      _result.message = message;
+    }
+    return _result;
+  }
+  factory OMEMOKeyExchange.fromBuffer($core.List<$core.int> i,
+          [$pb.ExtensionRegistry r = $pb.ExtensionRegistry.EMPTY]) =>
+      create()..mergeFromBuffer(i, r);
+  factory OMEMOKeyExchange.fromJson($core.String i,
+          [$pb.ExtensionRegistry r = $pb.ExtensionRegistry.EMPTY]) =>
+      create()..mergeFromJson(i, r);
+  @$core.Deprecated('Using this can add significant overhead to your binary. '
+      'Use [GeneratedMessageGenericExtensions.deepCopy] instead. '
+      'Will be removed in next major version')
+  OMEMOKeyExchange clone() => OMEMOKeyExchange()..mergeFromMessage(this);
+  @$core.Deprecated('Using this can add significant overhead to your binary. '
+      'Use [GeneratedMessageGenericExtensions.rebuild] instead. '
+      'Will be removed in next major version')
+  OMEMOKeyExchange copyWith(void Function(OMEMOKeyExchange) updates) =>
+      super.copyWith((message) => updates(message as OMEMOKeyExchange))
+          as OMEMOKeyExchange; // ignore: deprecated_member_use
+  $pb.BuilderInfo get info_ => _i;
+  @$core.pragma('dart2js:noInline')
+  static OMEMOKeyExchange create() => OMEMOKeyExchange._();
+  OMEMOKeyExchange createEmptyInstance() => create();
+  static $pb.PbList<OMEMOKeyExchange> createRepeated() =>
+      $pb.PbList<OMEMOKeyExchange>();
+  @$core.pragma('dart2js:noInline')
+  static OMEMOKeyExchange getDefault() => _defaultInstance ??=
+      $pb.GeneratedMessage.$_defaultFor<OMEMOKeyExchange>(create);
+  static OMEMOKeyExchange? _defaultInstance;
+
+  @$pb.TagNumber(1)
+  $core.int get pkId => $_getIZ(0);
+  @$pb.TagNumber(1)
+  set pkId($core.int v) {
+    $_setUnsignedInt32(0, v);
+  }
+
+  @$pb.TagNumber(1)
+  $core.bool hasPkId() => $_has(0);
+  @$pb.TagNumber(1)
+  void clearPkId() => clearField(1);
+
+  @$pb.TagNumber(2)
+  $core.int get spkId => $_getIZ(1);
+  @$pb.TagNumber(2)
+  set spkId($core.int v) {
+    $_setUnsignedInt32(1, v);
+  }
+
+  @$pb.TagNumber(2)
+  $core.bool hasSpkId() => $_has(1);
+  @$pb.TagNumber(2)
+  void clearSpkId() => clearField(2);
+
+  @$pb.TagNumber(3)
+  $core.List<$core.int> get ik => $_getN(2);
+  @$pb.TagNumber(3)
+  set ik($core.List<$core.int> v) {
+    $_setBytes(2, v);
+  }
+
+  @$pb.TagNumber(3)
+  $core.bool hasIk() => $_has(2);
+  @$pb.TagNumber(3)
+  void clearIk() => clearField(3);
+
+  @$pb.TagNumber(4)
+  $core.List<$core.int> get ek => $_getN(3);
+  @$pb.TagNumber(4)
+  set ek($core.List<$core.int> v) {
+    $_setBytes(3, v);
+  }
+
+  @$pb.TagNumber(4)
+  $core.bool hasEk() => $_has(3);
+  @$pb.TagNumber(4)
+  void clearEk() => clearField(4);
+
+  @$pb.TagNumber(5)
+  OMEMOAuthenticatedMessage get message => $_getN(4);
+  @$pb.TagNumber(5)
+  set message(OMEMOAuthenticatedMessage v) {
+    setField(5, v);
+  }
+
+  @$pb.TagNumber(5)
+  $core.bool hasMessage() => $_has(4);
+  @$pb.TagNumber(5)
+  void clearMessage() => clearField(5);
+  @$pb.TagNumber(5)
+  OMEMOAuthenticatedMessage ensureMessage() => $_ensure(4);
+}
--- a/lib/protobuf/schema.pbenum.dart
+++ b/lib/protobuf/schema.pbenum.dart
@@ -0,0 +1,6 @@
+///
+//  Generated code. Do not modify.
+//  source: schema.proto
+//
+// @dart = 2.12
+// ignore_for_file: annotate_overrides,camel_case_types,constant_identifier_names,directives_ordering,library_prefixes,non_constant_identifier_names,prefer_final_fields,return_of_invalid_type,unnecessary_const,unnecessary_import,unnecessary_this,unused_import,unused_shown_name
--- a/lib/protobuf/schema.pbjson.dart
+++ b/lib/protobuf/schema.pbjson.dart
@@ -0,0 +1,60 @@
+///
+//  Generated code. Do not modify.
+//  source: schema.proto
+//
+// @dart = 2.12
+// ignore_for_file: annotate_overrides,camel_case_types,constant_identifier_names,deprecated_member_use_from_same_package,directives_ordering,library_prefixes,non_constant_identifier_names,prefer_final_fields,return_of_invalid_type,unnecessary_const,unnecessary_import,unnecessary_this,unused_import,unused_shown_name
+
+import 'dart:core' as $core;
+import 'dart:convert' as $convert;
+import 'dart:typed_data' as $typed_data;
+
+@$core.Deprecated('Use oMEMOMessageDescriptor instead')
+const OMEMOMessage$json = {
+  '1': 'OMEMOMessage',
+  '2': [
+    {'1': 'n', '3': 1, '4': 2, '5': 13, '10': 'n'},
+    {'1': 'pn', '3': 2, '4': 2, '5': 13, '10': 'pn'},
+    {'1': 'dh_pub', '3': 3, '4': 2, '5': 12, '10': 'dhPub'},
+    {'1': 'ciphertext', '3': 4, '4': 1, '5': 12, '10': 'ciphertext'},
+  ],
+};
+
+/// Descriptor for `OMEMOMessage`. Decode as a `google.protobuf.DescriptorProto`.
+final $typed_data.Uint8List oMEMOMessageDescriptor = $convert.base64Decode(
+    'CgxPTUVNT01lc3NhZ2USDAoBbhgBIAIoDVIBbhIOCgJwbhgCIAIoDVICcG4SFQoGZGhfcHViGAMgAigMUgVkaFB1YhIeCgpjaXBoZXJ0ZXh0GAQgASgMUgpjaXBoZXJ0ZXh0');
+@$core.Deprecated('Use oMEMOAuthenticatedMessageDescriptor instead')
+const OMEMOAuthenticatedMessage$json = {
+  '1': 'OMEMOAuthenticatedMessage',
+  '2': [
+    {'1': 'mac', '3': 1, '4': 2, '5': 12, '10': 'mac'},
+    {'1': 'message', '3': 2, '4': 2, '5': 12, '10': 'message'},
+  ],
+};
+
+/// Descriptor for `OMEMOAuthenticatedMessage`. Decode as a `google.protobuf.DescriptorProto`.
+final $typed_data.Uint8List oMEMOAuthenticatedMessageDescriptor =
+    $convert.base64Decode(
+        'ChlPTUVNT0F1dGhlbnRpY2F0ZWRNZXNzYWdlEhAKA21hYxgBIAIoDFIDbWFjEhgKB21lc3NhZ2UYAiACKAxSB21lc3NhZ2U=');
+@$core.Deprecated('Use oMEMOKeyExchangeDescriptor instead')
+const OMEMOKeyExchange$json = {
+  '1': 'OMEMOKeyExchange',
+  '2': [
+    {'1': 'pk_id', '3': 1, '4': 2, '5': 13, '10': 'pkId'},
+    {'1': 'spk_id', '3': 2, '4': 2, '5': 13, '10': 'spkId'},
+    {'1': 'ik', '3': 3, '4': 2, '5': 12, '10': 'ik'},
+    {'1': 'ek', '3': 4, '4': 2, '5': 12, '10': 'ek'},
+    {
+      '1': 'message',
+      '3': 5,
+      '4': 2,
+      '5': 11,
+      '6': '.OMEMOAuthenticatedMessage',
+      '10': 'message'
+    },
+  ],
+};
+
+/// Descriptor for `OMEMOKeyExchange`. Decode as a `google.protobuf.DescriptorProto`.
+final $typed_data.Uint8List oMEMOKeyExchangeDescriptor = $convert.base64Decode(
+    'ChBPTUVNT0tleUV4Y2hhbmdlEhMKBXBrX2lkGAEgAigNUgRwa0lkEhUKBnNwa19pZBgCIAIoDVIFc3BrSWQSDgoCaWsYAyACKAxSAmlrEg4KAmVrGAQgAigMUgJlaxI0CgdtZXNzYWdlGAUgAigLMhouT01FTU9BdXRoZW50aWNhdGVkTWVzc2FnZVIHbWVzc2FnZQ==');
--- a/lib/protobuf/schema.pbserver.dart
+++ b/lib/protobuf/schema.pbserver.dart
@@ -0,0 +1,8 @@
+///
+//  Generated code. Do not modify.
+//  source: schema.proto
+//
+// @dart = 2.12
+// ignore_for_file: annotate_overrides,camel_case_types,constant_identifier_names,deprecated_member_use_from_same_package,directives_ordering,library_prefixes,non_constant_identifier_names,prefer_final_fields,return_of_invalid_type,unnecessary_const,unnecessary_import,unnecessary_this,unused_import,unused_shown_name
+
+export 'schema.pb.dart';
--- a/lib/src/crypto.dart
+++ b/lib/src/crypto.dart
@@ -6,7 +6,11 @@ import 'package:omemo_dart/src/keys.dart';
 /// it indicates which of [kp] ([identityKey] == 1) or [pk] ([identityKey] == 2)
 /// is the identity key. This is needed since the identity key pair/public key is
 /// an Ed25519 key, but we need them as X25519 keys for DH.
-Future<List<int>> omemoDH(OmemoKeyPair kp, OmemoPublicKey pk, int identityKey) async {
+Future<List<int>> omemoDH(
+  OmemoKeyPair kp,
+  OmemoPublicKey pk,
+  int identityKey,
+) async {
  var ckp = kp;
  var cpk = pk;

@@ -17,15 +21,14 @@ Future<List<int>> omemoDH(OmemoKeyPair kp, OmemoPublicKey pk, int identityKey) a
  }

  final shared = await Cryptography.instance.x25519().sharedSecretKey(
-    keyPair: await ckp.asKeyPair(),
-    remotePublicKey: cpk.asPublicKey(),
-  );
+        keyPair: await ckp.asKeyPair(),
+        remotePublicKey: cpk.asPublicKey(),
+      );

  return shared.extractBytes();
 }

 class HkdfKeyResult {
-
  const HkdfKeyResult(this.encryptionKey, this.authenticationKey, this.iv);
  final List<int> encryptionKey;
  final List<int> authenticationKey;
@@ -35,7 +38,8 @@ class HkdfKeyResult {
 /// cryptography _really_ wants to check the MAC output from AES-256-CBC. Since
 /// we don't have it, we need the MAC check to always "pass".
 class NoMacSecretBox extends SecretBox {
-  NoMacSecretBox(super.cipherText, { required super.nonce }) : super(mac: Mac.empty);
+  NoMacSecretBox(super.cipherText, {required super.nonce})
+      : super(mac: Mac.empty);

  @override
  Future<void> checkMac({
@@ -60,12 +64,20 @@ Future<HkdfKeyResult> deriveEncryptionKeys(List<int> input, String info) async {
  );
  final bytes = await result.extractBytes();

-  return HkdfKeyResult(bytes.sublist(0, 32), bytes.sublist(32, 64), bytes.sublist(64, 80));
+  return HkdfKeyResult(
+    bytes.sublist(0, 32),
+    bytes.sublist(32, 64),
+    bytes.sublist(64, 80),
+  );
 }

 /// A small helper function to make AES-256-CBC easier. Encrypt [plaintext] using [key] as
 /// the encryption key and [iv] as the IV. Returns the ciphertext.
-Future<List<int>> aes256CbcEncrypt(List<int> plaintext, List<int> key, List<int> iv) async {
+Future<List<int>> aes256CbcEncrypt(
+  List<int> plaintext,
+  List<int> key,
+  List<int> iv,
+) async {
  final algorithm = AesCbc.with256bits(
    macAlgorithm: MacAlgorithm.empty,
  );
@@ -80,7 +92,11 @@ Future<List<int>> aes256CbcEncrypt(List<int> plaintext, List<int> key, List<int>

 /// A small helper function to make AES-256-CBC easier. Decrypt [ciphertext] using [key] as
 /// the encryption key and [iv] as the IV. Returns the ciphertext.
-Future<List<int>> aes256CbcDecrypt(List<int> ciphertext, List<int> key, List<int> iv) async {
+Future<List<int>> aes256CbcDecrypt(
+  List<int> ciphertext,
+  List<int> key,
+  List<int> iv,
+) async {
  final algorithm = AesCbc.with256bits(
    macAlgorithm: MacAlgorithm.empty,
  );
--- a/lib/src/double_ratchet/crypto.dart
+++ b/lib/src/double_ratchet/crypto.dart
@@ -10,13 +10,20 @@ 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 {
+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 ciphertext =
+      await aes256CbcEncrypt(plaintext, keys.encryptionKey, keys.iv);

-  final header = OmemoMessage.fromBuffer(associatedData.sublist(sessionAd.length))
-    ..ciphertext = ciphertext;
+  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);
@@ -29,7 +36,12 @@ Future<List<int>> encrypt(List<int> mk, List<int> plaintext, List<int> associate
 /// 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 {
+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);

--- a/lib/src/double_ratchet/double_ratchet.dart
+++ b/lib/src/double_ratchet/double_ratchet.dart
@@ -1,5 +1,6 @@
 import 'dart:convert';
 import 'package:cryptography/cryptography.dart';
+import 'package:hex/hex.dart';
 import 'package:meta/meta.dart';
 import 'package:omemo_dart/src/crypto.dart';
 import 'package:omemo_dart/src/double_ratchet/crypto.dart';
@@ -13,7 +14,6 @@ import 'package:omemo_dart/src/protobuf/omemo_message.dart';
 const maxSkip = 1000;

 class RatchetStep {
-
  const RatchetStep(this.header, this.ciphertext);
  final OmemoMessage header;
  final List<int> ciphertext;
@@ -21,7 +21,6 @@ class RatchetStep {

@immutable
 class SkippedKey {
-
  const SkippedKey(this.dh, this.n);

  factory SkippedKey.fromJson(Map<String, dynamic> data) {
@@ -54,19 +53,21 @@ class SkippedKey {
 }

 class OmemoDoubleRatchet {
-
  OmemoDoubleRatchet(
    this.dhs, // DHs
    this.dhr, // DHr
-    this.rk,  // RK
+    this.rk, // RK
    this.cks, // CKs
    this.ckr, // CKr
-    this.ns,  // Ns
-    this.nr,  // Nr
-    this.pn,  // Pn
+    this.ns, // Ns
+    this.nr, // Nr
+    this.pn, // Pn
    this.ik,
    this.sessionAd,
    this.mkSkipped, // MKSKIPPED
+    this.acknowledged,
+    this.kexTimestamp,
+    this.kex,
  );

  factory OmemoDoubleRatchet.fromJson(Map<String, dynamic> data) {
@@ -81,8 +82,11 @@ class OmemoDoubleRatchet {
      'ns': 0,
      'nr': 0,
      'pn': 0,
-      'ik_pub': 'base/64/encoded',
+      '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',
@@ -92,11 +96,21 @@ class OmemoDoubleRatchet {
      ]
    }
    */
-    final mkSkipped = <SkippedKey, List<int>>{};
-    for (final entry in data['mkskipped']! as List<Map<String, dynamic>>) {
-      final key = SkippedKey.fromJson(entry);
-      mkSkipped[key] = base64.decode(entry['key']! as String);
-    }
+    // 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(
@@ -117,6 +131,9 @@ class OmemoDoubleRatchet {
      ),
      base64.decode(data['session_ad']! as String),
      mkSkipped,
+      data['acknowledged']! as bool,
+      data['kex_timestamp']! as int,
+      data['kex'] as String?,
    );
  }

@@ -148,13 +165,30 @@ class OmemoDoubleRatchet {

  final Map<SkippedKey, List<int>> mkSkipped;

+  /// The point in time at which we performed the kex exchange to create this ratchet.
+  /// Precision is milliseconds since epoch.
+  int kexTimestamp;
+
+  /// The key exchange that was used for initiating the session.
+  final String? kex;
+
+  /// Indicates whether we received an empty OMEMO message after building a session with
+  /// the device.
+  bool acknowledged;
+
  /// Create an OMEMO session using the Signed Pre Key [spk], the shared secret [sk] that
  /// was obtained using a X3DH and the associated data [ad] that was also obtained through
  /// a X3DH. [ik] refers to Bob's (the receiver's) IK public key.
-  static Future<OmemoDoubleRatchet> initiateNewSession(OmemoPublicKey spk, OmemoPublicKey ik, List<int> sk, List<int> ad) async {
+  static Future<OmemoDoubleRatchet> initiateNewSession(
+    OmemoPublicKey spk,
+    OmemoPublicKey ik,
+    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, dhr, 0));
+    final rk = await kdfRk(sk, await omemoDH(dhs, dhr, 0));
    final cks = rk;

    return OmemoDoubleRatchet(
@@ -169,6 +203,9 @@ class OmemoDoubleRatchet {
      ik,
      ad,
      {},
+      false,
+      timestamp,
+      '',
    );
  }

@@ -176,7 +213,13 @@ class OmemoDoubleRatchet {
  /// Pre Key keypair [spk], the shared secret [sk] that was obtained through a X3DH and
  /// the associated data [ad] that was also obtained through a X3DH. [ik] refers to
  /// Alice's (the initiator's) IK public key.
-  static Future<OmemoDoubleRatchet> acceptNewSession(OmemoKeyPair spk, OmemoPublicKey ik, List<int> sk, List<int> ad) async {
+  static Future<OmemoDoubleRatchet> acceptNewSession(
+    OmemoKeyPair spk,
+    OmemoPublicKey ik,
+    List<int> sk,
+    List<int> ad,
+    int kexTimestamp,
+  ) async {
    return OmemoDoubleRatchet(
      spk,
      null,
@@ -189,11 +232,15 @@ class OmemoDoubleRatchet {
      ik,
      ad,
      {},
+      true,
+      kexTimestamp,
+      null,
    );
  }

  Future<Map<String, dynamic>> toJson() async {
-    final mkSkippedSerialised = List<Map<String, dynamic>>.empty(growable: true);
+    final mkSkippedSerialised =
+        List<Map<String, dynamic>>.empty(growable: true);
    for (final entry in mkSkipped.entries) {
      final result = await entry.key.toJson();
      result['key'] = base64.encode(entry.value);
@@ -214,10 +261,22 @@ class OmemoDoubleRatchet {
      'ik_pub': base64.encode(await ik.getBytes()),
      'session_ad': base64.encode(sessionAd),
      'mkskipped': mkSkippedSerialised,
+      'acknowledged': acknowledged,
+      'kex_timestamp': kexTimestamp,
+      'kex': kex,
    };
  }

-  Future<List<int>?> _trySkippedMessageKeys(OmemoMessage header, List<int> ciphertext) async {
+  /// 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!,
@@ -226,7 +285,12 @@ class OmemoDoubleRatchet {
      final mk = mkSkipped[key]!;
      mkSkipped.remove(key);

-      return decrypt(mk, ciphertext, concat([sessionAd, header.writeToBuffer()]), sessionAd);
+      return decrypt(
+        mk,
+        ciphertext,
+        concat([sessionAd, header.writeToBuffer()]),
+        sessionAd,
+      );
    }

    return null;
@@ -249,18 +313,18 @@ class OmemoDoubleRatchet {
  }

  Future<void> _dhRatchet(OmemoMessage header) async {
-    pn = header.n!;
+    pn = ns;
    ns = 0;
    nr = 0;
    dhr = OmemoPublicKey.fromBytes(header.dhPub!, KeyPairType.x25519);

    final newRk = await kdfRk(rk, await omemoDH(dhs, dhr!, 0));
-    rk = newRk;
-    ckr = newRk;
+    rk = List.from(newRk);
+    ckr = List.from(newRk);
    dhs = await OmemoKeyPair.generateNewPair(KeyPairType.x25519);
    final newNewRk = await kdfRk(rk, await omemoDH(dhs, dhr!, 0));
-    rk = newNewRk;
-    cks = newNewRk;
+    rk = List.from(newNewRk);
+    cks = List.from(newNewRk);
  }

  /// Encrypt [plaintext] using the Double Ratchet.
@@ -278,7 +342,12 @@ class OmemoDoubleRatchet {

    return RatchetStep(
      header,
-      await encrypt(mk, plaintext, concat([sessionAd, header.writeToBuffer()]), sessionAd),
+      await encrypt(
+        mk,
+        plaintext,
+        concat([sessionAd, header.writeToBuffer()]),
+        sessionAd,
+      ),
    );
  }

@@ -286,14 +355,21 @@ class OmemoDoubleRatchet {
  /// 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 {
+  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;
    }

-    if (header.dhPub != await dhr?.getBytes()) {
+    final dhPubMatches = listsEqual(
+      header.dhPub!,
+      (await dhr?.getBytes()) ?? <int>[],
+    );
+    if (!dhPubMatches) {
      await _skipMessageKeys(header.pn!);
      await _dhRatchet(header);
    }
@@ -304,15 +380,65 @@ class OmemoDoubleRatchet {
    ckr = newCkr;
    nr++;

-    return decrypt(mk, ciphertext, concat([sessionAd, header.writeToBuffer()]), sessionAd);
+    return decrypt(
+      mk,
+      ciphertext,
+      concat([sessionAd, header.writeToBuffer()]),
+      sessionAd,
+    );
+  }
+
+  OmemoDoubleRatchet clone() {
+    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,
+    );
+  }
+
+  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 {
-    // ignore: invalid_use_of_visible_for_testing_member
-    final dhrMatch = dhr == null ? other.dhr == null : await dhr!.equals(other.dhr!);
-    final ckrMatch = ckr == null ? other.ckr == null : listsEqual(ckr!, other.ckr!);
-    final cksMatch = cks == null ? other.cks == null : listsEqual(cks!, other.cks!);
+    final dhrMatch = dhr == null
+        ? other.dhr == null
+        :
+        // ignore: invalid_use_of_visible_for_testing_member
+        other.dhr != null && await dhr!.equals(other.dhr!);
+    final ckrMatch = ckr == null
+        ? other.ckr == null
+        : other.ckr != null && listsEqual(ckr!, other.ckr!);
+    final cksMatch = cks == null
+        ? other.cks == null
+        : other.cks != null && listsEqual(cks!, other.cks!);

    // ignore: invalid_use_of_visible_for_testing_member
    final dhsMatch = await dhs.equals(other.dhs);
@@ -320,14 +446,15 @@ class OmemoDoubleRatchet {
    final ikMatch = await ik.equals(other.ik);

    return dhsMatch &&
-      ikMatch &&
-      dhrMatch &&
-      listsEqual(rk, other.rk) &&
-      cksMatch &&
-      ckrMatch &&
-      ns == other.ns &&
-      nr == other.nr &&
-      pn == other.pn &&
-      listsEqual(sessionAd, other.sessionAd);
+        ikMatch &&
+        dhrMatch &&
+        listsEqual(rk, other.rk) &&
+        cksMatch &&
+        ckrMatch &&
+        ns == other.ns &&
+        nr == other.nr &&
+        pn == other.pn &&
+        listsEqual(sessionAd, other.sessionAd) &&
+        kexTimestamp == other.kexTimestamp;
  }
 }
--- a/lib/src/errors.dart
+++ b/lib/src/errors.dart
@@ -1,32 +1,61 @@
+abstract class OmemoException {}
+
 /// Triggered during X3DH if the signature if the SPK does verify to the actual SPK.
-class InvalidSignatureException implements Exception {
-  String errMsg() => 'The signature of the SPK does not match the provided signature';
+class InvalidSignatureException extends OmemoException 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 InvalidMessageHMACException implements Exception {
+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 implements Exception {
+class SkippingTooManyMessagesException extends OmemoException
+    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 implements Exception {
+class NotEncryptedForDeviceException extends OmemoException
+    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 implements Exception {
+class NoDecryptionKeyException extends OmemoException 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 implements Exception {
+class UnknownSignedPrekeyException extends OmemoException implements Exception {
  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 {
+  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
+    implements Exception {
+  String errMsg() =>
+      'No key material available to create a ratchet session with';
+}
--- a/lib/src/helpers.dart
+++ b/lib/src/helpers.dart
@@ -43,7 +43,11 @@ int generateRandom32BitNumber() {
  return Random.secure().nextInt(4294967295 /*pow(2, 32) - 1*/);
 }

-OmemoPublicKey? decodeKeyIfNotNull(Map<String, dynamic> map, String key, KeyPairType type) {
+OmemoPublicKey? decodeKeyIfNotNull(
+  Map<String, dynamic> map,
+  String key,
+  KeyPairType type,
+) {
  if (map[key] == null) return null;

  return OmemoPublicKey.fromBytes(
@@ -73,3 +77,7 @@ OmemoKeyPair? decodeKeyPairIfNotNull(String? pk, String? sk, KeyPairType type) {
    type,
  );
 }
+
+int getTimestamp() {
+  return DateTime.now().millisecondsSinceEpoch;
+}
--- a/lib/src/keys.dart
+++ b/lib/src/keys.dart
@@ -9,7 +9,6 @@ const privateKeyLength = 32;
 const publicKeyLength = 32;

 class OmemoPublicKey {
-
  const OmemoPublicKey(this._pubkey);

  factory OmemoPublicKey.fromBytes(List<int> bytes, KeyPairType type) {
@@ -32,7 +31,10 @@ class OmemoPublicKey {
  Future<String> asBase64() async => base64Encode(_pubkey.bytes);

  Future<OmemoPublicKey> toCurve25519() async {
-    assert(type == KeyPairType.ed25519, 'Cannot convert non-Ed25519 public key to X25519');
+    assert(
+      type == KeyPairType.ed25519,
+      'Cannot convert non-Ed25519 public key to X25519',
+    );

    final pkc = Uint8List(publicKeyLength);
    TweetNaClExt.crypto_sign_ed25519_pk_to_x25519_pk(
@@ -40,22 +42,24 @@ class OmemoPublicKey {
      Uint8List.fromList(await getBytes()),
    );

-    return OmemoPublicKey(SimplePublicKey(List<int>.from(pkc), type: KeyPairType.x25519));
+    return OmemoPublicKey(
+      SimplePublicKey(List<int>.from(pkc), type: KeyPairType.x25519),
+    );
  }

  SimplePublicKey asPublicKey() => _pubkey;

  @visibleForTesting
  Future<bool> equals(OmemoPublicKey key) async {
-    return type == key.type && listsEqual(
-      await getBytes(),
-      await key.getBytes(),
-    );
+    return type == key.type &&
+        listsEqual(
+          await getBytes(),
+          await key.getBytes(),
+        );
  }
 }

 class OmemoPrivateKey {
-
  const OmemoPrivateKey(this._privkey, this.type);
  final List<int> _privkey;
  final KeyPairType type;
@@ -63,7 +67,10 @@ class OmemoPrivateKey {
  Future<List<int>> getBytes() async => _privkey;

  Future<OmemoPrivateKey> toCurve25519() async {
-    assert(type == KeyPairType.ed25519, 'Cannot convert non-Ed25519 private key to X25519');
+    assert(
+      type == KeyPairType.ed25519,
+      'Cannot convert non-Ed25519 private key to X25519',
+    );

    final skc = Uint8List(privateKeyLength);
    TweetNaClExt.crypto_sign_ed25519_sk_to_x25519_sk(
@@ -76,21 +83,25 @@ class OmemoPrivateKey {

  @visibleForTesting
  Future<bool> equals(OmemoPrivateKey key) async {
-    return type == key.type && listsEqual(
-      await getBytes(),
-      await key.getBytes(),
-    );
+    return type == key.type &&
+        listsEqual(
+          await getBytes(),
+          await key.getBytes(),
+        );
  }
 }

 /// A generic wrapper class for both Ed25519 and X25519 keypairs
 class OmemoKeyPair {
-
  const OmemoKeyPair(this.pk, this.sk, this.type);

  /// Create an OmemoKeyPair just from a [type] and the bytes of the private and public
  /// key.
-  factory OmemoKeyPair.fromBytes(List<int> publicKey, List<int> privateKey, KeyPairType type) {
+  factory OmemoKeyPair.fromBytes(
+    List<int> publicKey,
+    List<int> privateKey,
+    KeyPairType type,
+  ) {
    return OmemoKeyPair(
      OmemoPublicKey.fromBytes(
        publicKey,
@@ -107,7 +118,10 @@ class OmemoKeyPair {
  /// Generate a completely new random OmemoKeyPair of type [type]. [type] must be either
  /// KeyPairType.ed25519 or KeyPairType.x25519.
  static Future<OmemoKeyPair> generateNewPair(KeyPairType type) async {
-    assert(type == KeyPairType.ed25519 || type == KeyPairType.x25519, 'Keypair must be either Ed25519 or X25519');
+    assert(
+      type == KeyPairType.ed25519 || type == KeyPairType.x25519,
+      'Keypair must be either Ed25519 or X25519',
+    );

    SimpleKeyPair kp;
    if (type == KeyPairType.ed25519) {
@@ -136,7 +150,10 @@ class OmemoKeyPair {

  /// Return the bytes that comprise the public key.
  Future<OmemoKeyPair> toCurve25519() async {
-    assert(type == KeyPairType.ed25519, 'Cannot convert non-Ed25519 keypair to X25519');
+    assert(
+      type == KeyPairType.ed25519,
+      'Cannot convert non-Ed25519 keypair to X25519',
+    );

    return OmemoKeyPair(
      await pk.toCurve25519(),
@@ -156,7 +173,7 @@ class OmemoKeyPair {
  @visibleForTesting
  Future<bool> equals(OmemoKeyPair pair) async {
    return type == pair.type &&
-      await pk.equals(pair.pk) &&
-      await sk.equals(pair.sk);
+        await pk.equals(pair.pk) &&
+        await sk.equals(pair.sk);
  }
 }
--- a/lib/src/omemo/bundle.dart
+++ b/lib/src/omemo/bundle.dart
@@ -1,9 +1,9 @@
 import 'dart:convert';
 import 'package:cryptography/cryptography.dart';
+import 'package:hex/hex.dart';
 import 'package:omemo_dart/src/keys.dart';

 class OmemoBundle {
-
  const OmemoBundle(
    this.jid,
    this.id,
@@ -13,17 +13,23 @@ class OmemoBundle {
    this.ikEncoded,
    this.opksEncoded,
  );
+
  /// The bare Jid the Bundle belongs to
  final String jid;
+
  /// The device Id
  final int id;
+
  /// The SPK but base64 encoded
  final String spkEncoded;
  final int spkId;
+
  /// The SPK signature but base64 encoded
  final String spkSignatureEncoded;
+
  /// The IK but base64 encoded
  final String ikEncoded;
+
  /// The mapping of a OPK's id to the base64 encoded data
  final Map<int, String> opksEncoded;

@@ -43,4 +49,10 @@ class OmemoBundle {
  }

  List<int> get spkSignature => base64Decode(spkSignatureEncoded);
+
+  /// Calculates the fingerprint of the bundle (See
+  /// https://xmpp.org/extensions/xep-0384.html#security § 2).
+  Future<String> getFingerprint() async {
+    return HEX.encode(await ik.getBytes());
+  }
 }
--- a/lib/src/omemo/constants.dart
+++ b/lib/src/omemo/constants.dart
@@ -0,0 +1,2 @@
+/// The info used for when encrypting the AES key for the actual payload.
+const omemoPayloadInfoString = 'OMEMO Payload';
--- a/lib/src/omemo/decryption_result.dart
+++ b/lib/src/omemo/decryption_result.dart
@@ -0,0 +1,9 @@
+import 'package:meta/meta.dart';
+import 'package:omemo_dart/src/errors.dart';
+
+@immutable
+class DecryptionResult {
+  const DecryptionResult(this.payload, this.error);
+  final String? payload;
+  final OmemoException? error;
+}
--- a/lib/src/omemo/device.dart
+++ b/lib/src/omemo/device.dart
@@ -1,5 +1,6 @@
 import 'dart:convert';
 import 'package:cryptography/cryptography.dart';
+import 'package:hex/hex.dart';
 import 'package:meta/meta.dart';
 import 'package:omemo_dart/src/helpers.dart';
 import 'package:omemo_dart/src/keys.dart';
@@ -8,9 +9,8 @@ import 'package:omemo_dart/src/x3dh/x3dh.dart';

 /// This class represents an OmemoBundle but with all keypairs belonging to the keys
@immutable
-class Device {
-
-  const Device(
+class OmemoDevice {
+  const OmemoDevice(
    this.jid,
    this.id,
    this.ik,
@@ -23,7 +23,7 @@ class Device {
  );

  /// Deserialize the Device
-  factory Device.fromJson(Map<String, dynamic> data) {
+  factory OmemoDevice.fromJson(Map<String, dynamic> data) {
    // NOTE: We use the way OpenSSH names their keys, meaning that ik is the Identity
    //       Keypair's private key, while ik_pub refers to the Identity Keypair's public
    //       key.
@@ -49,16 +49,25 @@ class Device {
      ]
    }
    */
-    final opks = <int, OmemoKeyPair>{};
-    for (final opk in data['opks']! as List<Map<String, dynamic>>) {
-      opks[opk['id']! as int] = OmemoKeyPair.fromBytes(
-        base64.decode(opk['public']! as String),
-        base64.decode(opk['private']! as String),
-        KeyPairType.x25519,
-      );
-    }
+    // 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 opks = Map<int, OmemoKeyPair>.fromEntries(
+      (data['opks']! as List<dynamic>).map<MapEntry<int, OmemoKeyPair>>(
+        (opk) {
+          final map = opk as Map<String, dynamic>;
+          return MapEntry(
+            map['id']! as int,
+            OmemoKeyPair.fromBytes(
+              base64.decode(map['public']! as String),
+              base64.decode(map['private']! as String),
+              KeyPairType.x25519,
+            ),
+          );
+        },
+      ),
+    );

-    return Device(
+    return OmemoDevice(
      data['jid']! as String,
      data['id']! as int,
      OmemoKeyPair.fromBytes(
@@ -84,7 +93,10 @@ class Device {
  }

  /// Generate a completely new device, i.e. cryptographic identity.
-  static Future<Device> generateNewDevice(String jid, { int opkAmount = 100 }) async {
+  static Future<OmemoDevice> generateNewDevice(
+    String jid, {
+    int opkAmount = 100,
+  }) async {
    final id = generateRandom32BitNumber();
    final ik = await OmemoKeyPair.generateNewPair(KeyPairType.ed25519);
    final spk = await OmemoKeyPair.generateNewPair(KeyPairType.x25519);
@@ -96,7 +108,7 @@ class Device {
      opks[i] = await OmemoKeyPair.generateNewPair(KeyPairType.x25519);
    }

-    return Device(jid, id, ik, spk, spkId, signature, null, null, opks);
+    return OmemoDevice(jid, id, ik, spk, spkId, signature, null, null, opks);
  }

  /// Our bare Jid
@@ -110,13 +122,16 @@ class Device {

  /// The signed prekey...
  final OmemoKeyPair spk;
+
  /// ...its Id, ...
  final int spkId;
+
  /// ...and its signature
  final List<int> spkSignature;

  /// The old Signed Prekey...
  final OmemoKeyPair? oldSpk;
+
  /// ...and its Id
  final int? oldSpkId;

@@ -126,12 +141,12 @@ class Device {
  /// This replaces the Onetime-Prekey with id [id] with a completely new one. Returns
  /// a new Device object that copies over everything but replaces said key.
  @internal
-  Future<Device> replaceOnetimePrekey(int id) async {
+  Future<OmemoDevice> replaceOnetimePrekey(int id) async {
    opks[id] = await OmemoKeyPair.generateNewPair(KeyPairType.x25519);

-    return Device(
+    return OmemoDevice(
      jid,
-      id,
+      this.id,
      ik,
      spk,
      spkId,
@@ -145,12 +160,12 @@ class Device {
  /// This replaces the Signed-Prekey with a completely new one. Returns a new Device object
  /// that copies over everything but replaces the Signed-Prekey and its signature.
  @internal
-  Future<Device> replaceSignedPrekey() async {
+  Future<OmemoDevice> replaceSignedPrekey() async {
    final newSpk = await OmemoKeyPair.generateNewPair(KeyPairType.x25519);
    final newSpkId = generateRandom32BitNumber();
    final newSignature = await sig(ik, await newSpk.pk.getBytes());

-    return Device(
+    return OmemoDevice(
      jid,
      id,
      ik,
@@ -163,6 +178,23 @@ class Device {
    );
  }

+  /// Returns a new device that is equal to this one with the exception that the new
+  /// device's id is a new number between 0 and 2**32 - 1.
+  @internal
+  OmemoDevice withNewId() {
+    return OmemoDevice(
+      jid,
+      generateRandom32BitNumber(),
+      ik,
+      spk,
+      spkId,
+      spkSignature,
+      oldSpk,
+      oldSpkId,
+      opks,
+    );
+  }
+
  /// Converts this device into an OmemoBundle that could be used for publishing.
  Future<OmemoBundle> toBundle() async {
    final encodedOpks = <int, String>{};
@@ -182,6 +214,13 @@ class Device {
    );
  }

+  /// Returns the fingerprint of the current device
+  Future<String> getFingerprint() async {
+    // Since the local key is Ed25519, we must convert it to Curve25519 first
+    final curveKey = await ik.pk.toCurve25519();
+    return HEX.encode(await curveKey.getBytes());
+  }
+
  /// Serialise the device information.
  Future<Map<String, dynamic>> toJson() async {
    /// Serialise the OPKs
@@ -211,14 +250,16 @@ class Device {
  }

  @visibleForTesting
-  Future<bool> equals(Device other) async {
+  Future<bool> equals(OmemoDevice other) async {
    var opksMatch = true;
    if (opks.length != other.opks.length) {
      opksMatch = false;
    } else {
      for (final entry in opks.entries) {
        // ignore: invalid_use_of_visible_for_testing_member
-        final matches = await other.opks[entry.key]?.equals(entry.value) ?? false;
+        final matches =
+            // ignore: invalid_use_of_visible_for_testing_member
+            await other.opks[entry.key]?.equals(entry.value) ?? false;
        if (!matches) {
          opksMatch = false;
        }
@@ -230,15 +271,18 @@ class Device {
    // ignore: invalid_use_of_visible_for_testing_member
    final spkMatch = await spk.equals(other.spk);
    // ignore: invalid_use_of_visible_for_testing_member
-    final oldSpkMatch = oldSpk != null ? await oldSpk!.equals(other.oldSpk!) : other.oldSpk == null;
+    final oldSpkMatch = oldSpk != null
+        // ignore: invalid_use_of_visible_for_testing_member
+        ? await oldSpk!.equals(other.oldSpk!)
+        : other.oldSpk == null;
    return id == other.id &&
-      ikMatch &&
-      spkMatch &&
-      oldSpkMatch &&
-      jid == other.jid &&
-      listsEqual(spkSignature, other.spkSignature) &&
-      spkId == other.spkId &&
-      oldSpkId == other.oldSpkId &&
-      opksMatch;
+        ikMatch &&
+        spkMatch &&
+        oldSpkMatch &&
+        jid == other.jid &&
+        listsEqual(spkSignature, other.spkSignature) &&
+        spkId == other.spkId &&
+        oldSpkId == other.oldSpkId &&
+        opksMatch;
  }
 }
--- a/lib/src/omemo/encrypted_key.dart
+++ b/lib/src/omemo/encrypted_key.dart
@@ -4,7 +4,6 @@ import 'package:meta/meta.dart';
 /// <keys /> header.
@immutable
 class EncryptedKey {
-
  const EncryptedKey(this.jid, this.rid, this.value, this.kex);
  final String jid;
  final int rid;
--- a/lib/src/omemo/encryption_result.dart
+++ b/lib/src/omemo/encryption_result.dart
@@ -1,15 +1,33 @@
 import 'package:meta/meta.dart';
+import 'package:omemo_dart/src/errors.dart';
 import 'package:omemo_dart/src/omemo/encrypted_key.dart';
+import 'package:omemo_dart/src/omemo/ratchet_map_key.dart';

@immutable
 class EncryptionResult {
+  const EncryptionResult(
+    this.ciphertext,
+    this.encryptedKeys,
+    this.deviceEncryptionErrors,
+    this.jidEncryptionErrors,
+  );

-  const EncryptionResult(this.ciphertext, this.encryptedKeys);
-  
-  /// The actual message that was encrypted
+  /// The actual message that was encrypted.
  final List<int>? ciphertext;

  /// Mapping of the device Id to the key for decrypting ciphertext, encrypted
-  /// for the ratchet with said device Id
+  /// for the ratchet with said device Id.
  final List<EncryptedKey> encryptedKeys;
+
+  /// Mapping of a ratchet map keys to a possible exception.
+  final Map<RatchetMapKey, OmemoException> deviceEncryptionErrors;
+
+  /// Mapping of a JID to a possible exception.
+  final Map<String, OmemoException> jidEncryptionErrors;
+
+  /// True if the encryption was a success. This means that we could encrypt for
+  /// at least one ratchet.
+  bool isSuccess(int numberOfRecipients) =>
+      encryptedKeys.isNotEmpty &&
+      jidEncryptionErrors.length < numberOfRecipients;
 }
--- a/lib/src/omemo/events.dart
+++ b/lib/src/omemo/events.dart
@@ -5,24 +5,40 @@ abstract class OmemoEvent {}

 /// Triggered when a ratchet has been modified
 class RatchetModifiedEvent extends OmemoEvent {
-
-  RatchetModifiedEvent(this.jid, this.deviceId, this.ratchet);
+  RatchetModifiedEvent(
+    this.jid,
+    this.deviceId,
+    this.ratchet,
+    this.added,
+    this.replaced,
+  );
  final String jid;
  final int deviceId;
  final OmemoDoubleRatchet ratchet;
+
+  /// Indicates whether the ratchet has just been created (true) or just modified (false).
+  final bool added;
+
+  /// Indicates whether the ratchet has been replaced (true) or not.
+  final bool replaced;
+}
+
+/// Triggered when a ratchet has been removed and should be removed from storage.
+class RatchetRemovedEvent extends OmemoEvent {
+  RatchetRemovedEvent(this.jid, this.deviceId);
+  final String jid;
+  final int deviceId;
 }

 /// Triggered when the device map has been modified
-class DeviceMapModifiedEvent extends OmemoEvent {
-
-  DeviceMapModifiedEvent(this.map);
-  final Map<String, List<int>> map;
+class DeviceListModifiedEvent extends OmemoEvent {
+  DeviceListModifiedEvent(this.list);
+  final Map<String, List<int>> list;
 }

 /// Triggered by the OmemoSessionManager when our own device bundle was modified
 /// and thus should be republished.
 class DeviceModifiedEvent extends OmemoEvent {
-
  DeviceModifiedEvent(this.device);
-  final Device device;
+  final OmemoDevice device;
 }
--- a/lib/src/omemo/fingerprint.dart
+++ b/lib/src/omemo/fingerprint.dart
@@ -2,7 +2,6 @@ import 'package:meta/meta.dart';

@immutable
 class DeviceFingerprint {
-
  const DeviceFingerprint(this.deviceId, this.fingerprint);
  final String fingerprint;
  final int deviceId;
@@ -10,8 +9,8 @@ class DeviceFingerprint {
  @override
  bool operator ==(Object other) {
    return other is DeviceFingerprint &&
-      fingerprint == other.fingerprint &&
-      deviceId == other.deviceId;
+        fingerprint == other.fingerprint &&
+        deviceId == other.deviceId;
  }

  @override
--- a/lib/src/omemo/omemomanager.dart
+++ b/lib/src/omemo/omemomanager.dart
@@ -0,0 +1,854 @@
+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/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/omemo_authenticated_message.dart';
+import 'package:omemo_dart/src/protobuf/omemo_key_exchange.dart';
+import 'package:omemo_dart/src/protobuf/omemo_message.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;
+    });
+  }
+}
--- a/lib/src/omemo/ratchet_map_key.dart
+++ b/lib/src/omemo/ratchet_map_key.dart
@@ -2,14 +2,30 @@ import 'package:meta/meta.dart';

@immutable
 class RatchetMapKey {
-
  const RatchetMapKey(this.jid, this.deviceId);
+
+  factory RatchetMapKey.fromJsonKey(String key) {
+    final parts = key.split(':');
+    final deviceId = int.parse(parts.first);
+
+    return RatchetMapKey(
+      parts.sublist(1).join(':'),
+      deviceId,
+    );
+  }
+
  final String jid;
  final int deviceId;

+  String toJsonKey() {
+    return '$deviceId:$jid';
+  }
+
  @override
  bool operator ==(Object other) {
-    return other is RatchetMapKey && jid == other.jid && deviceId == other.deviceId;
+    return other is RatchetMapKey &&
+        jid == other.jid &&
+        deviceId == other.deviceId;
  }

  @override
--- a/lib/src/omemo/sessionmanager.dart
+++ b/lib/src/omemo/sessionmanager.dart
@@ -1,437 +0,0 @@
-import 'dart:async';
-import 'dart:convert';
-import 'package:collection/collection.dart';
-import 'package:cryptography/cryptography.dart';
-import 'package:hex/hex.dart';
-import 'package:meta/meta.dart';
-import 'package:omemo_dart/src/crypto.dart';
-import 'package:omemo_dart/src/double_ratchet/double_ratchet.dart';
-import 'package:omemo_dart/src/errors.dart';
-import 'package:omemo_dart/src/helpers.dart';
-import 'package:omemo_dart/src/keys.dart';
-import 'package:omemo_dart/src/omemo/bundle.dart';
-import 'package:omemo_dart/src/omemo/device.dart';
-import 'package:omemo_dart/src/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/protobuf/omemo_authenticated_message.dart';
-import 'package:omemo_dart/src/protobuf/omemo_key_exchange.dart';
-import 'package:omemo_dart/src/protobuf/omemo_message.dart';
-import 'package:omemo_dart/src/trust/base.dart';
-import 'package:omemo_dart/src/x3dh/x3dh.dart';
-import 'package:synchronized/synchronized.dart';
-
-/// The info used for when encrypting the AES key for the actual payload.
-const omemoPayloadInfoString = 'OMEMO Payload';
-
-class OmemoSessionManager {
-
-  OmemoSessionManager(this._device, this._deviceMap, this._ratchetMap, this._trustManager)
-    : _lock = Lock(),
-      _deviceLock = Lock(),
-      _eventStreamController = StreamController<OmemoEvent>.broadcast();
-
-  /// Deserialise the OmemoSessionManager from JSON data [data].
-  factory OmemoSessionManager.fromJson(Map<String, dynamic> data, TrustManager trustManager) {
-    final ratchetMap = <RatchetMapKey, OmemoDoubleRatchet>{};
-    for (final rawRatchet in data['sessions']! as List<Map<String, dynamic>>) {
-      final key = RatchetMapKey(rawRatchet['jid']! as String, rawRatchet['deviceId']! as int);
-      final ratchet = OmemoDoubleRatchet.fromJson(rawRatchet['ratchet']! as Map<String, dynamic>);
-      ratchetMap[key] = ratchet;
-    }
-
-    return OmemoSessionManager(
-      Device.fromJson(data['device']! as Map<String, dynamic>),
-      data['devices']! as Map<String, List<int>>,
-      ratchetMap,
-      trustManager,
-    );
-  }
-      
-  /// Generate a new cryptographic identity.
-  static Future<OmemoSessionManager> generateNewIdentity(String jid, TrustManager trustManager, { int opkAmount = 100 }) async {
-    assert(opkAmount > 0, 'opkAmount must be bigger than 0.');
-    final device = await Device.generateNewDevice(jid, opkAmount: opkAmount);
-
-    return OmemoSessionManager(device, {}, {}, trustManager);
-  }
-  
-  /// Lock for _ratchetMap and _bundleMap
-  final Lock _lock;
-
-  /// Mapping of the Device Id to its OMEMO session
-  final Map<RatchetMapKey, OmemoDoubleRatchet> _ratchetMap;
-
-  /// Mapping of a bare Jid to its Device Ids
-  final Map<String, List<int>> _deviceMap;
-
-  /// The event bus of the session manager
-  final StreamController<OmemoEvent> _eventStreamController;
-  
-  /// Our own keys...
-  // ignore: prefer_final_fields
-  Device _device;
-  /// and its lock
-  final Lock _deviceLock;
-
-  /// The trust manager
-  final TrustManager _trustManager;
-  TrustManager get trustManager => _trustManager;
-  
-  /// A stream that receives events regarding the session
-  Stream<OmemoEvent> get eventStream => _eventStreamController.stream;
-
-  Future<Device> getDevice() async {
-    Device? dev;
-    await _deviceLock.synchronized(() async {
-      dev = _device;
-    });
-
-    return dev!;
-  }
-  
-  /// Add a session [ratchet] with the [deviceId] to the internal tracking state.
-  Future<void> _addSession(String jid, int deviceId, OmemoDoubleRatchet ratchet) async {
-    await _lock.synchronized(() async {
-      // Add the bundle Id
-      if (!_deviceMap.containsKey(jid)) {
-        _deviceMap[jid] = [deviceId];
-      } else {
-        _deviceMap[jid]!.add(deviceId);
-      }
-
-      // Commit the device map
-      _eventStreamController.add(DeviceMapModifiedEvent(_deviceMap));
-
-      // Add the ratchet session
-      final key = RatchetMapKey(jid, deviceId);
-      if (!_ratchetMap.containsKey(key)) {
-        _ratchetMap[key] = ratchet;
-      } else {
-        // TODO(PapaTutuWawa): What do we do now?
-        throw Exception();
-      }
-
-      // Commit the ratchet
-      _eventStreamController.add(RatchetModifiedEvent(jid, deviceId, 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,
-    );
-
-    await _trustManager.onNewSession(jid, deviceId);
-    await _addSession(jid, deviceId, ratchet);
-
-    return OmemoKeyExchange()
-      ..pkId = kexResult.opkId
-      ..spkId = bundle.spkId
-      ..ik = await device.ik.pk.getBytes()
-      ..ek = await kexResult.ek.pk.getBytes();
-  }
-
-  /// Build a new session with the user at [jid] with the device [deviceId] using data
-  /// from the key exchange [kex]. In case [kex] contains an unknown Signed Prekey
-  /// identifier an UnknownSignedPrekeyException will be thrown.
-  Future<void> _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();
-    }
-
-    assert(spk != null, 'The used SPK must be found');
-    
-    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,
-    );
-
-    await _trustManager.onNewSession(jid, deviceId);
-    await _addSession(jid, deviceId, ratchet);
-  }
-
-  /// Like [encryptToJids] but only for one Jid [jid].
-  Future<EncryptionResult> encryptToJid(String jid, String? plaintext, { List<OmemoBundle>? newSessions }) {
-    return encryptToJids([jid], plaintext, newSessions: newSessions);
-  }
-  
-  /// 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.
-  Future<EncryptionResult> encryptToJids(List<String> jids, String? plaintext, { List<OmemoBundle>? newSessions }) 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 = <int, OmemoKeyExchange>{};
-    if (newSessions != null) {
-      for (final newSession in newSessions) {
-        kex[newSession.id] = await addSessionFromBundle(newSession.jid, newSession.id, newSession);
-      }
-    }
-    
-    await _lock.synchronized(() async {
-      // We assume that the user already checked if the session exists
-      for (final jid in jids) {
-        for (final deviceId in _deviceMap[jid]!) {
-          // 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;
-          }
-
-          final ratchetKey = RatchetMapKey(jid, deviceId);
-          final ratchet = _ratchetMap[ratchetKey]!;
-          final ciphertext = (await ratchet.ratchetEncrypt(keyPayload)).ciphertext;
-
-          // Commit the ratchet
-          _eventStreamController.add(RatchetModifiedEvent(jid, deviceId, ratchet));
-          
-          if (kex.isNotEmpty && kex.containsKey(deviceId)) {
-            final k = kex[deviceId]!
-              ..message = OmemoAuthenticatedMessage.fromBuffer(ciphertext);
-            encryptedKeys.add(
-              EncryptedKey(
-                jid,
-                deviceId,
-                base64.encode(k.writeToBuffer()),
-                true,
-              ),
-            );
-          } else {
-            encryptedKeys.add(
-              EncryptedKey(
-                jid,
-                deviceId,
-                base64.encode(ciphertext),
-                false,
-              ),
-            );
-          }
-        }
-      }
-    });
-
-    return EncryptionResult(
-      plaintext != null ? ciphertext : null,
-      encryptedKeys,
-    );
-  }
-
-  /// 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.
-  ///
-  /// 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<String?> decryptMessage(List<int>? ciphertext, String senderJid, int senderDeviceId, List<EncryptedKey> keys) 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);
-    OmemoAuthenticatedMessage authMessage;
-    if (rawKey.kex) {
-      // TODO(PapaTutuWawa): Only do this when we should
-      final kex = OmemoKeyExchange.fromBuffer(decodedRawKey);
-      await _addSessionFromKeyExchange(
-        senderJid,
-        senderDeviceId,
-        kex,
-      );
-
-      authMessage = kex.message!;
-
-      // Replace the OPK
-      await _deviceLock.synchronized(() async {
-        device = await device.replaceOnetimePrekey(kex.pkId!);
-
-        // Commit the device
-        _eventStreamController.add(DeviceModifiedEvent(device));
-      });
-    } else {
-      authMessage = OmemoAuthenticatedMessage.fromBuffer(decodedRawKey);
-    }
-    
-    final devices = _deviceMap[senderJid];
-    if (devices == null) {
-      throw NoDecryptionKeyException();
-    }
-    if (!devices.contains(senderDeviceId)) {
-      throw NoDecryptionKeyException();
-    }
-
-    final message = OmemoMessage.fromBuffer(authMessage.message!);
-    final ratchetKey = RatchetMapKey(senderJid, senderDeviceId);
-    List<int>? keyAndHmac;
-    await _lock.synchronized(() async {
-      final ratchet = _ratchetMap[ratchetKey]!;
-      if (rawKey.kex) {
-        keyAndHmac = await ratchet.ratchetDecrypt(message, authMessage.writeToBuffer());
-      } else {
-        keyAndHmac = await ratchet.ratchetDecrypt(message, decodedRawKey);
-      }
-
-      // Commit the ratchet
-      _eventStreamController.add(RatchetModifiedEvent(senderJid, senderDeviceId, ratchet));
-    });
-
-    // 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();
-    }
-    
-    final plaintext = await aes256CbcDecrypt(ciphertext, derivedKeys.encryptionKey, derivedKeys.iv);
-    return utf8.decode(plaintext);
-  }
-
-  /// Returns the list of hex-encoded fingerprints we have for sessions with [jid].
-  Future<List<DeviceFingerprint>> getHexFingerprintsForJid(String jid) async {
-    final fingerprints = List<DeviceFingerprint>.empty(growable: true);
-
-    await _lock.synchronized(() async {
-      // Get devices for jid
-      final devices = _deviceMap[jid]!;
-
-      for (final deviceId in devices) {
-        final ratchet = _ratchetMap[RatchetMapKey(jid, deviceId)]!;
-
-        fingerprints.add(
-          DeviceFingerprint(
-            deviceId,
-            HEX.encode(await ratchet.ik.getBytes()),
-          ),
-        );
-      }
-    });
-
-    return fingerprints;
-  }
-
-  /// Replaces the Signed Prekey and its signature in our own device bundle. Triggers
-  /// a DeviceModifiedEvent when done.
-  /// See https://xmpp.org/extensions/xep-0384.html#protocol-key_exchange under the point
-  /// "signed PreKey rotation period" for recommendations.
-  Future<void> rotateSignedPrekey() async {
-    await _deviceLock.synchronized(() async {
-      _device = await _device.replaceSignedPrekey();
-
-      // Commit the new device
-      _eventStreamController.add(DeviceModifiedEvent(_device));
-    });
-  }
-  
-  @visibleForTesting
-  OmemoDoubleRatchet getRatchet(String jid, int deviceId) => _ratchetMap[RatchetMapKey(jid, deviceId)]!;
-
-  @visibleForTesting
-  Map<String, List<int>> getDeviceMap() => _deviceMap;
-
-  @visibleForTesting
-  Map<RatchetMapKey, OmemoDoubleRatchet> getRatchetMap() => _ratchetMap;
-
-  /// Serialise the entire session manager into a JSON object.
-  Future<Map<String, dynamic>> toJson() async {
-    /*
-    {
-      'devices': {
-        'alice@...': [1, 2, ...],
-        'bob@...': [1],
-        ...
-      },
-      'device': { ... },
-      'sessions': [
-        {
-          'jid': 'alice@...',
-          'deviceId': 1,
-          'ratchet': { ... },
-        },
-        ...
-      ],
-      'trust': { ... }
-    }
-    */
-
-    final sessions = List<Map<String, dynamic>>.empty(growable: true);
-    for (final entry in _ratchetMap.entries) {
-      sessions.add({
-        'jid': entry.key.jid,
-        'deviceId': entry.key.deviceId,
-        'ratchet': await entry.value.toJson(),
-      });
-    }
-    return {
-      'devices': _deviceMap,
-      'device': await (await getDevice()).toJson(),
-      'sessions': sessions,
-      // TODO(PapaTutuWawa): Implement
-      'trust': <String, dynamic>{},
-    };
-  }
-}
--- a/lib/src/omemo/stanza.dart
+++ b/lib/src/omemo/stanza.dart
@@ -0,0 +1,41 @@
+import 'package:omemo_dart/src/omemo/encrypted_key.dart';
+
+/// Describes a stanza that was received by the underlying XMPP library.
+class OmemoIncomingStanza {
+  const OmemoIncomingStanza(
+    this.bareSenderJid,
+    this.senderDeviceId,
+    this.timestamp,
+    this.keys,
+    this.payload,
+  );
+
+  /// The bare JID of the sender of the stanza.
+  final String bareSenderJid;
+
+  /// The device ID of the sender.
+  final int senderDeviceId;
+
+  /// The timestamp when the stanza was received.
+  final int timestamp;
+
+  /// The included encrypted keys
+  final List<EncryptedKey> keys;
+
+  /// The string payload included in the <encrypted /> element.
+  final String? payload;
+}
+
+/// Describes a stanza that is to be sent out
+class OmemoOutgoingStanza {
+  const OmemoOutgoingStanza(
+    this.recipientJids,
+    this.payload,
+  );
+
+  /// The JIDs the stanza will be sent to.
+  final List<String> recipientJids;
+
+  /// The serialised XML data that should be encrypted.
+  final String payload;
+}
--- a/lib/src/protobuf/omemo_authenticated_message.dart
+++ b/lib/src/protobuf/omemo_authenticated_message.dart
@@ -2,7 +2,6 @@ import 'package:omemo_dart/src/helpers.dart';
 import 'package:omemo_dart/src/protobuf/protobuf.dart';

 class OmemoAuthenticatedMessage {
-
  OmemoAuthenticatedMessage();

  factory OmemoAuthenticatedMessage.fromBuffer(List<int> data) {
--- a/lib/src/protobuf/omemo_key_exchange.dart
+++ b/lib/src/protobuf/omemo_key_exchange.dart
@@ -3,7 +3,6 @@ import 'package:omemo_dart/src/protobuf/omemo_authenticated_message.dart';
 import 'package:omemo_dart/src/protobuf/protobuf.dart';

 class OmemoKeyExchange {
-
  OmemoKeyExchange();

  factory OmemoKeyExchange.fromBuffer(List<int> data) {
--- a/lib/src/protobuf/omemo_message.dart
+++ b/lib/src/protobuf/omemo_message.dart
@@ -2,7 +2,6 @@ import 'package:omemo_dart/src/helpers.dart';
 import 'package:omemo_dart/src/protobuf/protobuf.dart';

 class OmemoMessage {
-
  OmemoMessage();

  factory OmemoMessage.fromBuffer(List<int> data) {
--- a/lib/src/protobuf/protobuf.dart
+++ b/lib/src/protobuf/protobuf.dart
@@ -13,7 +13,6 @@ int fieldId(int number, int type) {
 }

 class VarintDecode {
-
  const VarintDecode(this.n, this.length);
  final int n;
  final int length;
@@ -46,21 +45,18 @@ List<int> encodeVarint(int i) {
  assert(i >= 0, "Two's complement is not implemented");
  final ret = List<int>.empty(growable: true);

-  var j = 0;
-  while (true) {
+  // Thanks to https://github.com/hathibelagal-dev/LEB128 for the trick with toRadixString!
+  final numSevenBlocks = (i.toRadixString(2).length / 7).ceil();
+  for (var j = 0; j < numSevenBlocks; j++) {
    // The 7 LSB of the byte we're creating
    final x = (i & (lsb7Mask << j * 7)) >> j * 7;
-    // The next bits
-    final next = i & (lsb7Mask << (j + 1) * 7);

-    if (next == 0) {
+    if (j == numSevenBlocks - 1) {
      // If we were to shift further, we only get zero, so we're at the end
      ret.add(x);
-      break;
    } else {
      // We still have at least one bit more to go, so set the MSB to 1
      ret.add(x + msb);
-      j++;
    }
  }

--- a/lib/src/trust/always.dart
+++ b/lib/src/trust/always.dart
@@ -11,4 +11,16 @@ class AlwaysTrustingTrustManager extends TrustManager {

  @override
  Future<void> onNewSession(String jid, int deviceId) async {}
+
+  @override
+  Future<bool> isEnabled(String jid, int deviceId) async => true;
+
+  @override
+  Future<void> setEnabled(String jid, int deviceId, bool enabled) async {}
+
+  @override
+  Future<void> removeTrustDecisionsForJid(String jid) async {}
+
+  @override
+  Future<Map<String, dynamic>> toJson() async => <String, dynamic>{};
 }
--- a/lib/src/trust/base.dart
+++ b/lib/src/trust/base.dart
@@ -8,4 +8,18 @@ abstract class TrustManager {
  /// Called by the OmemoSessionManager when a new session has been built. Should set
  /// a default trust state to [jid]'s device with identifier [deviceId].
  Future<void> onNewSession(String jid, int deviceId);
+
+  /// Return true if the device with id [deviceId] of Jid [jid] should be used for encryption.
+  /// If not, return false.
+  Future<bool> isEnabled(String jid, int deviceId);
+
+  /// Mark the device with id [deviceId] of Jid [jid] as enabled if [enabled] is true or as disabled
+  /// if [enabled] is false.
+  Future<void> setEnabled(String jid, int deviceId, bool enabled);
+
+  /// Serialize the trust manager to JSON.
+  Future<Map<String, dynamic>> toJson();
+
+  /// Removes all trust decisions for [jid].
+  Future<void> removeTrustDecisionsForJid(String jid);
 }
--- a/lib/src/trust/btbv.dart
+++ b/lib/src/trust/btbv.dart
@@ -8,24 +8,59 @@ import 'package:synchronized/synchronized.dart';
 /// - blindTrust: The fingerprint is not verified using OOB means
 /// - verified: The fingerprint has been verified using OOB means
 enum BTBVTrustState {
-  notTrusted,
-  blindTrust,
-  verified,
+  notTrusted, // = 1
+  blindTrust, // = 2
+  verified, // = 3
+}
+
+int _trustToInt(BTBVTrustState state) {
+  switch (state) {
+    case BTBVTrustState.notTrusted:
+      return 1;
+    case BTBVTrustState.blindTrust:
+      return 2;
+    case BTBVTrustState.verified:
+      return 3;
+  }
+}
+
+BTBVTrustState _trustFromInt(int i) {
+  switch (i) {
+    case 1:
+      return BTBVTrustState.notTrusted;
+    case 2:
+      return BTBVTrustState.blindTrust;
+    case 3:
+      return BTBVTrustState.verified;
+    default:
+      return BTBVTrustState.notTrusted;
+  }
 }

 /// A TrustManager that implements the idea of Blind Trust Before Verification.
 /// See https://gultsch.de/trust.html for more details.
 abstract class BlindTrustBeforeVerificationTrustManager extends TrustManager {
-  BlindTrustBeforeVerificationTrustManager()
-    : trustCache = {},
-      devices = {},
-      _lock = Lock();
+  BlindTrustBeforeVerificationTrustManager({
+    Map<RatchetMapKey, BTBVTrustState>? trustCache,
+    Map<RatchetMapKey, bool>? enablementCache,
+    Map<String, List<int>>? devices,
+  })  : trustCache = trustCache ?? {},
+        enablementCache = enablementCache ?? {},
+        devices = devices ?? {},
+        _lock = Lock();

-  /// The cache for Mapping a RatchetMapKey to its trust state
+  /// The cache for mapping a RatchetMapKey to its trust state
+  @visibleForTesting
  @protected
  final Map<RatchetMapKey, BTBVTrustState> trustCache;

+  /// The cache for mapping a RatchetMapKey to whether it is enabled or not
+  @visibleForTesting
+  @protected
+  final Map<RatchetMapKey, bool> enablementCache;
+
  /// Mapping of Jids to their device identifiers
+  @visibleForTesting
  @protected
  final Map<String, List<int>> devices;

@@ -74,10 +109,13 @@ abstract class BlindTrustBeforeVerificationTrustManager extends TrustManager {
  @override
  Future<void> onNewSession(String jid, int deviceId) async {
    await _lock.synchronized(() async {
+      final key = RatchetMapKey(jid, deviceId);
      if (_hasAtLeastOneVerifiedDevice(jid)) {
-        trustCache[RatchetMapKey(jid, deviceId)] = BTBVTrustState.notTrusted;
+        trustCache[key] = BTBVTrustState.notTrusted;
+        enablementCache[key] = false;
      } else {
-        trustCache[RatchetMapKey(jid, deviceId)] = BTBVTrustState.blindTrust;
+        trustCache[key] = BTBVTrustState.blindTrust;
+        enablementCache[key] = true;
      }

      if (devices.containsKey(jid)) {
@@ -91,21 +129,28 @@ abstract class BlindTrustBeforeVerificationTrustManager extends TrustManager {
    });
  }

-  /// Returns a mapping from the device identifiers of [jid] to their trust state.
+  /// Returns a mapping from the device identifiers of [jid] to their trust state. If
+  /// there are no devices known for [jid], then an empty map is returned.
  Future<Map<int, BTBVTrustState>> getDevicesTrust(String jid) async {
-    final map = <int, BTBVTrustState>{};
+    return _lock.synchronized(() async {
+      final map = <int, BTBVTrustState>{};
+
+      if (!devices.containsKey(jid)) return map;

-    await _lock.synchronized(() async {
      for (final deviceId in devices[jid]!) {
        map[deviceId] = trustCache[RatchetMapKey(jid, deviceId)]!;
      }
-    });

-    return map;
+      return map;
+    });
  }

  /// Sets the trust of [jid]'s device with identifier [deviceId] to [state].
-  Future<void> setDeviceTrust(String jid, int deviceId, BTBVTrustState state) async {
+  Future<void> setDeviceTrust(
+    String jid,
+    int deviceId,
+    BTBVTrustState state,
+  ) async {
    await _lock.synchronized(() async {
      trustCache[RatchetMapKey(jid, deviceId)] = state;

@@ -114,18 +159,95 @@ abstract class BlindTrustBeforeVerificationTrustManager extends TrustManager {
    });
  }

+  @override
+  Future<bool> isEnabled(String jid, int deviceId) async {
+    return _lock.synchronized(() async {
+      final value = enablementCache[RatchetMapKey(jid, deviceId)];
+
+      if (value == null) return false;
+      return value;
+    });
+  }
+
+  @override
+  Future<void> setEnabled(String jid, int deviceId, bool enabled) async {
+    await _lock.synchronized(() async {
+      enablementCache[RatchetMapKey(jid, deviceId)] = enabled;
+    });
+
+    // Commit the state
+    await commitState();
+  }
+
+  @override
+  Future<Map<String, dynamic>> toJson() async {
+    return {
+      'devices': devices,
+      'trust': trustCache.map(
+        (key, value) => MapEntry(
+          key.toJsonKey(),
+          _trustToInt(value),
+        ),
+      ),
+      'enable':
+          enablementCache.map((key, value) => MapEntry(key.toJsonKey(), value)),
+    };
+  }
+
+  /// From a serialized version of a BTBV trust manager, extract the device list.
+  /// NOTE: This is needed as Dart cannot just cast a List<dynamic> to List<int> and so on.
+  static Map<String, List<int>> deviceListFromJson(Map<String, dynamic> json) {
+    return (json['devices']! as Map<String, dynamic>).map<String, List<int>>(
+      (key, value) => MapEntry(
+        key,
+        (value as List<dynamic>).map<int>((i) => i as int).toList(),
+      ),
+    );
+  }
+
+  /// From a serialized version of a BTBV trust manager, extract the trust cache.
+  /// NOTE: This is needed as Dart cannot just cast a List<dynamic> to List<int> and so on.
+  static Map<RatchetMapKey, BTBVTrustState> trustCacheFromJson(
+    Map<String, dynamic> json,
+  ) {
+    return (json['trust']! as Map<String, dynamic>)
+        .map<RatchetMapKey, BTBVTrustState>(
+      (key, value) => MapEntry(
+        RatchetMapKey.fromJsonKey(key),
+        _trustFromInt(value as int),
+      ),
+    );
+  }
+
+  /// From a serialized version of a BTBV trust manager, extract the enable cache.
+  /// NOTE: This is needed as Dart cannot just cast a List<dynamic> to List<int> and so on.
+  static Map<RatchetMapKey, bool> enableCacheFromJson(
+    Map<String, dynamic> json,
+  ) {
+    return (json['enable']! as Map<String, dynamic>).map<RatchetMapKey, bool>(
+      (key, value) => MapEntry(
+        RatchetMapKey.fromJsonKey(key),
+        value as bool,
+      ),
+    );
+  }
+
+  @override
+  Future<void> removeTrustDecisionsForJid(String jid) async {
+    await _lock.synchronized(() async {
+      devices.remove(jid);
+      await commitState();
+    });
+  }
+
  /// Called when the state of the trust manager has been changed. Allows the user to
  /// commit the trust state to persistent storage.
  @visibleForOverriding
  Future<void> commitState();

-  /// Called when the user wants to restore the state of the trust manager. The format
-  /// and actual storage mechanism is left to the user.
-  @visibleForOverriding
-  Future<void> loadState();
-  
  @visibleForTesting
-  BTBVTrustState getDeviceTrust(String jid, int deviceId) => trustCache[RatchetMapKey(jid, deviceId)]!;
+  BTBVTrustState getDeviceTrust(String jid, int deviceId) =>
+      trustCache[RatchetMapKey(jid, deviceId)]!;
 }

 /// A BTBV TrustManager that does not commit its state to persistent storage. Well suited
@@ -133,7 +255,4 @@ abstract class BlindTrustBeforeVerificationTrustManager extends TrustManager {
 class MemoryBTBVTrustManager extends BlindTrustBeforeVerificationTrustManager {
  @override
  Future<void> commitState() async {}
-
-  @override
-  Future<void> loadState() async {}
 }
--- a/lib/src/trust/never.dart
+++ b/lib/src/trust/never.dart
@@ -11,4 +11,16 @@ class NeverTrustingTrustManager extends TrustManager {

  @override
  Future<void> onNewSession(String jid, int deviceId) async {}
+
+  @override
+  Future<bool> isEnabled(String jid, int deviceId) async => true;
+
+  @override
+  Future<void> setEnabled(String jid, int deviceId, bool enabled) async {}
+
+  @override
+  Future<void> removeTrustDecisionsForJid(String jid) async {}
+
+  @override
+  Future<Map<String, dynamic>> toJson() async => <String, dynamic>{};
 }
--- a/lib/src/x3dh/x3dh.dart
+++ b/lib/src/x3dh/x3dh.dart
@@ -12,7 +12,6 @@ const omemoX3DHInfoString = 'OMEMO X3DH';

 /// Performed by Alice
 class X3DHAliceResult {
-
  const X3DHAliceResult(this.ek, this.sk, this.opkId, this.ad);
  final OmemoKeyPair ek;
  final List<int> sk;
@@ -22,7 +21,6 @@ class X3DHAliceResult {

 /// Received by Bob
 class X3DHMessage {
-
  const X3DHMessage(this.ik, this.ek, this.opkId);
  final OmemoPublicKey ik;
  final OmemoPublicKey ek;
@@ -30,7 +28,6 @@ class X3DHMessage {
 }

 class X3DHBobResult {
-
  const X3DHBobResult(this.sk, this.ad);
  final List<int> sk;
  final List<int> ad;
@@ -39,7 +36,10 @@ class X3DHBobResult {
 /// 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, 'Signature keypair must be Ed25519');
+  assert(
+    keyPair.type == KeyPairType.ed25519,
+    'Signature keypair must be Ed25519',
+  );
  final signature = await Ed25519().sign(
    message,
    keyPair: await keyPair.asKeyPair(),
@@ -70,7 +70,10 @@ 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(OmemoBundle bundle, OmemoKeyPair ik) async {
+Future<X3DHAliceResult> x3dhFromBundle(
+  OmemoBundle bundle,
+  OmemoKeyPair ik,
+) async {
  // Check the signature first
  final signatureValue = await Ed25519().verify(
    await bundle.spk.getBytes(),
@@ -93,7 +96,7 @@ Future<X3DHAliceResult> x3dhFromBundle(OmemoBundle bundle, OmemoKeyPair ik) asyn
  final opk = bundle.getOpk(opkId);

  final dh1 = await omemoDH(ik, bundle.spk, 1);
-  final dh2 = await omemoDH(ek, bundle.ik,  2);
+  final dh2 = await omemoDH(ek, bundle.ik, 2);
  final dh3 = await omemoDH(ek, bundle.spk, 0);
  final dh4 = await omemoDH(ek, opk, 0);

@@ -108,9 +111,14 @@ Future<X3DHAliceResult> x3dhFromBundle(OmemoBundle bundle, OmemoKeyPair ik) asyn

 /// 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<X3DHBobResult> x3dhFromInitialMessage(X3DHMessage msg, OmemoKeyPair spk, OmemoKeyPair opk, OmemoKeyPair ik) async {
+Future<X3DHBobResult> x3dhFromInitialMessage(
+  X3DHMessage msg,
+  OmemoKeyPair spk,
+  OmemoKeyPair opk,
+  OmemoKeyPair ik,
+) async {
  final dh1 = await omemoDH(spk, msg.ik, 2);
-  final dh2 = await omemoDH(ik,  msg.ek, 1);
+  final dh2 = await omemoDH(ik, msg.ek, 1);
  final dh3 = await omemoDH(spk, msg.ek, 0);
  final dh4 = await omemoDH(opk, msg.ek, 0);

--- a/pubspec.yaml
+++ b/pubspec.yaml
@@ -1,6 +1,6 @@
 name: omemo_dart
 description: An XMPP library independent OMEMO library
-version: 0.1.0
+version: 0.4.3
 homepage: https://github.com/PapaTutuWawa/omemo_dart
 publish_to: https://git.polynom.me/api/packages/PapaTutuWawa/pub

@@ -11,7 +11,8 @@ dependencies:
  collection: ^1.16.0
  cryptography: ^2.0.5
  hex: ^0.2.0
-  meta: ^1.8.0
+  logging: ^1.0.2
+  meta: ^1.7.0
  pinenacl: ^0.5.1
  synchronized: ^3.0.0+2

--- a/test/double_ratchet_test.dart
+++ b/test/double_ratchet_test.dart
@@ -76,7 +76,6 @@ void main() {

    print('X3DH key exchange done');

-
    // Alice and Bob now share sk as a common secret and ad
    // Build a session
    final alicesRatchet = await OmemoDoubleRatchet.initiateNewSession(
@@ -84,12 +83,14 @@ void main() {
      ikBob.pk,
      resultAlice.sk,
      resultAlice.ad,
+      0,
    );
    final bobsRatchet = await OmemoDoubleRatchet.acceptNewSession(
      spkBob,
      ikAlice.pk,
      resultBob.sk,
      resultBob.ad,
+      0,
    );

    expect(alicesRatchet.sessionAd, bobsRatchet.sessionAd);
@@ -99,7 +100,8 @@ void main() {

      if (i.isEven) {
        // Alice encrypts a message
-        final aliceRatchetResult = await alicesRatchet.ratchetEncrypt(utf8.encode(messageText));
+        final aliceRatchetResult =
+            await alicesRatchet.ratchetEncrypt(utf8.encode(messageText));
        print('Alice sent the message');

        // Alice sends it to Bob
@@ -115,7 +117,8 @@ void main() {
        expect(utf8.encode(messageText), bobRatchetResult);
      } else {
        // Bob sends a message to Alice
-        final bobRatchetResult = await bobsRatchet.ratchetEncrypt(utf8.encode(messageText));
+        final bobRatchetResult =
+            await bobsRatchet.ratchetEncrypt(utf8.encode(messageText));
        print('Bob sent the message');

        // Bobs sends it to Alice
--- a/test/omemo_test.dart
+++ b/test/omemo_test.dart
@@ -1,13 +1,50 @@
+import 'package:logging/logging.dart';
 import 'package:omemo_dart/omemo_dart.dart';
 import 'package:omemo_dart/src/trust/always.dart';
 import 'package:omemo_dart/src/trust/never.dart';
 import 'package:test/test.dart';

 void main() {
+  Logger.root
+    ..level = Level.ALL
+    ..onRecord.listen((record) {
+      // ignore: avoid_print
+      print('${record.level.name}: ${record.message}');
+    });
+
+  test('Test replacing a onetime prekey', () async {
+    const aliceJid = 'alice@server.example';
+    final device = await OmemoDevice.generateNewDevice(aliceJid, opkAmount: 1);
+
+    final newDevice = await device.replaceOnetimePrekey(0);
+
+    expect(device.jid, newDevice.jid);
+    expect(device.id, newDevice.id);
+
+    var opksMatch = true;
+    if (newDevice.opks.length != device.opks.length) {
+      opksMatch = false;
+    } else {
+      for (final entry in device.opks.entries) {
+        final m = await newDevice.opks[entry.key]?.equals(entry.value) ?? false;
+        if (!m) opksMatch = false;
+      }
+    }
+
+    expect(opksMatch, true);
+    expect(await device.ik.equals(newDevice.ik), true);
+    expect(await device.spk.equals(newDevice.spk), true);
+
+    final oldSpkMatch = device.oldSpk != null
+        ? await device.oldSpk!.equals(newDevice.oldSpk!)
+        : newDevice.oldSpk == null;
+    expect(oldSpkMatch, true);
+    expect(listsEqual(device.spkSignature, newDevice.spkSignature), true);
+  });
+
  test('Test using OMEMO sessions with only one device per user', () async {
    const aliceJid = 'alice@server.example';
    const bobJid = 'bob@other.server.example';
-      
    // Alice and Bob generate their sessions
    var deviceModified = false;
    var ratchetModified = 0;
@@ -28,7 +65,7 @@ void main() {
        deviceModified = true;
      } else if (event is RatchetModifiedEvent) {
        ratchetModified++;
-      } else if (event is DeviceMapModifiedEvent) {
+      } else if (event is DeviceListModifiedEvent) {
        deviceMapModified++;
      }
    });
@@ -39,7 +76,7 @@ void main() {
      bobJid,
      messagePlaintext,
      newSessions: [
-        await (await bobSession.getDevice()).toBundle(),
+        await bobSession.getDeviceBundle(),
      ],
    );
    expect(aliceMessage.encryptedKeys.length, 1);
@@ -51,8 +88,9 @@ void main() {
    final bobMessage = await bobSession.decryptMessage(
      aliceMessage.ciphertext,
      aliceJid,
-      (await aliceSession.getDevice()).id,
+      await aliceSession.getDeviceId(),
      aliceMessage.encryptedKeys,
+      0,
    );
    expect(messagePlaintext, bobMessage);
    // The ratchet should be modified two times: Once for when the ratchet is created and
@@ -68,6 +106,12 @@ void main() {
      false,
    );

+    // Ratchets are acked
+    await aliceSession.ratchetAcknowledged(
+        bobJid, await bobSession.getDeviceId());
+    await bobSession.ratchetAcknowledged(
+        aliceJid, await aliceSession.getDeviceId());
+
    // Bob responds to Alice
    const bobResponseText = 'Oh, hello Alice!';
    final bobResponseMessage = await bobSession.encryptToJid(
@@ -82,13 +126,15 @@ void main() {
    final aliceReceivedMessage = await aliceSession.decryptMessage(
      bobResponseMessage.ciphertext,
      bobJid,
-      (await bobSession.getDevice()).id,
+      await bobSession.getDeviceId(),
      bobResponseMessage.encryptedKeys,
+      0,
    );
    expect(bobResponseText, aliceReceivedMessage);
  });

-  test('Test using OMEMO sessions with only two devices for the receiver', () async {
+  test('Test using OMEMO sessions with only two devices for the receiver',
+      () async {
    const aliceJid = 'alice@server.example';
    const bobJid = 'bob@other.server.example';

@@ -116,8 +162,8 @@ void main() {
      bobJid,
      messagePlaintext,
      newSessions: [
-        await (await bobSession.getDevice()).toBundle(),
-        await (await bobSession2.getDevice()).toBundle(),
+        await bobSession.getDeviceBundle(),
+        await bobSession2.getDeviceBundle(),
      ],
    );
    expect(aliceMessage.encryptedKeys.length, 2);
@@ -131,11 +177,18 @@ void main() {
    final bobMessage = await bobSession.decryptMessage(
      aliceMessage.ciphertext,
      aliceJid,
-      (await aliceSession.getDevice()).id,
+      await aliceSession.getDeviceId(),
      aliceMessage.encryptedKeys,
+      0,
    );
    expect(messagePlaintext, bobMessage);

+    // Ratchets are acked
+    await aliceSession.ratchetAcknowledged(
+        bobJid, await bobSession.getDeviceId());
+    await bobSession.ratchetAcknowledged(
+        aliceJid, await aliceSession.getDeviceId());
+
    // Bob responds to Alice
    const bobResponseText = 'Oh, hello Alice!';
    final bobResponseMessage = await bobSession.encryptToJid(
@@ -150,8 +203,9 @@ void main() {
    final aliceReceivedMessage = await aliceSession.decryptMessage(
      bobResponseMessage.ciphertext,
      bobJid,
-      (await bobSession.getDevice()).id,
+      await bobSession.getDeviceId(),
      bobResponseMessage.encryptedKeys,
+      0,
    );
    expect(bobResponseText, aliceReceivedMessage);

@@ -193,8 +247,8 @@ void main() {
      [bobJid, aliceJid],
      messagePlaintext,
      newSessions: [
-        await (await bobSession.getDevice()).toBundle(),
-        await (await aliceSession2.getDevice()).toBundle(),
+        await bobSession.getDeviceBundle(),
+        await aliceSession2.getDeviceBundle(),
      ],
    );
    expect(aliceMessage.encryptedKeys.length, 2);
@@ -206,8 +260,9 @@ void main() {
    final bobMessage = await bobSession.decryptMessage(
      aliceMessage.ciphertext,
      aliceJid,
-      (await aliceSession1.getDevice()).id,
+      await aliceSession1.getDeviceId(),
      aliceMessage.encryptedKeys,
+      0,
    );
    expect(messagePlaintext, bobMessage);

@@ -215,13 +270,14 @@ void main() {
    final aliceMessage2 = await aliceSession2.decryptMessage(
      aliceMessage.ciphertext,
      aliceJid,
-      (await aliceSession1.getDevice()).id,
+      await aliceSession1.getDeviceId(),
      aliceMessage.encryptedKeys,
+      0,
    );
    expect(messagePlaintext, aliceMessage2);
  });

-  test('Test using sending empty OMEMO messages', () async {
+  test('Test sending empty OMEMO messages', () async {
    const aliceJid = 'alice@server.example';
    const bobJid = 'bob@other.server.example';

@@ -242,7 +298,7 @@ void main() {
      bobJid,
      null,
      newSessions: [
-        await (await bobSession.getDevice()).toBundle(),
+        await bobSession.getDeviceBundle(),
      ],
    );
    expect(aliceMessage.encryptedKeys.length, 1);
@@ -255,13 +311,14 @@ void main() {
    final bobMessage = await bobSession.decryptMessage(
      aliceMessage.ciphertext,
      aliceJid,
-      (await aliceSession.getDevice()).id,
+      await aliceSession.getDeviceId(),
      aliceMessage.encryptedKeys,
+      0,
    );
    expect(bobMessage, null);

    // This call must not cause an exception
-    bobSession.getRatchet(aliceJid, (await aliceSession.getDevice()).id);
+    bobSession.getRatchet(aliceJid, await aliceSession.getDeviceId());
  });

  test('Test rotating the Signed Prekey', () async {
@@ -275,7 +332,7 @@ void main() {

    // Setup an event listener
    final oldDevice = await aliceSession.getDevice();
-    Device? newDevice;
+    OmemoDevice? newDevice;
    aliceSession.eventStream.listen((event) {
      if (event is DeviceModifiedEvent) {
        newDevice = event.device;
@@ -317,7 +374,7 @@ void main() {
      bobJid,
      messagePlaintext,
      newSessions: [
-        await (await bobSession.getDevice()).toBundle(),
+        await bobSession.getDeviceBundle(),
      ],
    );
    expect(aliceMessage.encryptedKeys.length, 1);
@@ -332,8 +389,9 @@ void main() {
    final bobMessage = await bobSession.decryptMessage(
      aliceMessage.ciphertext,
      aliceJid,
-      (await aliceSession.getDevice()).id,
+      await aliceSession.getDeviceId(),
      aliceMessage.encryptedKeys,
+      0,
    );
    expect(messagePlaintext, bobMessage);
  });
@@ -359,7 +417,7 @@ void main() {
      bobJid,
      null,
      newSessions: [
-        await (await bobSession.getDevice()).toBundle(),
+        await bobSession.getDeviceBundle(),
      ],
    );

@@ -367,4 +425,532 @@ void main() {
    // untrusted device.
    expect(aliceMessage.encryptedKeys.length, 1);
  });
+
+  test('Test by sending multiple messages back and forth', () async {
+    const aliceJid = 'alice@server.example';
+    const bobJid = 'bob@other.server.example';
+    // Alice and Bob generate their sessions
+    final aliceSession = await OmemoSessionManager.generateNewIdentity(
+      aliceJid,
+      AlwaysTrustingTrustManager(),
+      opkAmount: 1,
+    );
+    final bobSession = await OmemoSessionManager.generateNewIdentity(
+      bobJid,
+      AlwaysTrustingTrustManager(),
+      opkAmount: 1,
+    );
+
+    // Alice encrypts a message for Bob
+    final aliceMessage = await aliceSession.encryptToJid(
+      bobJid,
+      'Hello Bob!',
+      newSessions: [
+        await bobSession.getDeviceBundle(),
+      ],
+    );
+
+    // Alice sends the message to Bob
+    // ...
+
+    await bobSession.decryptMessage(
+      aliceMessage.ciphertext,
+      aliceJid,
+      await aliceSession.getDeviceId(),
+      aliceMessage.encryptedKeys,
+      0,
+    );
+
+    // Ratchets are acked
+    await aliceSession.ratchetAcknowledged(
+        bobJid, await bobSession.getDeviceId());
+    await bobSession.ratchetAcknowledged(
+        aliceJid, await aliceSession.getDeviceId());
+
+    for (var i = 0; i < 100; i++) {
+      final messageText = 'Test Message #$i';
+      // Bob responds to Alice
+      final bobResponseMessage = await bobSession.encryptToJid(
+        aliceJid,
+        messageText,
+      );
+
+      // Bob sends the message to Alice
+      // ...
+
+      // Alice decrypts it
+      final aliceReceivedMessage = await aliceSession.decryptMessage(
+        bobResponseMessage.ciphertext,
+        bobJid,
+        await bobSession.getDeviceId(),
+        bobResponseMessage.encryptedKeys,
+        0,
+      );
+      expect(messageText, aliceReceivedMessage);
+    }
+  });
+
+  group('Test removing a ratchet', () {
+    test('Test removing a ratchet when the user has multiple', () async {
+      const aliceJid = 'alice@server.local';
+      const bobJid = 'bob@some.server.local';
+      final aliceSession = await OmemoSessionManager.generateNewIdentity(
+        aliceJid,
+        AlwaysTrustingTrustManager(),
+        opkAmount: 1,
+      );
+      final bobSession1 = await OmemoSessionManager.generateNewIdentity(
+        bobJid,
+        AlwaysTrustingTrustManager(),
+        opkAmount: 1,
+      );
+      final bobSession2 = await OmemoSessionManager.generateNewIdentity(
+        bobJid,
+        AlwaysTrustingTrustManager(),
+        opkAmount: 1,
+      );
+
+      // Alice sends a message to those two Bobs
+      await aliceSession.encryptToJid(
+        bobJid,
+        'Hallo Welt',
+        newSessions: [
+          await bobSession1.getDeviceBundle(),
+          await bobSession2.getDeviceBundle(),
+        ],
+      );
+
+      // One of those two sessions is broken, so Alice removes the session2 ratchet
+      final id1 = await bobSession1.getDeviceId();
+      final id2 = await bobSession2.getDeviceId();
+      await aliceSession.removeRatchet(bobJid, id1);
+
+      final map = aliceSession.getRatchetMap();
+      expect(map.containsKey(RatchetMapKey(bobJid, id1)), false);
+      expect(map.containsKey(RatchetMapKey(bobJid, id2)), true);
+      final deviceMap = await aliceSession.getDeviceMap();
+      expect(deviceMap.containsKey(bobJid), true);
+      expect(deviceMap[bobJid], [id2]);
+    });
+
+    test('Test removing a ratchet when the user has only one', () async {
+      const aliceJid = 'alice@server.local';
+      const bobJid = 'bob@some.server.local';
+      final aliceSession = await OmemoSessionManager.generateNewIdentity(
+        aliceJid,
+        AlwaysTrustingTrustManager(),
+        opkAmount: 1,
+      );
+      final bobSession = await OmemoSessionManager.generateNewIdentity(
+        bobJid,
+        AlwaysTrustingTrustManager(),
+        opkAmount: 1,
+      );
+
+      // Alice sends a message to those two Bobs
+      await aliceSession.encryptToJid(
+        bobJid,
+        'Hallo Welt',
+        newSessions: [
+          await bobSession.getDeviceBundle(),
+        ],
+      );
+
+      // One of those two sessions is broken, so Alice removes the session2 ratchet
+      final id = await bobSession.getDeviceId();
+      await aliceSession.removeRatchet(bobJid, id);
+
+      final map = aliceSession.getRatchetMap();
+      expect(map.containsKey(RatchetMapKey(bobJid, id)), false);
+      final deviceMap = await aliceSession.getDeviceMap();
+      expect(deviceMap.containsKey(bobJid), false);
+    });
+  });
+
+  test('Test acknowledging a ratchet', () async {
+    const aliceJid = 'alice@server.example';
+    const bobJid = 'bob@other.server.example';
+    // Alice and Bob generate their sessions
+    final aliceSession = await OmemoSessionManager.generateNewIdentity(
+      aliceJid,
+      AlwaysTrustingTrustManager(),
+      opkAmount: 1,
+    );
+    final bobSession = await OmemoSessionManager.generateNewIdentity(
+      bobJid,
+      AlwaysTrustingTrustManager(),
+      opkAmount: 1,
+    );
+
+    // Alice sends Bob a message
+    await aliceSession.encryptToJid(
+      bobJid,
+      'Hallo Welt',
+      newSessions: [
+        await bobSession.getDeviceBundle(),
+      ],
+    );
+    expect(
+      await aliceSession.getUnacknowledgedRatchets(bobJid),
+      [
+        await bobSession.getDeviceId(),
+      ],
+    );
+
+    // Bob sends alice an empty message
+    // ...
+
+    // Alice decrypts it
+    // ...
+
+    // Alice marks the ratchet as acknowledged
+    await aliceSession.ratchetAcknowledged(
+        bobJid, await bobSession.getDeviceId());
+    expect(
+      (await aliceSession.getUnacknowledgedRatchets(bobJid))!.isEmpty,
+      true,
+    );
+  });
+
+  test('Test overwriting sessions', () async {
+    const aliceJid = 'alice@server.example';
+    const bobJid = 'bob@other.server.example';
+    // Alice and Bob generate their sessions
+    final aliceSession = await OmemoSessionManager.generateNewIdentity(
+      aliceJid,
+      AlwaysTrustingTrustManager(),
+      opkAmount: 1,
+    );
+    final bobSession = await OmemoSessionManager.generateNewIdentity(
+      bobJid,
+      AlwaysTrustingTrustManager(),
+      opkAmount: 2,
+    );
+
+    // Alice sends Bob a message
+    final msg1 = await aliceSession.encryptToJid(
+      bobJid,
+      'Hallo Welt',
+      newSessions: [
+        await bobSession.getDeviceBundle(),
+      ],
+    );
+    await bobSession.decryptMessage(
+      msg1.ciphertext,
+      aliceJid,
+      await aliceSession.getDeviceId(),
+      msg1.encryptedKeys,
+      0,
+    );
+    final aliceRatchet1 = aliceSession.getRatchet(
+      bobJid,
+      await bobSession.getDeviceId(),
+    );
+    final bobRatchet1 = bobSession.getRatchet(
+      aliceJid,
+      await aliceSession.getDeviceId(),
+    );
+
+    // Alice is impatient and immediately sends another message before the original one
+    // can be acknowledged by Bob
+    final msg2 = await aliceSession.encryptToJid(
+      bobJid,
+      "Why don't you answer?",
+      newSessions: [
+        await bobSession.getDeviceBundle(),
+      ],
+    );
+    await bobSession.decryptMessage(
+      msg2.ciphertext,
+      aliceJid,
+      await aliceSession.getDeviceId(),
+      msg2.encryptedKeys,
+      getTimestamp(),
+    );
+    final aliceRatchet2 = aliceSession.getRatchet(
+      bobJid,
+      await bobSession.getDeviceId(),
+    );
+    final bobRatchet2 = bobSession.getRatchet(
+      aliceJid,
+      await aliceSession.getDeviceId(),
+    );
+
+    // Both should only have one ratchet
+    expect(aliceSession.getRatchetMap().length, 1);
+    expect(bobSession.getRatchetMap().length, 1);
+
+    // The ratchets should both be different
+    expect(await aliceRatchet1.equals(aliceRatchet2), false);
+    expect(await bobRatchet1.equals(bobRatchet2), false);
+  });
+
+  test('Test resending key exchanges', () async {
+    const aliceJid = 'alice@server.example';
+    const bobJid = 'bob@other.server.example';
+    // Alice and Bob generate their sessions
+    final aliceSession = await OmemoSessionManager.generateNewIdentity(
+      aliceJid,
+      AlwaysTrustingTrustManager(),
+      opkAmount: 1,
+    );
+    final bobSession = await OmemoSessionManager.generateNewIdentity(
+      bobJid,
+      AlwaysTrustingTrustManager(),
+      opkAmount: 2,
+    );
+
+    // Alice sends Bob a message
+    final msg1 = await aliceSession.encryptToJid(
+      bobJid,
+      'Hallo Welt',
+      newSessions: [
+        await bobSession.getDeviceBundle(),
+      ],
+    );
+    // The first message should be a kex message
+    expect(msg1.encryptedKeys.first.kex, true);
+
+    await bobSession.decryptMessage(
+      msg1.ciphertext,
+      aliceJid,
+      await aliceSession.getDeviceId(),
+      msg1.encryptedKeys,
+      0,
+    );
+
+    // Alice is impatient and immediately sends another message before the original one
+    // can be acknowledged by Bob
+    final msg2 = await aliceSession.encryptToJid(
+      bobJid,
+      "Why don't you answer?",
+    );
+    expect(msg2.encryptedKeys.first.kex, true);
+
+    await bobSession.decryptMessage(
+      msg2.ciphertext,
+      aliceJid,
+      await aliceSession.getDeviceId(),
+      msg2.encryptedKeys,
+      getTimestamp(),
+    );
+  });
+
+  test('Test receiving old messages including a KEX', () async {
+    const aliceJid = 'alice@server.example';
+    const bobJid = 'bob@other.server.example';
+    // Alice and Bob generate their sessions
+    final aliceSession = await OmemoSessionManager.generateNewIdentity(
+      aliceJid,
+      AlwaysTrustingTrustManager(),
+      opkAmount: 1,
+    );
+    final bobSession = await OmemoSessionManager.generateNewIdentity(
+      bobJid,
+      AlwaysTrustingTrustManager(),
+      opkAmount: 2,
+    );
+
+    final bobsReceivedMessages = List<EncryptionResult>.empty(growable: true);
+    final bobsReceivedMessagesTimestamps = List<int>.empty(growable: true);
+
+    // Alice sends Bob a message
+    final msg1 = await aliceSession.encryptToJid(
+      bobJid,
+      'Hallo Welt',
+      newSessions: [
+        await bobSession.getDeviceBundle(),
+      ],
+    );
+    bobsReceivedMessages.add(msg1);
+    final t1 = getTimestamp();
+    bobsReceivedMessagesTimestamps.add(t1);
+
+    await bobSession.decryptMessage(
+      msg1.ciphertext,
+      aliceJid,
+      await aliceSession.getDeviceId(),
+      msg1.encryptedKeys,
+      t1,
+    );
+
+    // Ratchets are acked
+    await aliceSession.ratchetAcknowledged(
+        bobJid, await bobSession.getDeviceId());
+    await bobSession.ratchetAcknowledged(
+        aliceJid, await aliceSession.getDeviceId());
+
+    // Bob responds
+    final msg2 = await bobSession.encryptToJid(
+      aliceJid,
+      'Hello!',
+    );
+
+    await aliceSession.decryptMessage(
+      msg2.ciphertext,
+      bobJid,
+      await bobSession.getDeviceId(),
+      msg2.encryptedKeys,
+      getTimestamp(),
+    );
+
+    // Send some messages between the two
+    for (var i = 0; i < 100; i++) {
+      final msg = await aliceSession.encryptToJid(
+        bobJid,
+        'Hello $i',
+      );
+      bobsReceivedMessages.add(msg);
+      final t = getTimestamp();
+      bobsReceivedMessagesTimestamps.add(t);
+      final result = await bobSession.decryptMessage(
+        msg.ciphertext,
+        aliceJid,
+        await aliceSession.getDeviceId(),
+        msg.encryptedKeys,
+        t,
+      );
+
+      expect(result, 'Hello $i');
+    }
+
+    // Due to some issue with the transport protocol, the messages to Bob are received
+    // again.
+    final ratchetPreError = bobSession
+        .getRatchet(aliceJid, await aliceSession.getDeviceId())
+        .clone();
+    var invalidKex = 0;
+    var errorCounter = 0;
+    for (var i = 0; i < bobsReceivedMessages.length; i++) {
+      final msg = bobsReceivedMessages[i];
+      try {
+        await bobSession.decryptMessage(
+          msg.ciphertext,
+          aliceJid,
+          await aliceSession.getDeviceId(),
+          msg.encryptedKeys,
+          bobsReceivedMessagesTimestamps[i],
+        );
+        expect(true, false);
+      } on InvalidMessageHMACException catch (_) {
+        errorCounter++;
+      } on InvalidKeyExchangeException catch (_) {
+        invalidKex++;
+      }
+    }
+    final ratchetPostError = bobSession
+        .getRatchet(aliceJid, await aliceSession.getDeviceId())
+        .clone();
+
+    // The 100 messages including the initial KEX message
+    expect(invalidKex, 1);
+    expect(errorCounter, 100);
+    expect(await ratchetPreError.equals(ratchetPostError), true);
+
+    final msg3 = await aliceSession.encryptToJid(
+      bobJid,
+      'Are you okay?',
+    );
+    final result = await bobSession.decryptMessage(
+      msg3.ciphertext,
+      aliceJid,
+      await aliceSession.getDeviceId(),
+      msg3.encryptedKeys,
+      104,
+    );
+
+    expect(result, 'Are you okay?');
+  });
+
+  test("Test ignoring a new KEX when we haven't acket it yet", () async {
+    const aliceJid = 'alice@server.example';
+    const bobJid = 'bob@other.server.example';
+    // Alice and Bob generate their sessions
+    final aliceSession = await OmemoSessionManager.generateNewIdentity(
+      aliceJid,
+      AlwaysTrustingTrustManager(),
+      opkAmount: 1,
+    );
+    final bobSession = await OmemoSessionManager.generateNewIdentity(
+      bobJid,
+      AlwaysTrustingTrustManager(),
+      opkAmount: 1,
+    );
+
+    // Alice sends Bob a message
+    final msg1 = await aliceSession.encryptToJid(
+      bobJid,
+      'Hallo Welt',
+      newSessions: [
+        await bobSession.getDeviceBundle(),
+      ],
+    );
+    expect(msg1.encryptedKeys.first.kex, true);
+
+    await bobSession.decryptMessage(
+      msg1.ciphertext,
+      aliceJid,
+      await aliceSession.getDeviceId(),
+      msg1.encryptedKeys,
+      getTimestamp(),
+    );
+
+    // Alice sends another message before the ack can reach us
+    final msg2 = await aliceSession.encryptToJid(
+      bobJid,
+      'ANSWER ME!',
+    );
+    expect(msg2.encryptedKeys.first.kex, true);
+
+    await bobSession.decryptMessage(
+      msg2.ciphertext,
+      aliceJid,
+      await aliceSession.getDeviceId(),
+      msg2.encryptedKeys,
+      getTimestamp(),
+    );
+
+    // Now the acks reach us
+    await aliceSession.ratchetAcknowledged(
+        bobJid, await bobSession.getDeviceId());
+    await bobSession.ratchetAcknowledged(
+        aliceJid, await aliceSession.getDeviceId());
+
+    // Alice sends another message
+    final msg3 = await aliceSession.encryptToJid(
+      bobJid,
+      "You read the message, didn't you?",
+    );
+    expect(msg3.encryptedKeys.first.kex, false);
+
+    await bobSession.decryptMessage(
+      msg3.ciphertext,
+      aliceJid,
+      await aliceSession.getDeviceId(),
+      msg3.encryptedKeys,
+      getTimestamp(),
+    );
+
+    for (var i = 0; i < 100; i++) {
+      final messageText = 'Test Message #$i';
+      // Bob responds to Alice
+      final bobResponseMessage = await bobSession.encryptToJid(
+        aliceJid,
+        messageText,
+      );
+
+      // Bob sends the message to Alice
+      // ...
+
+      // Alice decrypts it
+      final aliceReceivedMessage = await aliceSession.decryptMessage(
+        bobResponseMessage.ciphertext,
+        bobJid,
+        await bobSession.getDeviceId(),
+        bobResponseMessage.encryptedKeys,
+        0,
+      );
+      expect(messageText, aliceReceivedMessage);
+    }
+  });
 }
--- a/test/omemomanager_test.dart
+++ b/test/omemomanager_test.dart
--- a/test/protobuf_test.dart
+++ b/test/protobuf_test.dart
@@ -48,6 +48,10 @@ void main() {
        <int>[172, 2],
      );
    });
+
+    test('Test some special cases', () {
+      expect(decodeVarint(encodeVarint(1042464893), 0).n, 1042464893);
+    });
  });

  group('OMEMOMessage', () {
@@ -165,10 +169,18 @@ void main() {
      expect(decoded.pkId, 698);
      expect(decoded.spkId, 245);
      expect(decoded.ik, <int>[1, 4, 6]);
-      expect(decoded.ek, <int>[4 ,6 ,7 , 80]);
+      expect(decoded.ek, <int>[4, 6, 7, 80]);

      expect(decoded.message!.mac, <int>[5, 6, 8, 0]);
      expect(decoded.message!.message, <int>[4, 5, 7, 3, 2]);
    });
+    test('Test decoding an issue', () {
+      /*
+      final data = 'CAAQfRogc2GwslU219dUkrMHNM4KdZRmuFnBTae+bQaJ+55IsAMiII7aZKj2sUpb6xR/3Ari7WZUmKFV0G6czUc4NMvjKDBaKnwKEM2ZpI8X3TgcxhxwENANnlsSaAgAEAAaICy8T9WPgLb7RdYd8/4JkrLF0RahEkC3ZaEfk5jw3dsLIkBMILzLyByweLgF4lCn0oNea+kbdrFr6rY7r/7WyI8hXEQz38QpnN+jyGGwC7Ga0dq70WuyqE7VpiFArQwqZh2G';
+      final kex = OmemoKeyExchange.fromBuffer(base64Decode(data));
+
+      expect(kex.spkId!, 1042464893);
+      */
+    });
  });
 }
--- a/test/serialisation_test.dart
+++ b/test/serialisation_test.dart
@@ -1,7 +1,12 @@
+import 'dart:convert';
 import 'package:omemo_dart/omemo_dart.dart';
 import 'package:omemo_dart/src/trust/always.dart';
 import 'package:test/test.dart';

+Map<String, dynamic> jsonify(Map<String, dynamic> map) {
+  return jsonDecode(jsonEncode(map)) as Map<String, dynamic>;
+}
+
 void main() {
  test('Test serialising and deserialising the Device', () async {
    // Generate a random session
@@ -11,23 +16,25 @@ void main() {
      opkAmount: 1,
    );
    final oldDevice = await oldSession.getDevice();
-    final serialised = await oldDevice.toJson();
+    final serialised = jsonify(await oldDevice.toJson());

-    final newDevice = Device.fromJson(serialised);
+    final newDevice = OmemoDevice.fromJson(serialised);
    expect(await oldDevice.equals(newDevice), true);
  });

-  test('Test serialising and deserialising the Device after rotating the SPK', () async {
+  test('Test serialising and deserialising the Device after rotating the SPK',
+      () async {
    // Generate a random session
    final oldSession = await OmemoSessionManager.generateNewIdentity(
      'user@test.server',
      AlwaysTrustingTrustManager(),
      opkAmount: 1,
    );
-    final oldDevice = await (await oldSession.getDevice()).replaceSignedPrekey();
-    final serialised = await oldDevice.toJson();
+    final oldDevice =
+        await (await oldSession.getDevice()).replaceSignedPrekey();
+    final serialised = jsonify(await oldDevice.toJson());

-    final newDevice = Device.fromJson(serialised);
+    final newDevice = OmemoDevice.fromJson(serialised);
    expect(await oldDevice.equals(newDevice), true);
  });

@@ -49,17 +56,19 @@ void main() {
      bobJid,
      'Hello Bob!',
      newSessions: [
-        await (await bobSession.getDevice()).toBundle(),
+        await bobSession.getDeviceBundle(),
      ],
    );
    await bobSession.decryptMessage(
      aliceMessage.ciphertext,
      aliceJid,
-      (await aliceSession.getDevice()).id,
+      await aliceSession.getDeviceId(),
      aliceMessage.encryptedKeys,
+      getTimestamp(),
    );
-    final aliceOld = aliceSession.getRatchet(bobJid, (await bobSession.getDevice()).id);
-    final aliceSerialised = await aliceOld.toJson();
+    final aliceOld =
+        aliceSession.getRatchet(bobJid, await bobSession.getDeviceId());
+    final aliceSerialised = jsonify(await aliceOld.toJson());
    final aliceNew = OmemoDoubleRatchet.fromJson(aliceSerialised);

    expect(await aliceOld.equals(aliceNew), true);
@@ -79,29 +88,68 @@ void main() {
    );
    await oldSession.addSessionFromBundle(
      'bob@localhost',
-      (await bobSession.getDevice()).id,
-      await (await bobSession.getDevice()).toBundle(),
+      await bobSession.getDeviceId(),
+      await bobSession.getDeviceBundle(),
    );

    // Serialise and deserialise
-    final serialised = await oldSession.toJson();
-    final newSession = OmemoSessionManager.fromJson(
+    final serialised = jsonify(await oldSession.toJsonWithoutSessions());
+    final newSession = OmemoSessionManager.fromJsonWithoutSessions(
      serialised,
+      // NOTE: At this point, we don't care about this attribute
+      {},
      AlwaysTrustingTrustManager(),
    );

    final oldDevice = await oldSession.getDevice();
    final newDevice = await newSession.getDevice();
    expect(await oldDevice.equals(newDevice), true);
-    expect(oldSession.getDeviceMap(), newSession.getDeviceMap());
+    expect(await oldSession.getDeviceMap(), await newSession.getDeviceMap());
+  });

-    expect(oldSession.getRatchetMap().length, newSession.getRatchetMap().length);
-    for (final session in oldSession.getRatchetMap().entries) {
-      expect(newSession.getRatchetMap().containsKey(session.key), true);
+  test('Test serializing and deserializing RatchetMapKey', () {
+    const test1 = RatchetMapKey('user@example.org', 1234);
+    final result1 = RatchetMapKey.fromJsonKey(test1.toJsonKey());
+    expect(result1.jid, test1.jid);
+    expect(result1.deviceId, test1.deviceId);

-      final oldRatchet = oldSession.getRatchetMap()[session.key]!;
-      final newRatchet = newSession.getRatchetMap()[session.key]!;
-      expect(await oldRatchet.equals(newRatchet), true);
-    }
+    const test2 = RatchetMapKey('user@example.org/hallo:welt', 3333);
+    final result2 = RatchetMapKey.fromJsonKey(test2.toJsonKey());
+    expect(result2.jid, test2.jid);
+    expect(result2.deviceId, test2.deviceId);
+  });
+
+  test('Test serializing and deserializing the components of the BTBV manager',
+      () async {
+    // Caroline's BTBV manager
+    final btbv = MemoryBTBVTrustManager();
+    // Example data
+    const aliceJid = 'alice@some.server';
+    const bobJid = 'bob@other.server';
+
+    await btbv.onNewSession(aliceJid, 1);
+    await btbv.setDeviceTrust(aliceJid, 1, BTBVTrustState.verified);
+    await btbv.onNewSession(aliceJid, 2);
+    await btbv.onNewSession(bobJid, 3);
+    await btbv.onNewSession(bobJid, 4);
+
+    final serialized = jsonify(await btbv.toJson());
+    final deviceList =
+        BlindTrustBeforeVerificationTrustManager.deviceListFromJson(
+      serialized,
+    );
+    expect(btbv.devices, deviceList);
+
+    final trustCache =
+        BlindTrustBeforeVerificationTrustManager.trustCacheFromJson(
+      serialized,
+    );
+    expect(btbv.trustCache, trustCache);
+
+    final enableCache =
+        BlindTrustBeforeVerificationTrustManager.enableCacheFromJson(
+      serialized,
+    );
+    expect(btbv.enablementCache, enableCache);
  });
 }
--- a/test/trust_test.dart
+++ b/test/trust_test.dart
@@ -12,6 +12,7 @@ void main() {
    // Caroline starts a chat a device from Alice
    await btbv.onNewSession(aliceJid, 1);
    expect(await btbv.isTrusted(aliceJid, 1), true);
+    expect(await btbv.isEnabled(aliceJid, 1), true);

    // Caroline meets with Alice and verifies her fingerprint
    await btbv.setDeviceTrust(aliceJid, 1, BTBVTrustState.verified);
@@ -21,6 +22,7 @@ void main() {
    await btbv.onNewSession(aliceJid, 2);
    expect(await btbv.isTrusted(aliceJid, 2), false);
    expect(btbv.getDeviceTrust(aliceJid, 2), BTBVTrustState.notTrusted);
+    expect(await btbv.isEnabled(aliceJid, 2), false);

    // Caronline starts a chat with Bob but since they live far apart, Caroline cannot
    // verify his fingerprint.
@@ -32,5 +34,7 @@ void main() {
    expect(await btbv.isTrusted(bobJid, 4), true);
    expect(btbv.getDeviceTrust(bobJid, 3), BTBVTrustState.blindTrust);
    expect(btbv.getDeviceTrust(bobJid, 4), BTBVTrustState.blindTrust);
+    expect(await btbv.isEnabled(bobJid, 3), true);
+    expect(await btbv.isEnabled(bobJid, 4), true);
  });
 }
--- a/test/x3dh_test.dart
+++ b/test/x3dh_test.dart
@@ -71,9 +71,13 @@ void main() {
    var exception = false;
    try {
      await x3dhFromBundle(bundleBob, ikAlice);
-    } catch(e) {
+    } catch (e) {
      exception = true;
-      expect(e is InvalidSignatureException, true, reason: 'Expected InvalidSignatureException, but got $e');
+      expect(
+        e is InvalidSignatureException,
+        true,
+        reason: 'Expected InvalidSignatureException, but got $e',
+      );
    }

    expect(exception, true, reason: 'Expected test failure');
Author	SHA1	Message	Date
Alexander "PapaTutuWawa	0ffc0b067a	docs: Migrate the example to the "new" OmemoManager Fixes #30.	2023-06-12 19:37:51 +02:00
Alexander "PapaTutuWawa	3376929c24	style: Formattiing issues	2023-06-12 19:37:45 +02:00
Alexander "PapaTutuWawa	3783ec6f13	feat: Remove OmemoSessionManager	2023-06-12 19:21:07 +02:00
Alexander "PapaTutuWawa	65f1daff55	style: Format using dart format	2023-06-12 19:20:43 +02:00
Alexander "PapaTutuWawa	f2ec7bd759	fix: Initial receiving ratchet has no trust data	2023-06-12 19:13:38 +02:00
Alexander "PapaTutuWawa	04bb70d657	release: Tag 0.4.3 All checks were successful ci/woodpecker/push/woodpecker Pipeline was successful Details ci/woodpecker/tag/woodpecker Pipeline was successful Details	2023-01-22 19:13:56 +01:00
Alexander "PapaTutuWawa	3fb5fd7eb8	style: Cleanup	2023-01-22 19:08:06 +01:00
Alexander "PapaTutuWawa	a97f8bc372	fix: Fix ratchet null issue Fixes #26.	2023-01-22 19:07:29 +01:00
Alexander "PapaTutuWawa	e29ee07015	release: Tag 0.4.2 All checks were successful ci/woodpecker/push/woodpecker Pipeline was successful Details ci/woodpecker/tag/woodpecker Pipeline was successful Details	2023-01-07 22:16:15 +01:00
Alexander "PapaTutuWawa	cc31475671	fix: Ratchets are not removed when using removeAllRatchets All checks were successful ci/woodpecker/push/woodpecker Pipeline was successful Details	2023-01-07 22:14:10 +01:00
Alexander "PapaTutuWawa	de85ab7955	chore: Bump version All checks were successful ci/woodpecker/push/woodpecker Pipeline was successful Details ci/woodpecker/tag/woodpecker Pipeline was successful Details	2023-01-02 13:57:15 +01:00
Alexander "PapaTutuWawa	a08c654295	feat: Do not fetch bundles for our own device	2023-01-02 13:56:22 +01:00
PapaTutuWawa	3825232ebe	Merge pull request #24 from PapaTutuWawa/feat/omemomanager Some checks are pending ci/woodpecker/push/woodpecker Pipeline is pending Details OmemoManager.	2023-01-01 16:59:00 +01:00
Alexander "PapaTutuWawa	32c0c5a4f0	feat: Update CHANGELOG and README	2023-01-01 16:55:15 +01:00
Alexander "PapaTutuWawa	092ce36410	feat: Add last features for use in moxxmpp	2023-01-01 16:49:19 +01:00
Alexander "PapaTutuWawa	4085631804	fix: Make logging less verbose	2022-12-27 12:49:11 +01:00
Alexander "PapaTutuWawa	6a9e98665d	fix: Track new KEX entries using a ratchet key	2022-12-27 12:47:45 +01:00
Alexander "PapaTutuWawa	5f844dafda	feat: Add getDeviceId	2022-12-27 12:43:50 +01:00
Alexander "PapaTutuWawa	480de1ce8f	test: Port the sending/receiving for the OmemoManager	2022-12-27 12:43:50 +01:00
Alexander "PapaTutuWawa	5e6b54aab5	feat: Better guard against failed lookups	2022-12-27 12:43:50 +01:00
Alexander "PapaTutuWawa	6c4dd62c5a	feat: _decryptMessage explicitly tells us that a ratchet was created	2022-12-27 12:43:50 +01:00
Alexander "PapaTutuWawa	54eeb816eb	feat: Export OmemoManager	2022-12-27 12:43:50 +01:00
Alexander "PapaTutuWawa	e35f65aff9	style: Fix style issues	2022-12-27 12:43:50 +01:00
Alexander "PapaTutuWawa	4dc3cfb2b1	feat: Deprecate OmemoSessionManager	2022-12-27 12:43:50 +01:00
Alexander "PapaTutuWawa	bca4840ca6	test: Add a test to ensure the library cannot get stuck	2022-12-27 12:43:50 +01:00
Alexander "PapaTutuWawa	caf841c53e	feat: Add more documentation	2022-12-27 12:43:50 +01:00
Alexander "PapaTutuWawa	5bc1136ec0	feat: Implement getting fingerprints	2022-12-27 12:43:50 +01:00
Alexander "PapaTutuWawa	d37a4bd719	chore: Bump version	2022-12-27 12:43:50 +01:00
Alexander "PapaTutuWawa	b48665c357	feat: Begin work on the OmemoManager interface	2022-12-27 12:43:50 +01:00
Alexander "PapaTutuWawa	06707d1a34	feat: Compute the fingerprint of a bundle	2022-12-27 12:43:18 +01:00
Alexander "PapaTutuWawa	4346b31637	feat: Add funding All checks were successful ci/woodpecker/push/woodpecker Pipeline was successful Details	2022-12-27 01:47:08 +01:00
Alexander "PapaTutuWawa	fe1ba99b14	feat: Indicate whether a ratchet has been created or not All checks were successful ci/woodpecker/push/woodpecker Pipeline was successful Details	2022-12-09 20:41:53 +01:00
Alexander "PapaTutuWawa	797bf69856	fix: Crash when calling getDevicesTrust for unknown Jid All checks were successful ci/woodpecker/push/woodpecker Pipeline was successful Details	2022-12-09 18:42:23 +01:00
Alexander "PapaTutuWawa	afad5056c0	release: Bump version to 0.3.2 All checks were successful ci/woodpecker/push/woodpecker Pipeline was successful Details ci/woodpecker/tag/woodpecker Pipeline was successful Details	2022-11-05 13:50:45 +01:00
Alexander "PapaTutuWawa	c68471349a	fix: Reuse old key exchange when the ratchet is unacked All checks were successful ci/woodpecker/push/woodpecker Pipeline was successful Details	2022-10-22 12:41:41 +02:00
Alexander "PapaTutuWawa	1472624b1d	fix: Use stanza receival timestamps to guard against stale kex messages All checks were successful ci/woodpecker/push/woodpecker Pipeline was successful Details	2022-10-02 19:23:58 +02:00
Alexander "PapaTutuWawa	0826d043d5	feat: Attempt to detect already decrypted messages	2022-10-02 17:03:39 +02:00
Alexander "PapaTutuWawa	2aa3674c4b	fix: Fix receiving an old key exchange breaking decryption This was mostly caused by Dart not copying values but referencing them. AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA. We know make some assumptions about received key exchanges, so this needs some field testing.	2022-10-02 14:56:20 +02:00
Alexander "PapaTutuWawa	7c3a9a75df	chore: Let pub ignore the protobuf build artifacts All checks were successful ci/woodpecker/push/woodpecker Pipeline was successful Details	2022-09-29 15:24:33 +02:00
Alexander "PapaTutuWawa	bc6f98bcd8	release: Version 0.3.1 All checks were successful ci/woodpecker/tag/woodpecker Pipeline was successful Details	2022-09-29 15:24:33 +02:00
Alexander "PapaTutuWawa	a107dfad87	ci: Fix issue with duplicate naming All checks were successful ci/woodpecker/push/woodpecker Pipeline was successful Details	2022-09-28 13:18:39 +02:00
Alexander "PapaTutuWawa	96771cf317	feat: Allow getting the device's fingerprint	2022-09-25 12:50:59 +02:00
Alexander "PapaTutuWawa	bea433e377	fix: React to all decryption errors with ratchet restoration	2022-09-25 11:35:12 +02:00
Alexander "PapaTutuWawa	47948fa6ea	build: Add Woodpecker CI	2022-09-22 21:33:22 +02:00
Alexander "PapaTutuWawa	a23dd30eee	chore: Commit the protobuf artifacts	2022-09-18 15:56:09 +02:00
Alexander "PapaTutuWawa	b4c14a9769	fix: Guard against a crash in the critical section	2022-09-15 23:33:07 +02:00
Alexander "PapaTutuWawa	b69acdd936	fix: Fix replaceOnetimePrekey mutating the device's id	2022-09-15 23:31:56 +02:00
Alexander "PapaTutuWawa	3d8c82fe5b	docs: Update README	2022-09-15 13:52:29 +02:00
Alexander "PapaTutuWawa	1e7f66ccc3	release: Bump to 0.3.0	2022-09-15 13:51:18 +02:00
Alexander "PapaTutuWawa	0480e9156f	fix: Fix occurence of not using synchronized's return	2022-09-15 13:46:52 +02:00
Alexander "PapaTutuWawa	96d9c55c87	fix: Make fromJson* functions work when reading JSON from a String	2022-09-15 13:41:33 +02:00
Alexander "PapaTutuWawa	49c847a96b	feat: Remove toJson and fromJson from OmemoSessionManager	2022-09-15 13:22:50 +02:00
Alexander "PapaTutuWawa	cf5331a026	feat: Introduce logging for logging purposes	2022-09-15 13:17:30 +02:00
Alexander "PapaTutuWawa	c1d8073af0	refactor: Remove BTBV's loadState method	2022-09-15 13:06:14 +02:00
Alexander "PapaTutuWawa	438012d8f8	fix: Hopefully fix all tests being flaky It seems that the varint encoding function would not work for some integers as input. This should in theory fix this issue. Since the SPK IDs are randomly between 0 and 2**32 - 1, it makes sense that the tests fail only sometimes.	2022-09-14 23:50:54 +02:00
Alexander "PapaTutuWawa	79704da99c	fix: Fix issues with the maps being unmodifiable	2022-09-14 22:02:50 +02:00
Alexander "PapaTutuWawa	4341797f14	fix: Commit the restored ratchet	2022-09-14 22:02:35 +02:00
Alexander "PapaTutuWawa	c5c579810e	fix: Restore the ratchets in case of an error This means that if the ratchet fails to decrypt a message, from the outside it will be as if that one message had never been received. Thus, the ratchet can be used normally. This is to guard against messages that are received again.	2022-09-14 21:58:41 +02:00
Alexander "PapaTutuWawa	8991599a0b	feat: Allow removing all ratchets for a given Jid	2022-09-11 17:26:54 +02:00
Alexander "PapaTutuWawa	dad938b0e1	feat: Allow initializing the BTBV trust manager in the constructor	2022-09-11 13:43:07 +02:00
Alexander "PapaTutuWawa	ff52c82039	feat: Help with serializing and deserializing the BTVT manager	2022-09-11 13:33:45 +02:00
Alexander "PapaTutuWawa	12e09947f6	feat: Implement enabling and disabling devices	2022-09-11 12:34:31 +02:00
Alexander "PapaTutuWawa	d530358359	chore: Lower meta's minimum version	2022-09-09 17:40:16 +02:00
Alexander "PapaTutuWawa	0e370a8e19	refactor: Use synchronized's return	2022-09-09 17:39:28 +02:00
Alexander "PapaTutuWawa	b6aa28e1ee	release: Bump version	2022-08-19 17:00:44 +02:00
Alexander "PapaTutuWawa	2e10842c54	feat: Make accepted ratchets unacknowledged by default	2022-08-19 16:59:24 +02:00
Alexander "PapaTutuWawa	0e2af1f2a3	feat: Add a function to check if a ratchet is acknowledged	2022-08-19 16:58:23 +02:00
Alexander "PapaTutuWawa	80e1b20f27	fix: Fix crash when calling getUnacknowledgedRatchets for a new Jid	2022-08-18 16:21:53 +02:00
Alexander "PapaTutuWawa	f68e45af26	docs: Update README	2022-08-18 15:35:07 +02:00
Alexander "PapaTutuWawa	345596923e	release: Bump version to 0.2.0	2022-08-18 15:33:49 +02:00
Alexander "PapaTutuWawa	d5d4aa9014	feat: Add getDeviceId and getDeviceBundle	2022-08-18 15:30:31 +02:00
Alexander "PapaTutuWawa	ee7b09bdb0	feat: Ratchets should overwrite each other	2022-08-18 15:20:32 +02:00
Alexander "PapaTutuWawa	73613e266f	feat: Allow regerating a device's id	2022-08-18 15:08:05 +02:00
Alexander "PapaTutuWawa	0a03483aaf	feat: Allow regenerating one's device identity	2022-08-18 15:02:17 +02:00
Alexander "PapaTutuWawa	fda06cef55	feat: Implement acknowledging ratchet sessions	2022-08-16 14:02:04 +02:00
Alexander "PapaTutuWawa	800b53b11f	style: Fix linter warnings	2022-08-16 13:46:08 +02:00
Alexander "PapaTutuWawa	5a097e4d2a	feat: Allow removing a ratchet session	2022-08-16 13:27:21 +02:00
Alexander "PapaTutuWawa	710b3c9497	feat: Allow serialising and deserialising without the ratchets	2022-08-16 12:57:16 +02:00
Alexander "PapaTutuWawa	f540a80ec2	docs: Remove trust manager serialization from OmemoSessionManager	2022-08-16 12:54:15 +02:00
Alexander "PapaTutuWawa	44ab31aebb	release: Bump version to 0.1.3	2022-08-11 12:06:02 +02:00
Alexander "PapaTutuWawa	e4f1d7d4b0	docs: Fix indent in README	2022-08-11 12:04:40 +02:00
Alexander "PapaTutuWawa	7600804aa1	test: Uncomment those tests	2022-08-11 12:02:33 +02:00
Alexander "PapaTutuWawa	a4589b6e09	feat: Allow access to the device map	2022-08-11 12:02:21 +02:00
Alexander "PapaTutuWawa	d0986a4608	fix: Fix ratchet only working for one message	2022-08-11 11:57:33 +02:00
Alexander "PapaTutuWawa	683a76cc80	test: Add test for sending multiple OMEMO messages	2022-08-10 17:24:35 +02:00
Alexander "PapaTutuWawa	dad707f71d	release: Sigh, bump to 0.1.2	2022-08-09 16:40:41 +02:00
Alexander "PapaTutuWawa	419be8af4d	Revert "meta: Lower Dart SDK requirement" This reverts commit `fafc4f2320`.	2022-08-09 16:40:19 +02:00
Alexander "PapaTutuWawa	fafc4f2320	meta: Lower Dart SDK requirement	2022-08-09 16:32:59 +02:00