feat: Append withheld info from room key bundle to store.

crates/matrix-sdk-crypto/CHANGELOG.md

@@ -6,6 +6,16 @@ All notable changes to this project will be documented in this file.
 
 ## [Unreleased] - ReleaseDate
 
+### Features
+
+- Improves feedback support for shared history when downloading room key bundles.
+  ([#5737](https://github.com/matrix-org/matrix-rust-sdk/pull/5737))
+  - Add `RoomKeyWithheldEntry` enum, wrapping either a received to-device `m.room_key.withheld` event or
+    its content, if derived from a downloaded room key bundle.
+  - `OlmMachine::receive_room_key_bundle` now appends withheld key information to the store.
+  - [**breaking**] `Changes::withheld_session_info` now stores a `RookKeyWithheldEntry` in each `room-id`-`session-id` entry.
+  - [**breaking**] `CryptoStore::get_withheld_info` now returns `Result<Option<RookKeyWithheldEntry>>`. This change also affects `MemoryStore`.
+
 ### Bug Fixes
 
 - Fix a bug introduced in 0.14.0 which meant that the serialization of the value returned by `OtherUserIdentity::verification_request_content` did not include a `msgtype` field.

crates/matrix-sdk-crypto/src/machine/mod.rs

@@ -998,7 +998,7 @@ impl OlmMachine {
         Ok(())
     }
 
-    fn add_withheld_info(&self, changes: &mut Changes, event: &RoomKeyWithheldEvent) {
+    pub(crate) fn add_withheld_info(&self, changes: &mut Changes, event: &RoomKeyWithheldEvent) {
         debug!(?event.content, "Processing `m.room_key.withheld` event");
 
         if let RoomKeyWithheldContent::MegolmV1AesSha2(
@@ -1010,7 +1010,7 @@ impl OlmMachine {
                 .withheld_session_info
                 .entry(c.room_id.to_owned())
                 .or_default()
-                .insert(c.session_id.to_owned(), event.to_owned());
+                .insert(c.session_id.to_owned(), event.to_owned().into());
         }
     }
 
@@ -2029,7 +2029,7 @@ impl OlmMachine {
                         .store
                         .get_withheld_info(room_id, content.session_id())
                         .await?
-                        .map(|e| e.content.withheld_code());
+                        .map(|e| e.content().withheld_code());
 
                     if withheld_code.is_some() {
                         // Partially withheld, report with a withheld code if we have one.
@@ -2145,7 +2145,7 @@ impl OlmMachine {
                     .store
                     .get_withheld_info(room_id, session_id)
                     .await?
-                    .map(|e| e.content.withheld_code());
+                    .map(|e| e.content().withheld_code());
                 Err(MegolmError::MissingRoomKey(withheld_code))
             }
         }

crates/matrix-sdk-crypto/src/machine/tests/mod.rs

@@ -1065,7 +1065,7 @@ async fn test_withheld_unverified() {
 
     assert_eq!(&withheld_received[0].room_id, room_id);
     assert_matches!(
-        &withheld_received[0].withheld_event.content,
+        &withheld_received[0].withheld_event.content(),
         RoomKeyWithheldContent::MegolmV1AesSha2(MegolmV1AesSha2WithheldContent::Unverified(
             unverified_withheld_content
         ))

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

@@ -66,6 +66,7 @@ macro_rules! cryptostore_integration_tests {
                         room_key_withheld::{
                             CommonWithheldCodeContent, MegolmV1AesSha2WithheldContent,
                             RoomKeyWithheldContent,
+                            RoomKeyWithheldEntry,
                         },
                         room_key_bundle::RoomKeyBundleContent,
                         secret_send::SecretSendContent,
@@ -1138,7 +1139,7 @@ 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<_, _>> = BTreeMap::new();
+                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");
@@ -1163,7 +1164,7 @@ macro_rules! cryptostore_integration_tests {
                 info_list
                     .entry(room_id.to_owned())
                     .or_default()
-                    .insert(session_id_1.to_owned(), event);
+                    .insert(session_id_1.to_owned(), event.into());
 
                 let content = RoomKeyWithheldContent::MegolmV1AesSha2(
                     MegolmV1AesSha2WithheldContent::BlackListed(
@@ -1183,7 +1184,7 @@ macro_rules! cryptostore_integration_tests {
                 info_list
                     .entry(room_id.to_owned())
                     .or_default()
-                    .insert(session_id_2.to_owned(), event);
+                    .insert(session_id_2.to_owned(), event.into());
 
                 let changes = Changes { withheld_session_info: info_list, ..Default::default() };
                 store.save_changes(changes).await.unwrap();
@@ -1192,16 +1193,16 @@ macro_rules! cryptostore_integration_tests {
 
                 assert_matches!(
                     is_withheld, Some(event)
-                    if event.content.algorithm() == EventEncryptionAlgorithm::MegolmV1AesSha2 &&
-                    event.content.withheld_code() == WithheldCode::Unverified
+                    if event.content().algorithm() == EventEncryptionAlgorithm::MegolmV1AesSha2 &&
+                    event.content().withheld_code() == WithheldCode::Unverified
                 );
 
                 let is_withheld = store.get_withheld_info(room_id, session_id_2).await.unwrap();
 
                 assert_matches!(
                     is_withheld, Some(event)
-                    if event.content.algorithm() == EventEncryptionAlgorithm::MegolmV1AesSha2 &&
-                    event.content.withheld_code() == WithheldCode::Blacklisted
+                    if event.content().algorithm() == EventEncryptionAlgorithm::MegolmV1AesSha2 &&
+                    event.content().withheld_code() == WithheldCode::Blacklisted
                 );
 
                 let other_room_id = room_id!("!nQRyiRFuyUhXeaQfiR:example.com");

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

@@ -45,7 +45,7 @@ use crate::{
         OutboundGroupSession, PickledAccount, PickledInboundGroupSession, PickledSession,
         PrivateCrossSigningIdentity, SenderDataType, StaticAccountData,
     },
-    types::events::room_key_withheld::RoomKeyWithheldEvent,
+    types::events::room_key_withheld::RoomKeyWithheldEntry,
 };
 
 fn encode_key_info(info: &SecretInfo) -> String {
@@ -97,7 +97,7 @@ pub struct MemoryStore {
     identities: StdRwLock<HashMap<OwnedUserId, String>>,
     outgoing_key_requests: StdRwLock<HashMap<OwnedTransactionId, GossipRequest>>,
     key_requests_by_info: StdRwLock<HashMap<String, OwnedTransactionId>>,
-    direct_withheld_info: StdRwLock<HashMap<OwnedRoomId, HashMap<String, RoomKeyWithheldEvent>>>,
+    direct_withheld_info: StdRwLock<HashMap<OwnedRoomId, HashMap<String, RoomKeyWithheldEntry>>>,
     custom_values: StdRwLock<HashMap<String, Vec<u8>>>,
     leases: StdRwLock<HashMap<String, (String, Instant)>>,
     secret_inbox: StdRwLock<HashMap<String, Vec<GossippedSecret>>>,
@@ -420,7 +420,7 @@ impl CryptoStore for MemoryStore {
         &self,
         room_id: &RoomId,
         session_id: &str,
-    ) -> Result<Option<RoomKeyWithheldEvent>> {
+    ) -> Result<Option<RoomKeyWithheldEntry>> {
         Ok(self
             .direct_withheld_info
             .read()
@@ -1276,7 +1276,7 @@ mod integration_tests {
             },
             CryptoStore,
         },
-        types::events::room_key_withheld::RoomKeyWithheldEvent,
+        types::events::room_key_withheld::RoomKeyWithheldEntry,
         Account, DeviceData, GossipRequest, GossippedSecret, SecretInfo, Session, UserIdentityData,
     };
 
@@ -1372,7 +1372,7 @@ mod integration_tests {
             &self,
             room_id: &RoomId,
             session_id: &str,
-        ) -> Result<Option<RoomKeyWithheldEvent>, Self::Error> {
+        ) -> Result<Option<RoomKeyWithheldEntry>, Self::Error> {
             self.0.get_withheld_info(room_id, session_id).await
         }
 

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

@@ -77,6 +77,7 @@ use crate::{
         Session, StaticAccountData,
     },
     types::{
+        events::room_key_withheld::{MegolmV1AesSha2WithheldContent, RoomKeyWithheldEntry},
         BackupSecrets, CrossSigningSecrets, MegolmBackupV1Curve25519AesSha2Secrets, RoomKeyExport,
         SecretsBundle,
     },
@@ -1632,12 +1633,11 @@ impl Store {
 
         tracing::Span::current().record("sender_data", tracing::field::debug(&sender_data));
 
-        match sender_data {
+        match &sender_data {
             SenderData::UnknownDevice { .. }
             | SenderData::VerificationViolation(_)
             | SenderData::DeviceInfo { .. } => {
                 warn!("Not accepting a historic room key bundle due to insufficient trust in the sender");
-                Ok(())
             }
             SenderData::SenderUnverified(_) | SenderData::SenderVerified(_) => {
                 let (good, bad): (Vec<_>, Vec<_>) = bundle.room_keys.iter().partition_map(|key| {
@@ -1681,10 +1681,30 @@ impl Store {
                         self.import_sessions_impl(good, None, progress_listener).await?;
                     }
                 }
+            }
+        }
 
-                Ok(())
+        let mut changes = Changes::default();
+        for withheld in &bundle.withheld {
+            if let RoomKeyWithheldContent::MegolmV1AesSha2(
+                MegolmV1AesSha2WithheldContent::BlackListed(c)
+                | MegolmV1AesSha2WithheldContent::Unverified(c)
+                | MegolmV1AesSha2WithheldContent::Unauthorised(c)
+                | MegolmV1AesSha2WithheldContent::Unavailable(c),
+            ) = withheld
+            {
+                changes.withheld_session_info.entry(c.room_id.to_owned()).or_default().insert(
+                    c.session_id.to_owned(),
+                    RoomKeyWithheldEntry::Bundle {
+                        sender: bundle_info.sender_user.clone(),
+                        content: withheld.to_owned(),
+                    },
+                );
             }
         }
+        self.save_changes(changes).await?;
+
+        Ok(())
     }
 }
 
@@ -1717,17 +1737,32 @@ impl matrix_sdk_common::cross_process_lock::TryLock for LockableCryptoStore {
 mod tests {
     use std::pin::pin;
 
+    use assert_matches2::assert_matches;
     use futures_util::StreamExt;
     use insta::{_macro_support::Content, assert_json_snapshot, internals::ContentPath};
     use matrix_sdk_test::async_test;
-    use ruma::{device_id, room_id, user_id, RoomId};
+    use ruma::{
+        device_id,
+        events::room::{EncryptedFileInit, JsonWebKeyInit},
+        owned_mxc_uri, room_id,
+        serde::Base64,
+        user_id, RoomId,
+    };
     use vodozemac::megolm::SessionKey;
 
     use crate::{
         machine::test_helpers::get_machine_pair,
         olm::{InboundGroupSession, SenderData},
-        store::types::DehydratedDeviceKey,
-        types::EventEncryptionAlgorithm,
+        store::types::{DehydratedDeviceKey, StoredRoomKeyBundleData},
+        types::{
+            events::{
+                room_key_bundle::RoomKeyBundleContent,
+                room_key_withheld::{
+                    MegolmV1AesSha2WithheldContent, RoomKeyWithheldContent, RoomKeyWithheldEntry,
+                },
+            },
+            EventEncryptionAlgorithm,
+        },
         OlmMachine,
     };
 
@@ -1969,6 +2004,93 @@ mod tests {
         });
     }
 
+    #[async_test]
+    async fn test_receive_room_key_bundle() {
+        let alice = OlmMachine::new(user_id!("@a:s.co"), device_id!("ALICE")).await;
+        let alice_key = alice.identity_keys().curve25519;
+        let bob = OlmMachine::new(user_id!("@b:s.co"), device_id!("BOB")).await;
+
+        let room_id = room_id!("!room1:localhost");
+
+        let session_key1 = "AgAAAAC2XHVzsMBKs4QCRElJ92CJKyGtknCSC8HY7cQ7UYwndMKLQAejXLh5UA0l6s736mgctcUMNvELScUWrObdflrHo+vth/gWreXOaCnaSxmyjjKErQwyIYTkUfqbHy40RJfEesLwnN23on9XAkch/iy8R2+Jz7B8zfG01f2Ow2SxPQFnAndcO1ZSD2GmXgedy6n4B20MWI1jGP2wiexOWbFSya8DO/VxC9m5+/mF+WwYqdpKn9g4Y05Yw4uz7cdjTc3rXm7xK+8E7hI//5QD1nHPvuKYbjjM9u2JSL+Bzp61Cw";
+        let session_key2 = "AgAAAAC1BXreFTUQQSBGekTEuYxhdytRKyv4JgDGcG+VOBYdPNGgs807SdibCGJky4lJ3I+7ZDGHoUzZPZP/4ogGu4kxni0PWdtWuN7+5zsuamgoFF/BkaGeUUGv6kgIkx8pyPpM5SASTUEP9bN2loDSpUPYwfiIqz74DgC4WQ4435sTBctYvKz8n+TDJwdLXpyT6zKljuqADAioud+s/iqx9LYn9HpbBfezZcvbg67GtE113pLrvde3IcPI5s6dNHK2onGO2B2eoaobcen18bbEDnlUGPeIivArLya7Da6us14jBQ";
+
+        let sessions = [
+            create_inbound_group_session_with_visibility(
+                &alice,
+                room_id,
+                &SessionKey::from_base64(session_key1).unwrap(),
+                true,
+            ),
+            create_inbound_group_session_with_visibility(
+                &alice,
+                room_id,
+                &SessionKey::from_base64(session_key2).unwrap(),
+                false,
+            ),
+        ];
+
+        alice.store().save_inbound_group_sessions(&sessions).await.unwrap();
+        let bundle = alice.store().build_room_key_bundle(room_id).await.unwrap();
+
+        bob.store()
+            .receive_room_key_bundle(
+                &StoredRoomKeyBundleData {
+                    sender_user: alice.user_id().to_owned(),
+                    sender_key: alice_key,
+                    sender_data: SenderData::sender_verified(
+                        alice.user_id(),
+                        device_id!("ALICE"),
+                        alice.identity_keys().ed25519,
+                    ),
+
+                    bundle_data: RoomKeyBundleContent {
+                        room_id: room_id.to_owned(),
+                        // This isn't used at all in the method call, so we can fill it with
+                        // garbage.
+                        file: EncryptedFileInit {
+                            url: owned_mxc_uri!("mxc://example.com/0"),
+                            key: JsonWebKeyInit {
+                                kty: "oct".to_owned(),
+                                key_ops: vec!["encrypt".to_owned(), "decrypt".to_owned()],
+                                alg: "A256CTR.".to_owned(),
+                                k: Base64::new(vec![0u8; 128]),
+                                ext: true,
+                            }
+                            .into(),
+                            iv: Base64::new(vec![0u8; 128]),
+                            hashes: vec![("sha256".to_owned(), Base64::new(vec![0u8; 128]))]
+                                .into_iter()
+                                .collect(),
+                            v: "v2".to_owned(),
+                        }
+                        .into(),
+                    },
+                },
+                bundle,
+                |_, _| {},
+            )
+            .await
+            .unwrap();
+
+        // The room key should be imported successfully
+        let imported_sessions =
+            bob.store().get_inbound_group_sessions_by_room_id(room_id).await.unwrap();
+
+        assert_eq!(imported_sessions.len(), 1);
+        assert_eq!(imported_sessions[0].room_id(), room_id);
+
+        assert_matches!(
+            bob.store().get_withheld_info(room_id, sessions[1].session_id()).await.unwrap(),
+            Some(RoomKeyWithheldEntry::Bundle {
+                content: RoomKeyWithheldContent::MegolmV1AesSha2(
+                    MegolmV1AesSha2WithheldContent::Unauthorised(_)
+                ),
+                ..
+            })
+        );
+    }
+
     /// Create an inbound Megolm session for the given room.
     ///
     /// `olm_machine` is used to set the `sender_key` and `signing_key`

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

@@ -35,7 +35,7 @@ use crate::{
         InboundGroupSession, OlmMessageHash, OutboundGroupSession, PrivateCrossSigningIdentity,
         SenderDataType, Session,
     },
-    types::events::room_key_withheld::RoomKeyWithheldEvent,
+    types::events::room_key_withheld::RoomKeyWithheldEntry,
     Account, DeviceData, GossipRequest, GossippedSecret, SecretInfo, UserIdentityData,
 };
 
@@ -116,7 +116,7 @@ pub trait CryptoStore: AsyncTraitDeps {
         &self,
         room_id: &RoomId,
         session_id: &str,
-    ) -> Result<Option<RoomKeyWithheldEvent>, Self::Error>;
+    ) -> Result<Option<RoomKeyWithheldEntry>, Self::Error>;
 
     /// Get all the inbound group sessions we have stored.
     async fn get_inbound_group_sessions(&self) -> Result<Vec<InboundGroupSession>, Self::Error>;
@@ -588,7 +588,7 @@ impl<T: CryptoStore> CryptoStore for EraseCryptoStoreError<T> {
         &self,
         room_id: &RoomId,
         session_id: &str,
-    ) -> Result<Option<RoomKeyWithheldEvent>, Self::Error> {
+    ) -> Result<Option<RoomKeyWithheldEntry>, Self::Error> {
         self.0.get_withheld_info(room_id, session_id).await.map_err(Into::into)
     }