From 4d6e1f5458625f386849fc757075c7c86b04b245 Mon Sep 17 00:00:00 2001
From: td <152161658+td-famedly@users.noreply.github.com>
Date: Wed, 3 Jul 2024 19:32:05 +0530
Subject: [PATCH] fix: ratchet on a single frame until ratchetWindowSize (#544)

---
 lib/src/e2ee/key_provider.dart |  4 ++-
 web/e2ee.cryptor.dart          | 51 +++++++++++++++++++---------------
 2 files changed, 32 insertions(+), 23 deletions(-)

diff --git a/lib/src/e2ee/key_provider.dart b/lib/src/e2ee/key_provider.dart
index 0c971d56..2dc738f6 100644
--- a/lib/src/e2ee/key_provider.dart
+++ b/lib/src/e2ee/key_provider.dart
@@ -50,10 +50,11 @@ abstract class KeyProvider {
 }
 
 class BaseKeyProvider implements KeyProvider {
+  final Map<String, int> _latestSetIndex = {};
   final Map<String, Map<int, Uint8List>> _keys = {};
 
   int getLatestIndex(String participantId) {
-    return _keys[participantId]?.keys.last ?? 0;
+    return _latestSetIndex[participantId] ?? 0;
   }
 
   Uint8List? _sharedKey;
@@ -152,6 +153,7 @@ class BaseKeyProvider implements KeyProvider {
     logger.info(
         '_setKey for ${keyInfo.participantId}, idx: ${keyInfo.keyIndex}, key: ${base64Encode(keyInfo.key)}');
     _keys[keyInfo.participantId]![keyInfo.keyIndex] = keyInfo.key;
+    _latestSetIndex[keyInfo.participantId] = keyInfo.keyIndex;
     await _keyProvider.setKey(
       participantId: keyInfo.participantId,
       index: keyInfo.keyIndex,
diff --git a/web/e2ee.cryptor.dart b/web/e2ee.cryptor.dart
index 2e697d98..70f486cf 100644
--- a/web/e2ee.cryptor.dart
+++ b/web/e2ee.cryptor.dart
@@ -314,15 +314,12 @@ class FrameCryptor {
   ) async {
     var buffer = frame.data.asUint8List();
 
-    if (!enabled ||
-        // skip for encryption for empty dtx frames
-        buffer.isEmpty) {
-      if (keyOptions.discardFrameWhenCryptorNotReady) {
-        return;
-      }
+    if (buffer.isEmpty) {
+      // skip for encryption for empty dtx frames
       controller.enqueue(frame);
       return;
     }
+    if (!enabled && keyOptions.discardFrameWhenCryptorNotReady) return;
 
     var secretKey = keyHandler.getKeySet(currentKeyIndex)?.encryptionKey;
     var keyIndex = currentKeyIndex;
@@ -381,7 +378,8 @@ class FrameCryptor {
       controller.enqueue(frame);
 
       if (lastError != CryptorError.kOk) {
-        logger.info('$participantIdentity trackId: $trackId kind: $kind cryptorState changed from $lastError to kOk');
+        logger.info(
+            '$participantIdentity trackId: $trackId kind: $kind cryptorState changed from $lastError to kOk');
         lastError = CryptorError.kOk;
         postMessage({
           'type': 'cryptorState',
@@ -425,15 +423,14 @@ class FrameCryptor {
     KeySet? initialKeySet;
     var initialKeyIndex = currentKeyIndex;
 
-    if (!enabled ||
-        // skip for encryption for empty dtx frames
-        buffer.isEmpty) {
-      sifGuard.recordUserFrame();
-      if (keyOptions.discardFrameWhenCryptorNotReady) return;
+    if (buffer.isEmpty) {
+      // skip for encryption for empty dtx frames
       logger.fine('enqueing empty frame');
+      sifGuard.recordUserFrame();
       controller.enqueue(frame);
       return;
     }
+    if (!enabled && keyOptions.discardFrameWhenCryptorNotReady) return;
 
     if (keyOptions.uncryptedMagicBytes != null) {
       var magicBytes = keyOptions.uncryptedMagicBytes!;
@@ -491,7 +488,8 @@ class FrameCryptor {
       /// to throw missingkeys faster lower your failureTolerance
       if (initialKeySet == null || !keyHandler.hasValidKey) {
         if (lastError != CryptorError.kMissingKey) {
-          logger.info('$participantIdentity trackId: $trackId kind: $kind cryptorState changed from $lastError to kMissingKey');
+          logger.info(
+              '$participantIdentity trackId: $trackId kind: $kind cryptorState changed from $lastError to kMissingKey');
           lastError = CryptorError.kMissingKey;
           postMessage({
             'type': 'cryptorState',
@@ -503,7 +501,6 @@ class FrameCryptor {
             'error': 'Missing key for track $trackId'
           });
         }
-        // controller.enqueue(frame);
         return;
       }
       var currentkeySet = initialKeySet;
@@ -528,7 +525,8 @@ class FrameCryptor {
         }
 
         if (currentkeySet != initialKeySet) {
-          logger.fine('ratchetKey: decryption ok, newState: kKeyRatcheted');
+          logger.fine(
+              'ratchetKey: $participantIdentity trackId: $trackId kind: $kind decryption ok, newState: kKeyRatcheted');
           await keyHandler.setKeySetFromMaterial(
               currentkeySet, initialKeyIndex);
         }
@@ -539,9 +537,10 @@ class FrameCryptor {
           logger.finer(
               'KeyRatcheted: ssrc ${metaData.synchronizationSource} timestamp ${frame.timestamp} ratchetCount $ratchetCount  participantId: $participantIdentity');
           logger.finer(
-              'ratchetKey: lastError != CryptorError.kKeyRatcheted, reset state to kKeyRatcheted');
+              'ratchetKey: $participantIdentity trackId: $trackId kind: $kind lastError != CryptorError.kKeyRatcheted, reset state to kKeyRatcheted');
 
-          logger.info('$participantIdentity trackId: $trackId kind: $kind cryptorState changed from $lastError to kKeyRatcheted');
+          logger.info(
+              '$participantIdentity trackId: $trackId kind: $kind cryptorState changed from $lastError to kKeyRatcheted');
           lastError = CryptorError.kKeyRatcheted;
           postMessage({
             'type': 'cryptorState',
@@ -556,9 +555,10 @@ class FrameCryptor {
       }
 
       Future<void> ratchedKeyInternal() async {
-        if (ratchetCount >= keyOptions.ratchetWindowSize ||
+        if (ratchetCount > keyOptions.ratchetWindowSize ||
             keyOptions.ratchetWindowSize <= 0) {
-          throw Exception('[ratchedKeyInternal] cannot ratchet anymore');
+          throw Exception(
+              '[ratchedKeyInternal] cannot ratchet anymore $participantIdentity trackId: $trackId kind: $kind');
         }
 
         var newKeyBuffer = crypto.jsArrayBufferFrom(await keyHandler.ratchet(
@@ -568,7 +568,13 @@ class FrameCryptor {
         currentkeySet =
             await keyHandler.deriveKeys(newMaterial, keyOptions.ratchetSalt);
         ratchetCount++;
-        await decryptFrameInternal();
+        try {
+          /// keep trying to ratchet until ratchetWindowSize
+          await decryptFrameInternal();
+        } catch (e) {
+          lastError = CryptorError.kInternalError;
+          await ratchedKeyInternal();
+        }
       }
 
       try {
@@ -583,7 +589,7 @@ class FrameCryptor {
 
       if (decrypted == null) {
         throw Exception(
-            '[decodeFunction] decryption failed even after ratchting');
+            '[decodeFunction] decryption failed even after ratchting $participantIdentity trackId: $trackId kind: $kind');
       }
 
       // we can now be sure that decryption was a success
@@ -600,7 +606,8 @@ class FrameCryptor {
       controller.enqueue(frame);
 
       if (lastError != CryptorError.kOk) {
-        logger.info('$participantIdentity trackId: $trackId kind: $kind cryptorState changed from $lastError to kOk');
+        logger.info(
+            '$participantIdentity trackId: $trackId kind: $kind cryptorState changed from $lastError to kOk');
         lastError = CryptorError.kOk;
         postMessage({
           'type': 'cryptorState',