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

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

@@ -998,12 +998,14 @@ impl OlmMachine {
         Ok(())
     }
 
-    pub(crate) fn add_withheld_info(&self, changes: &mut Changes, event: &RoomKeyWithheldEvent) {
+    fn add_withheld_info(&self, changes: &mut Changes, event: &RoomKeyWithheldEvent) {
         debug!(?event.content, "Processing `m.room_key.withheld` event");
 
         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};
 
@@ -78,8 +78,8 @@ use crate::{
     },
     store::types::RoomKeyWithheldEntry,
     types::{
-        events::room_key_withheld::MegolmV1AesSha2WithheldContent, BackupSecrets,
-        CrossSigningSecrets, MegolmBackupV1Curve25519AesSha2Secrets, RoomKeyExport, SecretsBundle,
+        BackupSecrets, CrossSigningSecrets, MegolmBackupV1Curve25519AesSha2Secrets, RoomKeyExport,
+        SecretsBundle,
     },
     verification::VerificationMachine,
     CrossSigningStatus, OwnUserIdentityData, RoomKeyImportResult,
@@ -1633,74 +1633,108 @@ impl Store {
 
         tracing::Span::current().record("sender_data", tracing::field::debug(&sender_data));
 
-        match &sender_data {
+        if matches!(
+            &sender_data,
             SenderData::UnknownDevice { .. }
-            | SenderData::VerificationViolation(_)
-            | SenderData::DeviceInfo { .. } => {
-                warn!("Not accepting a historic room key bundle due to insufficient trust in the sender");
+                | SenderData::VerificationViolation(_)
+                | SenderData::DeviceInfo { .. }
+        ) {
+            warn!(
+                "Not accepting a historic room key bundle due to insufficient trust in the sender"
+            );
+            return Ok(());
+        }
+
+        self.import_room_key_bundle_sessions(bundle_info, &bundle, progress_listener).await?;
+        self.import_room_key_bundle_withheld_info(bundle_info, &bundle).await?;
+
+        Ok(())
+    }
+
+    async fn import_room_key_bundle_sessions(
+        &self,
+        bundle_info: &StoredRoomKeyBundleData,
+        bundle: &RoomKeyBundle,
+        progress_listener: impl Fn(usize, usize),
+    ) -> Result<(), CryptoStoreError> {
+        let (good, bad): (Vec<_>, Vec<_>) = bundle.room_keys.iter().partition_map(|key| {
+            if key.room_id != bundle_info.bundle_data.room_id {
+                trace!("Ignoring key for incorrect room {} in bundle", key.room_id);
+                Either::Right(key)
+            } else {
+                Either::Left(key)
             }
-            SenderData::SenderUnverified(_) | SenderData::SenderVerified(_) => {
-                let (good, bad): (Vec<_>, Vec<_>) = bundle.room_keys.iter().partition_map(|key| {
-                    if key.room_id != bundle_info.bundle_data.room_id {
-                        trace!("Ignoring key for incorrect room {} in bundle", key.room_id);
-                        Either::Right(key)
-                    } else {
-                        Either::Left(key)
-                    }
-                });
+        });
 
-                match (bad.is_empty(), good.is_empty()) {
-                    // Case 1: Completely empty bundle.
-                    (true, true) => {
-                        warn!("Received a completely empty room key bundle");
-                    }
+        match (bad.is_empty(), good.is_empty()) {
+            // Case 1: Completely empty bundle.
+            (true, true) => {
+                warn!("Received a completely empty room key bundle");
+            }
 
-                    // Case 2: A bundle for the wrong room.
-                    (false, true) => {
-                        let bad_keys: Vec<_> =
-                            bad.iter().map(|&key| (&key.room_id, &key.session_id)).collect();
+            // Case 2: A bundle for the wrong room.
+            (false, true) => {
+                let bad_keys: Vec<_> =
+                    bad.iter().map(|&key| (&key.room_id, &key.session_id)).collect();
 
-                        warn!(
+                warn!(
                     ?bad_keys,
                     "Received a room key bundle for the wrong room, ignoring all room keys from the bundle"
                 );
-                    }
+            }
 
-                    // Case 3: A bundle containing useful room keys.
-                    (_, false) => {
-                        // We have at least some good keys, if we also have some bad ones let's
-                        // mention that here.
-                        if !bad.is_empty() {
-                            warn!(
-                                bad_key_count = bad.len(),
-                                "The room key bundle contained some room keys \
+            // Case 3: A bundle containing useful room keys.
+            (_, false) => {
+                // We have at least some good keys, if we also have some bad ones let's
+                // mention that here.
+                if !bad.is_empty() {
+                    warn!(
+                        bad_key_count = bad.len(),
+                        "The room key bundle contained some room keys \
                          that were meant for a different room"
-                            );
-                        }
-
-                        self.import_sessions_impl(good, None, progress_listener).await?;
-                    }
+                    );
                 }
+
+                self.import_sessions_impl(good, None, progress_listener).await?;
             }
         }
 
+        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 {
-            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 {
-                        sender: bundle_info.sender_user.clone(),
-                        content: withheld.to_owned(),
-                    },
-                );
+            let (room_id, session_id) = match withheld {
+                #[cfg(not(feature = "experimental-algorithms"))]
+                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,
+                },
+                _ => 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?;
 
@@ -1982,6 +2016,17 @@ mod tests {
         // We sort the sessions in the bundle, so that the snapshot is stable.
         bundle.room_keys.sort_by_key(|session| session.session_id.clone());
 
+        // We substitute the algorithm, since this changes based on feature flags.
+        let algorithm = if cfg!(feature = "experimental-algorithms") {
+            "m.megolm.v2.aes-sha2"
+        } else {
+            "m.megolm.v1.aes-sha2"
+        };
+        let map_algorithm = move |value: Content, _path: ContentPath<'_>| {
+            assert_eq!(value.as_str().unwrap(), algorithm);
+            "[algorithm]"
+        };
+
         // We also substitute alice's keys in the snapshot with placeholders
         let alice_curve_key = alice.identity_keys().curve25519.to_base64();
         let map_alice_curve_key = move |value: Content, _path: ContentPath<'_>| {
@@ -1996,6 +2041,8 @@ mod tests {
 
         insta::with_settings!({ sort_maps => true }, {
             assert_json_snapshot!(bundle, {
+                ".withheld[].algorithm" => insta::dynamic_redaction(map_algorithm),
+                ".room_keys[].algorithm" => insta::dynamic_redaction(map_algorithm),
                 ".room_keys[].sender_key" => insta::dynamic_redaction(map_alice_curve_key.clone()),
                 ".withheld[].sender_key" => insta::dynamic_redaction(map_alice_curve_key),
                 ".room_keys[].sender_claimed_keys.ed25519" => insta::dynamic_redaction(map_alice_ed25519_key),
@@ -2080,13 +2127,22 @@ mod tests {
         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 {
+            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(_)
+                ),
                 ..
-            })
+            }
         );
     }
 
@@ -2133,6 +2189,7 @@ mod tests {
         );
         assert_eq!(withheld_content.from_device, Some(owned_device_id!("ALICE")));
     }
+
     /// Create an inbound Megolm session for the given room.
     ///
     /// `olm_machine` is used to set the `sender_key` and `signing_key`
@@ -2150,7 +2207,10 @@ mod tests {
             room_id,
             session_key,
             SenderData::unknown(),
+            #[cfg(not(feature = "experimental-algorithms"))]
             EventEncryptionAlgorithm::MegolmV1AesSha2,
+            #[cfg(feature = "experimental-algorithms")]
+            EventEncryptionAlgorithm::MegolmV2AesSha2,
             None,
             shared_history,
         )

crates/matrix-sdk-crypto/src/store/snapshots/matrix_sdk_crypto__store__tests__build_room_key_bundle.snap

@@ -5,7 +5,7 @@ expression: bundle
 {
   "room_keys": [
     {
-      "algorithm": "m.megolm.v1.aes-sha2",
+      "algorithm": "[algorithm]",
       "room_id": "!room1:localhost",
       "sender_key": "[alice curve key]",
       "session_id": "AWcCd1w7VlIPYaZeB53LqfgHbQxYjWMY/bCJ7E5ZsVI",
@@ -15,7 +15,7 @@ expression: bundle
       }
     },
     {
-      "algorithm": "m.megolm.v1.aes-sha2",
+      "algorithm": "[algorithm]",
       "room_id": "!room1:localhost",
       "sender_key": "[alice curve key]",
       "session_id": "y1i8rPyf5MMnB0tenJPrMqWO6oAMCKi536z+KrH0tic",
@@ -27,7 +27,7 @@ expression: bundle
   ],
   "withheld": [
     {
-      "algorithm": "m.megolm.v1.aes-sha2",
+      "algorithm": "[algorithm]",
       "code": "m.unauthorised",
       "from_device": "BOB",
       "reason": "You are not authorised to read the message.",

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

@@ -487,18 +487,24 @@ pub struct RoomKeyWithheldInfo {
     /// 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.
+    /// The withheld entry from a `m.room_key.withheld` event or [MSC4268] room
+    /// key bundle.
+    ///
+    /// [MSC4268]: https://github.com/matrix-org/matrix-spec-proposals/pull/4362
     pub withheld_event: RoomKeyWithheldEntry,
 }
 
 /// Represents an entry for a withheld room key event, which can be either a
 /// to-device event or a bundle entry.
 #[derive(Clone, Debug, Serialize, Deserialize)]
 pub struct RoomKeyWithheldEntry {
-    /// The ID of the user who sent the bundle.
+    /// The user ID responsible for this entry, either from a
+    /// `m.room_key.withheld` to-device event or an [MSC4268] room key bundle.
+    ///
+    /// [MSC4268]: https://github.com/matrix-org/matrix-spec-proposals/pull/4362
     pub sender: OwnedUserId,
-    /// The contents of a single withheld entry in the bundle.
+    /// The content of the entry, which provides details about the reason the
+    /// key was withheld.
     pub content: RoomKeyWithheldContent,
 }
 

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 {