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

Author: kaylendog

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
+
+- Improve 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 `RoomKeyWithheldEntry` in each `room-id`-`session-id` entry.
+  - [**breaking**] `CryptoStore::get_withheld_info` now returns `Result<Option<RoomKeyWithheldEntry>>`. 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

@@ -1003,7 +1003,9 @@ impl OlmMachine {
 
         if let RoomKeyWithheldContent::MegolmV1AesSha2(
             MegolmV1AesSha2WithheldContent::BlackListed(c)
-            | MegolmV1AesSha2WithheldContent::Unverified(c),
+            | MegolmV1AesSha2WithheldContent::Unverified(c)
+            | MegolmV1AesSha2WithheldContent::Unauthorised(c)
+            | MegolmV1AesSha2WithheldContent::Unavailable(c),
         ) = &event.content
         {
             changes

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

@@ -59,7 +59,7 @@ use serde::{de::DeserializeOwned, Serialize};
 use thiserror::Error;
 use tokio::sync::{Mutex, Notify, OwnedRwLockWriteGuard, RwLock};
 use tokio_stream::wrappers::errors::BroadcastStreamRecvError;
-use tracing::{error, info, instrument, trace, warn};
+use tracing::{info, instrument, trace, warn};
 use types::{RoomKeyBundleInfo, StoredRoomKeyBundleData};
 use vodozemac::{megolm::SessionOrdering, Curve25519PublicKey};
 
@@ -76,6 +76,7 @@ use crate::{
         Account, ExportedRoomKey, InboundGroupSession, PrivateCrossSigningIdentity, SenderData,
         Session, StaticAccountData,
     },
+    store::types::RoomKeyWithheldEntry,
     types::{
         BackupSecrets, CrossSigningSecrets, MegolmBackupV1Curve25519AesSha2Secrets, RoomKeyExport,
         SecretsBundle,
@@ -1645,6 +1646,7 @@ impl Store {
         }
 
         self.import_room_key_bundle_sessions(bundle_info, &bundle, progress_listener).await?;
+        self.import_room_key_bundle_withheld_info(bundle_info, &bundle).await?;
 
         Ok(())
     }
@@ -1700,6 +1702,43 @@ impl Store {
         Ok(())
     }
 
+    async fn import_room_key_bundle_withheld_info(
+        &self,
+        bundle_info: &StoredRoomKeyBundleData,
+        bundle: &RoomKeyBundle,
+    ) -> Result<(), CryptoStoreError> {
+        let mut changes = Changes::default();
+        for withheld in &bundle.withheld {
+            let (room_id, session_id) = match withheld {
+                RoomKeyWithheldContent::MegolmV1AesSha2(c) => match (c.room_id(), c.session_id()) {
+                    (Some(room_id), Some(session_id)) => (room_id, session_id),
+                    _ => continue,
+                },
+                #[cfg(feature = "experimental-algorithms")]
+                RoomKeyWithheldContent::MegolmV2AesSha2(c) => match (c.room_id(), c.session_id()) {
+                    (Some(room_id), Some(session_id)) => (room_id, session_id),
+                    _ => continue,
+                },
+                RoomKeyWithheldContent::Unknown(_) => continue,
+            };
+
+            if room_id != bundle_info.bundle_data.room_id {
+                trace!("Ignoring withheld info for incorrect room {} in bundle", room_id);
+                continue;
+            }
+
+            changes.withheld_session_info.entry(room_id.to_owned()).or_default().insert(
+                session_id.to_owned(),
+                RoomKeyWithheldEntry {
+                    sender: bundle_info.sender_user.clone(),
+                    content: withheld.to_owned(),
+                },
+            );
+        }
+        self.save_changes(changes).await?;
+
+        Ok(())
+    }
 }
 
 impl Deref for Store {
@@ -1735,16 +1774,25 @@ mod tests {
     use futures_util::StreamExt;
     use insta::{_macro_support::Content, assert_json_snapshot, internals::ContentPath};
     use matrix_sdk_test::async_test;
-    use ruma::{device_id, owned_device_id, room_id, user_id, RoomId};
+    use ruma::{
+        device_id,
+        events::room::{EncryptedFileInit, JsonWebKeyInit},
+        owned_device_id, owned_mxc_uri, room_id,
+        serde::Base64,
+        user_id, RoomId,
+    };
     use serde_json::json;
     use vodozemac::megolm::SessionKey;
 
     use crate::{
         machine::test_helpers::get_machine_pair,
         olm::{InboundGroupSession, SenderData},
-        store::types::{DehydratedDeviceKey, RoomKeyWithheldEntry},
+        store::types::{DehydratedDeviceKey, RoomKeyWithheldEntry, StoredRoomKeyBundleData},
         types::{
-            events::room_key_withheld::{MegolmV1AesSha2WithheldContent, RoomKeyWithheldContent},
+            events::{
+                room_key_bundle::RoomKeyBundleContent,
+                room_key_withheld::{MegolmV1AesSha2WithheldContent, RoomKeyWithheldContent},
+            },
             EventEncryptionAlgorithm,
         },
         OlmMachine,
@@ -2001,6 +2049,102 @@ 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()
+                .expect("Withheld info should be present in the store."),
+            RoomKeyWithheldEntry {
+                #[cfg(not(feature = "experimental-algorithms"))]
+                content: RoomKeyWithheldContent::MegolmV1AesSha2(
+                    MegolmV1AesSha2WithheldContent::Unauthorised(_)
+                ),
+                #[cfg(feature = "experimental-algorithms")]
+                content: RoomKeyWithheldContent::MegolmV2AesSha2(
+                    MegolmV1AesSha2WithheldContent::Unauthorised(_)
+                ),
+                ..
+            }
+        );
+    }
+
     /// Tests that the new store format introduced in [#5737][#5737] does not
     /// conflict with items already in the store that were serialised with the
     /// older format.

crates/matrix-sdk-crypto/src/types/events/room_key_withheld.rs

@@ -17,7 +17,9 @@
 use std::collections::BTreeMap;
 
 use matrix_sdk_common::deserialized_responses::WithheldCode;
-use ruma::{events::AnyToDeviceEventContent, serde::JsonCastable, OwnedDeviceId, OwnedRoomId};
+use ruma::{
+    events::AnyToDeviceEventContent, serde::JsonCastable, OwnedDeviceId, OwnedRoomId, RoomId,
+};
 use serde::{Deserialize, Serialize};
 use serde_json::Value;
 use vodozemac::Curve25519PublicKey;
@@ -223,6 +225,28 @@ impl CommonWithheldCodeContent {
 }
 
 impl MegolmV1AesSha2WithheldContent {
+    /// Get the session ID for this content, if available.
+    pub fn session_id(&self) -> Option<&str> {
+        match self {
+            MegolmV1AesSha2WithheldContent::BlackListed(content)
+            | MegolmV1AesSha2WithheldContent::Unverified(content)
+            | MegolmV1AesSha2WithheldContent::Unauthorised(content)
+            | MegolmV1AesSha2WithheldContent::Unavailable(content) => Some(&content.session_id),
+            MegolmV1AesSha2WithheldContent::NoOlm(_) => None,
+        }
+    }
+
+    /// Get the room ID for this content, if available.
+    pub fn room_id(&self) -> Option<&RoomId> {
+        match self {
+            MegolmV1AesSha2WithheldContent::BlackListed(content)
+            | MegolmV1AesSha2WithheldContent::Unverified(content)
+            | MegolmV1AesSha2WithheldContent::Unauthorised(content)
+            | MegolmV1AesSha2WithheldContent::Unavailable(content) => Some(&content.room_id),
+            MegolmV1AesSha2WithheldContent::NoOlm(_) => None,
+        }
+    }
+
     /// Get the withheld code for this content
     pub fn withheld_code(&self) -> WithheldCode {
         match self {