Remove DirectWithheldCode

Cargo.toml

@@ -21,6 +21,7 @@ rust-version = "1.65"
 
 [workspace.dependencies]
 anyhow = "1.0.68"
+assert_matches = "1.5.0"
 async-stream = "0.3.3"
 async-trait = "0.1.60"
 base64 = "0.21.0"

crates/matrix-sdk-base/Cargo.toml

@@ -27,7 +27,7 @@ experimental-sliding-sync = ["ruma/unstable-msc3575"]
 testing = ["dep:http", "dep:matrix-sdk-test", "dep:assert_matches"]
 
 [dependencies]
-assert_matches = { version = "1.5.0", optional = true }
+assert_matches = { workspace = true, optional = true }
 async-stream = { workspace = true }
 async-trait = { workspace = true }
 dashmap = { workspace = true }
@@ -48,7 +48,7 @@ tracing = { workspace = true }
 zeroize = { workspace = true, features = ["zeroize_derive"] }
 
 [dev-dependencies]
-assert_matches = "1.5.0"
+assert_matches = { workspace = true }
 assign = "1.1.1"
 ctor = { workspace = true }
 futures = { version = "0.3.21", default-features = false, features = ["executor"] }

crates/matrix-sdk-crypto/Cargo.toml

@@ -63,7 +63,7 @@ tokio = { version = "1.24", default-features = false, features = ["time"] }
 
 [dev-dependencies]
 anyhow = { workspace = true }
-assert_matches = "1.5.0"
+assert_matches = { workspace = true }
 ctor.workspace = true
 futures = { version = "0.3.21", default-features = false, features = ["executor"] }
 http = { workspace = true }

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

