chore: Bump version

OmemoManager.

.github/FUNDING.yml

@@ -0,0 +1 @@
+ko_fi: papatutuwawa

.gitignore

@@ -12,6 +12,3 @@ pubspec.lock
 # NixOS
 .direnv
 .envrc
-
-# Protobuf build artifacts
-lib/protobuf/*.dart

.pubignore

@@ -0,0 +1 @@
+lib/protobuf

.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

CHANGELOG.md

@@ -9,7 +9,6 @@
 
 - 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`
@@ -20,3 +19,30 @@
 
 - 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

README.md

@@ -28,7 +28,7 @@ Include `omemo_dart` in your `pubspec.yaml` like this:
 dependencies:
   omemo_dart:
     hosted: https://git.polynom.me/api/packages/PapaTutuWawa/pub
-    version: ^0.2.0
+    version: ^0.4.0
   # [...]
 
 # [...]
@@ -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)

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"

example/omemo_dart_example.dart

@@ -119,6 +119,8 @@ void main() async {
     aliceDevice.id,
     // The deserialised keys
     keys,
+    // Since the message was not delayed, we use the current time
+    DateTime.now().millisecondsSinceEpoch,
   );
 
   // All Bob has to do now is replace the OMEMO wrapper element 

lib/omemo_dart.dart

@@ -10,8 +10,10 @@ 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';

lib/protobuf/schema.pb.dart

@@ -0,0 +1,263 @@
+///
+//  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);
+}
+

lib/protobuf/schema.pbenum.dart

@@ -0,0 +1,7 @@
+///
+//  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
+

lib/protobuf/schema.pbjson.dart

@@ -0,0 +1,48 @@
+///
+//  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 = const {
+  '1': 'OMEMOMessage',
+  '2': const [
+    const {'1': 'n', '3': 1, '4': 2, '5': 13, '10': 'n'},
+    const {'1': 'pn', '3': 2, '4': 2, '5': 13, '10': 'pn'},
+    const {'1': 'dh_pub', '3': 3, '4': 2, '5': 12, '10': 'dhPub'},
+    const {'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 = const {
+  '1': 'OMEMOAuthenticatedMessage',
+  '2': const [
+    const {'1': 'mac', '3': 1, '4': 2, '5': 12, '10': 'mac'},
+    const {'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 = const {
+  '1': 'OMEMOKeyExchange',
+  '2': const [
+    const {'1': 'pk_id', '3': 1, '4': 2, '5': 13, '10': 'pkId'},
+    const {'1': 'spk_id', '3': 2, '4': 2, '5': 13, '10': 'spkId'},
+    const {'1': 'ik', '3': 3, '4': 2, '5': 12, '10': 'ik'},
+    const {'1': 'ek', '3': 4, '4': 2, '5': 12, '10': 'ek'},
+    const {'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==');

lib/protobuf/schema.pbserver.dart

@@ -0,0 +1,9 @@
+///
+//  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';
+

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,7 +53,6 @@ class SkippedKey {
 }
 
 class OmemoDoubleRatchet {
-
   OmemoDoubleRatchet(
     this.dhs, // DHs
     this.dhr, // DHr
@@ -68,6 +66,8 @@ class OmemoDoubleRatchet {
     this.sessionAd,
     this.mkSkipped, // MKSKIPPED
     this.acknowledged,
+    this.kexTimestamp,
+    this.kex,
   );
   
   factory OmemoDoubleRatchet.fromJson(Map<String, dynamic> data) {
@@ -82,9 +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',
@@ -94,11 +96,20 @@ 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(
@@ -120,6 +131,8 @@ class OmemoDoubleRatchet {
       base64.decode(data['session_ad']! as String),
       mkSkipped,
       data['acknowledged']! as bool,
+      data['kex_timestamp']! as int,
+      data['kex'] as String?,
     );
   }
   
@@ -151,6 +164,13 @@ 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;
@@ -158,7 +178,7 @@ class OmemoDoubleRatchet {
   /// 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));
@@ -177,6 +197,8 @@ class OmemoDoubleRatchet {
       ad,
       {},
       false,
+      timestamp,
+      '',
     );
   }
 
@@ -184,7 +206,7 @@ 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,
@@ -197,7 +219,9 @@ class OmemoDoubleRatchet {
       ik,
       ad,
       {},
-      false,
+      true,
+      kexTimestamp,
+      null,
     );
   }
 
@@ -224,9 +248,17 @@ class OmemoDoubleRatchet {
       'session_ad': base64.encode(sessionAd),
       'mkskipped': mkSkippedSerialised,
       'acknowledged': acknowledged,
+      'kex_timestamp': kexTimestamp,
+      'kex': kex,
     };
   }
 
+  /// Returns the OMEMO compatible fingerprint of the ratchet session.
+  Future<String> getOmemoFingerprint() async {
+    final curveKey = await ik.toCurve25519();
+    return HEX.encode(await curveKey.getBytes());
+  }
+  
   Future<List<int>?> _trySkippedMessageKeys(OmemoMessage header, List<int> ciphertext) async {
     final key = SkippedKey(
       OmemoPublicKey.fromBytes(header.dhPub!, KeyPairType.x25519),
@@ -259,18 +291,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.
@@ -304,8 +336,8 @@ class OmemoDoubleRatchet {
     }
 
     final dhPubMatches = listsEqual(
-      header.dhPub ?? <int>[],
-      await dhr?.getBytes() ?? <int>[],
+      header.dhPub!,
+      (await dhr?.getBytes()) ?? <int>[],
     );
     if (!dhPubMatches) {
       await _skipMessageKeys(header.pn!);
@@ -321,12 +353,64 @@ class OmemoDoubleRatchet {
     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 {
+    final dhrMatch = dhr == null ?
+      other.dhr == null :
       // 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!);
+      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);
@@ -342,6 +426,7 @@ class OmemoDoubleRatchet {
       ns == other.ns &&
       nr == other.nr &&
       pn == other.pn &&
-      listsEqual(sessionAd, other.sessionAd);
+      listsEqual(sessionAd, other.sessionAd) &&
+      kexTimestamp == other.kexTimestamp;
   }
 }

lib/src/errors.dart

@@ -1,32 +1,54 @@
+abstract class OmemoException {}
+
 /// Triggered during X3DH if the signature if the SPK does verify to the actual SPK.
-class InvalidSignatureException implements Exception {
+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';
+}

lib/src/helpers.dart

@@ -73,3 +73,7 @@ OmemoKeyPair? decodeKeyPairIfNotNull(String? pk, String? sk, KeyPairType type) {
     type,
   );
 }
+
+int getTimestamp() {
+  return DateTime.now().millisecondsSinceEpoch;
+}

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) {
@@ -55,7 +54,6 @@ class OmemoPublicKey {
 }
 
 class OmemoPrivateKey {
-
   const OmemoPrivateKey(this._privkey, this.type);
   final List<int> _privkey;
   final KeyPairType type;
@@ -85,7 +83,6 @@ class OmemoPrivateKey {
 
 /// 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

lib/src/omemo/bundle.dart

@@ -1,5 +1,6 @@
 import 'dart:convert';
 import 'package:cryptography/cryptography.dart';
+import 'package:hex/hex.dart';
 import 'package:omemo_dart/src/keys.dart';
 
 class OmemoBundle {
@@ -43,4 +44,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());
+  }
 }

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';

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;
+}

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),
+    // 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,7 @@ 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 +105,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
@@ -126,12 +135,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 +154,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,
@@ -166,8 +175,8 @@ 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
-  Device withNewId() {
-    return Device(
+  OmemoDevice withNewId() {
+    return OmemoDevice(
       jid,
       generateRandom32BitNumber(),
       ik,
@@ -199,6 +208,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
@@ -228,7 +244,7 @@ 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;

lib/src/omemo/encryption_result.dart

@@ -1,15 +1,26 @@
 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;
 }

lib/src/omemo/events.dart

@@ -5,32 +5,31 @@ 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);
   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;
 }
 
 /// 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;
 }

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;

lib/src/omemo/omemomanager.dart

@@ -0,0 +1,786 @@
+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.payload);
+  final bool ratchetCreated;
+  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));
+  }
+
+  /// 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(),
+    );
+
+    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,
+      ),
+    );
+  }
+  
+  /// 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 ratchetKey = RatchetMapKey(senderJid, senderDeviceId);
+    final decodedRawKey = base64.decode(rawKey.value);
+    List<int>? keyAndHmac;
+    OmemoAuthenticatedMessage authMessage;
+    OmemoDoubleRatchet? oldRatchet;
+    OmemoMessage? message;
+    var ratchetCreated = false;
+    if (rawKey.kex) {
+      // If the ratchet already existed, we store it. If it didn't, oldRatchet will stay
+      // null.
+      final oldRatchet = _getRatchet(ratchetKey)?.clone();
+      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();
+        }
+        
+        // Try to decrypt it
+        try {
+          final decrypted = await oldRatchet.ratchetDecrypt(message, authMessage.writeToBuffer());
+
+          // Commit the ratchet
+          _eventStreamController.add(
+            RatchetModifiedEvent(
+              senderJid,
+              senderDeviceId,
+              oldRatchet,
+              false,
+            ),
+          );
+          
+          final plaintext = await _decryptAndVerifyHmac(
+            ciphertext,
+            decrypted,
+          );
+          _addSession(senderJid, senderDeviceId, oldRatchet);
+          return _InternalDecryptionResult(
+            true,
+            plaintext,
+          );
+        } catch (_) {
+          _log.finest('Failed to use old ratchet with KEX for existing ratchet');
+        }
+      }
+
+      final r = await _addSessionFromKeyExchange(senderJid, senderDeviceId, kex);
+      await _trustManager.onNewSession(senderJid, senderDeviceId);
+      _addSession(senderJid, senderDeviceId, r);
+      ratchetCreated = true;
+
+      // Replace the OPK
+      // TODO(PapaTutuWawa): Replace the OPK when we know that the KEX worked
+      await _deviceLock.synchronized(() async {
+        device = await device.replaceOnetimePrekey(kex.pkId!);
+
+        // Commit the device
+        _eventStreamController.add(DeviceModifiedEvent(device));
+      });
+    } else {
+      authMessage = OmemoAuthenticatedMessage.fromBuffer(decodedRawKey);
+      message = OmemoMessage.fromBuffer(authMessage.message!);
+    }
+    
+    final devices = _deviceList[senderJid];
+    if (devices?.contains(senderDeviceId) != true) {
+      throw NoDecryptionKeyException();
+    }
+
+    // We can guarantee that the ratchet exists at this point in time
+    final ratchet = _getRatchet(ratchetKey)!;
+    oldRatchet ??= ratchet.clone();
+
+    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,
+      ),
+    );
+
+    try {
+      return _InternalDecryptionResult(
+        ratchetCreated,
+        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.
+  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));
+      }
+    }
+
+    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) {
+        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));
+    } 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(jid);
+      _eventStreamController.add(RatchetRemovedEvent(jid, deviceId));
+    }
+
+    // Remove the devices from the device list cache and commit it
+    _deviceList.remove(jid);
+    _deviceListRequested.remove(jid);
+    _eventStreamController.add(DeviceListModifiedEvent(_deviceList));
+
+    // Remove trust decisions
+    await _trustManager.removeTrustDecisionsForJid(jid);
+    
+    await _leaveRatchetCriticalSection(jid);
+  }
+
+  /// Replaces the internal device with [newDevice]. Does not trigger an event.
+  Future<void> replaceDevice(OmemoDevice newDevice) async {
+    await _deviceLock.synchronized(() {
+      _device = newDevice;
+    });
+  }
+}

lib/src/omemo/ratchet_map_key.dart

@@ -2,7 +2,6 @@ import 'package:meta/meta.dart';
 
 @immutable
 class RatchetMapKey {
-
   const RatchetMapKey(this.jid, this.deviceId);
 
   factory RatchetMapKey.fromJsonKey(String key) {

lib/src/omemo/sessionmanager.dart

@@ -3,6 +3,7 @@ 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';
@@ -10,6 +11,7 @@ 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/device.dart';
 import 'package:omemo_dart/src/omemo/encrypted_key.dart';
 import 'package:omemo_dart/src/omemo/encryption_result.dart';
@@ -23,39 +25,35 @@ 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';
-
+@Deprecated('Use OmemoManager instead')
 class OmemoSessionManager {
-
+  @Deprecated('Use OmemoManager instead')
   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,
-    );
-  }
+      _eventStreamController = StreamController<OmemoEvent>.broadcast(),
+      _log = Logger('OmemoSessionManager');
 
   /// Deserialise the OmemoSessionManager from JSON data [data] that does not contain
   /// the ratchet sessions.
-  factory OmemoSessionManager.fromJsonWithoutSessions(Map<String, dynamic> data, Map<RatchetMapKey, OmemoDoubleRatchet> ratchetMap, TrustManager trustManager) {
+  @Deprecated('Use OmemoManager instead')
+  factory OmemoSessionManager.fromJsonWithoutSessions(
+    Map<String, dynamic> data,
+    Map<RatchetMapKey, OmemoDoubleRatchet> ratchetMap,
+    TrustManager trustManager,
+  ) {
+    // NOTE: Dart has some issues with just casting a List<dynamic> to List<Map<...>>, as
+    //       such we need to convert the items by hand.
     return OmemoSessionManager(
-      Device.fromJson(data['device']! as Map<String, dynamic>),
-      data['devices']! as Map<String, List<int>>,
+      OmemoDevice.fromJson(data['device']! as Map<String, dynamic>),
+      (data['devices']! as Map<String, dynamic>).map<String, List<int>>(
+        (key, value) {
+          return MapEntry(
+            key,
+            (value as List<dynamic>).map<int>((i) => i as int).toList(),
+          );
+        }
+      ),
       ratchetMap,
       trustManager,
     );
@@ -64,11 +62,14 @@ class OmemoSessionManager {
   /// 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);
+    final device = await OmemoDevice.generateNewDevice(jid, opkAmount: opkAmount);
 
     return OmemoSessionManager(device, {}, {}, trustManager);
   }
 
+  /// Logging
+  Logger _log;
+  
   /// Lock for _ratchetMap and _bundleMap
   final Lock _lock;
 
@@ -83,7 +84,7 @@ class OmemoSessionManager {
   
   /// Our own keys...
   // ignore: prefer_final_fields
-  Device _device;
+  OmemoDevice _device;
   /// and its lock
   final Lock _deviceLock;
 
@@ -95,13 +96,8 @@ class OmemoSessionManager {
   Stream<OmemoEvent> get eventStream => _eventStreamController.stream;
 
   /// Returns our own device.
-  Future<Device> getDevice() async {
-    Device? dev;
-    await _deviceLock.synchronized(() async {
-      dev = _device;
-    });
-
-    return dev!;
+  Future<OmemoDevice> getDevice() async {
+    return _deviceLock.synchronized(() => _device);
   }
 
   /// Returns the id attribute of our own device. This is just a short-hand for
@@ -124,14 +120,14 @@ class OmemoSessionManager {
         _deviceMap[jid] = [deviceId];
 
         // Commit the device map
-        _eventStreamController.add(DeviceMapModifiedEvent(_deviceMap));
+        _eventStreamController.add(DeviceListModifiedEvent(_deviceMap));
       } else {
         // Prevent having the same device multiple times in the list
         if (!_deviceMap[jid]!.contains(deviceId)) {
           _deviceMap[jid]!.add(deviceId);
 
           // Commit the device map
-          _eventStreamController.add(DeviceMapModifiedEvent(_deviceMap));
+          _eventStreamController.add(DeviceListModifiedEvent(_deviceMap));
         }
       }
 
@@ -140,7 +136,7 @@ class OmemoSessionManager {
       _ratchetMap[key] = ratchet;
 
       // Commit the ratchet
-      _eventStreamController.add(RatchetModifiedEvent(jid, deviceId, ratchet));
+      _eventStreamController.add(RatchetModifiedEvent(jid, deviceId, ratchet, true));
     });
   }
 
@@ -158,6 +154,7 @@ class OmemoSessionManager {
       bundle.ik,
       kexResult.sk,
       kexResult.ad,
+      getTimestamp(),
     );
 
     await _trustManager.onNewSession(jid, deviceId);
@@ -173,17 +170,17 @@ class OmemoSessionManager {
   /// 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 {
+  Future<OmemoDoubleRatchet> _addSessionFromKeyExchange(String jid, int deviceId, OmemoKeyExchange kex) async {
     // Pick the correct SPK
     final device = await getDevice();
-    OmemoKeyPair? spk;
-
-    await _lock.synchronized(() async {
+    final spk = await _lock.synchronized(() async {
       if (kex.spkId == _device.spkId) {
-        spk = _device.spk;
+        return _device.spk;
       } else if (kex.spkId == _device.oldSpkId) {
-        spk = _device.oldSpk;
+        return _device.oldSpk;
       }
+
+      return null;
     });
     if (spk == null) {
       throw UnknownSignedPrekeyException();
@@ -195,19 +192,19 @@ class OmemoSessionManager {
         OmemoPublicKey.fromBytes(kex.ek!, KeyPairType.x25519),
         kex.pkId!,
       ),
-      spk!,
+      spk,
       device.opks.values.elementAt(kex.pkId!),
       device.ik,
     );
     final ratchet = await OmemoDoubleRatchet.acceptNewSession(
-      spk!,
+      spk,
       OmemoPublicKey.fromBytes(kex.ik!, KeyPairType.ed25519),
       kexResult.sk,
       kexResult.ad,
+      getTimestamp(),
     );
 
-    await _trustManager.onNewSession(jid, deviceId);
-    await _addSession(jid, deviceId, ratchet);
+    return ratchet;
   }
 
   /// Like [encryptToJids] but only for one Jid [jid].
@@ -244,7 +241,11 @@ class OmemoSessionManager {
     final kex = <int, OmemoKeyExchange>{};
     if (newSessions != null) {
       for (final newSession in newSessions) {
-        kex[newSession.id] = await addSessionFromBundle(newSession.jid, newSession.id, newSession);
+        kex[newSession.id] = await addSessionFromBundle(
+          newSession.jid,
+          newSession.id,
+          newSession,
+        );
       }
     }
     
@@ -262,24 +263,42 @@ class OmemoSessionManager {
           }
 
           final ratchetKey = RatchetMapKey(jid, deviceId);
-          final ratchet = _ratchetMap[ratchetKey]!;
+          var 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)) {
+            // The ratchet did not exist
             final k = kex[deviceId]!
               ..message = OmemoAuthenticatedMessage.fromBuffer(ciphertext);
+            final buffer = base64.encode(k.writeToBuffer());
             encryptedKeys.add(
               EncryptedKey(
                 jid,
                 deviceId,
-                base64.encode(k.writeToBuffer()),
+                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,
@@ -289,6 +308,20 @@ class OmemoSessionManager {
                 ),
               );
             }
+          } 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));
         }
       }
     });
@@ -296,6 +329,8 @@ class OmemoSessionManager {
     return EncryptionResult(
       plaintext != null ? ciphertext : null,
       encryptedKeys,
+      const <RatchetMapKey, OmemoException>{},
+      const <String, OmemoException>{},
     );
   }
 
@@ -304,7 +339,7 @@ class OmemoSessionManager {
   /// [mapKey] with [oldRatchet].
   Future<void> _restoreRatchet(RatchetMapKey mapKey, OmemoDoubleRatchet oldRatchet) async {
     await _lock.synchronized(() {
-      print('RESTORING RATCHETS');
+      _log.finest('Restoring ratchet ${mapKey.jid}:${mapKey.deviceId} to ${oldRatchet.nr}');
       _ratchetMap[mapKey] = oldRatchet;
 
       // Commit the ratchet
@@ -313,21 +348,44 @@ class OmemoSessionManager {
           mapKey.jid,
           mapKey.deviceId,
           oldRatchet,
+          false,
         ),
       );
     });
   }
 
+  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),
+    );
+  }
+  
   /// 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<String?> decryptMessage(List<int>? ciphertext, String senderJid, int senderDeviceId, List<EncryptedKey> keys) async {
+  Future<String?> 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);
@@ -337,24 +395,56 @@ class OmemoSessionManager {
 
     final ratchetKey = RatchetMapKey(senderJid, senderDeviceId);
     final decodedRawKey = base64.decode(rawKey.value);
+    List<int>? keyAndHmac;
     OmemoAuthenticatedMessage authMessage;
     OmemoDoubleRatchet? oldRatchet;
+    OmemoMessage? message;
     if (rawKey.kex) {
       // If the ratchet already existed, we store it. If it didn't, oldRatchet will stay
       // null.
-      oldRatchet = await _getRatchet(ratchetKey);
-
-      // TODO(PapaTutuWawa): Only do this when we should
+      final oldRatchet = (await _getRatchet(ratchetKey))?.clone();
       final kex = OmemoKeyExchange.fromBuffer(decodedRawKey);
-      await _addSessionFromKeyExchange(
+      authMessage = kex.message!;
+      message = OmemoMessage.fromBuffer(authMessage.message!);
+
+      // Guard against old key exchanges
+      if (oldRatchet != null) {
+        _log.finest('KEX for existent ratchet. ${oldRatchet.pn}');
+        if (oldRatchet.kexTimestamp > timestamp) {
+          throw InvalidKeyExchangeException();
+        }
+        
+        // Try to decrypt it
+        try {
+          final decrypted = await oldRatchet.ratchetDecrypt(message, authMessage.writeToBuffer());
+
+          // Commit the ratchet
+          _eventStreamController.add(
+            RatchetModifiedEvent(
               senderJid,
               senderDeviceId,
-        kex,
+              oldRatchet,
+              false,
+            ),
           );
           
-      authMessage = kex.message!;
+          final plaintext = await _decryptAndVerifyHmac(
+            ciphertext,
+            decrypted,
+          );
+          await _addSession(senderJid, senderDeviceId, oldRatchet);
+          return plaintext;
+        } catch (_) {
+          _log.finest('Failed to use old ratchet with KEX for existing ratchet');
+        }
+      }
+
+      final r = await _addSessionFromKeyExchange(senderJid, senderDeviceId, kex);
+      await _trustManager.onNewSession(senderJid, senderDeviceId);
+      await _addSession(senderJid, senderDeviceId, r);
 
       // Replace the OPK
+      // TODO(PapaTutuWawa): Replace the OPK when we know that the KEX worked
       await _deviceLock.synchronized(() async {
         device = await device.replaceOnetimePrekey(kex.pkId!);
 
@@ -363,6 +453,7 @@ class OmemoSessionManager {
       });
     } else {
       authMessage = OmemoAuthenticatedMessage.fromBuffer(decodedRawKey);
+      message = OmemoMessage.fromBuffer(authMessage.message!);
     }
     
     final devices = _deviceMap[senderJid];
@@ -373,11 +464,9 @@ class OmemoSessionManager {
       throw NoDecryptionKeyException();
     }
 
-    final message = OmemoMessage.fromBuffer(authMessage.message!);
-    List<int>? keyAndHmac;
     // We can guarantee that the ratchet exists at this point in time
     final ratchet = (await _getRatchet(ratchetKey))!;
-    oldRatchet ??= ratchet ;
+    oldRatchet ??= ratchet.clone();
 
     try {
       if (rawKey.kex) {
@@ -385,32 +474,27 @@ class OmemoSessionManager {
       } else {
         keyAndHmac = await ratchet.ratchetDecrypt(message, decodedRawKey);
       }
-    } on InvalidMessageHMACException {
+    } catch (_) {
       await _restoreRatchet(ratchetKey, oldRatchet);
       rethrow;
     }
 
     // Commit the ratchet
-    _eventStreamController.add(RatchetModifiedEvent(senderJid, senderDeviceId, ratchet));
+    _eventStreamController.add(
+      RatchetModifiedEvent(
+        senderJid,
+        senderDeviceId,
+        ratchet,
+        false,
+      ),
+    );
 
-    // 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)) {
-      // TODO(PapaTutuWawa): I am unsure if we should restore the ratchet here
+    try {
+      return _decryptAndVerifyHmac(ciphertext, keyAndHmac);
+    } catch (_) { 
       await _restoreRatchet(ratchetKey, oldRatchet);
-      throw InvalidMessageHMACException();
+      rethrow;
     }
-    
-    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].
@@ -419,7 +503,7 @@ class OmemoSessionManager {
 
     await _lock.synchronized(() async {
       // Get devices for jid
-      final devices = _deviceMap[jid]!;
+      final devices = _deviceMap[jid] ?? [];
 
       for (final deviceId in devices) {
         final ratchet = _ratchetMap[RatchetMapKey(jid, deviceId)]!;
@@ -436,6 +520,16 @@ class OmemoSessionManager {
     return fingerprints;
   }
 
+  /// Returns the hex-encoded fingerprint of the current device.
+  Future<DeviceFingerprint> getHexFingerprintForDevice() async {
+    final device = await getDevice();
+
+    return DeviceFingerprint(
+      device.id,
+      HEX.encode(await device.ik.pk.getBytes()),
+    );
+  }
+
   /// 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
@@ -471,7 +565,7 @@ class OmemoSessionManager {
         _deviceMap.remove(jid);
       }
       // Commit it
-      _eventStreamController.add(DeviceMapModifiedEvent(_deviceMap));
+      _eventStreamController.add(DeviceListModifiedEvent(_deviceMap));
     });
   }
 
@@ -489,7 +583,7 @@ class OmemoSessionManager {
       // Remove the device from jid
       _deviceMap.remove(jid);
       // Commit it
-      _eventStreamController.add(DeviceMapModifiedEvent(_deviceMap));
+      _eventStreamController.add(DeviceListModifiedEvent(_deviceMap));
     });
   }
   
@@ -523,7 +617,7 @@ class OmemoSessionManager {
         ..acknowledged = true;
 
       // Commit it
-      _eventStreamController.add(RatchetModifiedEvent(jid, deviceId, ratchet));
+      _eventStreamController.add(RatchetModifiedEvent(jid, deviceId, ratchet, false));
     });
   }
 
@@ -531,7 +625,7 @@ class OmemoSessionManager {
   /// identity. Triggers an event to commit it to storage.
   Future<void> regenerateDevice({ int opkAmount = 100 }) async {
     await _deviceLock.synchronized(() async {
-      _device = await Device.generateNewDevice(_device.jid, opkAmount: opkAmount);
+      _device = await OmemoDevice.generateNewDevice(_device.jid, opkAmount: opkAmount);
 
       // Commit it
       _eventStreamController.add(DeviceModifiedEvent(_device));
@@ -561,42 +655,6 @@ class OmemoSessionManager {
   @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': { ... },
-        },
-        ...
-      ],
-    }
-    */
-
-    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,
-    };
-  }
-
   /// Serialise the entire session manager into a JSON object.
   Future<Map<String, dynamic>> toJsonWithoutSessions() async {
     /*

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;
+}

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) {

lib/src/trust/always.dart

@@ -18,6 +18,9 @@ class AlwaysTrustingTrustManager extends TrustManager {
   @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>{};
 }

lib/src/trust/base.dart

@@ -19,4 +19,7 @@ abstract class TrustManager {
 
   /// Serialize the trust manager to JSON.
   Future<Map<String, dynamic>> toJson();
+
+  /// Removes all trust decisions for [jid].
+  Future<void> removeTrustDecisionsForJid(String jid);
 }

lib/src/trust/btbv.dart

@@ -122,17 +122,20 @@ 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 {
+    return _lock.synchronized(() async {
       final map = <int, BTBVTrustState>{};
 
-    await _lock.synchronized(() async {
+      if (!devices.containsKey(jid)) return map;
+      
       for (final deviceId in devices[jid]!) {
         map[deviceId] = trustCache[RatchetMapKey(jid, deviceId)]!;
       }
-    });
 
       return map;
+    });
   }
 
   /// Sets the trust of [jid]'s device with identifier [deviceId] to [state].
@@ -209,16 +212,19 @@ abstract class BlindTrustBeforeVerificationTrustManager extends TrustManager {
     );
   }
 
+  @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)]!;
 }
@@ -228,7 +234,4 @@ abstract class BlindTrustBeforeVerificationTrustManager extends TrustManager {
 class MemoryBTBVTrustManager extends BlindTrustBeforeVerificationTrustManager {
   @override
   Future<void> commitState() async {}
-
-  @override
-  Future<void> loadState() async {}
 }

lib/src/trust/never.dart

@@ -18,6 +18,9 @@ class NeverTrustingTrustManager extends TrustManager {
   @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>{};
 }

pubspec.yaml

@@ -1,6 +1,6 @@
 name: omemo_dart
 description: An XMPP library independent OMEMO library
-version: 0.2.1
+version: 0.4.1
 homepage: https://github.com/PapaTutuWawa/omemo_dart
 publish_to: https://git.polynom.me/api/packages/PapaTutuWawa/pub
 
@@ -11,6 +11,7 @@ dependencies:
   collection: ^1.16.0
   cryptography: ^2.0.5
   hex: ^0.2.0
+  logging: ^1.0.2
   meta: ^1.7.0
   pinenacl: ^0.5.1
   synchronized: ^3.0.0+2

test/double_ratchet_test.dart

@@ -84,12 +84,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);

test/omemo_test.dart

@@ -1,9 +1,47 @@
+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';
@@ -27,7 +65,7 @@ void main() {
         deviceModified = true;
       } else if (event is RatchetModifiedEvent) {
         ratchetModified++;
-      } else if (event is DeviceMapModifiedEvent) {
+      } else if (event is DeviceListModifiedEvent) {
         deviceMapModified++;
       }
     });
@@ -52,6 +90,7 @@ void main() {
       aliceJid,
       await aliceSession.getDeviceId(),
       aliceMessage.encryptedKeys,
+      0,
     );
     expect(messagePlaintext, bobMessage);
     // The ratchet should be modified two times: Once for when the ratchet is created and
@@ -67,6 +106,10 @@ 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(
@@ -83,6 +126,7 @@ void main() {
       bobJid,
       await bobSession.getDeviceId(),
       bobResponseMessage.encryptedKeys,
+      0,
     );
     expect(bobResponseText, aliceReceivedMessage);
   });
@@ -132,9 +176,14 @@ void main() {
       aliceJid,
       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(
@@ -151,6 +200,7 @@ void main() {
       bobJid,
       await bobSession.getDeviceId(),
       bobResponseMessage.encryptedKeys,
+      0,
     );
     expect(bobResponseText, aliceReceivedMessage);
 
@@ -207,6 +257,7 @@ void main() {
       aliceJid,
       await aliceSession1.getDeviceId(),
       aliceMessage.encryptedKeys,
+      0,
     );
     expect(messagePlaintext, bobMessage);
 
@@ -216,6 +267,7 @@ void main() {
       aliceJid,
       await aliceSession1.getDeviceId(),
       aliceMessage.encryptedKeys,
+      0,
     );
     expect(messagePlaintext, aliceMessage2);
   });
@@ -256,6 +308,7 @@ void main() {
       aliceJid,
       await aliceSession.getDeviceId(),
       aliceMessage.encryptedKeys,
+      0,
     );
     expect(bobMessage, null);
 
@@ -274,7 +327,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;
@@ -333,6 +386,7 @@ void main() {
       aliceJid,
       await aliceSession.getDeviceId(),
       aliceMessage.encryptedKeys,
+      0,
     );
     expect(messagePlaintext, bobMessage);
   });
@@ -399,8 +453,13 @@ void main() {
       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
@@ -418,6 +477,7 @@ void main() {
         bobJid,
         await bobSession.getDeviceId(),
         bobResponseMessage.encryptedKeys,
+        0,
       );
       expect(messageText, aliceReceivedMessage);
     }
@@ -572,6 +632,7 @@ void main() {
       aliceJid,
       await aliceSession.getDeviceId(),
       msg1.encryptedKeys,
+      0,
     );
     final aliceRatchet1 = aliceSession.getRatchet(
       bobJid,
@@ -596,6 +657,7 @@ void main() {
       aliceJid,
       await aliceSession.getDeviceId(),
       msg2.encryptedKeys,
+      getTimestamp(),
     );
     final aliceRatchet2 = aliceSession.getRatchet(
       bobJid,
@@ -615,7 +677,7 @@ void main() {
     expect(await bobRatchet1.equals(bobRatchet2), false);
   });
 
-  test('Test receiving an old message that contains a KEX', () async {
+  test('Test resending key exchanges', () async {
     const aliceJid = 'alice@server.example';
     const bobJid = 'bob@other.server.example';
     // Alice and Bob generate their sessions
@@ -638,39 +700,144 @@ void main() {
         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(),
     );
     
-    // Due to some issue with the transport protocol, the first message Bob received is
-    // received again
-    try {
-      await bobSession.decryptMessage(
-        msg1.ciphertext,
+    // 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(),
-        msg1.encryptedKeys,
+        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 {
-      // NOOP
+      } 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,
@@ -681,12 +848,13 @@ void main() {
       aliceJid,
       await aliceSession.getDeviceId(),
       msg3.encryptedKeys,
+      104,
     );
 
     expect(result, 'Are you okay?');
   });
 
-  test('Test receiving an old message that does not contain a KEX', () async {
+  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
@@ -698,7 +866,7 @@ void main() {
     final bobSession = await OmemoSessionManager.generateNewIdentity(
       bobJid,
       AlwaysTrustingTrustManager(),
-      opkAmount: 2,
+      opkAmount: 1,
     );
 
     // Alice sends Bob a message
@@ -709,50 +877,70 @@ void main() {
         await bobSession.getDeviceBundle(),
       ],
     );
+    expect(msg1.encryptedKeys.first.kex, true);
+
     await bobSession.decryptMessage(
       msg1.ciphertext,
       aliceJid,
       await aliceSession.getDeviceId(),
       msg1.encryptedKeys,
+      getTimestamp(),
     );
 
-    // Bob responds
-    final msg2 = await bobSession.encryptToJid(
+    // 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,
-      'Hello!',
-    );
-    await aliceSession.decryptMessage(
-      msg2.ciphertext,
-      bobJid,
-      await bobSession.getDeviceId(),
+      await aliceSession.getDeviceId(),
       msg2.encryptedKeys,
+      getTimestamp(),
     );
 
-    // Due to some issue with the transport protocol, the first message Alice received is
-    // received again.
-    try {
-      await aliceSession.decryptMessage(
-        msg2.ciphertext,
-        bobJid,
-        await bobSession.getDeviceId(),
-        msg2.encryptedKeys,
-      );
-      expect(true, false);
-    } catch (_) {
-      // NOOP
-    }
+    // 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,
-      'Are you okay?',
+      "You read the message, didn't you?",
     );
-    final result = await bobSession.decryptMessage(
+    expect(msg3.encryptedKeys.first.kex, false);
+
+    await bobSession.decryptMessage(
       msg3.ciphertext,
       aliceJid,
       await aliceSession.getDeviceId(),
       msg3.encryptedKeys,
+      getTimestamp(),
     );
 
-    expect(result, 'Are you okay?');
+    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);
+    }
   });
 }

test/omemomanager_test.dart

@@ -0,0 +1,945 @@
+import 'dart:convert';
+import 'package:logging/logging.dart';
+import 'package:omemo_dart/omemo_dart.dart';
+import 'package:omemo_dart/src/trust/always.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 sending a message without the device list cache', () async {
+    const aliceJid = 'alice@server1';
+    const bobJid = 'bob@server2';
+    var aliceEmptyMessageSent = 0;
+    var bobEmptyMessageSent = 0;
+
+    final aliceDevice = await OmemoDevice.generateNewDevice(aliceJid, opkAmount: 1);
+    final bobDevice = await OmemoDevice.generateNewDevice(bobJid, opkAmount: 1);
+
+    final aliceManager = OmemoManager(
+      aliceDevice,
+      AlwaysTrustingTrustManager(),
+      (result, recipientJid) async {
+        aliceEmptyMessageSent++;
+      },
+      (jid) async {
+        expect(jid, bobJid);
+        return [ bobDevice.id ];
+      },
+      (jid, id) async {
+        expect(jid, bobJid);
+        return bobDevice.toBundle();
+      },
+      (jid) async {},
+    );
+    final bobManager = OmemoManager(
+      bobDevice,
+      AlwaysTrustingTrustManager(),
+      (result, recipientJid) async {
+        bobEmptyMessageSent++;
+      },
+      (jid) async {
+        expect(jid, aliceJid);
+        return [aliceDevice.id];
+      },
+      (jid, id) async {
+        expect(jid, aliceJid);
+        return aliceDevice.toBundle();
+      },
+      (jid) async {},
+    );
+
+    // Alice sends a message
+    final aliceResult = await aliceManager.onOutgoingStanza(
+      const OmemoOutgoingStanza(
+        [bobJid],
+        'Hello world',
+      ),
+    );
+
+    // Bob must be able to decrypt the message
+    final bobResult = await bobManager.onIncomingStanza(
+      OmemoIncomingStanza(
+        aliceJid,
+        aliceDevice.id,
+        DateTime.now().millisecondsSinceEpoch,
+        aliceResult.encryptedKeys,
+        base64.encode(aliceResult.ciphertext!),
+      ),
+    );
+
+    expect(bobResult.payload, 'Hello world');
+    expect(bobResult.error, null);
+    expect(aliceEmptyMessageSent, 0);
+    expect(bobEmptyMessageSent, 1);
+
+    // Alice receives the ack message
+    await aliceManager.ratchetAcknowledged(
+      bobJid,
+      bobDevice.id,
+    );
+    
+    // Bob now responds
+    final bobResult2 = await bobManager.onOutgoingStanza(
+      const OmemoOutgoingStanza(
+        [aliceJid],
+        'Hello world, Alice',
+      ),
+    );
+    final aliceResult2 = await aliceManager.onIncomingStanza(
+      OmemoIncomingStanza(
+        bobJid,
+        bobDevice.id,
+        DateTime.now().millisecondsSinceEpoch,
+        bobResult2.encryptedKeys,
+        base64.encode(bobResult2.ciphertext!),
+      ),
+    );
+
+    expect(aliceResult2.error, null);
+    expect(aliceEmptyMessageSent, 0);
+    expect(bobEmptyMessageSent, 1);
+    expect(aliceResult2.payload, 'Hello world, Alice');
+  });
+
+  test('Test triggering the heartbeat', () async {
+    const aliceJid = 'alice@server1';
+    const bobJid = 'bob@server2';
+    var aliceEmptyMessageSent = 0;
+    var bobEmptyMessageSent = 0;
+
+    final aliceDevice = await OmemoDevice.generateNewDevice(aliceJid, opkAmount: 1);
+    final bobDevice = await OmemoDevice.generateNewDevice(bobJid, opkAmount: 1);
+
+    final aliceManager = OmemoManager(
+      aliceDevice,
+      AlwaysTrustingTrustManager(),
+      (result, recipientJid) async {
+        aliceEmptyMessageSent++;
+      },
+      (jid) async {
+        expect(jid, bobJid);
+        return [ bobDevice.id ];
+      },
+      (jid, id) async {
+        expect(jid, bobJid);
+        return bobDevice.toBundle();
+      },
+      (jid) async {},
+    );
+    final bobManager = OmemoManager(
+      bobDevice,
+      AlwaysTrustingTrustManager(),
+      (result, recipientJid) async {
+        bobEmptyMessageSent++;
+      },
+      (jid) async {
+        expect(jid, aliceJid);
+        return [aliceDevice.id];
+      },
+      (jid, id) async {
+        expect(jid, aliceJid);
+        return aliceDevice.toBundle();
+      },
+      (jid) async {},
+    );
+
+    // Alice sends a message
+    final aliceResult = await aliceManager.onOutgoingStanza(
+      const OmemoOutgoingStanza(
+        [bobJid],
+        'Hello world',
+      ),
+    );
+
+    // Bob must be able to decrypt the message
+    final bobResult = await bobManager.onIncomingStanza(
+      OmemoIncomingStanza(
+        aliceJid,
+        aliceDevice.id,
+        DateTime.now().millisecondsSinceEpoch,
+        aliceResult.encryptedKeys,
+        base64.encode(aliceResult.ciphertext!),
+      ),
+    );
+
+    expect(aliceEmptyMessageSent, 0);
+    expect(bobEmptyMessageSent, 1);
+    expect(bobResult.payload, 'Hello world');
+
+    // Bob acknowledges the message
+    await aliceManager.ratchetAcknowledged(bobJid, bobDevice.id);
+    
+    // Alice now sends 52 messages that Bob decrypts
+    for (var i = 0; i <= 51; i++) {
+      final aliceResultLoop = await aliceManager.onOutgoingStanza(
+        OmemoOutgoingStanza(
+          [bobJid],
+          'Test message $i',
+        ),
+      );
+
+      final bobResultLoop = await bobManager.onIncomingStanza(
+        OmemoIncomingStanza(
+          aliceJid,
+          aliceDevice.id,
+          DateTime.now().millisecondsSinceEpoch,
+          aliceResultLoop.encryptedKeys,
+          base64.encode(aliceResultLoop.ciphertext!),
+        ),
+      );
+
+      expect(aliceEmptyMessageSent, 0);
+      expect(bobEmptyMessageSent, 1);
+      expect(bobResultLoop.payload, 'Test message $i');
+    }
+
+    // Alice sends a final message that triggers a heartbeat
+    final aliceResultFinal = await aliceManager.onOutgoingStanza(
+      const OmemoOutgoingStanza(
+        [bobJid],
+        'Test message last',
+      ),
+    );
+
+    final bobResultFinal = await bobManager.onIncomingStanza(
+      OmemoIncomingStanza(
+        aliceJid,
+        aliceDevice.id,
+        DateTime.now().millisecondsSinceEpoch,
+        aliceResultFinal.encryptedKeys,
+        base64.encode(aliceResultFinal.ciphertext!),
+      ),
+    );
+
+    expect(aliceEmptyMessageSent, 0);
+    expect(bobEmptyMessageSent, 2);
+    expect(bobResultFinal.payload, 'Test message last');
+  });
+
+  test('Test accessing data without it existing', () async {
+    const aliceJid = 'alice@server1';
+    const bobJid = 'bob@server2';
+    final aliceDevice = await OmemoDevice.generateNewDevice(aliceJid, opkAmount: 1);
+
+    final aliceManager = OmemoManager(
+      aliceDevice,
+      AlwaysTrustingTrustManager(),
+      (result, recipientJid) async {},
+      (jid) async => [],
+      (jid, id) async => null,
+      (jid) async {},
+    );
+
+    // Get non-existant fingerprints
+    expect(
+      await aliceManager.getFingerprintsForJid(bobJid),
+      null,
+    );
+
+    // Ack a non-existant ratchet
+    await aliceManager.ratchetAcknowledged(
+      bobJid,
+      42,
+    );
+  });
+
+  test('Test receiving a message encrypted for another device', () async {
+    const aliceJid = 'alice@server1';
+    const bobJid = 'bob@server2';
+    var oldDevice = true;
+
+    final aliceDevice = await OmemoDevice.generateNewDevice(aliceJid, opkAmount: 1);
+    final bobOldDevice = await OmemoDevice.generateNewDevice(bobJid, opkAmount: 1);
+    final bobCurrentDevice = await OmemoDevice.generateNewDevice(bobJid, opkAmount: 1);
+
+    final aliceManager = OmemoManager(
+      aliceDevice,
+      AlwaysTrustingTrustManager(),
+      (result, recipientJid) async {},
+      (jid) async {
+        expect(jid, bobJid);
+
+        return oldDevice ?
+          [ bobOldDevice.id ] :
+          [ bobCurrentDevice.id ];
+      },
+      (jid, id) async {
+        expect(jid, bobJid);
+        return oldDevice ?
+          bobOldDevice.toBundle() :
+          bobCurrentDevice.toBundle();
+      },
+      (jid) async {},
+    );
+    final bobManager = OmemoManager(
+      bobCurrentDevice,
+      AlwaysTrustingTrustManager(),
+      (result, recipientJid) async {},
+      (jid) async => [],
+      (jid, id) async => null,
+      (jid) async {},
+    );
+
+    // Alice encrypts a message to Bob
+    final aliceResult1 = await aliceManager.onOutgoingStanza(
+      const OmemoOutgoingStanza(
+        [bobJid],
+        'Hello Bob',
+      ),
+    );
+
+    // Bob's current device receives it
+    final bobResult1 = await bobManager.onIncomingStanza(
+      OmemoIncomingStanza(
+        aliceJid,
+        aliceDevice.id,
+        DateTime.now().millisecondsSinceEpoch,
+        aliceResult1.encryptedKeys,
+        base64.encode(aliceResult1.ciphertext!),
+      ),
+    );
+
+    expect(bobResult1.payload, null);
+    expect(bobResult1.error is NotEncryptedForDeviceException, true);
+
+    // Now Alice's client loses and regains the connection
+    aliceManager.onNewConnection();
+    oldDevice = false;
+
+    // And Alice sends a new message
+    final aliceResult2 = await aliceManager.onOutgoingStanza(
+      const OmemoOutgoingStanza(
+        [bobJid],
+        'Hello Bob x2',
+      ),
+    );
+    final bobResult2 = await bobManager.onIncomingStanza(
+      OmemoIncomingStanza(
+        aliceJid,
+        aliceDevice.id,
+        DateTime.now().millisecondsSinceEpoch,
+        aliceResult2.encryptedKeys,
+        base64.encode(aliceResult2.ciphertext!),
+      ),
+    );
+
+    expect(aliceResult2.encryptedKeys.length, 1);
+    expect(bobResult2.payload, 'Hello Bob x2');
+  });
+
+  test('Test receiving a response from a new device', () async {
+    const aliceJid = 'alice@server1';
+    const bobJid = 'bob@server2';
+    var bothDevices = false;
+
+    final aliceDevice = await OmemoDevice.generateNewDevice(aliceJid, opkAmount: 1);
+    final bobDevice1 = await OmemoDevice.generateNewDevice(bobJid, opkAmount: 1);
+    final bobDevice2 = await OmemoDevice.generateNewDevice(bobJid, opkAmount: 1);
+
+    final aliceManager = OmemoManager(
+      aliceDevice,
+      AlwaysTrustingTrustManager(),
+      (result, recipientJid) async {},
+      (jid) async {
+        expect(jid, bobJid);
+
+        return [
+          bobDevice1.id,
+
+          if (bothDevices)
+            bobDevice2.id,
+        ];
+      },
+      (jid, id) async {
+        expect(jid, bobJid);
+
+        if (bothDevices) {
+          if (id == bobDevice1.id) {
+            return bobDevice1.toBundle();
+          } else if (id == bobDevice2.id) {
+            return bobDevice2.toBundle();
+          }
+        } else {
+          if (id == bobDevice1.id) return bobDevice1.toBundle();
+        }
+
+        return null;
+      },
+      (jid) async {},
+    );
+    final bobManager1 = OmemoManager(
+      bobDevice1,
+      AlwaysTrustingTrustManager(),
+      (result, recipientJid) async {},
+      (jid) async => [],
+      (jid, id) async => null,
+      (jid) async {},
+    );
+    final bobManager2 = OmemoManager(
+      bobDevice2,
+      AlwaysTrustingTrustManager(),
+      (result, recipientJid) async {},
+      (jid) async {
+        expect(jid, aliceJid);
+        return [aliceDevice.id];
+      },
+      (jid, id) async {
+        expect(jid, aliceJid);
+        return aliceDevice.toBundle();
+      },
+      (jid) async {},
+    );
+
+    // Alice sends a message to Bob
+    final aliceResult1 = await aliceManager.onOutgoingStanza(
+      const OmemoOutgoingStanza(
+        [bobJid],
+        'Hello Bob!',
+      ),
+    );
+
+    // Bob decrypts it
+    final bobResult1 = await bobManager1.onIncomingStanza(
+      OmemoIncomingStanza(
+        aliceJid,
+        aliceDevice.id,
+        DateTime.now().millisecondsSinceEpoch,
+        aliceResult1.encryptedKeys,
+        base64.encode(aliceResult1.ciphertext!),
+      ),
+    );
+
+    expect(aliceResult1.encryptedKeys.length, 1);
+    expect(bobResult1.payload, 'Hello Bob!');
+
+    // Now Bob encrypts from his new device
+    bothDevices = true;
+    final bobResult2 = await bobManager2.onOutgoingStanza(
+      const OmemoOutgoingStanza(
+        [aliceJid],
+        'Hello from my new device',
+      ),
+    );
+
+    // And Alice decrypts it
+    final aliceResult2 = await aliceManager.onIncomingStanza(
+      OmemoIncomingStanza(
+        bobJid,
+        bobDevice2.id,
+        DateTime.now().millisecondsSinceEpoch,
+        bobResult2.encryptedKeys,
+        base64.encode(bobResult2.ciphertext!),
+      ),
+    );
+
+    expect(aliceResult2.payload, 'Hello from my new device');
+  });
+
+  test('Test receiving a device list update', () async {
+    const aliceJid = 'alice@server1';
+    const bobJid = 'bob@server2';
+    var bothDevices = false;
+
+    final aliceDevice = await OmemoDevice.generateNewDevice(aliceJid, opkAmount: 1);
+    final bobDevice1 = await OmemoDevice.generateNewDevice(bobJid, opkAmount: 1);
+    final bobDevice2 = await OmemoDevice.generateNewDevice(bobJid, opkAmount: 1);
+
+    final aliceManager = OmemoManager(
+      aliceDevice,
+      AlwaysTrustingTrustManager(),
+      (result, recipientJid) async {},
+      (jid) async {
+        expect(jid, bobJid);
+
+        return [
+          bobDevice1.id,
+
+          if (bothDevices)
+            bobDevice2.id,
+        ];
+      },
+      (jid, id) async {
+        expect(jid, bobJid);
+
+        if (bothDevices) {
+          if (id == bobDevice1.id) {
+            return bobDevice1.toBundle();
+          } else if (id == bobDevice2.id) {
+            return bobDevice2.toBundle();
+          }
+        } else {
+          if (id == bobDevice1.id) return bobDevice1.toBundle();
+        }
+
+        return null;
+      },
+      (jid) async {},
+    );
+    final bobManager1 = OmemoManager(
+      bobDevice1,
+      AlwaysTrustingTrustManager(),
+      (result, recipientJid) async {},
+      (jid) async => null,
+      (jid, id) async => null,
+      (jid) async {},
+    );
+    final bobManager2 = OmemoManager(
+      bobDevice2,
+      AlwaysTrustingTrustManager(),
+      (result, recipientJid) async {},
+      (jid) async => null,
+      (jid, id) async => null,
+      (jid) async {},
+    );
+
+    // Alice sends a message to Bob
+    final aliceResult1 = await aliceManager.onOutgoingStanza(
+      const OmemoOutgoingStanza(
+        [bobJid],
+        'Hello Bob!',
+      ),
+    );
+
+    // Bob decrypts it
+    final bobResult1 = await bobManager1.onIncomingStanza(
+      OmemoIncomingStanza(
+        aliceJid,
+        aliceDevice.id,
+        DateTime.now().millisecondsSinceEpoch,
+        aliceResult1.encryptedKeys,
+        base64.encode(aliceResult1.ciphertext!),
+      ),
+    );
+
+    expect(aliceResult1.encryptedKeys.length, 1);
+    expect(bobResult1.payload, 'Hello Bob!');
+
+    // Bob acks the ratchet session
+    await aliceManager.ratchetAcknowledged(bobJid, bobDevice1.id);
+    
+    // Bob now publishes a new device
+    bothDevices = true;
+    aliceManager.onDeviceListUpdate(
+      bobJid,
+      [
+        bobDevice1.id,
+        bobDevice2.id,
+      ],
+    );
+    
+    // Now Alice encrypts another message
+    final aliceResult2 = await aliceManager.onOutgoingStanza(
+      const OmemoOutgoingStanza(
+        [bobJid],
+        'Hello Bob! x2',
+      ),
+    );
+
+    expect(aliceResult2.encryptedKeys.length, 2);
+
+    // And Bob decrypts it
+    final bobResult21 = await bobManager1.onIncomingStanza(
+      OmemoIncomingStanza(
+        aliceJid,
+        aliceDevice.id,
+        DateTime.now().millisecondsSinceEpoch,
+        aliceResult2.encryptedKeys,
+        base64.encode(aliceResult2.ciphertext!),
+      ),
+    );
+    final bobResult22 = await bobManager2.onIncomingStanza(
+      OmemoIncomingStanza(
+        aliceJid,
+        aliceDevice.id,
+        DateTime.now().millisecondsSinceEpoch,
+        aliceResult2.encryptedKeys,
+        base64.encode(aliceResult2.ciphertext!),
+      ),
+    );
+
+    expect(bobResult21.payload, 'Hello Bob! x2');
+    expect(bobResult22.payload, 'Hello Bob! x2');
+
+    // Bob2 now responds
+    final bobResult32 = await bobManager2.onOutgoingStanza(
+      const OmemoOutgoingStanza(
+        [aliceJid],
+        'Hello Alice!',
+      ),
+    );
+
+    // And Alice responds
+    final aliceResult3 = await aliceManager.onIncomingStanza(
+      OmemoIncomingStanza(
+        bobJid,
+        bobDevice2.id,
+        DateTime.now().millisecondsSinceEpoch,
+        bobResult32.encryptedKeys,
+        base64.encode(bobResult32.ciphertext!),
+      ),
+    );
+
+    expect(aliceResult3.payload, 'Hello Alice!');
+  });
+
+  test('Test sending a message to two different JIDs', () async {
+    const aliceJid = 'alice@server1';
+    const bobJid = 'bob@server2';
+    const cocoJid = 'coco@server3';
+
+    final aliceDevice = await OmemoDevice.generateNewDevice(aliceJid, opkAmount: 1);
+    final bobDevice = await OmemoDevice.generateNewDevice(bobJid, opkAmount: 1);
+    final cocoDevice = await OmemoDevice.generateNewDevice(cocoJid, opkAmount: 1);
+
+    final aliceManager = OmemoManager(
+      aliceDevice,
+      AlwaysTrustingTrustManager(),
+      (result, recipientJid) async {},
+      (jid) async {
+        if (jid == bobJid) {
+          return [bobDevice.id];
+        } else if (jid == cocoJid) {
+          return [cocoDevice.id];
+        }
+
+        return null;
+      },
+      (jid, id) async {
+        if (jid == bobJid) {
+          return bobDevice.toBundle();
+        } else if (jid == cocoJid) {
+          return cocoDevice.toBundle();
+        }
+
+        return null;
+      },
+      (jid) async {},
+    );
+    final bobManager = OmemoManager(
+      bobDevice,
+      AlwaysTrustingTrustManager(),
+      (result, recipientJid) async {},
+      (jid) async => null,
+      (jid, id) async => null,
+      (jid) async {},
+    );
+    final cocoManager = OmemoManager(
+      cocoDevice,
+      AlwaysTrustingTrustManager(),
+      (result, recipientJid) async {},
+      (jid) async => null,
+      (jid, id) async => null,
+      (jid) async {},
+    );
+
+    // Alice sends a message to Bob and Coco
+    final aliceResult = await aliceManager.onOutgoingStanza(
+      const OmemoOutgoingStanza(
+        [bobJid, cocoJid],
+        'Hello Bob and Coco!',
+      ),
+    );
+
+    // Bob and Coco decrypt them
+    final bobResult = await bobManager.onIncomingStanza(
+      OmemoIncomingStanza(
+        aliceJid,
+        aliceDevice.id,
+        DateTime.now().millisecondsSinceEpoch,
+        aliceResult.encryptedKeys,
+        base64.encode(aliceResult.ciphertext!),
+      ),
+    );
+    final cocoResult = await cocoManager.onIncomingStanza(
+      OmemoIncomingStanza(
+        aliceJid,
+        aliceDevice.id,
+        DateTime.now().millisecondsSinceEpoch,
+        aliceResult.encryptedKeys,
+        base64.encode(aliceResult.ciphertext!),
+      ),
+    );
+
+    expect(bobResult.error, null);
+    expect(cocoResult.error, null);
+    expect(bobResult.payload, 'Hello Bob and Coco!');
+    expect(cocoResult.payload, 'Hello Bob and Coco!');
+  });
+
+  test('Test a fetch failure', () async {
+    const aliceJid = 'alice@server1';
+    const bobJid = 'bob@server2';
+    var failure = false;
+
+    final aliceDevice = await OmemoDevice.generateNewDevice(aliceJid, opkAmount: 1);
+    final bobDevice = await OmemoDevice.generateNewDevice(bobJid, opkAmount: 1);
+
+    final aliceManager = OmemoManager(
+      aliceDevice,
+      AlwaysTrustingTrustManager(),
+      (result, recipientJid) async {},
+      (jid) async {
+        expect(jid, bobJid);
+
+        return failure ?
+          null :
+          [bobDevice.id];
+      },
+      (jid, id) async {
+        expect(jid, bobJid);
+
+        return failure ?
+          null :
+          bobDevice.toBundle();
+      },
+      (jid) async {},
+    );
+    final bobManager = OmemoManager(
+      bobDevice,
+      AlwaysTrustingTrustManager(),
+      (result, recipientJid) async {},
+      (jid) async => null,
+      (jid, id) async => null,
+      (jid) async {},
+    );
+
+    // Alice sends a message to Bob and Coco
+    final aliceResult1 = await aliceManager.onOutgoingStanza(
+      const OmemoOutgoingStanza(
+        [bobJid],
+        'Hello Bob!',
+      ),
+    );
+
+    // Bob decrypts it
+    final bobResult1 = await bobManager.onIncomingStanza(
+      OmemoIncomingStanza(
+        aliceJid,
+        aliceDevice.id,
+        DateTime.now().millisecondsSinceEpoch,
+        aliceResult1.encryptedKeys,
+        base64.encode(aliceResult1.ciphertext!),
+      ),
+    );
+
+    expect(bobResult1.error, null);
+    expect(bobResult1.payload, 'Hello Bob!');
+
+    // Bob acks the message
+    await aliceManager.ratchetAcknowledged(
+      bobJid,
+      bobDevice.id,
+    );
+
+    // Alice has to reconnect but has no connection yet
+    failure = true;
+    aliceManager.onNewConnection();
+
+    // Alice sends another message to Bob
+    final aliceResult2 = await aliceManager.onOutgoingStanza(
+      const OmemoOutgoingStanza(
+        [bobJid],
+        'Hello Bob! x2',
+      ),
+    );
+
+    // And Bob decrypts it
+    final bobResult2 = await bobManager.onIncomingStanza(
+      OmemoIncomingStanza(
+        aliceJid,
+        aliceDevice.id,
+        DateTime.now().millisecondsSinceEpoch,
+        aliceResult2.encryptedKeys,
+        base64.encode(aliceResult2.ciphertext!),
+      ),
+    );
+
+    expect(bobResult2.error, null);
+    expect(bobResult2.payload, 'Hello Bob! x2');
+  });
+
+  test('Test sending a message with failed lookups', () async {
+    const aliceJid = 'alice@server1';
+    const bobJid = 'bob@server2';
+
+    final aliceDevice = await OmemoDevice.generateNewDevice(aliceJid, opkAmount: 1);
+
+    final aliceManager = OmemoManager(
+      aliceDevice,
+      AlwaysTrustingTrustManager(),
+      (result, recipientJid) async {},
+      (jid) async {
+        expect(jid, bobJid);
+
+        return null;
+      },
+      (jid, id) async {
+        expect(jid, bobJid);
+
+        return null;
+      },
+      (jid) async {},
+    );
+
+    // Alice sends a message to Bob
+    final aliceResult = await aliceManager.onOutgoingStanza(
+      const OmemoOutgoingStanza(
+        [bobJid],
+        'Hello Bob!',
+      ),
+    );
+
+    expect(aliceResult.isSuccess(1), false);
+    expect(aliceResult.jidEncryptionErrors[bobJid] is NoKeyMaterialAvailableException, true);
+  });
+
+  test('Test sending a message two two JIDs with failed lookups', () async {
+    const aliceJid = 'alice@server1';
+    const bobJid = 'bob@server2';
+    const cocoJid = 'coco@server3';
+
+    final aliceDevice = await OmemoDevice.generateNewDevice(aliceJid, opkAmount: 1);
+    final bobDevice = await OmemoDevice.generateNewDevice(bobJid, opkAmount: 1);
+
+    final aliceManager = OmemoManager(
+      aliceDevice,
+      AlwaysTrustingTrustManager(),
+      (result, recipientJid) async {},
+      (jid) async {
+        if (jid == bobJid) {
+          return [bobDevice.id];
+        }
+
+        return null;
+      },
+      (jid, id) async {
+        if (jid == bobJid) {
+          return bobDevice.toBundle();
+        }
+
+        return null;
+      },
+      (jid) async {},
+    );
+    final bobManager = OmemoManager(
+      bobDevice,
+      AlwaysTrustingTrustManager(),
+      (result, recipientJid) async {},
+      (jid) async => null,
+      (jid, id) async => null,
+      (jid) async {},
+    );
+
+    // Alice sends a message to Bob and Coco
+    final aliceResult = await aliceManager.onOutgoingStanza(
+      const OmemoOutgoingStanza(
+        [bobJid, cocoJid],
+        'Hello Bob and Coco!',
+      ),
+    );
+
+    expect(aliceResult.isSuccess(2), true);
+    expect(aliceResult.jidEncryptionErrors[cocoJid] is NoKeyMaterialAvailableException, true);
+
+    // Bob decrypts it
+    final bobResult = await bobManager.onIncomingStanza(
+      OmemoIncomingStanza(
+        aliceJid,
+        aliceDevice.id,
+        DateTime.now().millisecondsSinceEpoch,
+        aliceResult.encryptedKeys,
+        base64.encode(aliceResult.ciphertext!),
+      ),
+    );
+
+    expect(bobResult.payload, 'Hello Bob and Coco!');
+  });
+
+  test('Test sending multiple messages back and forth', () async {
+    const aliceJid = 'alice@server1';
+    const bobJid = 'bob@server2';
+
+    final aliceDevice = await OmemoDevice.generateNewDevice(aliceJid, opkAmount: 1);
+    final bobDevice = await OmemoDevice.generateNewDevice(bobJid, opkAmount: 1);
+
+    final aliceManager = OmemoManager(
+      aliceDevice,
+      AlwaysTrustingTrustManager(),
+      (result, recipientJid) async {},
+      (jid) async {
+        expect(jid, bobJid);
+
+        return [bobDevice.id];
+      },
+      (jid, id) async {
+        expect(jid, bobJid);
+
+        return bobDevice.toBundle();
+      },
+      (jid) async {},
+    );
+    final bobManager = OmemoManager(
+      bobDevice,
+      AlwaysTrustingTrustManager(),
+      (result, recipientJid) async {},
+      (jid) async => null,
+      (jid, id) async => null,
+      (jid) async {},
+    );
+
+    // Alice encrypts a message for Bob
+    final aliceMessage = await aliceManager.onOutgoingStanza(
+      const OmemoOutgoingStanza(
+        [bobJid],
+        'Hello Bob!',
+      ),
+    );
+
+    // And Bob decrypts it
+    await bobManager.onIncomingStanza(
+      OmemoIncomingStanza(
+        aliceJid,
+        aliceDevice.id,
+        DateTime.now().millisecondsSinceEpoch,
+        aliceMessage.encryptedKeys,
+        base64.encode(aliceMessage.ciphertext!),
+      ),
+
+    );
+
+    // Ratchets are acked
+    await aliceManager.ratchetAcknowledged(
+      bobJid,
+      bobDevice.id,
+    );
+    
+    for (var i = 0; i < 100; i++) {
+      final messageText = 'Test Message #$i';
+      // Bob responds to Alice
+      final bobResponseMessage = await bobManager.onOutgoingStanza(
+        OmemoOutgoingStanza(
+          [aliceJid],
+          messageText,
+        ),
+      );
+      expect(bobResponseMessage.isSuccess(1), true);
+      
+      final aliceReceivedMessage = await aliceManager.onIncomingStanza(
+        OmemoIncomingStanza(
+          bobJid,
+          bobDevice.id,
+          DateTime.now().millisecondsSinceEpoch,
+          bobResponseMessage.encryptedKeys,
+          base64.encode(bobResponseMessage.ciphertext!),
+        ),
+      );
+      expect(aliceReceivedMessage.payload, messageText);
+    }
+  });
+}

test/serialisation_test.dart

@@ -3,6 +3,10 @@ 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
@@ -12,9 +16,9 @@ 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);
   });
 
@@ -26,9 +30,9 @@ void main() {
       opkAmount: 1,
     );
     final oldDevice = await (await oldSession.getDevice()).replaceSignedPrekey();
-    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);
   });
   
@@ -58,9 +62,10 @@ void main() {
       aliceJid,
       await aliceSession.getDeviceId(),
       aliceMessage.encryptedKeys,
+      getTimestamp(),
     );
     final aliceOld = aliceSession.getRatchet(bobJid, await bobSession.getDeviceId());
-    final aliceSerialised = await aliceOld.toJson();
+    final aliceSerialised = jsonify(await aliceOld.toJson());
     final aliceNew = OmemoDoubleRatchet.fromJson(aliceSerialised);
 
     expect(await aliceOld.equals(aliceNew), true);
@@ -85,9 +90,11 @@ void main() {
     );
 
     // 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(),
     );
 
@@ -95,51 +102,6 @@ void main() {
     final newDevice = await newSession.getDevice();
     expect(await oldDevice.equals(newDevice), true);
     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);
-
-      final oldRatchet = oldSession.getRatchetMap()[session.key]!;
-      final newRatchet = newSession.getRatchetMap()[session.key]!;
-      expect(await oldRatchet.equals(newRatchet), true);
-    }
-  });
-
-  test('Test serialising and deserialising the BlindTrustBeforeVerificationTrustManager', () async {
-    // Caroline's BTBV manager
-    final btbv = MemoryBTBVTrustManager();
-    // Example data
-    const aliceJid = 'alice@some.server';
-    const bobJid = 'bob@other.server';
-    
-    // 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);
-    expect(await btbv.isTrusted(aliceJid, 1), true);
-
-    // Alice adds a new device
-    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.
-    await btbv.onNewSession(bobJid, 3);
-
-    // Bob adds a new device
-    await btbv.onNewSession(bobJid, 4);
-    expect(await btbv.isTrusted(bobJid, 3), true);
-    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);
   });
 
   test('Test serializing and deserializing RatchetMapKey', () {
@@ -167,21 +129,19 @@ void main() {
     await btbv.onNewSession(bobJid, 3);
     await btbv.onNewSession(bobJid, 4);
 
-    final managerJson = await btbv.toJson();
-    final managerString = jsonEncode(managerJson);
-    final managerPostJson = jsonDecode(managerString) as Map<String, dynamic>;
+    final serialized = jsonify(await btbv.toJson());
     final deviceList = BlindTrustBeforeVerificationTrustManager.deviceListFromJson(
-      managerPostJson,
+      serialized,
     );
     expect(btbv.devices, deviceList);
 
     final trustCache = BlindTrustBeforeVerificationTrustManager.trustCacheFromJson(
-      managerPostJson,
+      serialized,
     );
     expect(btbv.trustCache, trustCache);
 
     final enableCache = BlindTrustBeforeVerificationTrustManager.enableCacheFromJson(
-      managerPostJson,
+      serialized,
     );
     expect(btbv.enablementCache, enableCache);
   });