feat: Add RoomKeyWithheldEntry to wrap to-device and bundle payloads.

Author: kaylendog

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/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)
     }
 

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

@@ -34,7 +34,7 @@ use crate::{
         SenderData,
     },
     types::{
-        events::{room_key_bundle::RoomKeyBundleContent, room_key_withheld::RoomKeyWithheldEvent},
+        events::{room_key_bundle::RoomKeyBundleContent, room_key_withheld::RoomKeyWithheldEntry},
         EventEncryptionAlgorithm,
     },
     Account, Device, DeviceData, GossippedSecret, Session, UserIdentity, UserIdentityData,
@@ -75,7 +75,7 @@ pub struct Changes {
     pub identities: IdentityChanges,
     pub devices: DeviceChanges,
     /// Stores when a `m.room_key.withheld` is received
-    pub withheld_session_info: BTreeMap<OwnedRoomId, BTreeMap<String, RoomKeyWithheldEvent>>,
+    pub withheld_session_info: BTreeMap<OwnedRoomId, BTreeMap<String, RoomKeyWithheldEntry>>,
     pub room_settings: HashMap<OwnedRoomId, RoomSettings>,
     pub secrets: Vec<GossippedSecret>,
     pub next_batch_token: Option<String>,
@@ -486,7 +486,7 @@ pub struct RoomKeyWithheldInfo {
 
     /// The `m.room_key.withheld` event that notified us that the key is being
     /// withheld.
-    pub withheld_event: RoomKeyWithheldEvent,
+    pub withheld_event: RoomKeyWithheldEntry,
 }
 
 /// Information about a received historic room key bundle.

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

@@ -25,6 +25,39 @@ use vodozemac::Curve25519PublicKey;
 use super::{EventType, ToDeviceEvent};
 use crate::types::{deserialize_curve_key, serialize_curve_key, EventEncryptionAlgorithm};
 
+/// 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)]
+#[serde(untagged)]
+pub enum RoomKeyWithheldEntry {
+    /// A to-device event containing withheld room key information.
+    ToDevice(RoomKeyWithheldEvent),
+    /// Content held within a room key bundle.
+    Bundle(RoomKeyWithheldContent),
+}
+
+impl RoomKeyWithheldEntry {
+    /// Returns a reference to the underlying `RoomKeyWithheldContent`.
+    pub fn content(&self) -> &RoomKeyWithheldContent {
+        match self {
+            RoomKeyWithheldEntry::ToDevice(ev) => &ev.content,
+            RoomKeyWithheldEntry::Bundle(content) => content,
+        }
+    }
+}
+
+impl From<RoomKeyWithheldEvent> for RoomKeyWithheldEntry {
+    fn from(value: RoomKeyWithheldEvent) -> Self {
+        Self::ToDevice(value)
+    }
+}
+
+impl From<RoomKeyWithheldContent> for RoomKeyWithheldEntry {
+    fn from(value: RoomKeyWithheldContent) -> Self {
+        Self::Bundle(value)
+    }
+}
+
 /// The `m.room_key_request` to-device event.
 pub type RoomKeyWithheldEvent = ToDeviceEvent<RoomKeyWithheldContent>;
 

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

@@ -35,7 +35,7 @@ use matrix_sdk_crypto::{
         },
         CryptoStore, CryptoStoreError,
     },
-    types::events::room_key_withheld::RoomKeyWithheldEvent,
+    types::events::room_key_withheld::RoomKeyWithheldEntry,
     vodozemac::base64_encode,
     Account, DeviceData, GossipRequest, GossippedSecret, SecretInfo, TrackedUser, UserIdentityData,
 };
@@ -1386,7 +1386,7 @@ impl_crypto_store! {
         &self,
         room_id: &RoomId,
         session_id: &str,
-    ) -> Result<Option<RoomKeyWithheldEvent>> {
+    ) -> Result<Option<RoomKeyWithheldEntry>> {
         let key = self.serializer.encode_key(keys::DIRECT_WITHHELD_INFO, (session_id, room_id));
         if let Some(pickle) = self
             .inner

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

@@ -33,7 +33,7 @@ use matrix_sdk_crypto::{
         },
         CryptoStore,
     },
-    types::events::room_key_withheld::RoomKeyWithheldEvent,
+    types::events::room_key_withheld::RoomKeyWithheldEntry,
     Account, DeviceData, GossipRequest, GossippedSecret, SecretInfo, TrackedUser, UserIdentityData,
 };
 use matrix_sdk_store_encryption::StoreCipher;
@@ -1379,7 +1379,7 @@ impl CryptoStore for SqliteCryptoStore {
         &self,
         room_id: &RoomId,
         session_id: &str,
-    ) -> Result<Option<RoomKeyWithheldEvent>> {
+    ) -> Result<Option<RoomKeyWithheldEntry>> {
         let room_id = self.encode_key("direct_withheld_info", room_id);
         let session_id = self.encode_key("direct_withheld_info", session_id);
 
@@ -1388,7 +1388,7 @@ impl CryptoStore for SqliteCryptoStore {
             .get_direct_withheld_info(session_id, room_id)
             .await?
             .map(|value| {
-                let info = self.deserialize_json::<RoomKeyWithheldEvent>(&value)?;
+                let info = self.deserialize_json::<RoomKeyWithheldEntry>(&value)?;
                 Ok(info)
             })
             .transpose()
@@ -1894,7 +1894,7 @@ mod tests {
             .expect("This session should be withheld")
             .unwrap();
 
-        assert_eq!(withheld_info.content.withheld_code(), WithheldCode::Unverified);
+        assert_eq!(withheld_info.content().withheld_code(), WithheldCode::Unverified);
 
         let backup_keys = database.load_backup_keys().await.expect("backup key should be cached");
         assert_eq!(backup_keys.backup_version.unwrap(), "6");