@@ -69,8 +69,8 @@ use crate::{
     requests::{IncomingResponse, OutgoingRequest, UploadSigningKeysRequest},
     session_manager::{GroupSessionManager, SessionManager},
     store::{
-        withheld::DirectWithheldInfo, Changes, DeviceChanges, DynCryptoStore, IdentityChanges,
-        IntoCryptoStore, MemoryStore, Result as StoreResult, SecretImportError, Store,
+        Changes, DeviceChanges, DynCryptoStore, IdentityChanges, IntoCryptoStore, MemoryStore,
+        Result as StoreResult, SecretImportError, Store,
     },
     types::{
         events::{
@@ -80,7 +80,9 @@ use crate::{
                 RoomEventEncryptionScheme, SupportedEventEncryptionSchemes,
             },
             room_key::{MegolmV1AesSha2Content, RoomKeyContent},
-            room_key_withheld::RoomKeyWithheldEvent,
+            room_key_withheld::{
+                MegolmV1AesSha2WithheldContent, RoomKeyWithheldContent, RoomKeyWithheldEvent,
+            },
             ToDeviceEvents,
         },
         Signatures,
@@ -676,8 +678,19 @@ impl OlmMachine {
         changes: &mut Changes,
         event: &RoomKeyWithheldEvent,
     ) -> OlmResult<()> {
-        if let Some(info) = DirectWithheldInfo::from_event(event) {
-            changes.withheld_session_info.push(info)
+        match &event.content {
+            RoomKeyWithheldContent::MegolmV1AesSha2(c) => match c {
+                MegolmV1AesSha2WithheldContent::BlackListed(c)
+                | MegolmV1AesSha2WithheldContent::Unverified(c) => {
+                    changes
+                        .withheld_session_info
+                        .entry(c.room_id.to_owned())
+                        .or_insert_with(BTreeMap::default)
+                        .insert(c.session_id.to_owned(), event.to_owned());
+                }
+                _ => (),
+            },
+            _ => (),
         }
 
         Ok(())
@@ -1261,7 +1274,7 @@ impl OlmMachine {
                             .store
                             .get_withheld_info(room_id, content.session_id())
                             .await?
-                            .map(|i| i.withheld_code());
+                            .map(|e| e.content.withheld_code());
 
                         if withheld_code.is_some() {
                             // Partially withheld, report with a withheld code if we have one.
@@ -1279,7 +1292,7 @@ impl OlmMachine {
                 .store
                 .get_withheld_info(room_id, content.session_id())
                 .await?
-                .map(|i| i.withheld_code());
+                .map(|e| e.content.withheld_code());
 
             Err(MegolmError::MissingRoomKey(withheld_code))
         }

crates/matrix-sdk-crypto/src/session_manager/group_sessions.rs

@@ -34,10 +34,7 @@ use crate::{
     identities::device::MaybeEncryptedRoomKey,
     olm::{Account, InboundGroupSession, OutboundGroupSession, Session, ShareInfo, ShareState},
     store::{Changes, Result as StoreResult, Store},
-    types::events::{
-        room::encrypted::RoomEncryptedEventContent,
-        room_key_withheld::{RoomKeyWithheldContent, WithheldCode},
-    },
+    types::events::{room::encrypted::RoomEncryptedEventContent, room_key_withheld::WithheldCode},
     Device, EncryptionSettings, OlmError, ToDeviceRequest,
 };
 

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

@@ -5,6 +5,7 @@ macro_rules! cryptostore_integration_tests {
         mod cryptostore_integration_tests {
             use std::collections::{BTreeMap, HashMap};
 
+            use assert_matches::assert_matches;
             use matrix_sdk_test::async_test;
             use ruma::{
                 device_id, encryption::SignedKey, room_id, serde::Base64, user_id, DeviceId,
@@ -16,14 +17,18 @@ macro_rules! cryptostore_integration_tests {
                     PrivateCrossSigningIdentity, ReadOnlyAccount, Session,
                 },
                 store::{
-                    withheld::DirectWithheldInfo, Changes, CryptoStore, DeviceChanges,
+                    Changes, CryptoStore, DeviceChanges,
                     GossipRequest, IdentityChanges, RecoveryKey, RoomSettings,
                 },
                 testing::{get_device, get_other_identity, get_own_identity},
                 types::{
                     events::{
                         room_key_request::MegolmV1AesSha2Content,
-                        room_key_withheld::{CommonWithheldCodeContent, WithheldCode},
+                        room_key_withheld::{
+                            CommonWithheldCodeContent, MegolmV1AesSha2WithheldContent,
+                            RoomKeyWithheldContent, WithheldCode,
+                        },
+                        ToDeviceEvent,
                     },
                     EventEncryptionAlgorithm,
                 },
@@ -624,77 +629,80 @@ macro_rules! cryptostore_integration_tests {
             async fn withheld_info_storage() {
                 let (account, store) = get_loaded_store("withheld_info_storage").await;
 
-                let mut info_list: Vec<DirectWithheldInfo> = Vec::new();
+                let mut info_list: BTreeMap<_, BTreeMap<_, _>> = BTreeMap::new();
 
+                let user_id = account.user_id().to_owned();
                 let room_id = room_id!("!DwLygpkclUAfQNnfva:example.com");
                 let session_id_1 = "GBnDxGP9i3IkPsz3/ihNr6P7qjIXxSRVWZ1MYmSn09w";
                 let session_id_2 = "IDLtnNCH2kIr3xIf1B7JFkGpQmTjyMca2jww+X6zeOE";
 
-                let content = CommonWithheldCodeContent {
-                    room_id: room_id.to_owned(),
-                    session_id: session_id_1.into(),
-                    from_device: JsOption::Undefined,
-                    other: Default::default(),
-                    sender_key: Curve25519PublicKey::from_base64(
-                        "9n7mdWKOjr9c4NTlG6zV8dbFtNK79q9vZADoh7nMUwA",
-                    )
-                    .unwrap(),
-                };
-                let info = DirectWithheldInfo::new(
-                    EventEncryptionAlgorithm::MegolmV1AesSha2,
-                    WithheldCode::Unverified,
-                    content.to_owned(),
+                let content = RoomKeyWithheldContent::MegolmV1AesSha2(
+                    MegolmV1AesSha2WithheldContent::Unverified(
+                        CommonWithheldCodeContent {
+                            room_id: room_id.to_owned(),
+                            session_id: session_id_1.into(),
+                            from_device: JsOption::Undefined,
+                            other: Default::default(),
+                            sender_key: Curve25519PublicKey::from_base64(
+                                "9n7mdWKOjr9c4NTlG6zV8dbFtNK79q9vZADoh7nMUwA",
+                            )
+                            .unwrap(),
+                        }
+                        .into(),
+                    ),
                 );
-                info_list.push(info);
-
-                let content = CommonWithheldCodeContent {
-                    room_id: room_id.to_owned(),
-                    session_id: session_id_2.into(),
-                    from_device: JsOption::Undefined,
-                    other: Default::default(),
-                    sender_key: Curve25519PublicKey::from_base64(
-                        "9n7mdWKOjr9c4NTlG6zV8dbFtNK79q9vZADoh7nMUwA",
-                    )
-                    .unwrap(),
-                };
-
-                let info = DirectWithheldInfo::new(
-                    EventEncryptionAlgorithm::MegolmV1AesSha2,
-                    WithheldCode::Blacklisted,
-                    content,
+                let event = ToDeviceEvent::new(user_id.to_owned(), content);
+                info_list
+                    .entry(room_id.to_owned())
+                    .or_default()
+                    .insert(session_id_1.to_owned(), event);
+
+                let content = RoomKeyWithheldContent::MegolmV1AesSha2(
+                    MegolmV1AesSha2WithheldContent::BlackListed(
+                        CommonWithheldCodeContent {
+                            room_id: room_id.to_owned(),
+                            session_id: session_id_2.into(),
+                            from_device: JsOption::Undefined,
+                            other: Default::default(),
+                            sender_key: Curve25519PublicKey::from_base64(
+                                "9n7mdWKOjr9c4NTlG6zV8dbFtNK79q9vZADoh7nMUwA",
+                            )
+                            .unwrap(),
+                        }
+                        .into(),
+                    ),
                 );
-                info_list.push(info);
+                let event = ToDeviceEvent::new(user_id.to_owned(), content);
+                info_list
+                    .entry(room_id.to_owned())
+                    .or_default()
+                    .insert(session_id_2.to_owned(), event);
 
                 let changes = Changes { withheld_session_info: info_list, ..Default::default() };
-
                 store.save_changes(changes).await.unwrap();
 
                 let is_withheld = store.get_withheld_info(room_id, session_id_1).await.unwrap();
 
-                if let Some(info) = is_withheld {
-                    let actual_code = info.withheld_code();
-                    assert_eq!(EventEncryptionAlgorithm::MegolmV1AesSha2, info.algorithm());
-                    assert_eq!(room_id, info.room_id());
-                    assert_eq!(WithheldCode::Unverified, actual_code);
-                } else {
-                    panic!();
-                }
+                assert_matches!(
+                    is_withheld, Some(event)
+                    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();
 
-                if let Some(info) = is_withheld {
-                    let actual_code = info.withheld_code();
-                    assert_eq!(WithheldCode::Blacklisted, actual_code);
-                } else {
-                    panic!();
-                }
+                assert_matches!(
+                    is_withheld, Some(event)
+                    if event.content.algorithm() == EventEncryptionAlgorithm::MegolmV1AesSha2 &&
+                    event.content.withheld_code() == WithheldCode::Blacklisted
+                );
 
                 let other_room_id = room_id!("!nQRyiRFuyUhXeaQfiR:example.com");
 
                 let is_withheld =
                     store.get_withheld_info(other_room_id, session_id_2).await.unwrap();
 
-                assert_eq!(is_withheld, None);
+                assert!(is_withheld.is_none());
             }
 
             #[async_test]

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

@@ -15,10 +15,11 @@
 use std::{collections::HashMap, convert::Infallible, sync::Arc};
 
 use async_trait::async_trait;
-use dashmap::{mapref::entry::Entry, DashMap, DashSet};
+use dashmap::{DashMap, DashSet};
 use matrix_sdk_common::locks::Mutex;
 use ruma::{
-    DeviceId, OwnedDeviceId, OwnedTransactionId, OwnedUserId, RoomId, TransactionId, UserId,
+    DeviceId, OwnedDeviceId, OwnedRoomId, OwnedTransactionId, OwnedUserId, RoomId, TransactionId,
+    UserId,
 };
 use tracing::warn;
 
@@ -31,7 +32,7 @@ use crate::{
     gossiping::{GossipRequest, SecretInfo},
     identities::{ReadOnlyDevice, ReadOnlyUserIdentities},
     olm::{OutboundGroupSession, PrivateCrossSigningIdentity},
-    store::withheld::DirectWithheldInfo,
+    types::events::room_key_withheld::RoomKeyWithheldEvent,
     TrackedUser,
 };
 
@@ -54,7 +55,7 @@ pub struct MemoryStore {
     identities: Arc<DashMap<OwnedUserId, ReadOnlyUserIdentities>>,
     outgoing_key_requests: Arc<DashMap<OwnedTransactionId, GossipRequest>>,
     key_requests_by_info: Arc<DashMap<String, OwnedTransactionId>>,
-    direct_withheld_info: Arc<DashMap<String, DirectWithheldInfo>>,
+    direct_withheld_info: Arc<DashMap<OwnedRoomId, DashMap<String, RoomKeyWithheldEvent>>>,
 }
 
 impl Default for MemoryStore {
@@ -149,8 +150,13 @@ impl CryptoStore for MemoryStore {
             self.key_requests_by_info.insert(info_string, id);
         }
 
-        for info in changes.withheld_session_info {
-            self.direct_withheld_info.insert(info.session_id().to_owned(), info);
+        for (room_id, data) in changes.withheld_session_info {
+            for (session_id, event) in data {
+                self.direct_withheld_info
+                    .entry(room_id.to_owned())
+                    .or_insert_with(DashMap::new)
+                    .insert(session_id, event);
+            }
         }
 
         Ok(())
@@ -284,12 +290,11 @@ impl CryptoStore for MemoryStore {
         &self,
         room_id: &RoomId,
         session_id: &str,
-    ) -> Result<Option<DirectWithheldInfo>> {
+    ) -> Result<Option<RoomKeyWithheldEvent>> {
         Ok(self
             .direct_withheld_info
-            .get(session_id)
-            .filter(|e| e.value().room_id() == room_id)
-            .map(|e| e.value().to_owned()))
+            .get(room_id)
+            .and_then(|e| e.value().get(session_id).map(|v| v.value().to_owned())))
     }
 
     async fn get_room_settings(&self, _room_id: &RoomId) -> Result<Option<RoomSettings>> {

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

@@ -39,7 +39,7 @@
 //! [`CryptoStore`]: trait.Cryptostore.html
 
 use std::{
-    collections::{HashMap, HashSet},
+    collections::{BTreeMap, HashMap, HashSet},
     fmt::Debug,
     ops::Deref,
     sync::{atomic::AtomicBool, Arc},
@@ -68,7 +68,7 @@ use crate::{
         InboundGroupSession, OlmMessageHash, OutboundGroupSession, PrivateCrossSigningIdentity,
         ReadOnlyAccount, Session,
     },
-    types::EventEncryptionAlgorithm,
+    types::{events::room_key_withheld::RoomKeyWithheldEvent, EventEncryptionAlgorithm},
     utilities::encode,
     verification::VerificationMachine,
     CrossSigningStatus,
@@ -83,7 +83,6 @@ mod traits;
 #[macro_use]
 #[allow(missing_docs)]
 pub mod integration_tests;
-pub mod withheld;
 
 use caches::{SequenceNumber, UsersForKeyQuery};
 pub use error::{CryptoStoreError, Result};
@@ -92,7 +91,6 @@ pub use memorystore::MemoryStore;
 pub use traits::{CryptoStore, DynCryptoStore, IntoCryptoStore};
 
 pub use crate::gossiping::{GossipRequest, SecretInfo};
-use crate::store::withheld::DirectWithheldInfo;
 
 /// A wrapper for our CryptoStore trait object.
 ///
@@ -137,7 +135,7 @@ pub struct Changes {
     pub identities: IdentityChanges,
     pub devices: DeviceChanges,
     /// Stores when a `m.room_key.withheld` is received
-    pub withheld_session_info: Vec<DirectWithheldInfo>,
+    pub withheld_session_info: BTreeMap<OwnedRoomId, BTreeMap<String, RoomKeyWithheldEvent>>,
     pub room_settings: HashMap<OwnedRoomId, RoomSettings>,
 }