From e7f0a2b92c02d151dc753c1bb3886fff4ec11c81 Mon Sep 17 00:00:00 2001
From: Richard van der Hoff <richard@matrix.org>
Date: Tue, 9 Jul 2024 18:43:52 +0100
Subject: [PATCH] crypto: Pass room id and session id to
 `room_keys_withheld_received_stream`

---
 crates/matrix-sdk-crypto/CHANGELOG.md           |  3 ++-
 crates/matrix-sdk-crypto/src/machine.rs         |  4 +++-
 .../src/store/crypto_store_wrapper.rs           | 17 +++++++++++------
 crates/matrix-sdk-crypto/src/store/mod.rs       | 16 +++++++++++++++-
 4 files changed, 31 insertions(+), 9 deletions(-)

diff --git a/crates/matrix-sdk-crypto/CHANGELOG.md b/crates/matrix-sdk-crypto/CHANGELOG.md
index 7ced4886c82..fc8e737ab71 100644
--- a/crates/matrix-sdk-crypto/CHANGELOG.md
+++ b/crates/matrix-sdk-crypto/CHANGELOG.md
@@ -63,7 +63,8 @@ Additions:
 - Expose new method `OlmMachine::room_keys_withheld_received_stream`, to allow
   applications to receive notifications about received `m.room_key.withheld`
   events.
-  ([#3660](https://github.com/matrix-org/matrix-rust-sdk/pull/3660))
+  ([#3660](https://github.com/matrix-org/matrix-rust-sdk/pull/3660)),
+  ([#3673](https://github.com/matrix-org/matrix-rust-sdk/pull/3673))
 
 - Expose new method `OlmMachine::clear_crypto_cache()`, with FFI bindings
   ([#3462](https://github.com/matrix-org/matrix-rust-sdk/pull/3462))
diff --git a/crates/matrix-sdk-crypto/src/machine.rs b/crates/matrix-sdk-crypto/src/machine.rs
index 10dafe93c86..e290c799583 100644
--- a/crates/matrix-sdk-crypto/src/machine.rs
+++ b/crates/matrix-sdk-crypto/src/machine.rs
@@ -3292,8 +3292,10 @@ pub(crate) mod tests {
             .flatten()
             .expect("We should have received a notification of room key being withheld");
         assert_eq!(withheld_received.len(), 1);
+
+        assert_eq!(&withheld_received[0].room_id, room_id);
         assert_matches!(
-            &withheld_received[0].content,
+            &withheld_received[0].withheld_event.content,
             RoomKeyWithheldContent::MegolmV1AesSha2(MegolmV1AesSha2WithheldContent::Unverified(
                 unverified_withheld_content
             ))
diff --git a/crates/matrix-sdk-crypto/src/store/crypto_store_wrapper.rs b/crates/matrix-sdk-crypto/src/store/crypto_store_wrapper.rs
index da1884cd4de..da63f29832d 100644
--- a/crates/matrix-sdk-crypto/src/store/crypto_store_wrapper.rs
+++ b/crates/matrix-sdk-crypto/src/store/crypto_store_wrapper.rs
@@ -12,8 +12,7 @@ use super::{DeviceChanges, IdentityChanges, LockableCryptoStore};
 use crate::{
     olm::InboundGroupSession,
     store,
-    store::{Changes, DynCryptoStore, IntoCryptoStore, RoomKeyInfo},
-    types::events::room_key_withheld::RoomKeyWithheldEvent,
+    store::{Changes, DynCryptoStore, IntoCryptoStore, RoomKeyInfo, RoomKeyWithheldInfo},
     GossippedSecret, ReadOnlyOwnUserIdentity,
 };
 
@@ -32,7 +31,7 @@ pub(crate) struct CryptoStoreWrapper {
 
     /// The sender side of a broadcast stream that is notified whenever we
     /// receive an `m.room_key.withheld` message.
-    room_keys_withheld_received_sender: broadcast::Sender<Vec<RoomKeyWithheldEvent>>,
+    room_keys_withheld_received_sender: broadcast::Sender<Vec<RoomKeyWithheldInfo>>,
 
     /// The sender side of a broadcast channel which sends out secrets we
     /// received as a `m.secret.send` event.
@@ -77,8 +76,14 @@ impl CryptoStoreWrapper {
 
         let withheld_session_updates: Vec<_> = changes
             .withheld_session_info
-            .values()
-            .flat_map(|session_map| session_map.values().cloned())
+            .iter()
+            .flat_map(|(room_id, session_map)| {
+                session_map.iter().map(|(session_id, withheld_event)| RoomKeyWithheldInfo {
+                    room_id: room_id.to_owned(),
+                    session_id: session_id.to_owned(),
+                    withheld_event: withheld_event.clone(),
+                })
+            })
             .collect();
 
         let secrets = changes.secrets.to_owned();
@@ -161,7 +166,7 @@ impl CryptoStoreWrapper {
     /// logged and items will be dropped.
     pub fn room_keys_withheld_received_stream(
         &self,
-    ) -> impl Stream<Item = Vec<RoomKeyWithheldEvent>> {
+    ) -> impl Stream<Item = Vec<RoomKeyWithheldInfo>> {
         let stream = BroadcastStream::new(self.room_keys_withheld_received_sender.subscribe());
         Self::filter_errors_out_of_stream(stream, "room_keys_withheld_received_stream")
     }
diff --git a/crates/matrix-sdk-crypto/src/store/mod.rs b/crates/matrix-sdk-crypto/src/store/mod.rs
index 4ed122a132e..90f0ec23df8 100644
--- a/crates/matrix-sdk-crypto/src/store/mod.rs
+++ b/crates/matrix-sdk-crypto/src/store/mod.rs
@@ -913,6 +913,20 @@ impl From<&InboundGroupSession> for RoomKeyInfo {
     }
 }
 
+/// Information on a room key that has been withheld
+#[derive(Clone, Debug, Deserialize, Serialize)]
+pub struct RoomKeyWithheldInfo {
+    /// The room where the key is used.
+    pub room_id: OwnedRoomId,
+
+    /// The ID of the session that the key is for.
+    pub session_id: String,
+
+    /// The `m.room_key.withheld` event that notified us that the key is being
+    /// withheld.
+    pub withheld_event: RoomKeyWithheldEvent,
+}
+
 impl Store {
     /// Create a new Store.
     pub(crate) fn new(
@@ -1473,7 +1487,7 @@ impl Store {
     /// logged and items will be dropped.
     pub fn room_keys_withheld_received_stream(
         &self,
-    ) -> impl Stream<Item = Vec<RoomKeyWithheldEvent>> {
+    ) -> impl Stream<Item = Vec<RoomKeyWithheldInfo>> {
         self.inner.store.room_keys_withheld_received_stream()
     }