Merge pull request #5819 from matrix-org/rav/cryptostore_withheld_sessions_by_room_id

crypto: add new `CryptoStore` method `get_withheld_sessions_by_room_id`

Author: richvdh

crates/matrix-sdk-crypto/CHANGELOG.md

@@ -8,6 +8,8 @@ All notable changes to this project will be documented in this file.
 
 ### Features
 
+- Expose new method `CryptoStore::get_withheld_sessions_by_room_id`.
+  ([#5819](https://github.com/matrix-org/matrix-rust-sdk/pull/5819))
 - Use new withheld code in key bundles for sessions not marked as
   `shared_history`.
   ([#5807](https://github.com/matrix-org/matrix-rust-sdk/pull/5807)

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

@@ -1139,55 +1139,59 @@ macro_rules! cryptostore_integration_tests {
             async fn test_withheld_info_storage() {
                 let (account, store) = get_loaded_store("withheld_info_storage").await;
 
-                let mut info_list: BTreeMap<_, BTreeMap<_, RoomKeyWithheldEntry>> = BTreeMap::new();
-
                 let user_id = account.user_id().to_owned();
                 let room_id = room_id!("!DwLygpkclUAfQNnfva:example.com");
                 let session_id_1 = "GBnDxGP9i3IkPsz3/ihNr6P7qjIXxSRVWZ1MYmSn09w";
                 let session_id_2 = "IDLtnNCH2kIr3xIf1B7JFkGpQmTjyMca2jww+X6zeOE";
 
-                let content = RoomKeyWithheldContent::MegolmV1AesSha2(
-                    MegolmV1AesSha2WithheldContent::Unverified(
-                        CommonWithheldCodeContent::new(
-                            room_id.to_owned(),
-                            session_id_1.into(),
-                            Curve25519PublicKey::from_base64(
-                                "9n7mdWKOjr9c4NTlG6zV8dbFtNK79q9vZADoh7nMUwA",
+                {
+                    let mut info_list: BTreeMap<_, BTreeMap<_, RoomKeyWithheldEntry>> = BTreeMap::new();
+
+                    let content = RoomKeyWithheldContent::MegolmV1AesSha2(
+                        MegolmV1AesSha2WithheldContent::Unverified(
+                            CommonWithheldCodeContent::new(
+                                room_id.to_owned(),
+                                session_id_1.into(),
+                                Curve25519PublicKey::from_base64(
+                                    "9n7mdWKOjr9c4NTlG6zV8dbFtNK79q9vZADoh7nMUwA",
+                                )
+                                .unwrap(),
+                                "DEVICEID".into(),
                             )
-                            .unwrap(),
-                            "DEVICEID".into(),
-                        )
-                        .into(),
-                    ),
-                );
-                let event = ToDeviceEvent::new(user_id.to_owned(), content);
-                info_list
-                    .entry(room_id.to_owned())
-                    .or_default()
-                    .insert(session_id_1.to_owned(), event.into());
-
-                let content = RoomKeyWithheldContent::MegolmV1AesSha2(
-                    MegolmV1AesSha2WithheldContent::BlackListed(
-                        CommonWithheldCodeContent::new(
-                            room_id.to_owned(),
-                            session_id_2.into(),
-                            Curve25519PublicKey::from_base64(
-                                "9n7mdWKOjr9c4NTlG6zV8dbFtNK79q9vZADoh7nMUwA",
+                            .into(),
+                        ),
+                    );
+                    let event = ToDeviceEvent::new(user_id.to_owned(), content);
+                    info_list
+                        .entry(room_id.to_owned())
+                        .or_default()
+                        .insert(session_id_1.to_owned(), event.into());
+
+                    let content = RoomKeyWithheldContent::MegolmV1AesSha2(
+                        MegolmV1AesSha2WithheldContent::BlackListed(
+                            CommonWithheldCodeContent::new(
+                                room_id.to_owned(),
+                                session_id_2.into(),
+                                Curve25519PublicKey::from_base64(
+                                    "9n7mdWKOjr9c4NTlG6zV8dbFtNK79q9vZADoh7nMUwA",
+                                )
+                                .unwrap(),
+                                "DEVICEID".into(),
                             )
-                            .unwrap(),
-                            "DEVICEID".into(),
-                        )
-                        .into(),
-                    ),
-                );
-                let event = ToDeviceEvent::new(user_id.to_owned(), content);
-                info_list
-                    .entry(room_id.to_owned())
-                    .or_default()
-                    .insert(session_id_2.to_owned(), event.into());
+                            .into(),
+                        ),
+                    );
+                    let event = ToDeviceEvent::new(user_id.to_owned(), content);
+                    info_list
+                        .entry(room_id.to_owned())
+                        .or_default()
+                        .insert(session_id_2.to_owned(), event.into());
 
-                let changes = Changes { withheld_session_info: info_list, ..Default::default() };
-                store.save_changes(changes).await.unwrap();
+                    let changes = Changes { withheld_session_info: info_list, ..Default::default() };
+                    store.save_changes(changes).await.unwrap();
+                }
+
+                // Test `get_withheld_info`
 
                 let is_withheld = store.get_withheld_info(room_id, session_id_1).await.unwrap();
 
@@ -1211,6 +1215,12 @@ macro_rules! cryptostore_integration_tests {
                     store.get_withheld_info(other_room_id, session_id_2).await.unwrap();
 
                 assert!(is_withheld.is_none());
+
+                // Test `get_withheld_sessions_by_room_id`
+                let withhelds = store.get_withheld_sessions_by_room_id(room_id).await.expect("Error getting withheld sessions by room ID");
+                assert_eq!(withhelds.len(), 2);
+                let withheld1 = withhelds.iter().find(|entry| entry.content.megolm_session_id() == Some(session_id_1)).expect("Did not find session 1 in withhelds list");
+                assert_eq!(withheld1.content.withheld_code(), WithheldCode::Unverified)
             }
 
             #[async_test]

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-crypto/src/types/events/room_key_withheld.rs

@@ -154,6 +154,27 @@ impl RoomKeyWithheldContent {
             RoomKeyWithheldContent::Unknown(c) => c.algorithm.to_owned(),
         }
     }
+
+    /// Get the room ID of the withheld session, if known
+    pub fn room_id(&self) -> Option<&RoomId> {
+        match &self {
+            RoomKeyWithheldContent::MegolmV1AesSha2(c) => c.room_id(),
+            #[cfg(feature = "experimental-algorithms")]
+            RoomKeyWithheldContent::MegolmV2AesSha2(c) => c.room_id(),
+            RoomKeyWithheldContent::Unknown(_) => None,
+        }
+    }
+
+    /// Get the megolm session ID of the withheld session, if it is in fact a
+    /// megolm session.
+    pub fn megolm_session_id(&self) -> Option<&str> {
+        match &self {
+            RoomKeyWithheldContent::MegolmV1AesSha2(c) => c.session_id(),
+            #[cfg(feature = "experimental-algorithms")]
+            RoomKeyWithheldContent::MegolmV2AesSha2(c) => c.session_id(),
+            RoomKeyWithheldContent::Unknown(_) => None,
+        }
+    }
 }
 
 impl EventType for RoomKeyWithheldContent {

crates/matrix-sdk-indexeddb/CHANGELOG.md

@@ -8,7 +8,9 @@ All notable changes to this project will be documented in this file.
 
 ### Features
 
-- [**breaking**] `IndexeddbCryptoStore::get_withheld_info` now returns `Result<Option<RoomKeyWithheldEntry>, ...>`
+- Implement new method `CryptoStore::get_withheld_sessions_by_room_id`.
+  ([#5819](https://github.com/matrix-org/matrix-rust-sdk/pull/5819))
+- [**breaking**] `IndexeddbCryptoStore::get_withheld_info` now returns `Result<Option<RoomKeyWithheldEntry>, ...>`.
   ([#5737](https://github.com/matrix-org/matrix-rust-sdk/pull/5737))
 
 ### Performance

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

@@ -33,6 +33,7 @@ mod v10_to_v11;
 mod v11_to_v12;
 mod v12_to_v13;
 mod v13_to_v14;
+mod v14_to_v101;
 mod v5_to_v7;
 mod v7;
 mod v7_to_v8;
@@ -101,7 +102,7 @@ impl Drop for MigrationDb {
 /// as is version 200, but versions 101-199 are all backwards compatible with
 /// version 100. In other words, if you divide by 100, you get something
 /// approaching semver: version 200 is major version 2, minor version 0.
-const MAX_SUPPORTED_SCHEMA_VERSION: u32 = 99;
+const MAX_SUPPORTED_SCHEMA_VERSION: u32 = 199;
 
 /// Open the indexeddb with the given name, upgrading it to the latest version
 /// of the schema if necessary.
@@ -176,6 +177,15 @@ pub async fn open_and_upgrade_db(
         v13_to_v14::schema_bump(name).await?;
     }
 
+    if old_version < 100 {
+        v14_to_v101::schema_add(name).await?;
+    }
+
+    if old_version < 101 {
+        v14_to_v101::data_migrate(name, serializer).await?;
+        v14_to_v101::schema_delete(name).await?;
+    }
+
     // If you add more migrations here, you'll need to update
     // `tests::EXPECTED_SCHEMA_VERSION`.
 
@@ -254,16 +264,18 @@ mod tests {
     use indexed_db_futures::{
         database::VersionChangeEvent, prelude::*, transaction::TransactionMode,
     };
-    use matrix_sdk_common::js_tracing::make_tracing_subscriber;
+    use matrix_sdk_common::{
+        deserialized_responses::WithheldCode, js_tracing::make_tracing_subscriber,
+    };
     use matrix_sdk_crypto::{
         olm::{InboundGroupSession, SenderData, SessionKey},
-        store::CryptoStore,
-        types::EventEncryptionAlgorithm,
+        store::{types::RoomKeyWithheldEntry, CryptoStore},
+        types::{events::room_key_withheld::RoomKeyWithheldContent, EventEncryptionAlgorithm},
         vodozemac::{Curve25519PublicKey, Curve25519SecretKey, Ed25519PublicKey, Ed25519SecretKey},
     };
     use matrix_sdk_store_encryption::StoreCipher;
     use matrix_sdk_test::async_test;
-    use ruma::{room_id, OwnedRoomId, RoomId};
+    use ruma::{device_id, owned_user_id, room_id, OwnedRoomId, RoomId};
     use serde::Serialize;
     use tracing_subscriber::util::SubscriberInitExt;
     use wasm_bindgen::JsValue;
@@ -278,7 +290,7 @@ mod tests {
     wasm_bindgen_test::wasm_bindgen_test_configure!(run_in_browser);
 
     /// The schema version we expect after we open the store.
-    const EXPECTED_SCHEMA_VERSION: u32 = 14;
+    const EXPECTED_SCHEMA_VERSION: u32 = 101;
 
     /// Adjust this to test do a more comprehensive perf test
     const NUM_RECORDS_FOR_PERF: usize = 2_000;
@@ -786,6 +798,89 @@ mod tests {
         assert_eq!(backup_data.backup_version, Some("1".to_owned()));
     }
 
+    /// Test migrating `withheld_sessions` data from store v14 to latest,
+    /// on a store with encryption disabled.
+    #[async_test]
+    async fn test_v14_v101_migration_unencrypted() {
+        test_v14_v101_migration_with_cipher("test_v101_migration_unencrypted", None).await
+    }
+
+    /// Test migrating `withheld_sessions` data from store v14 to latest,
+    /// on a store with encryption enabled.
+    #[async_test]
+    async fn test_v14_v101_migration_encrypted() {
+        let cipher = StoreCipher::new().unwrap();
+        test_v14_v101_migration_with_cipher(
+            "test_v101_migration_encrypted",
+            Some(Arc::new(cipher)),
+        )
+        .await;
+    }
+
+    /// Helper function for `test_v14_v101_migration_{un,}encrypted`: test
+    /// migrating `withheld_sessions` data from store v14 to store v101.
+    async fn test_v14_v101_migration_with_cipher(
+        db_prefix: &str,
+        store_cipher: Option<Arc<StoreCipher>>,
+    ) {
+        let serializer = SafeEncodeSerializer::new(store_cipher.clone());
+
+        let _ = make_tracing_subscriber(None).try_init();
+        let db_name = format!("{db_prefix:0}::matrix-sdk-crypto");
+
+        // delete the db in case it was used in a previous run
+        let _ = Database::delete_by_name(&db_name).unwrap().await.unwrap();
+
+        let room_id = room_id!("!test:example.com");
+        let session_id = "12345";
+
+        // Given a DB with data in it as it was at v5
+        {
+            let db = create_v5_db(&db_name).await.unwrap();
+
+            let txn = db
+                .transaction(old_keys::DIRECT_WITHHELD_INFO)
+                .with_mode(TransactionMode::Readwrite)
+                .build()
+                .unwrap();
+            let store = txn.object_store(old_keys::DIRECT_WITHHELD_INFO).unwrap();
+
+            let sender_key =
+                Curve25519PublicKey::from_base64("9n7mdWKOjr9c4NTlG6zV8dbFtNK79q9vZADoh7nMUwA")
+                    .unwrap();
+
+            let withheld_entry = RoomKeyWithheldEntry {
+                sender: owned_user_id!("@alice:example.com"),
+                content: RoomKeyWithheldContent::new(
+                    EventEncryptionAlgorithm::MegolmV1AesSha2,
+                    WithheldCode::Blacklisted,
+                    room_id.to_owned(),
+                    session_id.to_owned(),
+                    sender_key,
+                    device_id!("ABC").to_owned(),
+                ),
+            };
+
+            let key = serializer.encode_key(old_keys::DIRECT_WITHHELD_INFO, (room_id, session_id));
+            let value = serializer.serialize_value(&withheld_entry).unwrap();
+            store.add(value).with_key(key).build().unwrap();
+            txn.commit().await.unwrap();
+            db.close();
+        }
+
+        // When I open a store based on that DB, triggering an upgrade
+        let store =
+            IndexeddbCryptoStore::open_with_store_cipher(&db_prefix, store_cipher).await.unwrap();
+
+        // Then I can read the withheld session settings
+        let withheld_entry = store
+            .get_withheld_info(room_id, session_id)
+            .await
+            .unwrap()
+            .expect("Should find a withheld entry in migrated data");
+        assert_eq!(withheld_entry.content.withheld_code(), WithheldCode::Blacklisted)
+    }
+
     async fn create_v5_db(name: &str) -> std::result::Result<Database, OpenDbError> {
         v0_to_v5::schema_add(name).await?;
         Database::open(name).with_version(5u32).build()?.await