stores: new method CryptoStore::get_withheld_sessions_by_room_id

Implement this across all the store implementations

crates/matrix-sdk-crypto/src/store/memorystore.rs

@@ -428,6 +428,18 @@ impl CryptoStore for MemoryStore {
             .and_then(|e| Some(e.get(session_id)?.to_owned())))
     }
 
+    async fn get_withheld_sessions_by_room_id(
+        &self,
+        room_id: &RoomId,
+    ) -> crate::store::Result<Vec<RoomKeyWithheldEntry>, Self::Error> {
+        Ok(self
+            .direct_withheld_info
+            .read()
+            .get(room_id)
+            .map(|e| e.values().cloned().collect())
+            .unwrap_or_default())
+    }
+
     async fn get_inbound_group_sessions(&self) -> Result<Vec<InboundGroupSession>> {
         let inbounds = self
             .inbound_group_sessions
@@ -1375,6 +1387,13 @@ mod integration_tests {
             self.0.get_withheld_info(room_id, session_id).await
         }
 
+        async fn get_withheld_sessions_by_room_id(
+            &self,
+            room_id: &RoomId,
+        ) -> Result<Vec<RoomKeyWithheldEntry>, Self::Error> {
+            self.0.get_withheld_sessions_by_room_id(room_id).await
+        }
+
         async fn get_inbound_group_sessions(
             &self,
         ) -> Result<Vec<InboundGroupSession>, Self::Error> {

crates/matrix-sdk-crypto/src/store/traits.rs

@@ -118,6 +118,19 @@ pub trait CryptoStore: AsyncTraitDeps {
         session_id: &str,
     ) -> Result<Option<RoomKeyWithheldEntry>, Self::Error>;
 
+    /// Get all the sessions where we have received an `m.room_key.withheld`
+    /// event (or, post-[MSC4268], where there was a `withheld` entry in the key
+    /// bundle).
+    ///
+    /// [MSC4268]: https://github.com/matrix-org/matrix-spec-proposals/pull/4268
+    ///
+    /// # Arguments
+    /// * `room_id` - The ID of the room to return withheld sessions for.
+    async fn get_withheld_sessions_by_room_id(
+        &self,
+        room_id: &RoomId,
+    ) -> Result<Vec<RoomKeyWithheldEntry>, Self::Error>;
+
     /// Get all the inbound group sessions we have stored.
     async fn get_inbound_group_sessions(&self) -> Result<Vec<InboundGroupSession>, Self::Error>;
 
@@ -592,6 +605,13 @@ impl<T: CryptoStore> CryptoStore for EraseCryptoStoreError<T> {
         self.0.get_withheld_info(room_id, session_id).await.map_err(Into::into)
     }
 
+    async fn get_withheld_sessions_by_room_id(
+        &self,
+        room_id: &RoomId,
+    ) -> Result<Vec<RoomKeyWithheldEntry>, Self::Error> {
+        self.0.get_withheld_sessions_by_room_id(room_id).await.map_err(Into::into)
+    }
+
     async fn get_room_settings(&self, room_id: &RoomId) -> Result<Option<RoomSettings>> {
         self.0.get_room_settings(room_id).await.map_err(Into::into)
     }

crates/matrix-sdk-indexeddb/src/crypto_store/mod.rs

@@ -1495,6 +1495,25 @@ impl_crypto_store! {
         }
     }
 
+    async fn get_withheld_sessions_by_room_id(
+        &self,
+        room_id: &RoomId,
+    ) -> Result<Vec<RoomKeyWithheldEntry>> {
+        let range = self.serializer.encode_to_range(keys::WITHHELD_SESSIONS, room_id);
+
+        self
+            .inner
+            .transaction(keys::WITHHELD_SESSIONS)
+            .with_mode(TransactionMode::Readonly)
+            .build()?
+            .object_store(keys::WITHHELD_SESSIONS)?
+            .get_all()
+            .with_query(&range)
+            .await?
+            .map(|val| self.serializer.deserialize_value(val?).map_err(Into::into))
+            .collect()
+    }
+
     async fn get_room_settings(&self, room_id: &RoomId) -> Result<Option<RoomSettings>> {
         let key = self.serializer.encode_key(keys::ROOM_SETTINGS, room_id);
         self.inner

crates/matrix-sdk-sqlite/src/crypto_store.rs

@@ -774,6 +774,14 @@ trait SqliteObjectCryptoStoreExt: SqliteAsyncConnExt {
             .optional()?)
     }
 
+    async fn get_withheld_sessions_by_room_id(&self, room_id: Key) -> Result<Vec<Vec<u8>>> {
+        Ok(self
+            .prepare("SELECT data FROM direct_withheld_info WHERE room_id = ?1", |mut stmt| {
+                stmt.query((room_id,))?.mapped(|row| row.get(0)).collect()
+            })
+            .await?)
+    }
+
     async fn get_room_settings(&self, room_id: Key) -> Result<Option<Vec<u8>>> {
         Ok(self
             .query_row("SELECT data FROM room_settings WHERE room_id = ?", (room_id,), |row| {
@@ -1403,6 +1411,21 @@ impl CryptoStore for SqliteCryptoStore {
             .transpose()
     }
 
+    async fn get_withheld_sessions_by_room_id(
+        &self,
+        room_id: &RoomId,
+    ) -> matrix_sdk_crypto::store::Result<Vec<RoomKeyWithheldEntry>, Self::Error> {
+        let room_id = self.encode_key("direct_withheld_info", room_id);
+
+        self.acquire()
+            .await?
+            .get_withheld_sessions_by_room_id(room_id)
+            .await?
+            .into_iter()
+            .map(|value| self.deserialize_json(&value))
+            .collect()
+    }
+
     async fn get_room_settings(&self, room_id: &RoomId) -> Result<Option<RoomSettings>> {
         let room_id = self.encode_key("room_settings", room_id.as_bytes());
         let Some(value) = self.acquire().await?.get_room_settings(room_id).await? else {