From 592be5bd579b27fefdecd2ce7c2ebcde7a91c185 Mon Sep 17 00:00:00 2001 From: valere Date: Fri, 10 Feb 2023 17:04:42 +0100 Subject: [PATCH] feat(withheld) basic no_olm --- .../src/session_manager/group_sessions.rs | 81 +++++++++++++++---- .../src/types/events/room_key_withheld.rs | 6 +- 2 files changed, 69 insertions(+), 18 deletions(-) diff --git a/crates/matrix-sdk-crypto/src/session_manager/group_sessions.rs b/crates/matrix-sdk-crypto/src/session_manager/group_sessions.rs index 77dd7d5fe..12a65daa4 100644 --- a/crates/matrix-sdk-crypto/src/session_manager/group_sessions.rs +++ b/crates/matrix-sdk-crypto/src/session_manager/group_sessions.rs @@ -228,6 +228,8 @@ impl GroupSessionManager { ToDeviceRequest, BTreeMap>, Vec, + // devices with no olm + BTreeMap>, )> { // Use a named type instead of a tuple with rather long type name struct EncryptResult { @@ -235,11 +237,13 @@ impl GroupSessionManager { share_info: BTreeMap>, message: BTreeMap>>, + device_info: (OwnedUserId, OwnedDeviceId), } let mut messages = BTreeMap::new(); let mut changed_sessions = Vec::new(); let mut share_infos = BTreeMap::new(); + let mut no_olm: BTreeMap> = BTreeMap::new(); let encrypt = |device: Device, session: OutboundGroupSession| async move { let mut message = BTreeMap::new(); @@ -280,7 +284,12 @@ impl GroupSessionManager { Err(e) => return Err(e), }; - Ok(EncryptResult { used_session, share_info, message }) + Ok(EncryptResult { + used_session, + share_info, + message, + device_info: (device.user_id().to_owned(), device.device_id().to_owned()), + }) }; let tasks: Vec<_> = @@ -289,11 +298,17 @@ impl GroupSessionManager { let results = join_all(tasks).await; for result in results { - let EncryptResult { used_session, share_info, message } = - result.expect("Encryption task panicked")?; + let EncryptResult { + used_session, + share_info, + message, + device_info: (user_id, device_id), + } = result.expect("Encryption task panicked")?; if let Some(session) = used_session { changed_sessions.push(session); + } else { + no_olm.entry(user_id).or_default().push(device_id) } for (user, device_messages) in message { @@ -318,7 +333,7 @@ impl GroupSessionManager { "Created a to-device request carrying a room_key" ); - Ok((txn_id, request, share_infos, changed_sessions)) + Ok((txn_id, request, share_infos, changed_sessions, no_olm)) } /// Given a list of user and an outbound session, return the list of users @@ -472,8 +487,8 @@ impl GroupSessionManager { outbound: OutboundGroupSession, message_index: u32, being_shared: Arc>, - ) -> OlmResult> { - let (id, request, share_infos, used_sessions) = + ) -> OlmResult<(Vec, BTreeMap>)> { + let (id, request, share_infos, used_sessions, no_olm) = Self::encrypt_session_for(outbound.clone(), chunk, message_index).await?; if !request.messages.is_empty() { @@ -481,7 +496,7 @@ impl GroupSessionManager { being_shared.insert(id, outbound.clone()); } - Ok(used_sessions) + Ok((used_sessions, no_olm)) } pub(crate) fn session_cache(&self) -> GroupSessionCache { @@ -508,6 +523,7 @@ impl GroupSessionManager { let encryption_settings = encryption_settings.into(); let mut changes = Changes::default(); + let mut no_olm: BTreeMap> = BTreeMap::new(); // Try to get an existing session or create a new one. let (outbound, inbound) = @@ -601,21 +617,38 @@ impl GroupSessionManager { // needed because each encryption step will mutate the Olm session, // ratcheting its state forward. for result in join_all(tasks).await { - let used_sessions: OlmResult> = result.expect("Encryption task panicked"); + let result: OlmResult<(Vec, BTreeMap>)> = + result.expect("Encryption task panicked"); - changes.sessions.extend(used_sessions?); + let (used_sessions, failed_no_olm) = result?; + changes.sessions.extend(used_sessions); + + no_olm.extend(failed_no_olm.into_iter()) } // Handle unverified and blacklisted withhelds - let flat_with_code: Vec<(Device, WithheldCode)> = - withheld_devices.into_iter().flat_map(|(_, list)| list.into_iter()).collect(); + let mut flat_with_code: Vec<(OwnedUserId, OwnedDeviceId, WithheldCode)> = withheld_devices + .into_iter() + .flat_map(|(_, list)| { + list.into_iter().map(|(d, c)| (d.user_id().to_owned(), d.device_id().to_owned(), c)) + }) + .collect(); + + no_olm.iter().for_each(|(u, list)| { + list.into_iter().for_each(|d| { + flat_with_code.push((u.to_owned(), d.to_owned(), WithheldCode::NoOlm)) + }); + }); let withheld_to_device_request: Vec> = flat_with_code .chunks(Self::MAX_TO_DEVICE_MESSAGES) .map(|chunk| { let mut messages = BTreeMap::new(); - chunk.into_iter().for_each(|(d, code)| { + // TODO do not resend code if already sent for that session (or target if + // no_olm) + chunk.into_iter().for_each(|(user_id, device_id, code)| { + // TODO use helper to create correct content depending on alg and code let content = MegolmV1AesSha2WithheldContent::new( room_id.to_owned(), Some(outbound.session_id().to_owned()), @@ -629,9 +662,9 @@ impl GroupSessionManager { .expect("We can always serialize a withheld content info") .cast(); messages - .entry(d.user_id().to_owned()) + .entry(user_id.to_owned()) .or_insert_with(BTreeMap::new) - .insert(DeviceIdOrAllDevices::DeviceId(d.device_id().into()), content); + .insert(DeviceIdOrAllDevices::DeviceId(device_id.to_owned()), content); }); Arc::new(ToDeviceRequest { @@ -897,12 +930,23 @@ mod tests { let requests = machine.share_room_key(room_id, users, EncryptionSettings::default()).await.unwrap(); - let event_count: usize = requests.iter().map(|r| r.message_count()).sum(); + let event_count: usize = requests + .iter() + .filter(|r| r.event_type == "m.room.encrypted".into()) + .map(|r| r.message_count()) + .sum(); // The keys claim response has a couple of one-time keys with invalid // signatures, thus only 148 sessions are actually created, we check // that all 148 valid sessions get an room key. assert_eq!(event_count, 148); + + let withheld_count: usize = requests + .iter() + .filter(|r| r.event_type == "m.room_key.withheld".into()) + .map(|r| r.message_count()) + .sum(); + assert_eq!(withheld_count, 2); } #[async_test] @@ -921,7 +965,12 @@ mod tests { .await .unwrap(); - let event_count: usize = requests.iter().map(|r| r.message_count()).sum(); + println!("{:?}", requests); + let event_count: usize = requests + .iter() + .filter(|r| r.event_type == "m.room.encrypted".into()) + .map(|r| r.message_count()) + .sum(); let outbound = machine.group_session_manager.get_outbound_group_session(room_id).unwrap(); assert_eq!(event_count, 1); 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 index 89324ca7a..41167f2ee 100644 --- a/crates/matrix-sdk-crypto/src/types/events/room_key_withheld.rs +++ b/crates/matrix-sdk-crypto/src/types/events/room_key_withheld.rs @@ -117,7 +117,9 @@ struct WithheldHelper { #[derive(Deserialize, Serialize)] pub struct MegolmV1AesSha2WithheldContent { /// The room where the key is used. - pub room_id: OwnedRoomId, + /// Required if code is not m.no_olm. + #[serde(default, skip_serializing_if = "JsOption::is_undefined")] + pub room_id: JsOption, /// Required if code is not m.no_olm. The ID of the session. #[serde(default, skip_serializing_if = "JsOption::is_undefined")] pub session_id: JsOption, @@ -152,7 +154,7 @@ impl MegolmV1AesSha2WithheldContent { from_device: Option, ) -> Self { Self { - room_id, + room_id: JsOption::Some(room_id), session_id: JsOption::from_implicit_option(session_id), sender_key, code,