feat: Use struct for RoomKeyWithheldEntry, move to store.

kaylendog · kaylendog · commit 24dcbf4b763e · 2025-10-02T15:23:08.000+01:00
diff --git a/crates/matrix-sdk-crypto/src/machine/mod.rs b/crates/matrix-sdk-crypto/src/machine/mod.rs
@@ -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))
             }
         }
diff --git a/crates/matrix-sdk-crypto/src/machine/tests/mod.rs b/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
         ))
diff --git a/crates/matrix-sdk-crypto/src/store/integration_tests.rs b/crates/matrix-sdk-crypto/src/store/integration_tests.rs
@@ -53,7 +53,8 @@ macro_rules! cryptostore_integration_tests {
                 store::{
                     types::{
                         BackupDecryptionKey, Changes, DehydratedDeviceKey, DeviceChanges,
-                        IdentityChanges, PendingChanges, StoredRoomKeyBundleData, RoomSettings,
+                        IdentityChanges, PendingChanges, StoredRoomKeyBundleData, RoomKeyWithheldEntry,
+                        RoomSettings
                     },
                     CryptoStore, GossipRequest,
                 },
@@ -66,7 +67,6 @@ macro_rules! cryptostore_integration_tests {
                         room_key_withheld::{
                             CommonWithheldCodeContent, MegolmV1AesSha2WithheldContent,
                             RoomKeyWithheldContent,
-                            RoomKeyWithheldEntry,
                         },
                         room_key_bundle::RoomKeyBundleContent,
                         secret_send::SecretSendContent,
@@ -1193,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");
diff --git a/crates/matrix-sdk-crypto/src/store/memorystore.rs b/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::RoomKeyWithheldEntry,
+    store::types::RoomKeyWithheldEntry,
 };
 
 fn encode_key_info(info: &SecretInfo) -> String {
@@ -1272,11 +1272,10 @@ mod integration_tests {
         store::{
             types::{
                 BackupKeys, Changes, DehydratedDeviceKey, PendingChanges, RoomKeyCounts,
-                RoomSettings, StoredRoomKeyBundleData, TrackedUser,
+                RoomKeyWithheldEntry, RoomSettings, StoredRoomKeyBundleData, TrackedUser,
             },
             CryptoStore,
         },
-        types::events::room_key_withheld::RoomKeyWithheldEntry,
         Account, DeviceData, GossipRequest, GossippedSecret, SecretInfo, Session, UserIdentityData,
     };
 
diff --git a/crates/matrix-sdk-crypto/src/store/mod.rs b/crates/matrix-sdk-crypto/src/store/mod.rs
@@ -76,10 +76,10 @@ use crate::{
         Account, ExportedRoomKey, InboundGroupSession, PrivateCrossSigningIdentity, SenderData,
         Session, StaticAccountData,
     },
+    store::types::RoomKeyWithheldEntry,
     types::{
-        events::room_key_withheld::{MegolmV1AesSha2WithheldContent, RoomKeyWithheldEntry},
-        BackupSecrets, CrossSigningSecrets, MegolmBackupV1Curve25519AesSha2Secrets, RoomKeyExport,
-        SecretsBundle,
+        events::room_key_withheld::MegolmV1AesSha2WithheldContent, BackupSecrets,
+        CrossSigningSecrets, MegolmBackupV1Curve25519AesSha2Secrets, RoomKeyExport, SecretsBundle,
     },
     verification::VerificationMachine,
     CrossSigningStatus, OwnUserIdentityData, RoomKeyImportResult,
@@ -1695,7 +1695,7 @@ impl Store {
             {
                 changes.withheld_session_info.entry(c.room_id.to_owned()).or_default().insert(
                     c.session_id.to_owned(),
-                    RoomKeyWithheldEntry::Bundle {
+                    RoomKeyWithheldEntry {
                         sender: bundle_info.sender_user.clone(),
                         content: withheld.to_owned(),
                     },
@@ -1753,13 +1753,11 @@ mod tests {
     use crate::{
         machine::test_helpers::get_machine_pair,
         olm::{InboundGroupSession, SenderData},
-        store::types::{DehydratedDeviceKey, StoredRoomKeyBundleData},
+        store::types::{DehydratedDeviceKey, RoomKeyWithheldEntry, StoredRoomKeyBundleData},
         types::{
             events::{
                 room_key_bundle::RoomKeyBundleContent,
-                room_key_withheld::{
-                    MegolmV1AesSha2WithheldContent, RoomKeyWithheldContent, RoomKeyWithheldEntry,
-                },
+                room_key_withheld::{MegolmV1AesSha2WithheldContent, RoomKeyWithheldContent},
             },
             EventEncryptionAlgorithm,
         },
@@ -2082,7 +2080,7 @@ mod tests {
 
         assert_matches!(
             bob.store().get_withheld_info(room_id, sessions[1].session_id()).await.unwrap(),
-            Some(RoomKeyWithheldEntry::Bundle {
+            Some(RoomKeyWithheldEntry {
                 content: RoomKeyWithheldContent::MegolmV1AesSha2(
                     MegolmV1AesSha2WithheldContent::Unauthorised(_)
                 ),
diff --git a/crates/matrix-sdk-crypto/src/store/traits.rs b/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::RoomKeyWithheldEntry,
+    store::types::RoomKeyWithheldEntry,
     Account, DeviceData, GossipRequest, GossippedSecret, SecretInfo, UserIdentityData,
 };
 
diff --git a/crates/matrix-sdk-crypto/src/store/types.rs b/crates/matrix-sdk-crypto/src/store/types.rs
@@ -34,7 +34,10 @@ use crate::{
         SenderData,
     },
     types::{
-        events::{room_key_bundle::RoomKeyBundleContent, room_key_withheld::RoomKeyWithheldEntry},
+        events::{
+            room_key_bundle::RoomKeyBundleContent,
+            room_key_withheld::{RoomKeyWithheldContent, RoomKeyWithheldEvent},
+        },
         EventEncryptionAlgorithm,
     },
     Account, Device, DeviceData, GossippedSecret, Session, UserIdentity, UserIdentityData,
@@ -489,6 +492,22 @@ pub struct RoomKeyWithheldInfo {
     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.
+    pub sender: OwnedUserId,
+    /// The contents of a single withheld entry in the bundle.
+    pub content: RoomKeyWithheldContent,
+}
+
+impl From<RoomKeyWithheldEvent> for RoomKeyWithheldEntry {
+    fn from(value: RoomKeyWithheldEvent) -> Self {
+        Self { sender: value.sender, content: value.content }
+    }
+}
+
 /// Information about a received historic room key bundle.
 ///
 /// This struct contains information needed to uniquely identify a room key
diff --git a/crates/matrix-sdk-crypto/src/types/events/room_key_withheld.rs b/crates/matrix-sdk-crypto/src/types/events/room_key_withheld.rs
@@ -28,46 +28,6 @@ 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 {
-        /// The ID of the user who sent the bundle.
-        sender: OwnedUserId,
-        /// The contents of a single withheld entry in the bundle.
-        content: 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,
-        }
-    }
-
-    /// Returns the sender's user ID, if available.
-    pub fn sender(&self) -> &UserId {
-        match self {
-            RoomKeyWithheldEntry::ToDevice(ev) => &ev.sender,
-            RoomKeyWithheldEntry::Bundle { sender, .. } => sender,
-        }
-    }
-}
-
-impl From<RoomKeyWithheldEvent> for RoomKeyWithheldEntry {
-    fn from(value: RoomKeyWithheldEvent) -> Self {
-        Self::ToDevice(value)
-    }
-}
-
 /// The `m.room_key_request` to-device event.
 pub type RoomKeyWithheldEvent = ToDeviceEvent<RoomKeyWithheldContent>;
 
diff --git a/crates/matrix-sdk-indexeddb/src/crypto_store/mod.rs b/crates/matrix-sdk-indexeddb/src/crypto_store/mod.rs
@@ -30,12 +30,11 @@ use matrix_sdk_crypto::{
     },
     store::{
         types::{
-            BackupKeys, Changes, DehydratedDeviceKey, PendingChanges, RoomKeyCounts, RoomSettings,
-            StoredRoomKeyBundleData,
+            BackupKeys, Changes, DehydratedDeviceKey, PendingChanges, RoomKeyCounts,
+            RoomKeyWithheldEntry, RoomSettings, StoredRoomKeyBundleData,
         },
         CryptoStore, CryptoStoreError,
     },
-    types::events::room_key_withheld::RoomKeyWithheldEntry,
     vodozemac::base64_encode,
     Account, DeviceData, GossipRequest, GossippedSecret, SecretInfo, TrackedUser, UserIdentityData,
 };
diff --git a/crates/matrix-sdk-sqlite/src/crypto_store.rs b/crates/matrix-sdk-sqlite/src/crypto_store.rs
@@ -28,12 +28,11 @@ use matrix_sdk_crypto::{
     },
     store::{
         types::{
-            BackupKeys, Changes, DehydratedDeviceKey, PendingChanges, RoomKeyCounts, RoomSettings,
-            StoredRoomKeyBundleData,
+            BackupKeys, Changes, DehydratedDeviceKey, PendingChanges, RoomKeyCounts,
+            RoomKeyWithheldEntry, RoomSettings, StoredRoomKeyBundleData,
         },
         CryptoStore,
     },
-    types::events::room_key_withheld::RoomKeyWithheldEntry,
     Account, DeviceData, GossipRequest, GossippedSecret, SecretInfo, TrackedUser, UserIdentityData,
 };
 use matrix_sdk_store_encryption::StoreCipher;
@@ -1894,7 +1893,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");
