.gitlint

@@ -5,6 +5,6 @@ line-length=72
 [title-trailing-punctuation]
 [title-hard-tab]
 [title-match-regex]
-regex=^(feat|fix|test|release|chore|security|docs|refactor|style|ci):.*$
+regex=^(feat|fix|test|release|chore|security|docs|refactor|style):.*$
 [body-trailing-whitespace]
 [body-first-line-empty]

.woodpecker.yml

@@ -1,22 +1,11 @@
 pipeline:
-  analysis:
+  lint:
-    image: dart:3.0.7
+    image: dart:2.18.1
     commands:
-      # Proxy requests to pub.dev using pubcached
+      - dart pub get
-      - PUB_HOSTED_URL=http://172.17.0.1:8000 dart pub get
       - dart analyze --fatal-infos --fatal-warnings
+  test:
+    image: dart:2.18.1
+    commands:
+      - dart pub get
       - dart test
-    when:
-      path:
-        includes: ['lib/**', 'test/**']
-  notify:
-    image: git.polynom.me/papatutuwawa/woodpecker-xmpp
-    settings:
-      xmpp_is_muc: 1
-      xmpp_tls: 1
-      xmpp_recipient: moxxy-build@muc.moxxy.org
-      xmpp_alias: 2Bot
-    secrets: [ xmpp_jid, xmpp_password, xmpp_server ]
-    when:
-      status:
-        - failure

CHANGELOG.md

@@ -50,28 +50,3 @@
 ## 0.4.2
 
 - Fix removeAllRatchets not removing, well, all ratchets. In fact, it did not remove any ratchet.
-
-## 0.4.3
-
-- Fix bug that causes ratchets to be unable to decrypt anything after receiving a heartbeat with a completely new session
-
-## 0.5.0
-
-This version is a complete rework of omemo_dart!
-
-- Removed events from `OmemoManager`
-- Removed `OmemoSessionManager`
-- Removed serialization/deserialization code
-- Replace exceptions with errors inside a result type
-- Ratchets and trust data is now loaded and cached on demand
-- Accessing the trust manager must happen via `withTrustManager`
-- Overriding the base implementations is replaced by providing callback functions
-
-## 0.5.1
-
-- Remove `added` and `replaced` from the data passed to the `CommitRatchetsCallback`
-- Added a list of newly added and replaced ratchets to the encryption and decryption results. This is useful for displaying messages like "Contact added a new device"
-
-## 0.6.0
-
-- Bump dependencies to fix running with never version of Dart

README.md

@@ -1,7 +1,5 @@
 # omemo_dart
 
