From fda06cef55909b02806366e103f58a38c5e957bc Mon Sep 17 00:00:00 2001
From: "Alexander \"PapaTutuWawa" <papatutuwawa@polynom.me>
Date: Tue, 16 Aug 2022 14:02:04 +0200
Subject: [PATCH] feat: Implement acknowledging ratchet sessions

---
 lib/src/double_ratchet/double_ratchet.dart | 10 +++++
 lib/src/omemo/sessionmanager.dart          | 27 +++++++++++++
 test/omemo_test.dart                       | 44 ++++++++++++++++++++++
 3 files changed, 81 insertions(+)

diff --git a/lib/src/double_ratchet/double_ratchet.dart b/lib/src/double_ratchet/double_ratchet.dart
index 7bb86da..bf73810 100644
--- a/lib/src/double_ratchet/double_ratchet.dart
+++ b/lib/src/double_ratchet/double_ratchet.dart
@@ -67,6 +67,7 @@ class OmemoDoubleRatchet {
     this.ik,
     this.sessionAd,
     this.mkSkipped, // MKSKIPPED
+    this.acknowledged,
   );
 
   factory OmemoDoubleRatchet.fromJson(Map<String, dynamic> data) {
@@ -83,6 +84,7 @@ class OmemoDoubleRatchet {
       'pn': 0,
       'ik_pub': 'base/64/encoded',
       'session_ad': 'base/64/encoded',
+      'acknowledged': true | false,
       'mkskipped': [
         {
           'key': 'base/64/encoded',
@@ -117,6 +119,7 @@ class OmemoDoubleRatchet {
       ),
       base64.decode(data['session_ad']! as String),
       mkSkipped,
+      data['acknowledged']! as bool,
     );
   }
   
@@ -148,6 +151,10 @@ class OmemoDoubleRatchet {
 
   final Map<SkippedKey, List<int>> mkSkipped;
 
+  /// Indicates whether we received an empty OMEMO message after building a session with
+  /// the device.
+  bool acknowledged;
+
   /// Create an OMEMO session using the Signed Pre Key [spk], the shared secret [sk] that
   /// was obtained using a X3DH and the associated data [ad] that was also obtained through
   /// a X3DH. [ik] refers to Bob's (the receiver's) IK public key.
@@ -169,6 +176,7 @@ class OmemoDoubleRatchet {
       ik,
       ad,
       {},
+      false,
     );
   }
 
@@ -189,6 +197,7 @@ class OmemoDoubleRatchet {
       ik,
       ad,
       {},
+      true,
     );
   }
 
@@ -214,6 +223,7 @@ class OmemoDoubleRatchet {
       'ik_pub': base64.encode(await ik.getBytes()),
       'session_ad': base64.encode(sessionAd),
       'mkskipped': mkSkippedSerialised,
+      'acknowledged': acknowledged,
     };
   }
   
diff --git a/lib/src/omemo/sessionmanager.dart b/lib/src/omemo/sessionmanager.dart
index b995347..35b2a02 100644
--- a/lib/src/omemo/sessionmanager.dart
+++ b/lib/src/omemo/sessionmanager.dart
@@ -430,6 +430,33 @@ class OmemoSessionManager {
       _eventStreamController.add(DeviceMapModifiedEvent(_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 {
+    final ret = List<int>.empty(growable: true);
+
+    await _lock.synchronized(() async {
+      final devices = _deviceMap[jid]!;
+      for (final device in devices) {
+        final ratchet = _ratchetMap[RatchetMapKey(jid, device)]!;
+        if (!ratchet.acknowledged) ret.add(device);
+      }
+    });
+
+    return ret;
+  }
+
+  /// 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));
+    });
+  }
   
   @visibleForTesting
   OmemoDoubleRatchet getRatchet(String jid, int deviceId) => _ratchetMap[RatchetMapKey(jid, deviceId)]!;
diff --git a/test/omemo_test.dart b/test/omemo_test.dart
index f5f1171..a9cd102 100644
--- a/test/omemo_test.dart
+++ b/test/omemo_test.dart
@@ -499,4 +499,48 @@ void main() {
       expect(deviceMap.containsKey(bobJid), false);
     });
   });
+
+  test('Test acknowledging a ratchet', () async {
+    const aliceJid = 'alice@server.example';
+    const bobJid = 'bob@other.server.example';
+    // Alice and Bob generate their sessions
+    final aliceSession = await OmemoSessionManager.generateNewIdentity(
+      aliceJid,
+      AlwaysTrustingTrustManager(),
+      opkAmount: 1,
+    );
+    final bobSession = await OmemoSessionManager.generateNewIdentity(
+      bobJid,
+      AlwaysTrustingTrustManager(),
+      opkAmount: 1,
+    );
+
+    // Alice sends Bob a message
+    await aliceSession.encryptToJid(
+      bobJid,
+      'Hallo Welt',
+      newSessions: [
+        await (await bobSession.getDevice()).toBundle(),
+      ],
+    );
+    expect(
+      await aliceSession.getUnacknowledgedRatchets(bobJid),
+      [
+        (await bobSession.getDevice()).id,
+      ],
+    );
+
+    // Bob sends alice an empty message
+    // ...
+
+    // Alice decrypts it
+    // ...
+
+    // Alice marks the ratchet as acknowledged
+    await aliceSession.ratchetAcknowledged(bobJid, (await bobSession.getDevice()).id);
+    expect(
+      (await aliceSession.getUnacknowledgedRatchets(bobJid)).isEmpty,
+      true,
+    );
+  });
 }