From a85752d8b2a74e5f1578263f19f93cccb2bfa5eb Mon Sep 17 00:00:00 2001
From: Valere <bill.carson@valrsoft.com>
Date: Fri, 20 Dec 2024 15:29:39 +0100
Subject: [PATCH] feat(room_key_stream): Add support for lagging stream error

---
 CHANGELOG.md          |  6 ++++
 src/machine.rs        | 73 +++++++++++++++++++++++++++++++++++--------
 tests/machine.test.ts | 16 ++++++++++
 3 files changed, 82 insertions(+), 13 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index aab9d31e9..4d753624d 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -8,6 +8,12 @@
     `DehydratedDevices.keysForUpload` and `DehydratedDevices.rehydrate` now use the `DehydratedDeviceKey` as parameter
     instead of a raw UInt8Array.Use `DehydratedDeviceKey::createKeyFromArray` to migrate.
 
+**Other changes**
+
+-   Add `OlmMachine.registerRoomKeysWithheldCallbacks` that supports now a success and error callback.
+    An error will occur when the stream missed some updates; in that case decryption could be retried
+    for all current failures.
+
 # matrix-sdk-crypto-wasm v12.1.0
 
 -   Update matrix-rusk-sdk to `37c17cf854a70f` for the fix for
diff --git a/src/machine.rs b/src/machine.rs
index ba2efc45c..0b3764347 100644
--- a/src/machine.rs
+++ b/src/machine.rs
@@ -1320,22 +1320,69 @@ impl OlmMachine {
     /// of {@link RoomKeyInfo}) and returns a Promise.
     #[wasm_bindgen(js_name = "registerRoomKeyUpdatedCallback")]
     pub async fn register_room_key_updated_callback(&self, callback: Function) {
+        self.register_room_key_updated_callbacks(callback, None).await
+    }
+
+    /// Register a success and an error callback which will be called whenever
+    /// there is an update to room keys.
+    ///
+    /// `success_callback` should be a function that takes a single argument (an
+    /// array of {@link RoomKeyInfo}) and returns a Promise.
+    ///
+    /// `error_callback` should be a function that takes a single argument (the
+    /// error) and returns a Promise. When such an error happens that means
+    /// that update stream lost track and that all current decryption failures
+    /// can be retried as the key might as been imported without notice.
+    #[wasm_bindgen(js_name = "registerRoomKeyUpdatedCallbacks")]
+    pub async fn register_room_key_updated_callbacks(
+        &self,
+        success_callback: Function,
+        error_callback: Option<Function>,
+    ) {
         let stream = self.inner.store().room_keys_received_stream();
 
-        copy_stream_to_callback(
-            stream,
-            |input| match input {
-                Ok(keys) => iter::once(
-                    keys.into_iter().map(RoomKeyInfo::from).map(JsValue::from).collect::<Array>(),
-                ),
-                Err(e) => {
-                    warn!("Error reading  room_keys_received_stream {:?}", e);
-                    iter::once(Array::new())
+        spawn_local(async move {
+            pin_mut!(stream);
+
+            while let Some(item) = stream.next().await {
+                match item {
+                    Ok(input) => {
+                        let js_array = input
+                            .into_iter()
+                            .map(RoomKeyInfo::from)
+                            .map(JsValue::from)
+                            .collect::<Array>();
+                        match promise_result_to_future(
+                            success_callback.call1(&JsValue::NULL, &js_array.into()),
+                        )
+                        .await
+                        {
+                            Ok(_) => (),
+                            Err(e) => {
+                                warn!("Error calling registerRoomKeyUpdatedCallback success callback: {:?}", e);
+                            }
+                        }
+                    }
+                    Err(e) => {
+                        if let Some(ref error_callback) = error_callback {
+                            let js_error = JsError::new(&e.to_string());
+                            match promise_result_to_future(
+                                error_callback.call1(&JsValue::NULL, &js_error.into()),
+                            )
+                            .await
+                            {
+                                Ok(_) => (),
+                                Err(e) => {
+                                    warn!("Error calling registerRoomKeyUpdatedCallback error callback: {:?}", e);
+                                }
+                            }
+                        } else {
+                            warn!("Error reading room_keys_received_stream {:?}, no callback specified", e);
+                        }
+                    }
                 }
-            },
-            callback,
-            "room-key-received",
-        );
+            }
+        })
     }
 
     /// Register a callback which will be called whenever we receive a
diff --git a/tests/machine.test.ts b/tests/machine.test.ts
index 0c070219e..a9302082e 100644
--- a/tests/machine.test.ts
+++ b/tests/machine.test.ts
@@ -929,6 +929,22 @@ describe(OlmMachine.name, () => {
             expect(keyInfoList.length).toEqual(1);
             expect(keyInfoList[0].roomId.toString()).toStrictEqual(room.toString());
         });
+
+        test("importing room keys calls RoomKeyUpdatedCallbacks", async () => {
+            const success = jest.fn();
+            success.mockImplementation(() => Promise.resolve(undefined));
+            const error = jest.fn();
+            error.mockImplementation(() => Promise.resolve(undefined));
+            let m = await machine();
+            m.registerRoomKeyUpdatedCallbacks(success, error);
+            await m.importExportedRoomKeys(exportedRoomKeys, () => undefined);
+            expect(success).toHaveBeenCalledTimes(1);
+            let keyInfoList = success.mock.calls[0][0];
+            expect(keyInfoList.length).toEqual(1);
+            expect(keyInfoList[0].roomId.toString()).toStrictEqual(room.toString());
+
+            expect(error).toHaveBeenCalledTimes(0);
+        });
     });
 
     describe("can do in-room verification", () => {