-[![status-badge](https://ci.polynom.me/api/badges/16/status.svg)](https://ci.polynom.me/repos/16)
-
 `omemo_dart` is a Dart library to help developers of Dart/Flutter XMPP clients to implement
 [OMEMO](https://xmpp.org/extensions/xep-0384.html) in its newest version - currently 0.8.3.
 
@@ -30,33 +28,23 @@ Include `omemo_dart` in your `pubspec.yaml` like this:
 dependencies:
   omemo_dart:
     hosted: https://git.polynom.me/api/packages/PapaTutuWawa/pub
-    version: ^0.5.0
+    version: ^0.4.0
   # [...]
 
 # [...]
 ```
 
-### Example
-
-This repository includes a documented ["example"](./example/omemo_dart_example.dart) that explains the basic usage of the library while
-leaving out the XMPP specific bits. For a more functional and integrated example, see the `omemo_client.dart` example from
-[moxxmpp](https://codeberg.org/moxxy/moxxmpp).
-
-### Persistence
-
-By default, `omemo_dart` uses in-memory implementations for everything. For a real-world application, this is unsuitable as OMEMO devices would be constantly added.
-In order to allow persistence, your application needs to keep track of the following:
-
-- The `OmemoDevice` assigned to the `OmemoManager`
-- `JID -> [int]`: The device list for each JID
-- `(JID, device) -> Ratchet`: The actual ratchet
-
-If you also use the `BlindTrustBeforeVerificationTrustManager`, you additionally need to keep track of:
-
-- `(JID, device) -> (int, bool)`: The trust level and the enablement state
-
 ## Contributing
 
+Due to issues with `protobuf`, `omemo_dart` reimplements the Protobuf encoding for the required
+OMEMO messages. As such, `protobuf` is only a dependency for testing that the serialisation and
+deserialisation of the custom implementation. In order to run tests, you need the Protbuf
+compiler. After that, making sure that
+the [Dart Protobuf compiler addon](https://pub.dev/packages/protoc_plugin) and the
+Protobuf compiler itself is in your PATH,
+run `protoc -I=./protobuf/ --dart_out=lib/protobuf/ ./protobuf/schema.proto` in the
+repository's root to generate the real Protobuf bindings.
+
 When submitting a PR, please run the linter using `dart analyze` and make sure that all
 tests still pass using `dart test`.
 

analysis_options.yaml

@@ -9,5 +9,8 @@ linter:
 
 analyzer:
   exclude:
-    - "lib/src/protobuf/*.dart"
+    - "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

@@ -8,54 +8,32 @@ void main() async {
   const aliceJid = 'alice@some.server';
   const bobJid = 'bob@other.serve';
 
-  // You are Alice and want to begin using OMEMO, so you first create an OmemoManager.
+  // You are Alice and want to begin using OMEMO, so you first create a SessionManager
-  final aliceManager = OmemoManager(
+  final aliceSession = await OmemoSessionManager.generateNewIdentity(
-    // Generate Alice's OMEMO device bundle. We can specify how many One-time Prekeys we want, but
+    // The bare Jid of Alice as a String
-    // per default, omemo_dart generates 100 (recommended by XEP-0384).
+    aliceJid,
-    await OmemoDevice.generateNewDevice(aliceJid),
     // The trust manager we want to use. In this case, we use the provided one that
     // implements "Blind Trust Before Verification". To make things simpler, we keep
     // no persistent data and can thus use the MemoryBTBVTrustManager. If we wanted to keep
     // the state, we would have to override BlindTrustBeforeVerificationTrustManager.
-    BlindTrustBeforeVerificationTrustManager(),
+    MemoryBTBVTrustManager(),
-    // This function is called whenever we need to send an OMEMO heartbeat to [recipient].
+    // Here we specify how many Onetime Prekeys we want to have. XEP-0384 recommends around
-    // [result] is the encryted data to include. This needs to be wired into your XMPP library's
+    // 100 OPKs, so let's generate 100. The parameter defaults to 100.
-    // OMEMO implementation.
+    //opkAmount: 100,
-    // For simplicity, we use an empty function and imagine it works.
-    (result, recipient) async => {},
-    // This function is called whenever we need to fetch the device list for [jid].
-    // This needs to be wired into your XMPP library's OMEMO implementation.
-    // For simplicity, we use an empty function and imagine it works.
-    (jid) async => [],
-    // This function is called whenever we need to fetch the device bundle with id [id] from [jid].
-    // This needs to be wired into your XMPP library's OMEMO implementation.
-    // For simplicity, we use an empty function and imagine it works.
-    (jid, id) async => null,
-    // This function is called whenever we need to subscribe to [jid]'s device list PubSub node.
-    // This needs to be wired into your XMPP library's OMEMO implementation.
-    // For simplicity, we use an empty function and imagine it works.
-    (jid) async {},
-    // This function is called whenever our own device bundle has to be republished to our PEP node.
-    // This needs to be wired into your XMPP library's OMEMO implementation.
-    // For simplicity, we use an empty function and imagine it works.
-    (device) async {},
   );
 
-  // Bob, on his side, also creates an [OmemoManager] similar to Alice.
+  // Alice now wants to chat with Bob at his bare Jid "bob@other.server". To make things
-  final bobManager = OmemoManager(
+  // simple, we just generate the identity bundle ourselves. In the real world, we would
-    await OmemoDevice.generateNewDevice(bobJid),
+  // request it using PEP and then convert the device bundle into a OmemoBundle object.
-    BlindTrustBeforeVerificationTrustManager(),
+  final bobSession = await OmemoSessionManager.generateNewIdentity(
-    (result, recipient) async => {},
+    bobJid,
-    (jid) async => [],
+    MemoryBTBVTrustManager(),
-    (jid, id) async => null,
+    // Just for illustrative purposes
-    (jid) async {},
+    opkAmount: 1,
-    (device) async {},
   );
 
   // Alice prepares to send the message to Bob, so she builds the message stanza and
   // collects all the children of the stanza that should be encrypted into a string.
-  // Note that this leaves out the wrapping stanza, i.e. if we want to send a <message />
-  // we only include the <message />'s children.
   const aliceMessageStanzaBody = '''
   <body>Hello Bob, it's me, Alice!</body>
   <super-secret-element xmlns='super-secret-element' />
@@ -75,28 +53,31 @@ void main() async {
   <time stamp='1969-07-20T21:56:15-05:00' />
 </envelope>
 ''';
-
+ 
-  // Next, we encrypt the envelope element using Alice's [OmemoManager]. It will
+  // Since Alice has no open session with Bob, we need to tell the session manager to build
-  // automatically attempt to fetch the device bundles of Bob.
+  // it when sending the message.
-  final message = await aliceManager.onOutgoingStanza(
+  final message = await aliceSession.encryptToJid(
-    const OmemoOutgoingStanza(
+    // The bare receiver Jid
-      // The bare receiver Jid
+    bobJid,
-      [bobJid],
+    // The envelope we want to encrypt
-
+    envelope,
-      // The payload we want to encrypt, i.e. the envelope.
+    // Since this is the first time Alice contacts Bob from this device, we need to create
-      envelope,
+    // a new session. Let's also assume that Bob only has one device. We may, however,
-    ),
+    // add more bundles to newSessions, if we know of more.
+    newSessions: [
+      await bobSession.getDeviceBundle(),
+    ],
   );
 
-  // In a proper implementation, we would also do some error checking here.
-
   // Alice now builds the actual message stanza for Bob
   final payload = base64.encode(message.ciphertext!);
-  final aliceDevice = await aliceManager.getDevice();
+  final aliceDevice = await aliceSession.getDevice();
+  // ignore: unused_local_variable
+  final bobDevice = await bobSession.getDevice();
   // Since we know we have just one key for Bob, we take a shortcut. However, in the real
   // world, we have to serialise every EncryptedKey to a <key /> element and group them
   // per Jid.
-  final key = message.encryptedKeys[bobJid]![0];
+  final key = message.encryptedKeys[0];
 
   // Note that the key's "kex" attribute refers to key.kex. It just means that the
   // encrypted key also contains the required data for Bob to build a session with Alice.
@@ -106,7 +87,7 @@ void main() async {
   <encrypted xmlns='urn:xmpp:omemo:2'>
     <header sid='${aliceDevice.id}'>
       <keys jid='$bobJid'>
-        <key rid='${key.rid} kex='${key.kex}'>
+        <key rid='${key.rid} kex='true'>
           ${key.value}
         </key>
       </keys>
@@ -124,34 +105,25 @@ void main() async {
   // Bob now receives an OMEMO encrypted message from Alice and wants to decrypt it.
   // Since we have just one key, let's just deserialise the one key by hand.
   final keys = [
-    EncryptedKey(key.rid, key.value, true),
+    EncryptedKey(bobJid, key.rid, key.value, true),
   ];
 
   // Bob extracts the payload and attempts to decrypt it.
   // ignore: unused_local_variable
-  final bobMessage = await bobManager.onIncomingStanza(
+  final bobMessage = await bobSession.decryptMessage(
-    OmemoIncomingStanza(
+    // base64 decode the payload
-      // The bare sender JID of the message. In this case, it's Alice's.
+    base64.decode(payload),
-      aliceJid,
+    // Specify the Jid of the sender
-      // The 'sid' attribute of the <header /> element. Here, we know that Alice only has one device.
+    aliceJid,
-      aliceDevice.id,
+    // Specify the device identifier of the sender (the "sid" attribute of <header />)
-
+    aliceDevice.id,
-      /// The decoded <key /> elements. from the header. Note that we only include the ones
+    // The deserialised keys
-      /// relevant for Bob, so all children of <keys jid='$bobJid' />.
+    keys,
-      keys,
+    // Since the message was not delayed, we use the current time
-
+    DateTime.now().millisecondsSinceEpoch,
-      /// The text of the <payload /> element, if it exists. If not, then the message might be
-      /// a hearbeat, where no payload is sent. In that case, use null.
-      payload,
-
-      /// Since we did not receive this message due to a catch-up mechanism, like MAM, we
-      /// set this to false. If we, however, did use a catch-up mechanism, we must set this
-      /// to true to prevent the OPKs from being replaced.
-      false,
-    ),
   );
 
-  // All Bob has to do now is replace the OMEMO wrapper element
+  // All Bob has to do now is replace the OMEMO wrapper element 
   // <encrypted xmlns='urn:xmpp:omemo:2' />) with the content of the <content /> element
   // of the envelope we just decrypted.
 

flake.lock

@@ -1,15 +1,12 @@
 {
   "nodes": {
     "flake-utils": {
-      "inputs": {
-        "systems": "systems"
-      },
       "locked": {
-        "lastModified": 1692799911,
+        "lastModified": 1656065134,
-        "narHash": "sha256-3eihraek4qL744EvQXsK1Ha6C3CR7nnT8X2qWap4RNk=",
+        "narHash": "sha256-oc6E6ByIw3oJaIyc67maaFcnjYOz1mMcOtHxbEf9NwQ=",
         "owner": "numtide",
         "repo": "flake-utils",
-        "rev": "f9e7cf818399d17d347f847525c5a5a8032e4e44",
+        "rev": "bee6a7250dd1b01844a2de7e02e4df7d8a0a206c",
         "type": "github"
       },
       "original": {
@@ -20,16 +17,16 @@
     },
     "nixpkgs": {
       "locked": {
-        "lastModified": 1727586919,
+        "lastModified": 1657540956,
-        "narHash": "sha256-e/YXG0tO5GWHDS8QQauj8aj4HhXEm602q9swrrlTlKQ=",
+        "narHash": "sha256-ihGbOFWtAkENwxBE5kV/yWt2MncvW+BObLDsmxCLo/Q=",
-        "owner": "NixOS",
+        "owner": "NANASHI0X74",
         "repo": "nixpkgs",
-        "rev": "2dcd9c55e8914017226f5948ac22c53872a13ee2",
+        "rev": "043de04db8a6b0391b3fefaaade160514d866946",
         "type": "github"
       },
       "original": {
-        "owner": "NixOS",
+        "owner": "NANASHI0X74",
-        "ref": "nixpkgs-unstable",
+        "ref": "flutter-3-0-0",
         "repo": "nixpkgs",
         "type": "github"
       }
@@ -39,21 +36,6 @@
         "flake-utils": "flake-utils",
         "nixpkgs": "nixpkgs"
       }
-    },
-    "systems": {
-      "locked": {
-        "lastModified": 1681028828,
-        "narHash": "sha256-Vy1rq5AaRuLzOxct8nz4T6wlgyUR7zLU309k9mBC768=",
-        "owner": "nix-systems",
-        "repo": "default",
-        "rev": "da67096a3b9bf56a91d16901293e51ba5b49a27e",
-        "type": "github"
-      },
-      "original": {
-        "owner": "nix-systems",
-        "repo": "default",
-        "type": "github"
-      }
     }
   },
   "root": "root",

flake.nix

@@ -1,7 +1,7 @@
 {
   description = "omemo_dart";
   inputs = {
-    nixpkgs.url = "github:NixOS/nixpkgs/nixpkgs-unstable";
+    nixpkgs.url = "github:NANASHI0X74/nixpkgs/flutter-3-0-0";
     flake-utils.url = "github:numtide/flake-utils";
   };
   outputs = { self, nixpkgs, flake-utils }: flake-utils.lib.eachDefaultSystem (system: let

lib/omemo_dart.dart

@@ -8,11 +8,11 @@ export 'src/omemo/bundle.dart';
 export 'src/omemo/device.dart';
 export 'src/omemo/encrypted_key.dart';
 export 'src/omemo/encryption_result.dart';
-export 'src/omemo/errors.dart';
+export 'src/omemo/events.dart';
 export 'src/omemo/fingerprint.dart';
-export 'src/omemo/omemo.dart';
+export 'src/omemo/omemomanager.dart';
-export 'src/omemo/ratchet_data.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';

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

@@ -4,3 +4,4 @@
 //
 // @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

@@ -6,3 +6,4 @@
 // 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/common/constants.dart

@@ -1,11 +0,0 @@
-/// The overarching assumption is that we use Ed25519 keys for the identity keys
-const omemoX3DHInfoString = 'OMEMO X3DH';
-
-/// The info used for when encrypting the AES key for the actual payload.
-const omemoPayloadInfoString = 'OMEMO Payload';
-
-/// Info string for ENCRYPT
-const encryptHkdfInfoString = 'OMEMO Message Key Material';
-
-/// Amount of messages we may skip per session
-const maxSkip = 1000;

lib/src/crypto.dart

@@ -1,18 +1,12 @@
 import 'dart:convert';
 import 'package:cryptography/cryptography.dart';
-import 'package:moxlib/moxlib.dart';
-import 'package:omemo_dart/src/errors.dart';
 import 'package:omemo_dart/src/keys.dart';
 
 /// Performs X25519 with [kp] and [pk]. If [identityKey] is set, then
 /// it indicates which of [kp] ([identityKey] == 1) or [pk] ([identityKey] == 2)
 /// is the identity key. This is needed since the identity key pair/public key is
 /// an Ed25519 key, but we need them as X25519 keys for DH.
-Future<List<int>> omemoDH(
+Future<List<int>> omemoDH(OmemoKeyPair kp, OmemoPublicKey pk, int identityKey) async {
-  OmemoKeyPair kp,
-  OmemoPublicKey pk,
-  int identityKey,
-) async {
   var ckp = kp;
   var cpk = pk;
 
@@ -23,14 +17,15 @@ Future<List<int>> omemoDH(
   }
 
   final shared = await Cryptography.instance.x25519().sharedSecretKey(
-        keyPair: await ckp.asKeyPair(),
+    keyPair: await ckp.asKeyPair(),
-        remotePublicKey: cpk.asPublicKey(),
+    remotePublicKey: cpk.asPublicKey(),
-      );
+  );
 
   return shared.extractBytes();
 }
 
 class HkdfKeyResult {
+
   const HkdfKeyResult(this.encryptionKey, this.authenticationKey, this.iv);
   final List<int> encryptionKey;
   final List<int> authenticationKey;
@@ -40,8 +35,7 @@ class HkdfKeyResult {
 /// cryptography _really_ wants to check the MAC output from AES-256-CBC. Since
 /// we don't have it, we need the MAC check to always "pass".
 class NoMacSecretBox extends SecretBox {
-  NoMacSecretBox(super.cipherText, {required super.nonce})
+  NoMacSecretBox(super.cipherText, { required super.nonce }) : super(mac: Mac.empty);
-      : super(mac: Mac.empty);
 
   @override
   Future<void> checkMac({
@@ -65,21 +59,13 @@ Future<HkdfKeyResult> deriveEncryptionKeys(List<int> input, String info) async {
     info: utf8.encode(info),
   );
   final bytes = await result.extractBytes();
-
+  
-  return HkdfKeyResult(
+  return HkdfKeyResult(bytes.sublist(0, 32), bytes.sublist(32, 64), bytes.sublist(64, 80));
-    bytes.sublist(0, 32),
-    bytes.sublist(32, 64),
-    bytes.sublist(64, 80),
-  );
 }
 
 /// A small helper function to make AES-256-CBC easier. Encrypt [plaintext] using [key] as
 /// the encryption key and [iv] as the IV. Returns the ciphertext.
-Future<List<int>> aes256CbcEncrypt(
+Future<List<int>> aes256CbcEncrypt(List<int> plaintext, List<int> key, List<int> iv) async {
-  List<int> plaintext,
-  List<int> key,
-  List<int> iv,
-) async {
   final algorithm = AesCbc.with256bits(
     macAlgorithm: MacAlgorithm.empty,
   );
@@ -94,27 +80,17 @@ Future<List<int>> aes256CbcEncrypt(
 
 /// A small helper function to make AES-256-CBC easier. Decrypt [ciphertext] using [key] as
 /// the encryption key and [iv] as the IV. Returns the ciphertext.
-Future<Result<MalformedCiphertextError, List<int>>> aes256CbcDecrypt(
+Future<List<int>> aes256CbcDecrypt(List<int> ciphertext, List<int> key, List<int> iv) async {
-  List<int> ciphertext,
-  List<int> key,
-  List<int> iv,
-) async {
   final algorithm = AesCbc.with256bits(
     macAlgorithm: MacAlgorithm.empty,
   );
-  try {
+  return algorithm.decrypt(
-    return Result(
+    NoMacSecretBox(
-      await algorithm.decrypt(
+      ciphertext,
-        NoMacSecretBox(
+      nonce: iv,
-          ciphertext,
+    ),
-          nonce: iv,
+    secretKey: SecretKey(key),
-        ),
+  );
-        secretKey: SecretKey(key),
-      ),
-    );
-  } catch (ex) {
-    return Result(MalformedCiphertextError(ex));
-  }
 }
 
 /// OMEMO often uses the output of a HMAC-SHA-256 truncated to its first 16 bytes.

lib/src/double_ratchet/crypto.dart

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

lib/src/double_ratchet/double_ratchet.dart

@@ -1,25 +1,48 @@
+import 'dart:convert';
 import 'package:cryptography/cryptography.dart';
 import 'package:hex/hex.dart';
 import 'package:meta/meta.dart';
-import 'package:moxlib/moxlib.dart';
-import 'package:omemo_dart/src/common/constants.dart';
 import 'package:omemo_dart/src/crypto.dart';
+import 'package:omemo_dart/src/double_ratchet/crypto.dart';
 import 'package:omemo_dart/src/double_ratchet/kdf.dart';
 import 'package:omemo_dart/src/errors.dart';
 import 'package:omemo_dart/src/helpers.dart';
 import 'package:omemo_dart/src/keys.dart';
-import 'package:omemo_dart/src/protobuf/schema.pb.dart';
+import 'package:omemo_dart/src/protobuf/omemo_message.dart';
+
+/// Amount of messages we may skip per session
+const maxSkip = 1000;
+
+class RatchetStep {
+  const RatchetStep(this.header, this.ciphertext);
+  final OmemoMessage header;
+  final List<int> ciphertext;
+}
 
 @immutable
 class SkippedKey {
   const SkippedKey(this.dh, this.n);
 
-  /// The DH public key for which we skipped a message key.
+  factory SkippedKey.fromJson(Map<String, dynamic> data) {
-  final OmemoPublicKey dh;
+    return SkippedKey(
+      OmemoPublicKey.fromBytes(
+        base64.decode(data['public']! as String),
+        KeyPairType.x25519,
+      ),
+      data['n']! as int,
+    );
+  }
 
-  /// The associated number of the message key we skipped.
+  final OmemoPublicKey dh;
   final int n;
 
+  Future<Map<String, dynamic>> toJson() async {
+    return {
+      'public': base64.encode(await dh.getBytes()),
+      'n': n,
+    };
+  }
+  
   @override
   bool operator ==(Object other) {
     return other is SkippedKey && other.dh == dh && other.n == n;
@@ -29,45 +52,90 @@ class SkippedKey {
   int get hashCode => dh.hashCode ^ n.hashCode;
 }
 
-@immutable
-class KeyExchangeData {
-  const KeyExchangeData(
-    this.pkId,
-    this.spkId,
-    this.ik,
-    this.ek,
-  );
-
-  /// The id of the used OPK.
-  final int pkId;
-
-  /// The id of the used SPK.
-  final int spkId;
-
-  /// The ephemeral key used while the key exchange.
-  final OmemoPublicKey ek;
-
-  /// The identity key used in the key exchange.
-  final OmemoPublicKey ik;
-}
-
 class OmemoDoubleRatchet {
   OmemoDoubleRatchet(
     this.dhs, // DHs
     this.dhr, // DHr
-    this.rk, // RK
+    this.rk,  // RK
     this.cks, // CKs
     this.ckr, // CKr
-    this.ns, // Ns
+    this.ns,  // Ns
-    this.nr, // Nr
+    this.nr,  // Nr
-    this.pn, // Pn
+    this.pn,  // Pn
     this.ik,
     this.sessionAd,
     this.mkSkipped, // MKSKIPPED
     this.acknowledged,
+    this.kexTimestamp,
     this.kex,
   );
-
+  
+  factory OmemoDoubleRatchet.fromJson(Map<String, dynamic> data) {
+    /*
+    {
+      'dhs': 'base/64/encoded',
+      'dhs_pub': 'base/64/encoded',
+      'dhr': null | 'base/64/encoded',
+      'rk': 'base/64/encoded',
+      'cks': null | 'base/64/encoded',
+      'ckr': null | 'base/64/encoded',
+      'ns': 0,
+      'nr': 0,
+      'pn': 0,
+      'ik_pub': null | 'base/64/encoded',
+      'session_ad': 'base/64/encoded',
+      'acknowledged': true | false,
+      'kex_timestamp': int,
+      'kex': 'base/64/encoded',
+      'mkskipped': [
+        {
+          'key': 'base/64/encoded',
+          'public': 'base/64/encoded',
+          'n': 0
+        }, ...
+      ]
+    }
+    */
+    // NOTE: Dart has some issues with just casting a List<dynamic> to List<Map<...>>, as
+    //       such we need to convert the items by hand.
+    final mkSkipped = Map<SkippedKey, List<int>>.fromEntries(
+      (data['mkskipped']! as List<dynamic>).map<MapEntry<SkippedKey, List<int>>>(
+        (entry) {
+          final map = entry as Map<String, dynamic>;
+          final key = SkippedKey.fromJson(map);
+          return MapEntry(
+            key,
+            base64.decode(map['key']! as String),
+          );
+        },
+      ),
+    );
+    
+    return OmemoDoubleRatchet(
+      OmemoKeyPair.fromBytes(
+        base64.decode(data['dhs_pub']! as String),
+        base64.decode(data['dhs']! as String),
+        KeyPairType.x25519,
+      ),
+      decodeKeyIfNotNull(data, 'dhr', KeyPairType.x25519),
+      base64.decode(data['rk']! as String),
+      base64DecodeIfNotNull(data, 'cks'),
+      base64DecodeIfNotNull(data, 'ckr'),
+      data['ns']! as int,
+      data['nr']! as int,
+      data['pn']! as int,
+      OmemoPublicKey.fromBytes(
+        base64.decode(data['ik_pub']! as String),
+        KeyPairType.ed25519,
+      ),
+      base64.decode(data['session_ad']! as String),
+      mkSkipped,
+      data['acknowledged']! as bool,
+      data['kex_timestamp']! as int,
+      data['kex'] as String?,
+    );
+  }
+  
   /// Sending DH keypair
   OmemoKeyPair dhs;
 
@@ -91,41 +159,36 @@ class OmemoDoubleRatchet {
   /// The IK public key from the chat partner. Not used for the actual encryption but
   /// for verification purposes
   final OmemoPublicKey ik;
-
+  
-  /// Associated data for this ratchet.
   final List<int> sessionAd;
 
-  /// List of skipped message keys.
   final Map<SkippedKey, List<int>> mkSkipped;
 
-  /// The key exchange that was used for initiating the session.
+  /// The point in time at which we performed the kex exchange to create this ratchet.
-  final KeyExchangeData kex;
+  /// 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.
+  /// the device. 
   bool acknowledged;
 
   /// Create an OMEMO session using the Signed Pre Key [spk], the shared secret [sk] that
   /// was obtained using a X3DH and the associated data [ad] that was also obtained through
   /// a X3DH. [ik] refers to Bob's (the receiver's) IK public key.
-  static Future<OmemoDoubleRatchet> initiateNewSession(
+  static Future<OmemoDoubleRatchet> initiateNewSession(OmemoPublicKey spk, OmemoPublicKey ik, List<int> sk, List<int> ad, int timestamp) async {
-    OmemoPublicKey spk,
-    int spkId,
-    OmemoPublicKey ik,
-    OmemoPublicKey ownIk,
-    OmemoPublicKey ek,
-    List<int> sk,
-    List<int> ad,
-    int pkId,
-  ) async {
     final dhs = await OmemoKeyPair.generateNewPair(KeyPairType.x25519);
-    final rk = await kdfRk(sk, await omemoDH(dhs, spk, 0));
+    final dhr = spk;
+    final rk  = await kdfRk(sk, await omemoDH(dhs, dhr, 0));
+    final cks = rk;
 
     return OmemoDoubleRatchet(
       dhs,
-      spk,
+      dhr,
-      List.from(rk),
+      rk,
-      List.from(rk),
+      cks,
       null,
       0,
       0,
@@ -134,12 +197,8 @@ class OmemoDoubleRatchet {
       ad,
       {},
       false,
-      KeyExchangeData(
+      timestamp,
-        pkId,
+      '',
-        spkId,
-        ownIk,
-        ek,
-      ),
     );
   }
 
@@ -147,15 +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(
+  static Future<OmemoDoubleRatchet> acceptNewSession(OmemoKeyPair spk, OmemoPublicKey ik, List<int> sk, List<int> ad, int kexTimestamp) async {
-    OmemoKeyPair spk,
-    int spkId,
-    OmemoPublicKey ik,
-    int pkId,
-    OmemoPublicKey ek,
-    List<int> sk,
-    List<int> ad,
-  ) async {
     return OmemoDoubleRatchet(
       spk,
       null,
@@ -169,52 +220,63 @@ class OmemoDoubleRatchet {
       ad,
       {},
       true,
-      KeyExchangeData(
+      kexTimestamp,
-        pkId,
+      null,
-        spkId,
-        ik,
-        ek,
-      ),
     );
   }
 
-  /// Performs a single ratchet step in case we received a new
+  Future<Map<String, dynamic>> toJson() async {
-  /// public key in [header].
+    final mkSkippedSerialised = List<Map<String, dynamic>>.empty(growable: true);
-  Future<void> _dhRatchet(OMEMOMessage header) async {
+    for (final entry in mkSkipped.entries) {
-    pn = ns;
+      final result = await entry.key.toJson();
-    ns = 0;
+      result['key'] = base64.encode(entry.value);
-    nr = 0;
-    dhr = OmemoPublicKey.fromBytes(header.dhPub, KeyPairType.x25519);
-    final newRk1 = await kdfRk(
-      rk,
-      await omemoDH(
-        dhs,
-        dhr!,
-        0,
-      ),
-    );
-    rk = List.from(newRk1);
-    ckr = List.from(newRk1);
 
-    dhs = await OmemoKeyPair.generateNewPair(KeyPairType.x25519);
+      mkSkippedSerialised.add(result);
-    final newRk2 = await kdfRk(
+    }
-      rk,
+    
-      await omemoDH(
+    return {
-        dhs,
+      'dhs': base64.encode(await dhs.sk.getBytes()),
-        dhr!,
+      'dhs_pub': base64.encode(await dhs.pk.getBytes()),
-        0,
+      'dhr': dhr != null ? base64.encode(await dhr!.getBytes()) : null,
-      ),
+      'rk': base64.encode(rk),
-    );
+      'cks': cks != null ? base64.encode(cks!) : null,
-    rk = List.from(newRk2);
+      'ckr': ckr != null ? base64.encode(ckr!) : null,
-    cks = List.from(newRk2);
+      'ns': ns,
+      'nr': nr,
+      'pn': pn,
+      'ik_pub': base64.encode(await ik.getBytes()),
+      'session_ad': base64.encode(sessionAd),
+      'mkskipped': mkSkippedSerialised,
+      'acknowledged': acknowledged,
+      'kex_timestamp': kexTimestamp,
+      'kex': kex,
+    };
   }
 
-  /// Skip (and keep track of) message keys until our receive counter is
+  /// Returns the OMEMO compatible fingerprint of the ratchet session.
-  /// equal to [until]. If we would skip too many messages, returns
+  Future<String> getOmemoFingerprint() async {
-  /// a [SkippingTooManyKeysError]. If not, returns null.
+    final curveKey = await ik.toCurve25519();
-  Future<OmemoError?> _skipMessageKeys(int until) async {
+    return HEX.encode(await curveKey.getBytes());
+  }
+  
+  Future<List<int>?> _trySkippedMessageKeys(OmemoMessage header, List<int> ciphertext) async {
+    final key = SkippedKey(
+      OmemoPublicKey.fromBytes(header.dhPub!, KeyPairType.x25519),
+      header.n!,
+    );
+    if (mkSkipped.containsKey(key)) {
+      final mk = mkSkipped[key]!;
+      mkSkipped.remove(key);
+
+      return decrypt(mk, ciphertext, concat([sessionAd, header.writeToBuffer()]), sessionAd);
+    }
+
+    return null;
+  }
+
+  Future<void> _skipMessageKeys(int until) async {
     if (nr + maxSkip < until) {
-      return SkippingTooManyKeysError();
+      throw SkippingTooManyMessagesException();
     }
 
     if (ckr != null) {
@@ -222,147 +284,86 @@ class OmemoDoubleRatchet {
         final newCkr = await kdfCk(ckr!, kdfCkNextChainKey);
         final mk = await kdfCk(ckr!, kdfCkNextMessageKey);
         ckr = newCkr;
-
         mkSkipped[SkippedKey(dhr!, nr)] = mk;
         nr++;
       }
     }
-
-    return null;
   }
 
-  /// Decrypt [ciphertext] using keys derived from the message key [mk]. Also computes the
+  Future<void> _dhRatchet(OmemoMessage header) async {
-  /// HMAC from the [OMEMOMessage] embedded in [message].
+    pn = ns;
-  ///
+    ns = 0;
-  /// If the computed HMAC does not match the HMAC in [message], returns
+    nr = 0;
-  /// [InvalidMessageHMACError]. If it matches, returns the decrypted
+    dhr = OmemoPublicKey.fromBytes(header.dhPub!, KeyPairType.x25519);
-  /// payload.
-  Future<Result<OmemoError, List<int>>> _decrypt(
-    OMEMOAuthenticatedMessage message,
-    List<int> ciphertext,
-    List<int> mk,
-  ) async {
-    final keys = await deriveEncryptionKeys(mk, encryptHkdfInfoString);
 
-    final hmacInput = concat([sessionAd, message.message]);
+    final newRk = await kdfRk(rk, await omemoDH(dhs, dhr!, 0));
-    final hmacResult = await truncatedHmac(hmacInput, keys.authenticationKey);
+    rk = List.from(newRk);
-    if (!listsEqual(hmacResult, message.mac)) {
+    ckr = List.from(newRk);
-      return Result(InvalidMessageHMACError());
+    dhs = await OmemoKeyPair.generateNewPair(KeyPairType.x25519);
-    }
+    final newNewRk = await kdfRk(rk, await omemoDH(dhs, dhr!, 0));
-
+    rk = List.from(newNewRk);
-    final plaintext =
+    cks = List.from(newNewRk);
-        await aes256CbcDecrypt(ciphertext, keys.encryptionKey, keys.iv);
-    if (plaintext.isType<MalformedCiphertextError>()) {
-      return Result(plaintext.get<MalformedCiphertextError>());
-    }
-
-    return Result(plaintext.get<List<int>>());
   }
 
-  /// Checks whether we could decrypt the payload in [header] with a skipped key. If yes,
+  /// Encrypt [plaintext] using the Double Ratchet.
-  /// attempts to decrypt it. If not, returns null.
+  Future<RatchetStep> ratchetEncrypt(List<int> plaintext) async {
-  ///
+    final newCks = await kdfCk(cks!, kdfCkNextChainKey);
-  /// If the decryption is successful, returns the plaintext payload. If an error occurs, like
-  /// an [InvalidMessageHMACError], that is returned instead.
-  Future<Result<OmemoError, List<int>?>> _trySkippedMessageKeys(
-    OMEMOAuthenticatedMessage message,
-    OMEMOMessage header,
-  ) async {
-    final key = SkippedKey(
-      OmemoPublicKey.fromBytes(header.dhPub, KeyPairType.x25519),
-      header.n,
-    );
-    if (mkSkipped.containsKey(key)) {
-      final mk = mkSkipped[key]!;
-      mkSkipped.remove(key);
-
-      return _decrypt(message, header.ciphertext, mk);
-    }
-
-    return const Result(null);
-  }
-
-  /// Decrypt the payload (deeply) embedded in [message].
-  ///
-  /// If everything goes well, returns the plaintext payload. If an error occurs, that
-  /// is returned instead.
-  Future<Result<OmemoError, List<int>>> ratchetDecrypt(
-    OMEMOAuthenticatedMessage message,
-  ) async {
-    final header = OMEMOMessage.fromBuffer(message.message);
-
-    // Try skipped keys
-    final plaintextRaw = await _trySkippedMessageKeys(message, header);
-    if (plaintextRaw.isType<OmemoError>()) {
-      // Propagate the error
-      return Result(plaintextRaw.get<OmemoError>());
-    }
-
-    final plaintext = plaintextRaw.get<List<int>?>();
-    if (plaintext != null) {
-      return Result(plaintext);
-    }
-
-    if (dhr == null || !listsEqual(header.dhPub, await dhr!.getBytes())) {
-      final skipResult1 = await _skipMessageKeys(header.pn);
-      if (skipResult1 != null) {
-        return Result(skipResult1);
-      }
-
-      await _dhRatchet(header);
-    }
-
-    final skipResult2 = await _skipMessageKeys(header.n);
-    if (skipResult2 != null) {
-      return Result(skipResult2);
-    }
-
-    final ck = await kdfCk(ckr!, kdfCkNextChainKey);
-    final mk = await kdfCk(ckr!, kdfCkNextMessageKey);
-    ckr = ck;
-    nr++;
-
-    return _decrypt(message, header.ciphertext, mk);
-  }
-
-  /// Encrypt the payload [plaintext] using the double ratchet session.
-  Future<OMEMOAuthenticatedMessage> ratchetEncrypt(List<int> plaintext) async {
-    // Advance the ratchet
-    final ck = await kdfCk(cks!, kdfCkNextChainKey);
     final mk = await kdfCk(cks!, kdfCkNextMessageKey);
-    cks = ck;
 
-    // Generate encryption, authentication key and IV
+    cks = newCks;
-    final keys = await deriveEncryptionKeys(mk, encryptHkdfInfoString);
+    final header = OmemoMessage()
-    final ciphertext =
-        await aes256CbcEncrypt(plaintext, keys.encryptionKey, keys.iv);
-
-    // Fill-in the header and serialize it here so we do it only once
-    final header = OMEMOMessage()
       ..dhPub = await dhs.pk.getBytes()
       ..pn = pn
-      ..n = ns
+      ..n = ns;
-      ..ciphertext = ciphertext;
-    final headerBytes = header.writeToBuffer();
 
-    // Increment the send counter
     ns++;
 
-    final newAd = concat([sessionAd, headerBytes]);
+    return RatchetStep(
-    final hmac = await truncatedHmac(newAd, keys.authenticationKey);
+      header,
-    return OMEMOAuthenticatedMessage()
+      await encrypt(mk, plaintext, concat([sessionAd, header.writeToBuffer()]), sessionAd),
-      ..mac = hmac
+    );
-      ..message = headerBytes;
+  }
+
+  /// Decrypt a [ciphertext] that was sent with the header [header] using the Double
+  /// Ratchet. Returns the decrypted (raw) plaintext.
+  ///
+  /// Throws an SkippingTooManyMessagesException if too many messages were to be skipped.
+  Future<List<int>> ratchetDecrypt(OmemoMessage header, List<int> ciphertext) async {
+    // Check if we skipped too many messages
+    final plaintext = await _trySkippedMessageKeys(header, ciphertext);
+    if (plaintext != null) {
+      return plaintext;
+    }
+
+    final dhPubMatches = listsEqual(
+      header.dhPub!,
+      (await dhr?.getBytes()) ?? <int>[],
+    );
+    if (!dhPubMatches) {
+      await _skipMessageKeys(header.pn!);
+      await _dhRatchet(header);
+    }
+    
+    await _skipMessageKeys(header.n!);
+    final newCkr = await kdfCk(ckr!, kdfCkNextChainKey);
+    final mk = await kdfCk(ckr!, kdfCkNextMessageKey);
+    ckr = newCkr;
+    nr++;
+
+    return decrypt(mk, ciphertext, concat([sessionAd, header.writeToBuffer()]), sessionAd);
   }
 
-  /// Returns a copy of the ratchet.
   OmemoDoubleRatchet clone() {
     return OmemoDoubleRatchet(
       dhs,
       dhr,
       rk,
-      cks != null ? List<int>.from(cks!) : null,
+      cks != null ?
-      ckr != null ? List<int>.from(ckr!) : null,
+        List<int>.from(cks!) :
+        null,
+      ckr != null ?
+        List<int>.from(ckr!) :
+        null,
       ns,
       nr,
       pn,
@@ -370,47 +371,62 @@ class OmemoDoubleRatchet {
       sessionAd,
       Map<SkippedKey, List<int>>.from(mkSkipped),
       acknowledged,
+      kexTimestamp,
       kex,
     );
   }
 
-  /// Computes the fingerprint of the double ratchet, according to
+  OmemoDoubleRatchet cloneWithKex(String kex) {
-  /// XEP-0384.
+    return OmemoDoubleRatchet(
-  Future<String> get fingerprint async {
+      dhs,
-    final curveKey = await ik.toCurve25519();
+      dhr,
-    return HEX.encode(
+      rk,
-      await curveKey.getBytes(),
+      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
+    final dhrMatch = dhr == null ?
-        ? other.dhr == null
+      other.dhr == null :
-        :
+      // ignore: invalid_use_of_visible_for_testing_member
-        // ignore: invalid_use_of_visible_for_testing_member
+      other.dhr != null && await dhr!.equals(other.dhr!);
-        other.dhr != null && await dhr!.equals(other.dhr!);
+    final ckrMatch = ckr == null ?
-    final ckrMatch = ckr == null
+      other.ckr == null :
-        ? other.ckr == null
+      other.ckr != null && listsEqual(ckr!, other.ckr!);
-        : other.ckr != null && listsEqual(ckr!, other.ckr!);
+    final cksMatch = cks == null ?
-    final cksMatch = cks == null
+      other.cks == null :
-        ? other.cks == null
+      other.cks != null && listsEqual(cks!, other.cks!);
-        : other.cks != null && listsEqual(cks!, other.cks!);
+ 
-
     // ignore: invalid_use_of_visible_for_testing_member
     final dhsMatch = await dhs.equals(other.dhs);
     // ignore: invalid_use_of_visible_for_testing_member
     final ikMatch = await ik.equals(other.ik);
 
     return dhsMatch &&
-        ikMatch &&
+      ikMatch &&
-        dhrMatch &&
+      dhrMatch &&
-        listsEqual(rk, other.rk) &&
+      listsEqual(rk, other.rk) &&
-        cksMatch &&
+      cksMatch &&
-        ckrMatch &&
+      ckrMatch &&
-        ns == other.ns &&
+      ns == other.ns &&
-        nr == other.nr &&
+      nr == other.nr &&
-        pn == other.pn &&
+      pn == other.pn &&
-        listsEqual(sessionAd, other.sessionAd);
+      listsEqual(sessionAd, other.sessionAd) &&
+      kexTimestamp == other.kexTimestamp;
   }
 }

lib/src/double_ratchet/kdf.dart

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

lib/src/errors.dart

@@ -1,39 +1,54 @@
-abstract class OmemoError {}
+abstract class OmemoException {}
 
 /// Triggered during X3DH if the signature if the SPK does verify to the actual SPK.
-class InvalidKeyExchangeSignatureError extends OmemoError {}
+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.
-class InvalidMessageHMACError extends OmemoError {}
+/// Triggered by the Session Manager if the computed HMAC does not match the attached HMAC.
+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 SkippingTooManyKeysError extends OmemoError {}
+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 NotEncryptedForDeviceError extends OmemoError {}
+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 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 UnknownSignedPrekeyError extends OmemoError {}
+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 NoKeyMaterialAvailableError extends OmemoError {}
+class NoKeyMaterialAvailableException extends OmemoException implements Exception {
-
+  String errMsg() => 'No key material available to create a ratchet session with';
-/// A non-key-exchange message was received that was encrypted for our device, but we have no ratchet with
-/// the device that sent the message.
-class NoSessionWithDeviceError extends OmemoError {}
-
-/// Caused when the AES-256 CBC decryption failed.
-class MalformedCiphertextError extends OmemoError {
-  MalformedCiphertextError(this.ex);
-
-  /// The exception that was raised while decryption.
-  final Object ex;
 }
-
-/// Caused by an empty <key /> element
-class MalformedEncryptedKeyError extends OmemoError {}

lib/src/helpers.dart

@@ -1,5 +1,7 @@
 import 'dart:convert';
 import 'dart:math';
+import 'package:cryptography/cryptography.dart';
+import 'package:omemo_dart/src/keys.dart';
 
 /// Flattens [inputs] and concatenates the elements.
 List<int> concat(List<List<int>> inputs) {
@@ -41,35 +43,37 @@ int generateRandom32BitNumber() {
   return Random.secure().nextInt(4294967295 /*pow(2, 32) - 1*/);
 }
 
-/// Describes the differences between two lists in terms of its items.
+OmemoPublicKey? decodeKeyIfNotNull(Map<String, dynamic> map, String key, KeyPairType type) {
-class ListDiff<T> {
+  if (map[key] == null) return null;
-  ListDiff(this.added, this.removed);
 
-  /// The items that were added.
+  return OmemoPublicKey.fromBytes(
-  final List<T> added;
+    base64.decode(map[key]! as String),
-
+    type,
-  /// The items that were removed.
+  );
-  final List<T> removed;
 }
 
-extension AppendToListOrCreateExtension<K, V> on Map<K, List<V>> {
+List<int>? base64DecodeIfNotNull(Map<String, dynamic> map, String key) {
-  /// Create or append [value] to the list identified with key [key].
+  if (map[key] == null) return null;
-  void appendOrCreate(K key, V value, {bool checkExistence = false}) {
+
-    if (containsKey(key)) {
+  return base64.decode(map[key]! as String);
-      if (!checkExistence) {
-        this[key]!.add(value);
-      }
-      if (!this[key]!.contains(value)) {
-        this[key]!.add(value);
-      }
-    } else {
-      this[key] = [value];
-    }
-  }
 }
 
-extension StringFromBase64Extension on String {
+String? base64EncodeIfNotNull(List<int>? bytes) {
-  /// Base64-decode this string. Useful for doing `someString?.fromBase64()` instead
+  if (bytes == null) return null;
-  /// of `someString != null ? base64Decode(someString) : null`.
+
-  List<int> fromBase64() => base64Decode(this);
+  return base64.encode(bytes);
+}
+
+OmemoKeyPair? decodeKeyPairIfNotNull(String? pk, String? sk, KeyPairType type) {
+  if (pk == null || sk == null) return null;
+
+  return OmemoKeyPair.fromBytes(
+    base64.decode(pk),
+    base64.decode(sk),
+    type,
+  );
+}
+
+int getTimestamp() {
+  return DateTime.now().millisecondsSinceEpoch;
 }

lib/src/keys.dart

@@ -21,7 +21,7 @@ class OmemoPublicKey {
   }
 
   final SimplePublicKey _pubkey;
-
+  
   KeyPairType get type => _pubkey.type;
 
   /// Return the bytes that comprise the public key.
@@ -31,10 +31,7 @@ class OmemoPublicKey {
   Future<String> asBase64() async => base64Encode(_pubkey.bytes);
 
   Future<OmemoPublicKey> toCurve25519() async {
-    assert(
+    assert(type == KeyPairType.ed25519, 'Cannot convert non-Ed25519 public key to X25519');
-      type == KeyPairType.ed25519,
-      'Cannot convert non-Ed25519 public key to X25519',
-    );
 
     final pkc = Uint8List(publicKeyLength);
     TweetNaClExt.crypto_sign_ed25519_pk_to_x25519_pk(
@@ -42,20 +39,17 @@ class OmemoPublicKey {
       Uint8List.fromList(await getBytes()),
     );
 
-    return OmemoPublicKey(
+    return OmemoPublicKey(SimplePublicKey(List<int>.from(pkc), type: KeyPairType.x25519));
-      SimplePublicKey(List<int>.from(pkc), type: KeyPairType.x25519),
-    );
   }
 
   SimplePublicKey asPublicKey() => _pubkey;
 
   @visibleForTesting
   Future<bool> equals(OmemoPublicKey key) async {
-    return type == key.type &&
+    return type == key.type && listsEqual(
-        listsEqual(
+      await getBytes(),
-          await getBytes(),
+      await key.getBytes(),
-          await key.getBytes(),
+    );
-        );
   }
 }
 
@@ -65,12 +59,9 @@ class OmemoPrivateKey {
   final KeyPairType type;
 
   Future<List<int>> getBytes() async => _privkey;
-
+  
   Future<OmemoPrivateKey> toCurve25519() async {
-    assert(
+    assert(type == KeyPairType.ed25519, 'Cannot convert non-Ed25519 private key to X25519');
-      type == KeyPairType.ed25519,
-      'Cannot convert non-Ed25519 private key to X25519',
-    );
 
     final skc = Uint8List(privateKeyLength);
     TweetNaClExt.crypto_sign_ed25519_sk_to_x25519_sk(
@@ -83,11 +74,10 @@ class OmemoPrivateKey {
 
   @visibleForTesting
   Future<bool> equals(OmemoPrivateKey key) async {
-    return type == key.type &&
+    return type == key.type && listsEqual(
-        listsEqual(
+      await getBytes(),
-          await getBytes(),
+      await key.getBytes(),
-          await key.getBytes(),
+    );
-        );
   }
 }
 
@@ -97,11 +87,7 @@ class OmemoKeyPair {
 
   /// Create an OmemoKeyPair just from a [type] and the bytes of the private and public
   /// key.
-  factory OmemoKeyPair.fromBytes(
+  factory OmemoKeyPair.fromBytes(List<int> publicKey, List<int> privateKey, KeyPairType type) {
-    List<int> publicKey,
-    List<int> privateKey,
-    KeyPairType type,
-  ) {
     return OmemoKeyPair(
       OmemoPublicKey.fromBytes(
         publicKey,
@@ -118,10 +104,7 @@ class OmemoKeyPair {
   /// Generate a completely new random OmemoKeyPair of type [type]. [type] must be either
   /// KeyPairType.ed25519 or KeyPairType.x25519.
   static Future<OmemoKeyPair> generateNewPair(KeyPairType type) async {
-    assert(
+    assert(type == KeyPairType.ed25519 || type == KeyPairType.x25519, 'Keypair must be either Ed25519 or X25519');
-      type == KeyPairType.ed25519 || type == KeyPairType.x25519,
-      'Keypair must be either Ed25519 or X25519',
-    );
 
     SimpleKeyPair kp;
     if (type == KeyPairType.ed25519) {
@@ -136,7 +119,7 @@ class OmemoKeyPair {
     }
 
     final kpd = await kp.extract();
-
+    
     return OmemoKeyPair(
       OmemoPublicKey(await kp.extractPublicKey()),
       OmemoPrivateKey(await kpd.extractPrivateKeyBytes(), type),
@@ -147,13 +130,10 @@ class OmemoKeyPair {
   final KeyPairType type;
   final OmemoPublicKey pk;
   final OmemoPrivateKey sk;
-
+  
   /// Return the bytes that comprise the public key.
   Future<OmemoKeyPair> toCurve25519() async {
-    assert(
+    assert(type == KeyPairType.ed25519, 'Cannot convert non-Ed25519 keypair to X25519');
-      type == KeyPairType.ed25519,
-      'Cannot convert non-Ed25519 keypair to X25519',
-    );
 
     return OmemoKeyPair(
       await pk.toCurve25519(),
@@ -173,7 +153,7 @@ class OmemoKeyPair {
   @visibleForTesting
   Future<bool> equals(OmemoKeyPair pair) async {
     return type == pair.type &&
-        await pk.equals(pair.pk) &&
+      await pk.equals(pair.pk) &&
-        await sk.equals(pair.sk);
+      await sk.equals(pair.sk);
   }
 }

lib/src/omemo/bundle.dart

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

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

@@ -3,28 +3,7 @@ import 'package:omemo_dart/src/errors.dart';
 
 @immutable
 class DecryptionResult {
-  const DecryptionResult(
+  const DecryptionResult(this.payload, this.error);
-    this.payload,
-    this.usedOpkId,
-    this.newRatchets,
-    this.replacedRatchets,
-    this.error,
-  );
-
-  /// The decrypted payload or null, if it was an empty OMEMO message.
   final String? payload;
-
+  final OmemoException? error;
-  /// In case a key exchange has been performed: The id of the used OPK. Useful for
-  /// replacing the OPK after a message catch-up.
-  final int? usedOpkId;
-
-  /// Mapping of JIDs to a list of device ids for which we created a new ratchet session.
-  final Map<String, List<int>> newRatchets;
-
-  /// Similar to [newRatchets], but the ratchets listed in [replacedRatchets] where also existent before
-  /// and replaced with the new ratchet.
-  final Map<String, List<int>> replacedRatchets;
-
-  /// The error that occurred during decryption or null, if no error occurred.
-  final OmemoError? error;
 }

lib/src/omemo/device.dart

@@ -22,11 +22,78 @@ class OmemoDevice {
     this.opks,
   );
 
+  /// Deserialize the Device
+  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.
+    /*
+    {
+      'jid': 'alice@...',
+      'id': 123,
+      'ik': 'base/64/encoded',
+      'ik_pub': 'base/64/encoded',
+      'spk': 'base/64/encoded',
+      'spk_pub': 'base/64/encoded',
+      'spk_id': 123,
+      'spk_sig': 'base/64/encoded',
+      'old_spk': 'base/64/encoded',
+      'old_spk_pub': 'base/64/encoded',
+      'old_spk_id': 122,
+      'opks': [
+        {
+          'id': 0,
+          'public': 'base/64/encoded',
+          'private': 'base/64/encoded'
+        }, ...
+      ]
+    }
+    */
+    // 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 OmemoDevice(
+      data['jid']! as String,
+      data['id']! as int,
+      OmemoKeyPair.fromBytes(
+        base64.decode(data['ik_pub']! as String),
+        base64.decode(data['ik']! as String),
+        KeyPairType.ed25519,
+      ),
+      OmemoKeyPair.fromBytes(
+        base64.decode(data['spk_pub']! as String),
+        base64.decode(data['spk']! as String),
+        KeyPairType.x25519,
+      ),
+      data['spk_id']! as int,
+      base64.decode(data['spk_sig']! as String),
+      decodeKeyPairIfNotNull(
+        data['old_spk_pub'] as String?,
+        data['old_spk'] as String?,
+        KeyPairType.x25519,
+      ),
+      data['old_spk_id'] as int?,
+      opks,
+    );
+  }
+  
   /// Generate a completely new device, i.e. cryptographic identity.
-  static Future<OmemoDevice> generateNewDevice(
+  static Future<OmemoDevice> generateNewDevice(String jid, { int opkAmount = 100 }) async {
-    String jid, {
-    int opkAmount = 100,
-  }) async {
     final id = generateRandom32BitNumber();
     final ik = await OmemoKeyPair.generateNewPair(KeyPairType.ed25519);
     final spk = await OmemoKeyPair.generateNewPair(KeyPairType.x25519);
@@ -35,16 +102,7 @@ class OmemoDevice {
 
     final opks = <int, OmemoKeyPair>{};
     for (var i = 0; i < opkAmount; i++) {
-      // Generate unique ids for each key
+      opks[i] = await OmemoKeyPair.generateNewPair(KeyPairType.x25519);
-      while (true) {
-        final opkId = generateRandom32BitNumber();
-        if (opks.containsKey(opkId)) {
-          continue;
-        }
-
-        opks[opkId] = await OmemoKeyPair.generateNewPair(KeyPairType.x25519);
-        break;
-      }
     }
 
     return OmemoDevice(jid, id, ik, spk, spkId, signature, null, null, opks);
@@ -52,7 +110,7 @@ class OmemoDevice {
 
   /// Our bare Jid
   final String jid;
-
+  
   /// The device Id
   final int id;
 
@@ -61,16 +119,13 @@ class OmemoDevice {
 
   /// The signed prekey...
   final OmemoKeyPair spk;
-
   /// ...its Id, ...
   final int spkId;
-
   /// ...and its signature
   final List<int> spkSignature;
 
   /// The old Signed Prekey...
   final OmemoKeyPair? oldSpk;
-
   /// ...and its Id
   final int? oldSpkId;
 
@@ -81,19 +136,8 @@ class OmemoDevice {
   /// a new Device object that copies over everything but replaces said key.
   @internal
   Future<OmemoDevice> replaceOnetimePrekey(int id) async {
-    opks.remove(id);
+    opks[id] = await OmemoKeyPair.generateNewPair(KeyPairType.x25519);
-
+    
-    // Generate a new unique id for the OPK.
-    while (true) {
-      final newId = generateRandom32BitNumber();
-      if (opks.containsKey(newId)) {
-        continue;
-      }
-
-      opks[newId] = await OmemoKeyPair.generateNewPair(KeyPairType.x25519);
-      break;
-    }
-
     return OmemoDevice(
       jid,
       this.id,
@@ -144,7 +188,7 @@ class OmemoDevice {
       opks,
     );
   }
-
+  
   /// Converts this device into an OmemoBundle that could be used for publishing.
   Future<OmemoBundle> toBundle() async {
     final encodedOpks = <int, String>{};
@@ -170,6 +214,34 @@ class OmemoDevice {
     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
+    final serialisedOpks = List<Map<String, dynamic>>.empty(growable: true);
+    for (final entry in opks.entries) {
+      serialisedOpks.add({
+        'id': entry.key,
+        'public': base64.encode(await entry.value.pk.getBytes()),
+        'private': base64.encode(await entry.value.sk.getBytes()),
+      });
+    }
+    
+    return {
+      'jid': jid,
+      'id': id,
+      'ik': base64.encode(await ik.sk.getBytes()),
+      'ik_pub': base64.encode(await ik.pk.getBytes()),
+      'spk': base64.encode(await spk.sk.getBytes()),
+      'spk_pub': base64.encode(await spk.pk.getBytes()),
+      'spk_id': spkId,
+      'spk_sig': base64.encode(spkSignature),
+      'old_spk': base64EncodeIfNotNull(await oldSpk?.sk.getBytes()),
+      'old_spk_pub': base64EncodeIfNotNull(await oldSpk?.pk.getBytes()),
+      'old_spk_id': oldSpkId,
+      'opks': serialisedOpks,
+    };
+  }
 
   @visibleForTesting
   Future<bool> equals(OmemoDevice other) async {
@@ -179,9 +251,7 @@ class OmemoDevice {
     } else {
       for (final entry in opks.entries) {
         // ignore: invalid_use_of_visible_for_testing_member
-        final matches =
+        final matches = await other.opks[entry.key]?.equals(entry.value) ?? false;
-            // ignore: invalid_use_of_visible_for_testing_member
-            await other.opks[entry.key]?.equals(entry.value) ?? false;
         if (!matches) {
           opksMatch = false;
         }
@@ -193,18 +263,15 @@ class OmemoDevice {
     // ignore: invalid_use_of_visible_for_testing_member
     final spkMatch = await spk.equals(other.spk);
     // ignore: invalid_use_of_visible_for_testing_member
-    final oldSpkMatch = oldSpk != null
+    final oldSpkMatch = oldSpk != null ? await oldSpk!.equals(other.oldSpk!) : other.oldSpk == null;
-        // ignore: invalid_use_of_visible_for_testing_member
-        ? await oldSpk!.equals(other.oldSpk!)
-        : other.oldSpk == null;
     return id == other.id &&
-        ikMatch &&
+      ikMatch &&
-        spkMatch &&
+      spkMatch &&
-        oldSpkMatch &&
+      oldSpkMatch &&
-        jid == other.jid &&
+      jid == other.jid &&
-        listsEqual(spkSignature, other.spkSignature) &&
+      listsEqual(spkSignature, other.spkSignature) &&
-        spkId == other.spkId &&
+      spkId == other.spkId &&
-        oldSpkId == other.oldSpkId &&
+      oldSpkId == other.oldSpkId &&
-        opksMatch;
+      opksMatch;
   }
 }

lib/src/omemo/encrypted_key.dart

@@ -1,23 +1,13 @@
-import 'dart:convert';
-
 import 'package:meta/meta.dart';
 
 /// EncryptedKey is the intermediary format of a <key /> element in the OMEMO message's
 /// <keys /> header.
 @immutable
 class EncryptedKey {
-  const EncryptedKey(this.rid, this.value, this.kex);
 
-  /// The id of the device the key is encrypted for.
+  const EncryptedKey(this.jid, this.rid, this.value, this.kex);
+  final String jid;
   final int rid;
-
-  /// The base64-encoded payload.
   final String value;
-
-  /// Flag indicating whether the payload is a OMEMOKeyExchange (true) or
-  /// an OMEMOAuthenticatedMessage (false).
   final bool kex;
-
-  /// The base64-decoded payload.
-  List<int> get data => base64Decode(value);
 }

lib/src/omemo/encryption_result.dart

@@ -1,36 +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/errors.dart';
+import 'package:omemo_dart/src/omemo/ratchet_map_key.dart';
 
 @immutable
 class EncryptionResult {
-  const EncryptionResult(
+  const EncryptionResult(this.ciphertext, this.encryptedKeys, this.deviceEncryptionErrors, this.jidEncryptionErrors);
-    this.ciphertext,
+  
-    this.encryptedKeys,
-    this.deviceEncryptionErrors,
-    this.newRatchets,
-    this.replacedRatchets,
-    this.canSend,
-  );
-
   /// 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.
-  final Map<String, List<EncryptedKey>> encryptedKeys;
+  final List<EncryptedKey> encryptedKeys;
 
-  /// Mapping of a JID to
+  /// Mapping of a ratchet map keys to a possible exception.
-  final Map<String, List<EncryptToJidError>> deviceEncryptionErrors;
+  final Map<RatchetMapKey, OmemoException> deviceEncryptionErrors;
 
-  /// Mapping of JIDs to a list of device ids for which we created a new ratchet session.
+  /// Mapping of a JID to a possible exception.
-  final Map<String, List<int>> newRatchets;
+  final Map<String, OmemoException> jidEncryptionErrors;
 
-  /// Similar to [newRatchets], but the ratchets listed in [replacedRatchets] where also existent before
+  /// True if the encryption was a success. This means that we could encrypt for
-  /// and replaced with the new ratchet.
+  /// at least one ratchet.
-  final Map<String, List<int>> replacedRatchets;
+  bool isSuccess(int numberOfRecipients) => encryptedKeys.isNotEmpty && jidEncryptionErrors.length < numberOfRecipients;
-
-  /// A flag indicating that the message could be sent like that, i.e. we were able
-  /// to encrypt to at-least one device per recipient.
-  final bool canSend;
 }

lib/src/omemo/errors.dart

@@ -1,12 +0,0 @@
-import 'package:omemo_dart/src/errors.dart';
-
-/// Returned on encryption, if encryption failed for some reason.
-class EncryptToJidError extends OmemoError {
-  EncryptToJidError(this.device, this.error);
-
-  /// The device the error occurred with
-  final int? device;
-
-  /// The actual error.
-  final OmemoError error;
-}

lib/src/omemo/events.dart

@@ -0,0 +1,35 @@
+import 'package:omemo_dart/src/double_ratchet/double_ratchet.dart';
+import 'package:omemo_dart/src/omemo/device.dart';
+
+abstract class OmemoEvent {}
+
+/// Triggered when a ratchet has been modified
+class RatchetModifiedEvent extends OmemoEvent {
+  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 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 OmemoDevice device;
+}

lib/src/omemo/fingerprint.dart

@@ -9,8 +9,8 @@ class DeviceFingerprint {
   @override
   bool operator ==(Object other) {
     return other is DeviceFingerprint &&
-        fingerprint == other.fingerprint &&
+      fingerprint == other.fingerprint &&
-        deviceId == other.deviceId;
+      deviceId == other.deviceId;
   }
 
   @override

lib/src/omemo/omemomanager.dart

@@ -0,0 +1,787 @@
+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.
+  @visibleForTesting
+  OmemoDoubleRatchet? getRatchet(RatchetMapKey key) => _ratchetMap[key];
+
+  /// Figure out what bundles we have to still build a session with.
+  Future<List<OmemoBundle>> _fetchNewBundles(String jid) async {
+    // Check if we already requested the device list for [jid]
+    List<int> bundlesToFetch;
+    if (!_deviceListRequested.containsKey(jid) || !_deviceList.containsKey(jid)) {
+      // We don't have an up-to-date version of the device list
+      final newDeviceList = await fetchDeviceListImpl(jid);
+      if (newDeviceList == null) return [];
+
+      _deviceList[jid] = newDeviceList;
+      bundlesToFetch = newDeviceList
+        .where((id) {
+          return !_ratchetMap.containsKey(RatchetMapKey(jid, id)) ||
+                 _deviceList[jid]?.contains(id) == false;
+        }).toList();
+
+      // Trigger an event with the new device list
+      _eventStreamController.add(DeviceListModifiedEvent(_deviceList));
+    } else {
+      // We already have an up-to-date version of the device list
+      bundlesToFetch = _deviceList[jid]!
+        .where((id) => !_ratchetMap.containsKey(RatchetMapKey(jid, id)))
+        .toList();
+    }
+
+    if (bundlesToFetch.isNotEmpty) {
+      _log.finest('Fetching bundles $bundlesToFetch for $jid');
+    }
+
+    final device = await getDevice();
+    final newBundles = List<OmemoBundle>.empty(growable: true);
+    for (final id in bundlesToFetch) {
+      if (jid == device.jid && id == device.id) continue;
+
+      final bundle = await fetchDeviceBundleImpl(jid, id);
+      if (bundle != null) newBundles.add(bundle);
+    }
+
+    return newBundles;
+  }
+  
+  /// Encrypt the key [plaintext] for all known bundles of the Jids in [jids]. Returns a
+  /// map that maps the device Id to the ciphertext of [plaintext].
+  ///
+  /// If [plaintext] is null, then the result will be an empty OMEMO message, i.e. one that
+  /// does not contain a <payload /> element. This means that the ciphertext attribute of
+  /// the result will be null as well.
+  /// NOTE: Must be called within the ratchet critical section
+  Future<EncryptionResult> _encryptToJids(List<String> jids, String? plaintext) async {
+    final encryptedKeys = List<EncryptedKey>.empty(growable: true);
+
+    var ciphertext = const <int>[];
+    var keyPayload = const <int>[];
+    if (plaintext != null) {
+      // Generate the key and encrypt the plaintext
+      final key = generateRandomBytes(32);
+      final keys = await deriveEncryptionKeys(key, omemoPayloadInfoString);
+      ciphertext = await aes256CbcEncrypt(
+        utf8.encode(plaintext),
+        keys.encryptionKey,
+        keys.iv,
+      );
+      final hmac = await truncatedHmac(ciphertext, keys.authenticationKey);
+      keyPayload = concat([key, hmac]);
+    } else {
+      keyPayload = List<int>.filled(32, 0x0);
+    }
+
+    final kex = <RatchetMapKey, OmemoKeyExchange>{};
+    for (final jid in jids) {
+      for (final newSession in await _fetchNewBundles(jid)) {
+        kex[RatchetMapKey(jid, newSession.id)] = await addSessionFromBundle(
+          newSession.jid,
+          newSession.id,
+          newSession,
+        );
+      }
+    }
+
+    // We assume that the user already checked if the session exists
+    final deviceEncryptionErrors = <RatchetMapKey, OmemoException>{};
+    final jidEncryptionErrors = <String, OmemoException>{};
+    for (final jid in jids) {
+      final devices = _deviceList[jid];
+      if (devices == null) {
+        _log.severe('Device list does not exist for $jid.');
+        jidEncryptionErrors[jid] = NoKeyMaterialAvailableException();
+        continue;
+      }
+
+      if (!_subscriptionMap.containsKey(jid)) {
+        unawaited(subscribeToDeviceListNodeImpl(jid));
+        _subscriptionMap[jid] = true;
+      }
+
+      for (final deviceId in devices) {
+        // Empty OMEMO messages are allowed to bypass trust
+        if (plaintext != null) {
+          // Only encrypt to devices that are trusted
+          if (!(await _trustManager.isTrusted(jid, deviceId))) continue;
+
+          // Only encrypt to devices that are enabled
+          if (!(await _trustManager.isEnabled(jid, deviceId))) continue;
+        }
+
+        final ratchetKey = RatchetMapKey(jid, deviceId);
+        var ratchet = _ratchetMap[ratchetKey];
+        if (ratchet == null) {
+          _log.severe('Ratchet ${ratchetKey.toJsonKey()} does not exist.');
+          deviceEncryptionErrors[ratchetKey] = NoKeyMaterialAvailableException();
+          continue;
+        }
+
+        final ciphertext = (await ratchet.ratchetEncrypt(keyPayload)).ciphertext;
+ 
+        if (kex.containsKey(ratchetKey)) {
+          // The ratchet did not exist
+          final k = kex[ratchetKey]!
+            ..message = OmemoAuthenticatedMessage.fromBuffer(ciphertext);
+          final buffer = base64.encode(k.writeToBuffer());
+          encryptedKeys.add(
+            EncryptedKey(
+              jid,
+              deviceId,
+              buffer,
+              true,
+            ),
+          );
+
+          ratchet = ratchet.cloneWithKex(buffer);
+          _ratchetMap[ratchetKey] = ratchet;
+        } else if (!ratchet.acknowledged) {
+          // The ratchet exists but is not acked
+          if (ratchet.kex != null) {
+            final oldKex = OmemoKeyExchange.fromBuffer(base64.decode(ratchet.kex!))
+              ..message = OmemoAuthenticatedMessage.fromBuffer(ciphertext);
+            
+            encryptedKeys.add(
+              EncryptedKey(
+                jid,
+                deviceId,
+                base64.encode(oldKex.writeToBuffer()),
+                true,
+              ),
+            );
+          } else {
+            // The ratchet is not acked but we don't have the old key exchange
+            _log.warning('Ratchet for $jid:$deviceId is not acked but the kex attribute is null');
+            encryptedKeys.add(
+              EncryptedKey(
+                jid,
+                deviceId,
+                base64.encode(ciphertext),
+                false,
+              ),
+            );
+          }
+        } else {
+          // The ratchet exists and is acked
+          encryptedKeys.add(
+            EncryptedKey(
+              jid,
+              deviceId,
+              base64.encode(ciphertext),
+              false,
+            ),
+          );
+        }
+
+        // Commit the ratchet
+        _eventStreamController.add(RatchetModifiedEvent(jid, deviceId, ratchet, false));
+      }
+    }
+
+    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(RatchetMapKey(jid, deviceId));
+      _eventStreamController.add(RatchetRemovedEvent(jid, deviceId));
+    }
+
+    // Remove the devices from the device list cache and commit it
+    _deviceList.remove(jid);
+    _deviceListRequested.remove(jid);
+    _eventStreamController.add(DeviceListModifiedEvent(_deviceList));
+
+    // Remove trust decisions
+    await _trustManager.removeTrustDecisionsForJid(jid);
+    
+    await _leaveRatchetCriticalSection(jid);
+  }
+
+  /// Replaces the internal device with [newDevice]. Does not trigger an event.
+  Future<void> replaceDevice(OmemoDevice newDevice) async {
+    await _deviceLock.synchronized(() {
+      _device = newDevice;
+    });
+  }
+}

lib/src/omemo/queue.dart

@@ -1,100 +0,0 @@
-import 'dart:async';
-import 'dart:collection';
-
-import 'package:meta/meta.dart';
-import 'package:synchronized/synchronized.dart';
-
-extension UtilAllMethodsList<T> on List<T> {
-  void removeAll(List<T> values) {
-    for (final value in values) {
-      remove(value);
-    }
-  }
-
-  bool containsAll(List<T> values) {
-    for (final value in values) {
-      if (!contains(value)) {
-        return false;
-      }
-    }
-
-    return true;
-  }
-}
-
-class _RatchetAccessQueueEntry {
-  _RatchetAccessQueueEntry(
-    this.jids,
-    this.completer,
-  );
-
-  final List<String> jids;
-  final Completer<void> completer;
-}
-
-class RatchetAccessQueue {
-  final Queue<_RatchetAccessQueueEntry> _queue = Queue();
-
-  @visibleForTesting
-  final List<String> runningOperations = List<String>.empty(growable: true);
-
-  final Lock lock = Lock();
-
-  bool canBypass(List<String> jids) {
-    for (final jid in jids) {
-      if (runningOperations.contains(jid)) {
-        return false;
-      }
-    }
-
-    return true;
-  }
-
-  Future<void> enterCriticalSection(List<String> jids) async {
-    final completer = await lock.synchronized<Completer<void>?>(() {
-      if (canBypass(jids)) {
-        runningOperations.addAll(jids);
-        return null;
-      }
-
-      final completer = Completer<void>();
-      _queue.add(
-        _RatchetAccessQueueEntry(
-          jids,
-          completer,
-        ),
-      );
-
-      return completer;
-    });
-
-    await completer?.future;
-  }
-
-  Future<void> leaveCriticalSection(List<String> jids) async {
-    await lock.synchronized(() {
-      runningOperations.removeAll(jids);
-
-      while (_queue.isNotEmpty) {
-        if (canBypass(_queue.first.jids)) {
-          final head = _queue.removeFirst();
-          runningOperations.addAll(head.jids);
-          head.completer.complete();
-        } else {
-          break;
-        }
-      }
-    });
-  }
-
-  Future<T> synchronized<T>(
-    List<String> jids,
-    Future<T> Function() function,
-  ) async {
-    await enterCriticalSection(jids);
-    final result = await function();
-    await leaveCriticalSection(jids);
-
-    return result;
-  }
-}

lib/src/omemo/ratchet_data.dart

@@ -1,18 +0,0 @@
-import 'package:omemo_dart/src/double_ratchet/double_ratchet.dart';
-
-class OmemoRatchetData {
-  const OmemoRatchetData(
-    this.jid,
-    this.id,
-    this.ratchet,
-  );
-
-  /// The JID we have the ratchet with.
-  final String jid;
-
-  /// The device id we have the ratchet with.
-  final int id;
-
-  /// The actual double ratchet to commit.
-  final OmemoDoubleRatchet ratchet;
-}

lib/src/omemo/ratchet_map_key.dart

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

lib/src/omemo/sessionmanager.dart

@@ -0,0 +1,676 @@
+import 'dart:async';
+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/device.dart';
+import 'package:omemo_dart/src/omemo/encrypted_key.dart';
+import 'package:omemo_dart/src/omemo/encryption_result.dart';
+import 'package:omemo_dart/src/omemo/events.dart';
+import 'package:omemo_dart/src/omemo/fingerprint.dart';
+import 'package:omemo_dart/src/omemo/ratchet_map_key.dart';
+import 'package:omemo_dart/src/protobuf/omemo_authenticated_message.dart';
+import 'package:omemo_dart/src/protobuf/omemo_key_exchange.dart';
+import 'package:omemo_dart/src/protobuf/omemo_message.dart';
+import 'package:omemo_dart/src/trust/base.dart';
+import 'package:omemo_dart/src/x3dh/x3dh.dart';
+import 'package:synchronized/synchronized.dart';
+
+@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(),
+      _log = Logger('OmemoSessionManager');
+
+  /// Deserialise the OmemoSessionManager from JSON data [data] that does not contain
+  /// the ratchet sessions.
+  @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(
+      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,
+    );
+  }
+  
+  /// 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 OmemoDevice.generateNewDevice(jid, opkAmount: opkAmount);
+
+    return OmemoSessionManager(device, {}, {}, trustManager);
+  }
+
+  /// Logging
+  Logger _log;
+  
+  /// Lock for _ratchetMap and _bundleMap
+  final Lock _lock;
+
+  /// Mapping of the Device Id to its OMEMO session
+  final Map<RatchetMapKey, OmemoDoubleRatchet> _ratchetMap;
+
+  /// Mapping of a bare Jid to its Device Ids
+  final Map<String, List<int>> _deviceMap;
+
+  /// The event bus of the session manager
+  final StreamController<OmemoEvent> _eventStreamController;
+  
+  /// Our own keys...
+  // ignore: prefer_final_fields
+  OmemoDevice _device;
+  /// and its lock
+  final Lock _deviceLock;
+
+  /// The trust manager
+  final TrustManager _trustManager;
+  TrustManager get trustManager => _trustManager;
+  
+  /// A stream that receives events regarding the session
+  Stream<OmemoEvent> get eventStream => _eventStreamController.stream;
+
+  /// Returns our own device.
+  Future<OmemoDevice> getDevice() async {
+    return _deviceLock.synchronized(() => _device);
+  }
+
+  /// Returns the id attribute of our own device. This is just a short-hand for
+  /// ```await (session.getDevice()).id```.
+  Future<int> getDeviceId() async {
+    return _deviceLock.synchronized(() => _device.id);
+  }
+
+  /// Returns the device as an OmemoBundle. This is just a short-hand for
+  /// ```await (await session.getDevice()).toBundle()```.
+  Future<OmemoBundle> getDeviceBundle() async {
+    return _deviceLock.synchronized(() async => _device.toBundle());
+  }
+  
+  /// Add a session [ratchet] with the [deviceId] to the internal tracking state.
+  Future<void> _addSession(String jid, int deviceId, OmemoDoubleRatchet ratchet) async {
+    await _lock.synchronized(() async {
+      // Add the bundle Id
+      if (!_deviceMap.containsKey(jid)) {
+        _deviceMap[jid] = [deviceId];
+
+        // Commit the device map
+        _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(DeviceListModifiedEvent(_deviceMap));
+        }
+      }
+
+      // Add the ratchet session
+      final key = RatchetMapKey(jid, deviceId);
+      _ratchetMap[key] = ratchet;
+
+      // Commit the ratchet
+      _eventStreamController.add(RatchetModifiedEvent(jid, deviceId, ratchet, true));
+    });
+  }
+
+  /// 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);
+    await _addSession(jid, deviceId, ratchet);
+
+    return OmemoKeyExchange()
+      ..pkId = kexResult.opkId
+      ..spkId = bundle.spkId
+      ..ik = await device.ik.pk.getBytes()
+      ..ek = await kexResult.ek.pk.getBytes();
+  }
+
+  /// Build a new session with the user at [jid] with the device [deviceId] using data
+  /// from the key exchange [kex]. In case [kex] contains an unknown Signed Prekey
+  /// identifier an UnknownSignedPrekeyException will be thrown.
+  Future<OmemoDoubleRatchet> _addSessionFromKeyExchange(String jid, int deviceId, OmemoKeyExchange kex) async {
+    // Pick the correct SPK
+    final device = await getDevice();
+    final spk = await _lock.synchronized(() async {
+      if (kex.spkId == _device.spkId) {
+        return _device.spk;
+      } else if (kex.spkId == _device.oldSpkId) {
+        return _device.oldSpk;
+      }
+
+      return null;
+    });
+    if (spk == null) {
+      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;
+  }
+
+  /// Like [encryptToJids] but only for one Jid [jid].
+  Future<EncryptionResult> encryptToJid(String jid, String? plaintext, { List<OmemoBundle>? newSessions }) {
+    return encryptToJids([jid], plaintext, newSessions: newSessions);
+  }
+  
+  /// Encrypt the key [plaintext] for all known bundles of the Jids in [jids]. Returns a
+  /// map that maps the device Id to the ciphertext of [plaintext].
+  ///
+  /// If [plaintext] is null, then the result will be an empty OMEMO message, i.e. one that
+  /// does not contain a <payload /> element. This means that the ciphertext attribute of
+  /// the result will be null as well.
+  Future<EncryptionResult> encryptToJids(List<String> jids, String? plaintext, { List<OmemoBundle>? newSessions }) async {
+    final encryptedKeys = List<EncryptedKey>.empty(growable: true);
+
+    var ciphertext = const <int>[];
+    var keyPayload = const <int>[];
+    if (plaintext != null) {
+      // Generate the key and encrypt the plaintext
+      final key = generateRandomBytes(32);
+      final keys = await deriveEncryptionKeys(key, omemoPayloadInfoString);
+      ciphertext = await aes256CbcEncrypt(
+        utf8.encode(plaintext),
+        keys.encryptionKey,
+        keys.iv,
+      );
+      final hmac = await truncatedHmac(ciphertext, keys.authenticationKey);
+      keyPayload = concat([key, hmac]);
+    } else {
+      keyPayload = List<int>.filled(32, 0x0);
+    }
+
+    final kex = <int, OmemoKeyExchange>{};
+    if (newSessions != null) {
+      for (final newSession in newSessions) {
+        kex[newSession.id] = await addSessionFromBundle(
+          newSession.jid,
+          newSession.id,
+          newSession,
+        );
+      }
+    }
+    
+    await _lock.synchronized(() async {
+      // We assume that the user already checked if the session exists
+      for (final jid in jids) {
+        for (final deviceId in _deviceMap[jid]!) {
+          // Empty OMEMO messages are allowed to bypass trust
+          if (plaintext != null) {
+            // Only encrypt to devices that are trusted
+            if (!(await _trustManager.isTrusted(jid, deviceId))) continue;
+
+            // Onyl encrypt to devices that are enabled
+            if (!(await _trustManager.isEnabled(jid, deviceId))) continue;
+          }
+
+          final ratchetKey = RatchetMapKey(jid, deviceId);
+          var ratchet = _ratchetMap[ratchetKey]!;
+          final ciphertext = (await ratchet.ratchetEncrypt(keyPayload)).ciphertext;
+ 
+          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,
+                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,
+      const <RatchetMapKey, OmemoException>{},
+      const <String, OmemoException>{},
+    );
+  }
+
+  /// 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].
+  Future<void> _restoreRatchet(RatchetMapKey mapKey, OmemoDoubleRatchet oldRatchet) async {
+    await _lock.synchronized(() {
+      _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,
+        ),
+      );
+    });
+  }
+
+  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, 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;
+    if (rawKey.kex) {
+      // If the ratchet already existed, we store it. If it didn't, oldRatchet will stay
+      // null.
+      final oldRatchet = (await _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. ${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,
+              oldRatchet,
+              false,
+            ),
+          );
+          
+          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!);
+
+        // Commit the device
+        _eventStreamController.add(DeviceModifiedEvent(device));
+      });
+    } else {
+      authMessage = OmemoAuthenticatedMessage.fromBuffer(decodedRawKey);
+      message = OmemoMessage.fromBuffer(authMessage.message!);
+    }
+    
+    final devices = _deviceMap[senderJid];
+    if (devices == null) {
+      throw NoDecryptionKeyException();
+    }
+    if (!devices.contains(senderDeviceId)) {
+      throw NoDecryptionKeyException();
+    }
+
+    // We can guarantee that the ratchet exists at this point in time
+    final ratchet = (await _getRatchet(ratchetKey))!;
+    oldRatchet ??= ratchet.clone();
+
+    try {
+      if (rawKey.kex) {
+        keyAndHmac = await ratchet.ratchetDecrypt(message, authMessage.writeToBuffer());
+      } else {
+        keyAndHmac = await ratchet.ratchetDecrypt(message, decodedRawKey);
+      }
+    } catch (_) {
+      await _restoreRatchet(ratchetKey, oldRatchet);
+      rethrow;
+    }
+
+    // Commit the ratchet
+    _eventStreamController.add(
+      RatchetModifiedEvent(
+        senderJid,
+        senderDeviceId,
+        ratchet,
+        false,
+      ),
+    );
+
+    try {
+      return _decryptAndVerifyHmac(ciphertext, keyAndHmac);
+    } catch (_) { 
+      await _restoreRatchet(ratchetKey, oldRatchet);
+      rethrow;
+    }
+  }
+
+  /// Returns the list of hex-encoded fingerprints we have for sessions with [jid].
+  Future<List<DeviceFingerprint>> getHexFingerprintsForJid(String jid) async {
+    final fingerprints = List<DeviceFingerprint>.empty(growable: true);
+
+    await _lock.synchronized(() async {
+      // Get devices for jid
+      final devices = _deviceMap[jid] ?? [];
+
+      for (final deviceId in devices) {
+        final ratchet = _ratchetMap[RatchetMapKey(jid, deviceId)]!;
+
+        fingerprints.add(
+          DeviceFingerprint(
+            deviceId,
+            HEX.encode(await ratchet.ik.getBytes()),
+          ),
+        );
+      }
+    });
+
+    return fingerprints;
+  }
+
+  /// 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
+  /// "signed PreKey rotation period" for recommendations.
+  Future<void> rotateSignedPrekey() async {
+    await _deviceLock.synchronized(() async {
+      _device = await _device.replaceSignedPrekey();
+
+      // Commit the new device
+      _eventStreamController.add(DeviceModifiedEvent(_device));
+    });
+  }
+
+  /// Returns the device map, i.e. the mapping of bare Jid to its device identifiers
+  /// we have built sessions with.
+  Future<Map<String, List<int>>> getDeviceMap() async {
+    return _lock.synchronized(() => _deviceMap);
+  }
+
+  /// Removes the ratchet identified by [jid] and [deviceId] from the session manager.
+  /// Also triggers events for commiting the new device map to storage and removing
+  /// the old ratchet.
+  Future<void> removeRatchet(String jid, int deviceId) async {
+    await _lock.synchronized(() async {
+      // Remove the ratchet
+      _ratchetMap.remove(RatchetMapKey(jid, deviceId));
+      // Commit it
+      _eventStreamController.add(RatchetRemovedEvent(jid, deviceId));
+
+      // Remove the device from jid
+      _deviceMap[jid]!.remove(deviceId);
+      if (_deviceMap[jid]!.isEmpty) {
+        _deviceMap.remove(jid);
+      }
+      // Commit it
+      _eventStreamController.add(DeviceListModifiedEvent(_deviceMap));
+    });
+  }
+
+  /// Removes all ratchets for Jid [jid]. Triggers a DeviceMapModified event at the end and an
+  /// RatchetRemovedEvent for each ratchet.
+  Future<void> removeAllRatchets(String jid) async {
+    await _lock.synchronized(() async {
+      for (final deviceId in _deviceMap[jid]!) {
+        // Remove the ratchet
+        _ratchetMap.remove(RatchetMapKey(jid, deviceId));
+        // Commit it
+        _eventStreamController.add(RatchetRemovedEvent(jid, deviceId));
+      }
+
+      // Remove the device from jid
+      _deviceMap.remove(jid);
+      // Commit it
+      _eventStreamController.add(DeviceListModifiedEvent(_deviceMap));
+    });
+  }
+  
+  /// Returns the list of device identifiers belonging to [jid] that are yet unacked, i.e.
+  /// we have not yet received an empty OMEMO message from.
+  Future<List<int>?> getUnacknowledgedRatchets(String jid) async {
+    return _lock.synchronized(() async {
+      final ret = List<int>.empty(growable: true);
+      final devices = _deviceMap[jid];
+      if (devices == null) return null;
+
+      for (final device in devices) {
+        final ratchet = _ratchetMap[RatchetMapKey(jid, device)]!;
+        if (!ratchet.acknowledged) ret.add(device);
+      }
+
+      return ret;
+    });
+  }
+
+  /// Returns true if the ratchet for [jid] with device identifier [deviceId] is
+  /// acknowledged. Returns false if not.
+  Future<bool> isRatchetAcknowledged(String jid, int deviceId) async {
+    return _lock.synchronized(() => _ratchetMap[RatchetMapKey(jid, deviceId)]!.acknowledged);
+  }
+  
+  /// Mark the ratchet for device [deviceId] from [jid] as acked.
+  Future<void> ratchetAcknowledged(String jid, int deviceId) async {
+    await _lock.synchronized(() async {
+      final ratchet = _ratchetMap[RatchetMapKey(jid, deviceId)]!
+        ..acknowledged = true;
+
+      // Commit it
+      _eventStreamController.add(RatchetModifiedEvent(jid, deviceId, ratchet, false));
+    });
+  }
+
+  /// 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({ int opkAmount = 100 }) async {
+    await _deviceLock.synchronized(() async {
+      _device = await OmemoDevice.generateNewDevice(_device.jid, opkAmount: opkAmount);
+
+      // Commit it
+      _eventStreamController.add(DeviceModifiedEvent(_device));
+    });
+  }
+
+  /// Make our device have a new identifier. Only useful before publishing it as a bundle
+  /// to make sure that our device has a id that is account unique.
+  Future<void> regenerateDeviceId() async {
+    await _deviceLock.synchronized(() async {
+      _device = _device.withNewId();
+
+      // Commit it
+      _eventStreamController.add(DeviceModifiedEvent(_device));
+    });
+  }
+
+  Future<OmemoDoubleRatchet?> _getRatchet(RatchetMapKey key) async {
+    return _lock.synchronized(() async {
+      return _ratchetMap[key];
+    });
+  }
+  
+  @visibleForTesting
+  OmemoDoubleRatchet getRatchet(String jid, int deviceId) => _ratchetMap[RatchetMapKey(jid, deviceId)]!;
+
+  @visibleForTesting
+  Map<RatchetMapKey, OmemoDoubleRatchet> getRatchetMap() => _ratchetMap;
+
+  /// Serialise the entire session manager into a JSON object.
+  Future<Map<String, dynamic>> toJsonWithoutSessions() async {
+    /*
+    {
+      'devices': {
+        'alice@...': [1, 2, ...],
+        'bob@...': [1],
+        ...
+      },
+      'device': { ... },
+    }
+    */
+
+    return {
+      'devices': _deviceMap,
+      'device': await (await getDevice()).toJson(),
+    };
+  }
+}

lib/src/omemo/stanza.dart

@@ -5,9 +5,9 @@ class OmemoIncomingStanza {
   const OmemoIncomingStanza(
     this.bareSenderJid,
     this.senderDeviceId,
+    this.timestamp,
     this.keys,
     this.payload,
-    this.isCatchup,
   );
 
   /// The bare JID of the sender of the stanza.
@@ -16,14 +16,14 @@ class OmemoIncomingStanza {
   /// The device ID of the sender.
   final int senderDeviceId;
 
-  /// The included encrypted keys for our own JID
+  /// 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;
-
-  /// Flag indicating whether the message was received due to a catchup.
-  final bool isCatchup;
 }
 
 /// Describes a stanza that is to be sent out
@@ -37,5 +37,5 @@ class OmemoOutgoingStanza {
   final List<String> recipientJids;
 
   /// The serialised XML data that should be encrypted.
-  final String? payload;
+  final String payload;
 }

lib/src/protobuf/omemo_authenticated_message.dart

@@ -0,0 +1,39 @@
+import 'package:omemo_dart/src/helpers.dart';
+import 'package:omemo_dart/src/protobuf/protobuf.dart';
+
+class OmemoAuthenticatedMessage {
+
+  OmemoAuthenticatedMessage();
+
+  factory OmemoAuthenticatedMessage.fromBuffer(List<int> data) {
+    var i = 0;
+    
+    // required bytes mac     = 1;
+    if (data[0] != fieldId(1, fieldTypeByteArray)) {
+      throw Exception();
+    }
+    final mac = data.sublist(2, i + 2 + data[1]);
+    i += data[1] + 2;
+
+    if (data[i] != fieldId(2, fieldTypeByteArray)) {
+      throw Exception();
+    }
+    final message = data.sublist(i + 2, i + 2 + data[i + 1]);
+
+    return OmemoAuthenticatedMessage()
+      ..mac = mac
+      ..message = message;
+  }
+  
+  List<int>? mac;
+  List<int>? message;
+
+  List<int> writeToBuffer() {
+    return concat([
+      [fieldId(1, fieldTypeByteArray), mac!.length],
+      mac!,
+      [fieldId(2, fieldTypeByteArray), message!.length],
+      message!,
+    ]);
+  }
+}

lib/src/protobuf/omemo_key_exchange.dart

@@ -0,0 +1,71 @@
+import 'package:omemo_dart/src/helpers.dart';
+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) {
+    var i = 0;
+
+    if (data[i] != fieldId(1, fieldTypeUint32)) {
+      throw Exception();
+    }
+    var decoded = decodeVarint(data, 1);
+    final pkId = decoded.n;
+    i += decoded.length + 1;
+
+    if (data[i] != fieldId(2, fieldTypeUint32)) {
+      throw Exception();
+    }
+    decoded = decodeVarint(data, i + 1);
+    final spkId = decoded.n;
+    i += decoded.length + 1;
+
+    if (data[i] != fieldId(3, fieldTypeByteArray)) {
+      throw Exception();
+    }
+    final ik = data.sublist(i + 2, i + 2 + data[i + 1]);
+    i += 2 + data[i + 1];
+
+    if (data[i] != fieldId(4, fieldTypeByteArray)) {
+      throw Exception();
+    }
+    final ek = data.sublist(i + 2, i + 2 + data[i + 1]);
+    i += 2 + data[i + 1];
+
+    if (data[i] != fieldId(5, fieldTypeByteArray)) {
+      throw Exception();
+    }
+    final message = OmemoAuthenticatedMessage.fromBuffer(data.sublist(i + 2));
+
+    return OmemoKeyExchange()
+      ..pkId = pkId
+      ..spkId = spkId
+      ..ik = ik
+      ..ek = ek
+      ..message = message;
+  }
+  
+  int? pkId;
+  int? spkId;
+  List<int>? ik;
+  List<int>? ek;
+  OmemoAuthenticatedMessage? message;
+
+  List<int> writeToBuffer() {
+    final msg = message!.writeToBuffer();
+    return concat([
+      [fieldId(1, fieldTypeUint32)],
+      encodeVarint(pkId!),
+      [fieldId(2, fieldTypeUint32)],
+      encodeVarint(spkId!),
+      [fieldId(3, fieldTypeByteArray), ik!.length],
+      ik!,
+      [fieldId(4, fieldTypeByteArray), ek!.length],
+      ek!,
+      [fieldId(5, fieldTypeByteArray), msg.length],
+      msg,
+    ]);
+  }
+}

lib/src/protobuf/omemo_message.dart

@@ -0,0 +1,76 @@
+import 'package:omemo_dart/src/helpers.dart';
+import 'package:omemo_dart/src/protobuf/protobuf.dart';
+
+class OmemoMessage {
+
+  OmemoMessage();
+
+  factory OmemoMessage.fromBuffer(List<int> data) {
+    var i = 0;
+    
+    // required uint32 n          = 1;
+    if (data[0] != fieldId(1, fieldTypeUint32)) {
+      throw Exception();
+    }
+    var decode = decodeVarint(data, 1);
+    final n = decode.n;
+    i += decode.length + 1;
+
+    // required uint32 pn         = 2;
+    if (data[i] != fieldId(2, fieldTypeUint32)) {
+      throw Exception();
+    }
+    decode = decodeVarint(data, i + 1);
+    final pn = decode.n;
+    i += decode.length + 1;
+
+    // required bytes  dh_pub     = 3;
+    if (data[i] != fieldId(3, fieldTypeByteArray)) {
+      throw Exception();
+    }
+    final dhPub = data.sublist(i + 2, i + 2 + data[i + 1]);
+    i += 2 + data[i + 1];
+
+    // optional bytes  ciphertext = 4;
+    List<int>? ciphertext;
+    if (i < data.length) {
+      if (data[i] != fieldId(4, fieldTypeByteArray)) {
+        throw Exception();
+      }
+
+      ciphertext = data.sublist(i + 2, i + 2 + data[i + 1]);
+    }
+
+    return OmemoMessage()
+      ..n = n
+      ..pn = pn
+      ..dhPub = dhPub
+      ..ciphertext = ciphertext;
+  }
+
+  int? n;
+  int? pn;
+  List<int>? dhPub;
+  List<int>? ciphertext;
+
+  List<int> writeToBuffer() {
+    final data = concat([
+      [fieldId(1, fieldTypeUint32)],
+      encodeVarint(n!),
+      [fieldId(2, fieldTypeUint32)],
+      encodeVarint(pn!),
+      [fieldId(3, fieldTypeByteArray), dhPub!.length],
+      dhPub!,
+    ]);
+
+    if (ciphertext != null) {
+      return concat([
+        data,
+        [fieldId(4, fieldTypeByteArray), ciphertext!.length],
+        ciphertext!,
+      ]);
+    }
+
+    return data;
+  }
+}

lib/src/protobuf/protobuf.dart

@@ -0,0 +1,65 @@
+/// Masks the 7 LSB
+const lsb7Mask = 0x7F;
+
+/// Constant for setting the MSB
+const msb = 1 << 7;
+
+/// Field types
+const fieldTypeUint32 = 0;
+const fieldTypeByteArray = 2;
+
+int fieldId(int number, int type) {
+  return (number << 3) | type;
+}
+
+class VarintDecode {
+
+  const VarintDecode(this.n, this.length);
+  final int n;
+  final int length;
+}
+
+/// Decode a Varint that begins at [input]'s index [offset].
+VarintDecode decodeVarint(List<int> input, int offset) {
+  // The return value
+  var n = 0;
+  // The byte offset counter
+  var i = 0;
+
+  // Iterate until the MSB of the byte is 0
+  while (true) {
+    // Mask only the 7 LSB and "move" them accordingly
+    n += (input[offset + i] & lsb7Mask) << (7 * i);
+
+    // Break if we reached the end
+    if (input[offset + i] & 1 << 7 == 0) {
+      break;
+    }
+    i++;
+  }
+
+  return VarintDecode(n, i + 1);
+}
+
+// Encodes the integer [i] into a Varint.
+List<int> encodeVarint(int i) {
+  assert(i >= 0, "Two's complement is not implemented");
+  final ret = List<int>.empty(growable: true);
+
+  // Thanks to https://github.com/hathibelagal-dev/LEB128 for the trick with toRadixString!
+  final numSevenBlocks = (i.toRadixString(2).length / 7).ceil();
+  for (var j = 0; j < numSevenBlocks; j++) {
+    // The 7 LSB of the byte we're creating
+    final x = (i & (lsb7Mask << j * 7)) >> j * 7;
+
+    if (j == numSevenBlocks - 1) {
+      // If we were to shift further, we only get zero, so we're at the end
+      ret.add(x);
+    } else {
+      // We still have at least one bit more to go, so set the MSB to 1
+      ret.add(x + msb);
+    }
+  }
+
+  return ret;
+}

lib/src/protobuf/schema.pb.dart

@@ -1,379 +0,0 @@
-///
-//  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/src/protobuf/schema.pbjson.dart

@@ -1,60 +0,0 @@
-///
-//  Generated code. Do not modify.
-//  source: schema.proto
-//
-// @dart = 2.12
-// ignore_for_file: annotate_overrides,camel_case_types,constant_identifier_names,deprecated_member_use_from_same_package,directives_ordering,library_prefixes,non_constant_identifier_names,prefer_final_fields,return_of_invalid_type,unnecessary_const,unnecessary_import,unnecessary_this,unused_import,unused_shown_name
-
-import 'dart:core' as $core;
-import 'dart:convert' as $convert;
-import 'dart:typed_data' as $typed_data;
-
-@$core.Deprecated('Use oMEMOMessageDescriptor instead')
-const OMEMOMessage$json = {
-  '1': 'OMEMOMessage',
-  '2': [
-    {'1': 'n', '3': 1, '4': 2, '5': 13, '10': 'n'},
-    {'1': 'pn', '3': 2, '4': 2, '5': 13, '10': 'pn'},
-    {'1': 'dh_pub', '3': 3, '4': 2, '5': 12, '10': 'dhPub'},
-    {'1': 'ciphertext', '3': 4, '4': 1, '5': 12, '10': 'ciphertext'},
-  ],
-};
-
-/// Descriptor for `OMEMOMessage`. Decode as a `google.protobuf.DescriptorProto`.
-final $typed_data.Uint8List oMEMOMessageDescriptor = $convert.base64Decode(
-    'CgxPTUVNT01lc3NhZ2USDAoBbhgBIAIoDVIBbhIOCgJwbhgCIAIoDVICcG4SFQoGZGhfcHViGAMgAigMUgVkaFB1YhIeCgpjaXBoZXJ0ZXh0GAQgASgMUgpjaXBoZXJ0ZXh0');
-@$core.Deprecated('Use oMEMOAuthenticatedMessageDescriptor instead')
-const OMEMOAuthenticatedMessage$json = {
-  '1': 'OMEMOAuthenticatedMessage',
-  '2': [
-    {'1': 'mac', '3': 1, '4': 2, '5': 12, '10': 'mac'},
-    {'1': 'message', '3': 2, '4': 2, '5': 12, '10': 'message'},
-  ],
-};
-
-/// Descriptor for `OMEMOAuthenticatedMessage`. Decode as a `google.protobuf.DescriptorProto`.
-final $typed_data.Uint8List oMEMOAuthenticatedMessageDescriptor =
-    $convert.base64Decode(
-        'ChlPTUVNT0F1dGhlbnRpY2F0ZWRNZXNzYWdlEhAKA21hYxgBIAIoDFIDbWFjEhgKB21lc3NhZ2UYAiACKAxSB21lc3NhZ2U=');
-@$core.Deprecated('Use oMEMOKeyExchangeDescriptor instead')
-const OMEMOKeyExchange$json = {
-  '1': 'OMEMOKeyExchange',
-  '2': [
-    {'1': 'pk_id', '3': 1, '4': 2, '5': 13, '10': 'pkId'},
-    {'1': 'spk_id', '3': 2, '4': 2, '5': 13, '10': 'spkId'},
-    {'1': 'ik', '3': 3, '4': 2, '5': 12, '10': 'ik'},
-    {'1': 'ek', '3': 4, '4': 2, '5': 12, '10': 'ek'},
-    {
-      '1': 'message',
-      '3': 5,
-      '4': 2,
-      '5': 11,
-      '6': '.OMEMOAuthenticatedMessage',
-      '10': 'message'
-    },
-  ],
-};
-
-/// Descriptor for `OMEMOKeyExchange`. Decode as a `google.protobuf.DescriptorProto`.
-final $typed_data.Uint8List oMEMOKeyExchangeDescriptor = $convert.base64Decode(
-    'ChBPTUVNT0tleUV4Y2hhbmdlEhMKBXBrX2lkGAEgAigNUgRwa0lkEhUKBnNwa19pZBgCIAIoDVIFc3BrSWQSDgoCaWsYAyACKAxSAmlrEg4KAmVrGAQgAigMUgJlaxI0CgdtZXNzYWdlGAUgAigLMhouT01FTU9BdXRoZW50aWNhdGVkTWVzc2FnZVIHbWVzc2FnZQ==');

lib/src/trust/always.dart

@@ -18,9 +18,9 @@ class AlwaysTrustingTrustManager extends TrustManager {
   @override
   Future<void> setEnabled(String jid, int deviceId, bool enabled) async {}
 
-  @override
+  @override 
   Future<void> removeTrustDecisionsForJid(String jid) async {}
-
+  
   @override
-  Future<void> loadTrustData(String jid) async {}
+  Future<Map<String, dynamic>> toJson() async => <String, dynamic>{};
 }

lib/src/trust/base.dart

@@ -1,5 +1,3 @@
-import 'package:meta/meta.dart';
-
 /// The base class for managing trust in OMEMO sessions.
 // ignore: one_member_abstracts
 abstract class TrustManager {
@@ -9,7 +7,6 @@ abstract class TrustManager {
 
   /// Called by the OmemoSessionManager when a new session has been built. Should set
   /// a default trust state to [jid]'s device with identifier [deviceId].
-  @internal
   Future<void> onNewSession(String jid, int deviceId);
 
   /// Return true if the device with id [deviceId] of Jid [jid] should be used for encryption.
@@ -20,14 +17,9 @@ abstract class TrustManager {
   /// if [enabled] is false.
   Future<void> setEnabled(String jid, int deviceId, bool enabled);
 
-  /// Removes all trust decisions for [jid].
+  /// Serialize the trust manager to JSON.
-  @internal
+  Future<Map<String, dynamic>> toJson();
-  Future<void> removeTrustDecisionsForJid(String jid);
 
-  // ignore: comment_references
+  /// Removes all trust decisions for [jid].
-  /// Called from within the [OmemoManager].
+  Future<void> removeTrustDecisionsForJid(String jid);
-  /// Loads the trust data for the JID [jid] from persistent storage
-  /// into the internal cache, if applicable.
-  @internal
-  Future<void> loadTrustData(String jid);
 }

lib/src/trust/btbv.dart

@@ -1,116 +1,64 @@
 import 'package:meta/meta.dart';
-import 'package:omemo_dart/src/helpers.dart';
 import 'package:omemo_dart/src/omemo/ratchet_map_key.dart';
 import 'package:omemo_dart/src/trust/base.dart';
-
+import 'package:synchronized/synchronized.dart';
-@immutable
-class BTBVTrustData {
-  const BTBVTrustData(
-    this.jid,
-    this.device,
-    this.state,
-    this.enabled,
-    this.trusted,
-  );
-
-  /// The JID in question.
-  final String jid;
-
-  /// The device (ratchet) in question.
-  final int device;
-
-  /// The trust state of the ratchet.
-  final BTBVTrustState state;
-
-  /// Flag indicating whether the ratchet is enabled (true) or not (false).
-  final bool enabled;
-
-  /// Flag indicating whether the ratchet is trusted. For loading and commiting a ratchet, this field
-  /// contains an arbitrary value.
-  /// When using [BlindTrustBeforeVerificationTrustManager.getDevicesTrust], this flag will be true if
-  /// the ratchet is trusted and false if not.
-  final bool trusted;
-}
-
-/// A callback for when a trust decision is to be commited to persistent storage.
-typedef BTBVTrustCommitCallback = Future<void> Function(BTBVTrustData data);
-
-/// A stub-implementation of [BTBVTrustCommitCallback].
-Future<void> btbvCommitStub(BTBVTrustData _) async {}
-
-/// A callback for when all trust decisions for a JID should be removed from persistent storage.
-typedef BTBVRemoveTrustForJidCallback = Future<void> Function(String jid);
-
-/// A stub-implementation of [BTBVRemoveTrustForJidCallback].
-Future<void> btbvRemoveTrustStub(String _) async {}
-
-/// A callback for when trust data should be loaded.
-typedef BTBVLoadDataCallback = Future<List<BTBVTrustData>> Function(String jid);
-
-/// A stub-implementation for [BTBVLoadDataCallback].
-Future<List<BTBVTrustData>> btbvLoadDataStub(String _) async => [];
 
 /// Every device is in either of those two trust states:
 /// - notTrusted: The device is absolutely not trusted
 /// - blindTrust: The fingerprint is not verified using OOB means
 /// - verified: The fingerprint has been verified using OOB means
 enum BTBVTrustState {
-  notTrusted(1),
+  notTrusted, // = 1
-  blindTrust(2),
+  blindTrust, // = 2
-  verified(3);
+  verified,   // = 3
+}
 
-  const BTBVTrustState(this.value);
+int _trustToInt(BTBVTrustState state) {
-
+  switch (state) {
-  factory BTBVTrustState.fromInt(int value) {
+    case BTBVTrustState.notTrusted: return 1;
-    switch (value) {
+    case BTBVTrustState.blindTrust: return 2;
-      case 1:
+    case BTBVTrustState.verified:   return 3;
-        return BTBVTrustState.notTrusted;
-      case 2:
-        return BTBVTrustState.blindTrust;
-      case 3:
-        return BTBVTrustState.verified;
-      // TODO(Unknown): Should we handle this better?
-      default:
-        return BTBVTrustState.notTrusted;
-    }
   }
+}
 
-  /// The value backing the trust state.
+BTBVTrustState _trustFromInt(int i) {
-  final int value;
+  switch (i) {
+    case 1: return BTBVTrustState.notTrusted;
+    case 2: return BTBVTrustState.blindTrust;
+    case 3: return BTBVTrustState.verified;
+    default: return BTBVTrustState.notTrusted;
+  }
 }
 
 /// A TrustManager that implements the idea of Blind Trust Before Verification.
 /// See https://gultsch.de/trust.html for more details.
-class BlindTrustBeforeVerificationTrustManager extends TrustManager {
+abstract class BlindTrustBeforeVerificationTrustManager extends TrustManager {
   BlindTrustBeforeVerificationTrustManager({
-    this.loadData = btbvLoadDataStub,
+    Map<RatchetMapKey, BTBVTrustState>? trustCache,
-    this.commit = btbvCommitStub,
+    Map<RatchetMapKey, bool>? enablementCache,
-    this.removeTrust = btbvRemoveTrustStub,
+    Map<String, List<int>>? devices,
-  });
+  }) : trustCache = trustCache ?? {},
+       enablementCache = enablementCache ?? {},
+       devices = devices ?? {},
+       _lock = Lock();
 
   /// The cache for mapping a RatchetMapKey to its trust state
   @visibleForTesting
   @protected
-  final Map<RatchetMapKey, BTBVTrustState> trustCache = {};
+  final Map<RatchetMapKey, BTBVTrustState> trustCache;
 
   /// The cache for mapping a RatchetMapKey to whether it is enabled or not
   @visibleForTesting
   @protected
-  final Map<RatchetMapKey, bool> enablementCache = {};
+  final Map<RatchetMapKey, bool> enablementCache;
-
+  
   /// Mapping of Jids to their device identifiers
   @visibleForTesting
   @protected
-  final Map<String, List<int>> devices = {};
+  final Map<String, List<int>> devices;
 
-  /// Callback for loading trust data.
+  /// The lock for devices and trustCache
-  final BTBVLoadDataCallback loadData;
+  final Lock _lock;
-
-  /// Callback for commiting trust data to persistent storage.
-  final BTBVTrustCommitCallback commit;
-
-  /// Callback for removing trust data for a JID.
-  final BTBVRemoveTrustForJidCallback removeTrust;
 
   /// Returns true if [jid] has at least one device that is verified. If not, returns false.
   /// Note that this function accesses devices and trustCache, which requires that the
@@ -122,148 +70,168 @@ class BlindTrustBeforeVerificationTrustManager extends TrustManager {
       return trustCache[RatchetMapKey(jid, id)]! == BTBVTrustState.verified;
     });
   }
-
+  
   @override
   Future<bool> isTrusted(String jid, int deviceId) async {
-    final trustCacheValue = trustCache[RatchetMapKey(jid, deviceId)];
+    var returnValue = false;
-    if (trustCacheValue == BTBVTrustState.notTrusted) {
+    await _lock.synchronized(() async {
-      return false;
+      final trustCacheValue = trustCache[RatchetMapKey(jid, deviceId)];
-    } else if (trustCacheValue == BTBVTrustState.verified) {
+      if (trustCacheValue == BTBVTrustState.notTrusted) {
-      // The key is verified, so it's safe.
+        returnValue = false;
-      return true;
+        return;
-    } else {
+      } else if (trustCacheValue == BTBVTrustState.verified) {
-      if (_hasAtLeastOneVerifiedDevice(jid)) {
+        // The key is verified, so it's safe.
-        // Do not trust if there is at least one device with full trust
+        returnValue = true;
-        return false;
+        return;
       } else {
-        // We have not verified a key from [jid], so it is blind trust all the way.
+        if (_hasAtLeastOneVerifiedDevice(jid)) {
-        return true;
+          // Do not trust if there is at least one device with full trust
+          returnValue = false;
+          return;
+        } else {
+          // We have not verified a key from [jid], so it is blind trust all the way.
+          returnValue = true;
+          return;
+        }
       }
-    }
+    });
+
+    return returnValue;
   }
 
   @override
   Future<void> onNewSession(String jid, int deviceId) async {
-    final key = RatchetMapKey(jid, deviceId);
+    await _lock.synchronized(() async {
-    if (_hasAtLeastOneVerifiedDevice(jid)) {
-      trustCache[key] = BTBVTrustState.notTrusted;
-      enablementCache[key] = false;
-    } else {
-      trustCache[key] = BTBVTrustState.blindTrust;
-      enablementCache[key] = true;
-    }
-
-    // Append to the device list
-    devices.appendOrCreate(jid, deviceId, checkExistence: true);
-
-    // Commit the state
-    await commit(
-      BTBVTrustData(
-        jid,
-        deviceId,
-        trustCache[key]!,
-        enablementCache[key]!,
-        false,
-      ),
-    );
-  }
-
-  /// 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, BTBVTrustData>> getDevicesTrust(String jid) async {
-    final map = <int, BTBVTrustData>{};
-
-    if (!devices.containsKey(jid)) return map;
-
-    for (final deviceId in devices[jid]!) {
       final key = RatchetMapKey(jid, deviceId);
-      if (!trustCache.containsKey(key) || !enablementCache.containsKey(key)) {
+      if (_hasAtLeastOneVerifiedDevice(jid)) {
-        continue;
+        trustCache[key] = BTBVTrustState.notTrusted;
+        enablementCache[key] = false;
+      } else {
+        trustCache[key] = BTBVTrustState.blindTrust;
+        enablementCache[key] = true;
       }
 
-      map[deviceId] = BTBVTrustData(
+      if (devices.containsKey(jid)) {
-        jid,
+        devices[jid]!.add(deviceId);
-        deviceId,
+      } else {
-        trustCache[key]!,
+        devices[jid] = List<int>.from([deviceId]);
-        enablementCache[key]!,
+      }
-        await isTrusted(jid, deviceId),
+      
-      );
+      // Commit the state
-    }
+      await commitState();
+    });
+  }
+  
+  /// 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>{};
 
-    return map;
+      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].
-  Future<void> setDeviceTrust(
+  Future<void> setDeviceTrust(String jid, int deviceId, BTBVTrustState state) async {
-    String jid,
+    await _lock.synchronized(() async {
-    int deviceId,
+      trustCache[RatchetMapKey(jid, deviceId)] = state;
-    BTBVTrustState state,
-  ) async {
-    final key = RatchetMapKey(jid, deviceId);
-    trustCache[key] = state;
 
-    // Commit the state
+      // Commit the state
-    await commit(
+      await commitState();
-      BTBVTrustData(
+    });
-        jid,
-        deviceId,
-        state,
-        enablementCache[key]!,
-        false,
-      ),
-    );
   }
 
   @override
   Future<bool> isEnabled(String jid, int deviceId) async {
-    final value = enablementCache[RatchetMapKey(jid, deviceId)];
+    return _lock.synchronized(() async {
+      final value = enablementCache[RatchetMapKey(jid, deviceId)];
 
-    if (value == null) return false;
+      if (value == null) return false;
-    return value;
+      return value;
+    });
   }
 
   @override
   Future<void> setEnabled(String jid, int deviceId, bool enabled) async {
-    final key = RatchetMapKey(jid, deviceId);
+    await _lock.synchronized(() async {
-    enablementCache[key] = enabled;
+      enablementCache[RatchetMapKey(jid, deviceId)] = enabled;
+    });
 
     // Commit the state
-    await commit(
+    await commitState();
-      BTBVTrustData(
+  }
-        jid,
+
-        deviceId,
+  @override
-        trustCache[key]!,
+  Future<Map<String, dynamic>> toJson() async {
-        enabled,
+    return {
-        false,
+      'devices': devices,
+      'trust': trustCache.map((key, value) => MapEntry(
+        key.toJsonKey(), _trustToInt(value),
+      ),),
+      'enable': enablementCache.map((key, value) => MapEntry(key.toJsonKey(), value)),
+    };
+  }
+
+  /// From a serialized version of a BTBV trust manager, extract the device list.
+  /// NOTE: This is needed as Dart cannot just cast a List<dynamic> to List<int> and so on.
+  static Map<String, List<int>> deviceListFromJson(Map<String, dynamic> json) {
+    return (json['devices']! as Map<String, dynamic>).map<String, List<int>>(
+      (key, value) => MapEntry(
+        key,
+        (value as List<dynamic>).map<int>((i) => i as int).toList(),
       ),
     );
   }
 
-  @override
+  /// From a serialized version of a BTBV trust manager, extract the trust cache.
+  /// NOTE: This is needed as Dart cannot just cast a List<dynamic> to List<int> and so on.
+  static Map<RatchetMapKey, BTBVTrustState> trustCacheFromJson(Map<String, dynamic> json) {
+    return (json['trust']! as Map<String, dynamic>).map<RatchetMapKey, BTBVTrustState>(
+      (key, value) => MapEntry(
+        RatchetMapKey.fromJsonKey(key),
+        _trustFromInt(value as int),
+      ),
+    );
+  }
+
+  /// From a serialized version of a BTBV trust manager, extract the enable cache.
+  /// NOTE: This is needed as Dart cannot just cast a List<dynamic> to List<int> and so on.
+  static Map<RatchetMapKey, bool> enableCacheFromJson(Map<String, dynamic> json) {
+    return (json['enable']! as Map<String, dynamic>).map<RatchetMapKey, bool>(
+      (key, value) => MapEntry(
+        RatchetMapKey.fromJsonKey(key),
+        value as bool,
+      ),
+    );
+  }
+
+  @override 
   Future<void> removeTrustDecisionsForJid(String jid) async {
-    // Clear the caches
+    await _lock.synchronized(() async {
-    for (final device in devices[jid]!) {
+      devices.remove(jid);
-      final key = RatchetMapKey(jid, device);
+      await commitState();
-      trustCache.remove(key);
+    });
-      enablementCache.remove(key);
-    }
-    devices.remove(jid);
-
-    // Commit the state
-    await removeTrust(jid);
-  }
-
-  @override
-  Future<void> loadTrustData(String jid) async {
-    for (final result in await loadData(jid)) {
-      final key = RatchetMapKey(jid, result.device);
-      trustCache[key] = result.state;
-      enablementCache[key] = result.enabled;
-      devices.appendOrCreate(jid, result.device, checkExistence: true);
-    }
   }
+  
+  /// 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();
 
   @visibleForTesting
-  BTBVTrustState getDeviceTrust(String jid, int deviceId) =>
+  BTBVTrustState getDeviceTrust(String jid, int deviceId) => trustCache[RatchetMapKey(jid, deviceId)]!;
-      trustCache[RatchetMapKey(jid, deviceId)]!;
+}
+
+/// A BTBV TrustManager that does not commit its state to persistent storage. Well suited
+/// for testing.
+class MemoryBTBVTrustManager extends BlindTrustBeforeVerificationTrustManager {
+  @override
+  Future<void> commitState() async {}
 }

lib/src/trust/never.dart

@@ -18,9 +18,9 @@ class NeverTrustingTrustManager extends TrustManager {
   @override
   Future<void> setEnabled(String jid, int deviceId, bool enabled) async {}
 
-  @override
+  @override 
   Future<void> removeTrustDecisionsForJid(String jid) async {}
 
   @override
-  Future<void> loadTrustData(String jid) async {}
+  Future<Map<String, dynamic>> toJson() async => <String, dynamic>{};
 }

lib/src/x3dh/x3dh.dart

@@ -1,16 +1,18 @@
 import 'dart:convert';
 import 'dart:math';
 import 'package:cryptography/cryptography.dart';
-import 'package:moxlib/moxlib.dart';
-import 'package:omemo_dart/src/common/constants.dart';
 import 'package:omemo_dart/src/crypto.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';
 
+/// The overarching assumption is that we use Ed25519 keys for the identity keys
+const omemoX3DHInfoString = 'OMEMO X3DH';
+
 /// Performed by Alice
 class X3DHAliceResult {
+
   const X3DHAliceResult(this.ek, this.sk, this.opkId, this.ad);
   final OmemoKeyPair ek;
   final List<int> sk;
@@ -20,6 +22,7 @@ class X3DHAliceResult {
 
 /// Received by Bob
 class X3DHMessage {
+
   const X3DHMessage(this.ik, this.ek, this.opkId);
   final OmemoPublicKey ik;
   final OmemoPublicKey ek;
@@ -27,6 +30,7 @@ class X3DHMessage {
 }
 
 class X3DHBobResult {
+
   const X3DHBobResult(this.sk, this.ad);
   final List<int> sk;
   final List<int> ad;
@@ -35,10 +39,7 @@ class X3DHBobResult {
 /// Sign [message] using the keypair [keyPair]. Note that [keyPair] must be
 /// a Ed25519 keypair.
 Future<List<int>> sig(OmemoKeyPair keyPair, List<int> message) async {
-  assert(
+  assert(keyPair.type == KeyPairType.ed25519, 'Signature keypair must be Ed25519');
-    keyPair.type == KeyPairType.ed25519,
-    'Signature keypair must be Ed25519',
-  );
   final signature = await Ed25519().sign(
     message,
     keyPair: await keyPair.asKeyPair(),
@@ -69,11 +70,7 @@ Future<List<int>> kdf(List<int> km) async {
 
 /// Alice builds a session with Bob using his bundle [bundle] and Alice's identity key
 /// pair [ik].
-Future<Result<InvalidKeyExchangeSignatureError, X3DHAliceResult>>
+Future<X3DHAliceResult> x3dhFromBundle(OmemoBundle bundle, OmemoKeyPair ik) async {
-    x3dhFromBundle(
-  OmemoBundle bundle,
-  OmemoKeyPair ik,
-) async {
   // Check the signature first
   final signatureValue = await Ed25519().verify(
     await bundle.spk.getBytes(),
@@ -84,7 +81,7 @@ Future<Result<InvalidKeyExchangeSignatureError, X3DHAliceResult>>
   );
 
   if (!signatureValue) {
-    return Result(InvalidKeyExchangeSignatureError());
+    throw InvalidSignatureException();
   }
 
   // Generate EK
@@ -94,9 +91,9 @@ Future<Result<InvalidKeyExchangeSignatureError, X3DHAliceResult>>
   final opkIndex = random.nextInt(bundle.opksEncoded.length);
   final opkId = bundle.opksEncoded.keys.elementAt(opkIndex);
   final opk = bundle.getOpk(opkId);
-
+  
   final dh1 = await omemoDH(ik, bundle.spk, 1);
-  final dh2 = await omemoDH(ek, bundle.ik, 2);
+  final dh2 = await omemoDH(ek, bundle.ik,  2);
   final dh3 = await omemoDH(ek, bundle.spk, 0);
   final dh4 = await omemoDH(ek, opk, 0);
 
@@ -106,19 +103,14 @@ Future<Result<InvalidKeyExchangeSignatureError, X3DHAliceResult>>
     await bundle.ik.getBytes(),
   ]);
 
-  return Result(X3DHAliceResult(ek, sk, opkId, ad));
+  return X3DHAliceResult(ek, sk, opkId, ad);
 }
 
 /// Bob builds the X3DH shared secret from the inital message [msg], the SPK [spk], the
 /// OPK [opk] that was selected by Alice and our IK [ik]. Returns the shared secret.
-Future<X3DHBobResult> x3dhFromInitialMessage(
+Future<X3DHBobResult> x3dhFromInitialMessage(X3DHMessage msg, OmemoKeyPair spk, OmemoKeyPair opk, OmemoKeyPair ik) async {
-  X3DHMessage msg,
-  OmemoKeyPair spk,
-  OmemoKeyPair opk,
-  OmemoKeyPair ik,
-) async {
   final dh1 = await omemoDH(spk, msg.ik, 2);
-  final dh2 = await omemoDH(ik, msg.ek, 1);
+  final dh2 = await omemoDH(ik,  msg.ek, 1);
   final dh3 = await omemoDH(spk, msg.ek, 0);
   final dh4 = await omemoDH(opk, msg.ek, 0);
 

omemo_dart.doap

@@ -1,53 +0,0 @@
-<?xml version='1.0' encoding='UTF-8'?>
-<rdf:RDF xmlns:rdf='http://www.w3.org/1999/02/22-rdf-syntax-ns#'
-         xmlns="http://usefulinc.com/ns/doap#"
-         xmlns:foaf="http://xmlns.com/foaf/0.1/"
-         xmlns:xmpp="https://linkmauve.fr/ns/xmpp-doap#"
-         xmlns:schema="https://schema.org/">
-  <Project xml:lang='en'>
-    <!-- Moxxy information -->
-    <name>omemo_dart</name>
-    <created>2022-06-30</created>
-    <homepage rdf:resource="https://github.com/PapaTutuWawa/omemo_dart" />
-    <os>Linux</os>
-    <os>Windows</os>
-    <os>macOS</os>
-    <os>Android</os>
-    <os>iOS</os>
-
-    <!-- Description -->
-    <shortdesc xml:lang="de">A Dart implementation of the cryptography needed for OMEMO 0.8.3</shortdesc>
-    <description xml:lang="en">omemo_dart is a pure Dart implementation of the low-level components (X3DH, Double Ratchet) with a high-level library-agnostic interface for implementing OMEMO 0.8.3.</description>
-
-    <!-- Maintainer information -->
-    <maintainer>
-      <foaf:Person>
-      	<foaf:name>Alexander "Polynomdivision"</foaf:name>
-      	<foaf:homepage rdf:resource="https://polynom.me" />
-      </foaf:Person>
-    </maintainer>
-
-    <!-- Channel list -->
-    <developer-forum rdf:resource='xmpp:dev@muc.moxxy.org?join'/>
-
-    <!-- Repository information -->
-    <programming-language>Dart</programming-language>
-    <bug-database rdf:resource="https://github.com/PapaTutuWawa/omemo_dart/issues" />
-    <license rdf:resource="https://github.com/PapaTutuWawa/omemo_dart/blob/master/LICENSE" />
-    <repository>
-      <GitRepository>
-        <browse rdf:resource="https://github.com/PapaTutuWawa/omemo_dart" />
-        <location rdf:resource="https://github.com/PapaTutuWawa/omemo_dart.git" />
-      </GitRepository>
-    </repository>
-
-    <!-- XEP list -->
-    <implements>
-      <xmpp:SupportedXep>
-        <xmpp:xep rdf:resource="https://xmpp.org/extensions/xep-0384.html"/>
-        <xmpp:status>partial</xmpp:status>
-        <xmpp:version>0.8.3</xmpp:version>
-      </xmpp:SupportedXep>
-    </implements>
-  </Project>
-</rdf:RDF>

pubspec.yaml

@@ -1,27 +1,24 @@
 name: omemo_dart
 description: An XMPP library independent OMEMO library
-version: 0.6.0
+version: 0.4.2
 homepage: https://github.com/PapaTutuWawa/omemo_dart
 publish_to: https://git.polynom.me/api/packages/PapaTutuWawa/pub
 
 environment:
-  sdk: ">=3.0.0 <4.0.0"
+  sdk: '>=2.17.0 <3.0.0'
 
 dependencies:
-  collection: ^1.18.0
+  collection: ^1.16.0
-  cryptography: ^2.7.0
+  cryptography: ^2.0.5
   hex: ^0.2.0
-  logging: ^1.2.0
+  logging: ^1.0.2
-  meta: ^1.15.0
+  meta: ^1.7.0
-  moxlib:
+  pinenacl: ^0.5.1
-    version: ^0.2.0
+  synchronized: ^3.0.0+2
-    hosted: https://git.polynom.me/api/packages/Moxxy/pub
-  pinenacl: ^0.6.0
-  protobuf: ^3.1.0
-  protoc_plugin: ^21.1.2
-  synchronized: ^3.3.0+3
 
 dev_dependencies:
-  lints: ^5.0.0
+  lints: ^2.0.0
-  test: ^1.25.8
+  protobuf: ^2.1.0
-  very_good_analysis: ^6.0.0
+  protoc_plugin: ^20.0.1
+  test: ^1.21.0
+  very_good_analysis: ^3.0.1

test/double_ratchet_test.dart

@@ -1,10 +1,39 @@
+// ignore_for_file: avoid_print
 import 'dart:convert';
-import 'dart:developer';
 import 'package:cryptography/cryptography.dart';
 import 'package:omemo_dart/omemo_dart.dart';
+import 'package:omemo_dart/protobuf/schema.pb.dart';
+import 'package:omemo_dart/src/double_ratchet/crypto.dart';
 import 'package:test/test.dart';
 
 void main() {
+  test('Test encrypting and decrypting', () async {
+    final sessionAd = List<int>.filled(32, 0x0);
+    final mk = List<int>.filled(32, 0x1);
+    final plaintext = utf8.encode('Hallo');
+    final header = OMEMOMessage()
+      ..n = 0
+      ..pn = 0
+      ..dhPub = List<int>.empty();
+    final asd = concat([sessionAd, header.writeToBuffer()]);
+      
+    final ciphertext = await encrypt(
+      mk,
+      plaintext,
+      asd,
+      sessionAd,
+    );
+
+    final decrypted = await decrypt(
+      mk,
+      ciphertext,
+      asd,
+      sessionAd,
+    );
+
+    expect(decrypted, plaintext);
+  });
+
   test('Test the Double Ratchet', () async {
     // Generate keys
     const bobJid = 'bob@other.example.server';
@@ -26,14 +55,13 @@ void main() {
         2: await opkBob.pk.asBase64(),
       },
     );
-
+    
     // Alice does X3DH
-    final resultAliceRaw = await x3dhFromBundle(bundleBob, ikAlice);
+    final resultAlice = await x3dhFromBundle(bundleBob, ikAlice);
-    final resultAlice = resultAliceRaw.get<X3DHAliceResult>();
 
     // Alice sends the inital message to Bob
     // ...
-
+    
     // Bob does X3DH
     final resultBob = await x3dhFromInitialMessage(
       X3DHMessage(
@@ -46,71 +74,63 @@ void main() {
       ikBob,
     );
 
-    log('X3DH key exchange done');
+    print('X3DH key exchange done');
+
 
     // Alice and Bob now share sk as a common secret and ad
     // Build a session
     final alicesRatchet = await OmemoDoubleRatchet.initiateNewSession(
       spkBob.pk,
-      bundleBob.spkId,
       ikBob.pk,
-      ikAlice.pk,
-      resultAlice.ek.pk,
       resultAlice.sk,
       resultAlice.ad,
-      resultAlice.opkId,
+      0,
     );
     final bobsRatchet = await OmemoDoubleRatchet.acceptNewSession(
       spkBob,
-      bundleBob.spkId,
       ikAlice.pk,
-      2,
-      resultAlice.ek.pk,
       resultBob.sk,
       resultBob.ad,
+      0,
     );
 
     expect(alicesRatchet.sessionAd, bobsRatchet.sessionAd);
 
     for (var i = 0; i < 100; i++) {
       final messageText = 'Hello, dear $i';
-
+      
-      log('${i + 1}/100');
       if (i.isEven) {
         // Alice encrypts a message
-        final aliceRatchetResult =
+        final aliceRatchetResult = await alicesRatchet.ratchetEncrypt(utf8.encode(messageText));
-            await alicesRatchet.ratchetEncrypt(utf8.encode(messageText));
+        print('Alice sent the message');
-        log('Alice sent the message');
 
         // Alice sends it to Bob
         // ...
 
         // Bob tries to decrypt it
         final bobRatchetResult = await bobsRatchet.ratchetDecrypt(
-          aliceRatchetResult,
+          aliceRatchetResult.header,
+          aliceRatchetResult.ciphertext,
         );
-        log('Bob decrypted the message');
+        print('Bob decrypted the message');
 
-        expect(bobRatchetResult.isType<List<int>>(), true);
+        expect(utf8.encode(messageText), bobRatchetResult);
-        expect(bobRatchetResult.get<List<int>>(), utf8.encode(messageText));
       } else {
         // Bob sends a message to Alice
-        final bobRatchetResult =
+        final bobRatchetResult = await bobsRatchet.ratchetEncrypt(utf8.encode(messageText));
-            await bobsRatchet.ratchetEncrypt(utf8.encode(messageText));
+        print('Bob sent the message');
-        log('Bob sent the message');
 
         // Bobs sends it to Alice
         // ...
 
         // Alice tries to decrypt it
         final aliceRatchetResult = await alicesRatchet.ratchetDecrypt(
-          bobRatchetResult,
+          bobRatchetResult.header,
+          bobRatchetResult.ciphertext,
         );
-        log('Alice decrypted the message');
+        print('Alice decrypted the message');
 
-        expect(aliceRatchetResult.isType<List<int>>(), true);
+        expect(utf8.encode(messageText), aliceRatchetResult);
-        expect(aliceRatchetResult.get<List<int>>(), utf8.encode(messageText));
-        expect(utf8.encode(messageText), aliceRatchetResult.get<List<int>>());
       }
     }
   });

test/protobuf_test.dart

@@ -0,0 +1,186 @@
+import 'package:omemo_dart/protobuf/schema.pb.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/protobuf/protobuf.dart';
+import 'package:test/test.dart';
+
+void main() {
+  group('Base 128 Varints', () {
+    test('Test simple parsing of Varints', () {
+      expect(
+        decodeVarint(<int>[1], 0).n,
+        1,
+      );
+      expect(
+        decodeVarint(<int>[1], 0).length,
+        1,
+      );
+      expect(
+        decodeVarint(<int>[0x96, 0x01, 0x00], 0).n,
+        150,
+      );
+      expect(
+        decodeVarint(<int>[0x96, 0x01, 0x00], 0).length,
+        2,
+      );
+      expect(
+        decodeVarint(<int>[172, 2, 0x8], 0).n,
+        300,
+      );
+      expect(
+        decodeVarint(<int>[172, 2, 0x8], 0).length,
+        2,
+      );
+    });
+
+    test('Test encoding Varints', () {
+      expect(
+        encodeVarint(1),
+        <int>[1],
+      );
+      expect(
+        encodeVarint(150),
+        <int>[0x96, 0x01],
+      );
+      expect(
+        encodeVarint(300),
+        <int>[172, 2],
+      );
+    });
+
+    test('Test some special cases', () {
+      expect(decodeVarint(encodeVarint(1042464893), 0).n, 1042464893);
+    });
+  });
+
+  group('OMEMOMessage', () {
+    test('Decode a OMEMOMessage', () {
+      final pbMessage = OMEMOMessage()
+        ..n = 1
+        ..pn = 5
+        ..dhPub = <int>[1, 2, 3]
+        ..ciphertext = <int>[4, 5, 6];
+      final serial = pbMessage.writeToBuffer();
+      final msg = OmemoMessage.fromBuffer(serial);
+
+      expect(msg.n, 1);
+      expect(msg.pn, 5);
+      expect(msg.dhPub, <int>[1, 2, 3]);
+      expect(msg.ciphertext, <int>[4, 5, 6]);
+    });
+    test('Decode a OMEMOMessage without ciphertext', () {
+      final pbMessage = OMEMOMessage()
+        ..n = 1
+        ..pn = 5
+        ..dhPub = <int>[1, 2, 3];
+      final serial = pbMessage.writeToBuffer();
+      final msg = OmemoMessage.fromBuffer(serial);
+      
+      expect(msg.n, 1);
+      expect(msg.pn, 5);
+      expect(msg.dhPub, <int>[1, 2, 3]);
+      expect(msg.ciphertext, null);
+    });
+    test('Encode a OMEMOMessage', () {
+      final m = OmemoMessage()
+        ..n = 1
+        ..pn = 5
+        ..dhPub = <int>[1, 2, 3]
+        ..ciphertext = <int>[4, 5, 6];
+      final serial = m.writeToBuffer();
+      final msg = OMEMOMessage.fromBuffer(serial);
+      
+      expect(msg.n, 1);
+      expect(msg.pn, 5);
+      expect(msg.dhPub, <int>[1, 2, 3]);
+      expect(msg.ciphertext, <int>[4, 5, 6]);
+    });
+    test('Encode a OMEMOMessage without ciphertext', () {
+      final m = OmemoMessage()
+        ..n = 1
+        ..pn = 5
+        ..dhPub = <int>[1, 2, 3];
+      final serial = m.writeToBuffer();
+      final msg = OMEMOMessage.fromBuffer(serial);
+
+      expect(msg.n, 1);
+      expect(msg.pn, 5);
+      expect(msg.dhPub, <int>[1, 2, 3]);
+      expect(msg.ciphertext, <int>[]);
+    });
+  });
+
+  group('OMEMOAuthenticatedMessage', () {
+    test('Test encoding a message', () {
+      final msg = OmemoAuthenticatedMessage()
+        ..mac = <int>[1, 2, 3]
+        ..message = <int>[4, 5, 6];
+      final decoded = OMEMOAuthenticatedMessage.fromBuffer(msg.writeToBuffer());
+
+      expect(decoded.mac, <int>[1, 2, 3]);
+      expect(decoded.message, <int>[4, 5, 6]);
+    });
+    test('Test decoding a message', () {
+      final msg = OMEMOAuthenticatedMessage()
+        ..mac = <int>[1, 2, 3]
+        ..message = <int>[4, 5, 6];
+      final bytes = msg.writeToBuffer();
+      final decoded = OmemoAuthenticatedMessage.fromBuffer(bytes);
+
+      expect(decoded.mac, <int>[1, 2, 3]);
+      expect(decoded.message, <int>[4, 5, 6]);
+    });
+  });
+
+  group('OMEMOKeyExchange', () {
+    test('Test encoding a message', () {
+      final authMessage = OmemoAuthenticatedMessage()
+        ..mac = <int>[5, 6, 8, 0]
+        ..message = <int>[4, 5, 7, 3, 2];
+      final message = OmemoKeyExchange()
+        ..pkId = 698
+        ..spkId = 245
+        ..ik = <int>[1, 4, 6]
+        ..ek = <int>[4, 6, 7, 80]
+        ..message = authMessage;
+      final kex = OMEMOKeyExchange.fromBuffer(message.writeToBuffer());
+
+      expect(kex.pkId, 698);
+      expect(kex.spkId, 245);
+      expect(kex.ik, <int>[1, 4, 6]);
+      expect(kex.ek, <int>[4, 6, 7, 80]);
+
+      expect(kex.message.mac, <int>[5, 6, 8, 0]);
+      expect(kex.message.message, <int>[4, 5, 7, 3, 2]);
+    });
+    test('Test decoding a message', () {
+      final message = OMEMOAuthenticatedMessage()
+        ..mac = <int>[5, 6, 8, 0]
+        ..message = <int>[4, 5, 7, 3, 2];
+      final kex = OMEMOKeyExchange()
+        ..pkId = 698
+        ..spkId = 245
+        ..ik = <int>[1, 4, 6]
+        ..ek = <int>[4, 6, 7, 80]
+        ..message = message;
+      final decoded = OmemoKeyExchange.fromBuffer(kex.writeToBuffer());
+
+      expect(decoded.pkId, 698);
+      expect(decoded.spkId, 245);
+      expect(decoded.ik, <int>[1, 4, 6]);
+      expect(decoded.ek, <int>[4 ,6 ,7 , 80]);
+
+      expect(decoded.message!.mac, <int>[5, 6, 8, 0]);
+      expect(decoded.message!.message, <int>[4, 5, 7, 3, 2]);
+    });
+    test('Test decoding an issue', () {
+      /*
+      final data = 'CAAQfRogc2GwslU219dUkrMHNM4KdZRmuFnBTae+bQaJ+55IsAMiII7aZKj2sUpb6xR/3Ari7WZUmKFV0G6czUc4NMvjKDBaKnwKEM2ZpI8X3TgcxhxwENANnlsSaAgAEAAaICy8T9WPgLb7RdYd8/4JkrLF0RahEkC3ZaEfk5jw3dsLIkBMILzLyByweLgF4lCn0oNea+kbdrFr6rY7r/7WyI8hXEQz38QpnN+jyGGwC7Ga0dq70WuyqE7VpiFArQwqZh2G';
+      final kex = OmemoKeyExchange.fromBuffer(base64Decode(data));
+
+      expect(kex.spkId!, 1042464893);
+      */
+    });
+  });
+}

test/queue_test.dart

@@ -1,62 +0,0 @@
-import 'dart:async';
-import 'package:omemo_dart/src/omemo/queue.dart';
-import 'package:test/test.dart';
-
-Future<void> testMethod(
-  RatchetAccessQueue queue,
-  List<String> data,
-  int duration,
-) async {
-  await queue.enterCriticalSection(data);
-
-  await Future<void>.delayed(Duration(seconds: duration));
-
-  await queue.leaveCriticalSection(data);
-}
-
-void main() {
-  test('Test blocking due to conflicts', () async {
-    final queue = RatchetAccessQueue();
-
-    unawaited(testMethod(queue, ['a', 'b', 'c'], 5));
-    unawaited(testMethod(queue, ['a'], 4));
-
-    await Future<void>.delayed(const Duration(seconds: 1));
-    expect(
-      queue.runningOperations.containsAll(['a', 'b', 'c']),
-      isTrue,
-    );
-    expect(queue.runningOperations.length, 3);
-
-    await Future<void>.delayed(const Duration(seconds: 4));
-
-    expect(
-      queue.runningOperations.containsAll(['a']),
-      isTrue,
-    );
-    expect(queue.runningOperations.length, 1);
-
-    await Future<void>.delayed(const Duration(seconds: 4));
-    expect(queue.runningOperations.length, 0);
-  });
-
-  test('Test not blocking due to no conflicts', () async {
-    final queue = RatchetAccessQueue();
-
-    unawaited(testMethod(queue, ['a', 'b'], 5));
-    unawaited(testMethod(queue, ['c'], 5));
-    unawaited(testMethod(queue, ['d'], 5));
-
-    await Future<void>.delayed(const Duration(seconds: 1));
-    expect(queue.runningOperations.length, 4);
-    expect(
-      queue.runningOperations.containsAll([
-        'a',
-        'b',
-        'c',
-        'd',
-      ]),
-      isTrue,
-    );
-  });
-}

test/serialisation_test.dart

@@ -0,0 +1,148 @@
+import 'dart:convert';
+import 'package:omemo_dart/omemo_dart.dart';
+import 'package:omemo_dart/src/trust/always.dart';
+import 'package:test/test.dart';
+
+Map<String, dynamic> jsonify(Map<String, dynamic> map) {
+  return jsonDecode(jsonEncode(map)) as Map<String, dynamic>;
+}
+
+void main() {
+  test('Test serialising and deserialising the Device', () async {
+    // Generate a random session
+    final oldSession = await OmemoSessionManager.generateNewIdentity(
+      'user@test.server',
+      AlwaysTrustingTrustManager(),
+      opkAmount: 1,
+    );
+    final oldDevice = await oldSession.getDevice();
+    final serialised = jsonify(await oldDevice.toJson());
+
+    final newDevice = OmemoDevice.fromJson(serialised);
+    expect(await oldDevice.equals(newDevice), true);
+  });
+
+  test('Test serialising and deserialising the Device after rotating the SPK', () async {
+    // Generate a random session
+    final oldSession = await OmemoSessionManager.generateNewIdentity(
+      'user@test.server',
+      AlwaysTrustingTrustManager(),
+      opkAmount: 1,
+    );
+    final oldDevice = await (await oldSession.getDevice()).replaceSignedPrekey();
+    final serialised = jsonify(await oldDevice.toJson());
+
+    final newDevice = OmemoDevice.fromJson(serialised);
+    expect(await oldDevice.equals(newDevice), true);
+  });
+  
+  test('Test serialising and deserialising the OmemoDoubleRatchet', () async {
+    // Generate a random ratchet
+    const aliceJid = 'alice@server.example';
+    const bobJid = 'bob@other.server.example';
+    final aliceSession = await OmemoSessionManager.generateNewIdentity(
+      aliceJid,
+      AlwaysTrustingTrustManager(),
+      opkAmount: 1,
+    );
+    final bobSession = await OmemoSessionManager.generateNewIdentity(
+      bobJid,
+      AlwaysTrustingTrustManager(),
+      opkAmount: 1,
+    );
+    final aliceMessage = await aliceSession.encryptToJid(
+      bobJid,
+      'Hello Bob!',
+      newSessions: [
+        await bobSession.getDeviceBundle(),
+      ],
+    );
+    await bobSession.decryptMessage(
+      aliceMessage.ciphertext,
+      aliceJid,
+      await aliceSession.getDeviceId(),
+      aliceMessage.encryptedKeys,
+      getTimestamp(),
+    );
+    final aliceOld = aliceSession.getRatchet(bobJid, await bobSession.getDeviceId());
+    final aliceSerialised = jsonify(await aliceOld.toJson());
+    final aliceNew = OmemoDoubleRatchet.fromJson(aliceSerialised);
+
+    expect(await aliceOld.equals(aliceNew), true);
+  });
+
+  test('Test serialising and deserialising the OmemoSessionManager', () async {
+    // Generate a random session
+    final oldSession = await OmemoSessionManager.generateNewIdentity(
+      'a@server',
+      AlwaysTrustingTrustManager(),
+      opkAmount: 4,
+    );
+    final bobSession = await OmemoSessionManager.generateNewIdentity(
+      'b@other.server',
+      AlwaysTrustingTrustManager(),
+      opkAmount: 4,
+    );
+    await oldSession.addSessionFromBundle(
+      'bob@localhost',
+      await bobSession.getDeviceId(),
+      await bobSession.getDeviceBundle(),
+    );
+
+    // Serialise and deserialise
+    final serialised = jsonify(await oldSession.toJsonWithoutSessions());
+    final newSession = OmemoSessionManager.fromJsonWithoutSessions(
+      serialised,
+      // NOTE: At this point, we don't care about this attribute
+      {},
+      AlwaysTrustingTrustManager(),
+    );
+
+    final oldDevice = await oldSession.getDevice();
+    final newDevice = await newSession.getDevice();
+    expect(await oldDevice.equals(newDevice), true);
+    expect(await oldSession.getDeviceMap(), await newSession.getDeviceMap());
+  });
+
+  test('Test serializing and deserializing RatchetMapKey', () {
+    const test1 = RatchetMapKey('user@example.org', 1234);
+    final result1 = RatchetMapKey.fromJsonKey(test1.toJsonKey());
+    expect(result1.jid, test1.jid);
+    expect(result1.deviceId, test1.deviceId);
+
+    const test2 = RatchetMapKey('user@example.org/hallo:welt', 3333);
+    final result2 = RatchetMapKey.fromJsonKey(test2.toJsonKey());
+    expect(result2.jid, test2.jid);
+    expect(result2.deviceId, test2.deviceId);
+  });
+
+  test('Test serializing and deserializing the components of the BTBV manager', () async {
+    // Caroline's BTBV manager
+    final btbv = MemoryBTBVTrustManager();
+    // Example data
+    const aliceJid = 'alice@some.server';
+    const bobJid = 'bob@other.server';
+    
+    await btbv.onNewSession(aliceJid, 1);
+    await btbv.setDeviceTrust(aliceJid, 1, BTBVTrustState.verified);
+    await btbv.onNewSession(aliceJid, 2);
+    await btbv.onNewSession(bobJid, 3);
+    await btbv.onNewSession(bobJid, 4);
+
+    final serialized = jsonify(await btbv.toJson());
+    final deviceList = BlindTrustBeforeVerificationTrustManager.deviceListFromJson(
+      serialized,
+    );
+    expect(btbv.devices, deviceList);
+
+    final trustCache = BlindTrustBeforeVerificationTrustManager.trustCacheFromJson(
+      serialized,
+    );
+    expect(btbv.trustCache, trustCache);
+
+    final enableCache = BlindTrustBeforeVerificationTrustManager.enableCacheFromJson(
+      serialized,
+    );
+    expect(btbv.enablementCache, enableCache);
+  });
+}

test/trust_test.dart

@@ -4,11 +4,11 @@ import 'package:test/test.dart';
 void main() {
   test('Test the Blind Trust Before Verification TrustManager', () async {
     // Caroline's BTBV manager
-    final btbv = BlindTrustBeforeVerificationTrustManager();
+    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);

test/x3dh_test.dart

@@ -24,14 +24,13 @@ void main() {
         2: await opkBob.pk.asBase64(),
       },
     );
-
+    
     // Alice does X3DH
-    final resultAliceRaw = await x3dhFromBundle(bundleBob, ikAlice);
+    final resultAlice = await x3dhFromBundle(bundleBob, ikAlice);
-    final resultAlice = resultAliceRaw.get<X3DHAliceResult>();
 
     // Alice sends the inital message to Bob
     // ...
-
+    
     // Bob does X3DH
     final resultBob = await x3dhFromInitialMessage(
       X3DHMessage(
@@ -43,7 +42,7 @@ void main() {
       opkBob,
       ikBob,
     );
-
+    
     expect(resultAlice.sk, resultBob.sk);
     expect(resultAlice.ad, resultBob.ad);
   });
@@ -67,9 +66,16 @@ void main() {
         2: await opkBob.pk.asBase64(),
       },
     );
-
+    
     // Alice does X3DH
-    final result = await x3dhFromBundle(bundleBob, ikAlice);
+    var exception = false;
-    expect(result.isType<InvalidKeyExchangeSignatureError>(), isTrue);
+    try {
+      await x3dhFromBundle(bundleBob, ikAlice);
+    } catch(e) {
+      exception = true;
+      expect(e is InvalidSignatureException, true, reason: 'Expected InvalidSignatureException, but got $e');
+    }
+
+    expect(exception, true, reason: 'Expected test failure');
   });
 }