From 96144404ed1434317822fc7561ad99ca43e2d7b3 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Damir=20Jeli=C4=87?= Date: Mon, 27 Apr 2020 16:26:03 +0200 Subject: [PATCH 1/2] crypto: Return a result when searching for missing sessions. --- src/async_client.rs | 2 +- src/base_client.rs | 6 +++--- src/crypto/machine.rs | 8 ++++---- 3 files changed, 8 insertions(+), 8 deletions(-) diff --git a/src/async_client.rs b/src/async_client.rs index 54098166..3021bf3b 100644 --- a/src/async_client.rs +++ b/src/async_client.rs @@ -898,7 +898,7 @@ impl AsyncClient { .read() .await .get_missing_sessions(users) - .await + .await? }; if !missing_sessions.is_empty() { diff --git a/src/base_client.rs b/src/base_client.rs index b3e2f913..16fea864 100644 --- a/src/base_client.rs +++ b/src/base_client.rs @@ -435,12 +435,12 @@ impl Client { pub async fn get_missing_sessions( &self, users: impl Iterator, - ) -> BTreeMap> { + ) -> Result>> { let mut olm = self.olm.lock().await; match &mut *olm { - Some(o) => o.get_missing_sessions(users).await, - None => BTreeMap::new(), + Some(o) => Ok(o.get_missing_sessions(users).await?), + None => Ok(BTreeMap::new()), } } diff --git a/src/crypto/machine.rs b/src/crypto/machine.rs index 357975b2..cc8c9d34 100644 --- a/src/crypto/machine.rs +++ b/src/crypto/machine.rs @@ -202,11 +202,11 @@ impl OlmMachine { pub async fn get_missing_sessions( &mut self, users: impl Iterator, - ) -> BTreeMap> { + ) -> Result>> { let mut missing = BTreeMap::new(); for user_id in users { - let user_devices = self.store.get_user_devices(user_id).await.unwrap(); + let user_devices = self.store.get_user_devices(user_id).await?; for device in user_devices.devices() { let sender_key = if let Some(k) = device.get_key(&KeyAlgorithm::Curve25519) { @@ -215,7 +215,7 @@ impl OlmMachine { continue; }; - let sessions = self.store.get_sessions(sender_key).await.unwrap(); + let sessions = self.store.get_sessions(sender_key).await?; let is_missing = if let Some(sessions) = sessions { sessions.lock().await.is_empty() @@ -237,7 +237,7 @@ impl OlmMachine { } } - missing + Ok(missing) } pub async fn receive_keys_claim_response( From b6c9358d148605d36fae8341d163e8c405d8d909 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Damir=20Jeli=C4=87?= Date: Mon, 27 Apr 2020 16:31:28 +0200 Subject: [PATCH 2/2] crypto: Test the creation of an Olm session. --- src/crypto/device.rs | 34 +++++++++++++ src/crypto/machine.rs | 109 ++++++++++++++++++++++++++++++++++++++---- 2 files changed, 133 insertions(+), 10 deletions(-) diff --git a/src/crypto/device.rs b/src/crypto/device.rs index 6eb69b70..ad0ad82d 100644 --- a/src/crypto/device.rs +++ b/src/crypto/device.rs @@ -13,12 +13,16 @@ // limitations under the License. use std::collections::BTreeMap; +#[cfg(test)] +use std::convert::TryFrom; use std::mem; use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::Arc; use atomic::Atomic; +#[cfg(test)] +use super::OlmMachine; use crate::api::r0::keys::{DeviceKeys, KeyAlgorithm}; use crate::events::Algorithm; use crate::identifiers::{DeviceId, UserId}; @@ -151,6 +155,36 @@ impl Device { } } +#[cfg(test)] +impl From<&OlmMachine> for Device { + fn from(machine: &OlmMachine) -> Self { + Device { + user_id: Arc::new(machine.user_id.clone()), + device_id: Arc::new(machine.device_id.clone()), + algorithms: Arc::new(vec![ + Algorithm::MegolmV1AesSha2, + Algorithm::OlmV1Curve25519AesSha2, + ]), + keys: Arc::new( + machine + .account + .identity_keys() + .iter() + .map(|(key, value)| { + ( + KeyAlgorithm::try_from(key.as_ref()).unwrap(), + value.to_owned(), + ) + }) + .collect(), + ), + display_name: Arc::new(None), + deleted: Arc::new(AtomicBool::new(false)), + trust_state: Arc::new(Atomic::new(TrustState::Unset)), + } + } +} + impl From<&DeviceKeys> for Device { fn from(device_keys: &DeviceKeys) -> Self { let mut keys = BTreeMap::new(); diff --git a/src/crypto/machine.rs b/src/crypto/machine.rs index cc8c9d34..ff03f22d 100644 --- a/src/crypto/machine.rs +++ b/src/crypto/machine.rs @@ -61,11 +61,11 @@ pub type OneTimeKeys = BTreeMap; pub struct OlmMachine { /// The unique user id that owns this account. - user_id: UserId, + pub(crate) user_id: UserId, /// The unique device id of the device that holds this account. - device_id: DeviceId, + pub(crate) device_id: DeviceId, /// Our underlying Olm Account holding our identity keys. - account: Account, + pub(crate) account: Account, /// The number of signed one-time keys we have uploaded to the server. If /// this is None, no action will be taken. After a sync request the client /// needs to set this for us, depending on the count we will suggest the @@ -1366,21 +1366,32 @@ impl OlmMachine { #[cfg(test)] mod test { - static USER_ID: &str = "@test:example.org"; - const DEVICE_ID: &str = "DEVICEID"; + static USER_ID: &str = "@bob:example.org"; + static DEVICE_ID: &str = "DEVICEID"; use js_int::UInt; + use std::collections::BTreeMap; use std::convert::TryFrom; use std::fs::File; use std::io::prelude::*; - use ruma_identifiers::UserId; use serde_json::json; use crate::api::r0::keys; - use crate::crypto::machine::OlmMachine; + use crate::crypto::machine::{OlmMachine, OneTimeKeys}; + use crate::crypto::Device; + use crate::identifiers::{DeviceId, UserId}; + use http::Response; + fn alice_id() -> UserId { + UserId::try_from("@alice:example.org").unwrap() + } + + fn alice_device_id() -> DeviceId { + "JLAFKJWSCS".to_string() + } + fn user_id() -> UserId { UserId::try_from(USER_ID).unwrap() } @@ -1405,10 +1416,10 @@ mod test { keys::get_keys::Response::try_from(data).expect("Can't parse the keys upload response") } - async fn get_prepared_machine() -> OlmMachine { + async fn get_prepared_machine() -> (OlmMachine, OneTimeKeys) { let mut machine = OlmMachine::new(&user_id(), DEVICE_ID).unwrap(); machine.uploaded_signed_key_count = Some(0); - let (_, _) = machine + let (_, otk) = machine .keys_for_upload() .await .expect("Can't prepare initial key upload"); @@ -1418,7 +1429,34 @@ mod test { .await .unwrap(); + (machine, otk.unwrap()) + } + + async fn get_machine_after_query() -> (OlmMachine, OneTimeKeys) { + let (mut machine, otk) = get_prepared_machine().await; + let response = keys_query_response(); + machine + .receive_keys_query_response(&response) + .await + .unwrap(); + + (machine, otk) + } + + async fn get_machine_pair() -> (OlmMachine, OlmMachine, OneTimeKeys) { + let (bob, otk) = get_prepared_machine().await; + + let alice_id = alice_id(); + let alice_device = alice_device_id(); + let alice = OlmMachine::new(&alice_id, &alice_device).unwrap(); + + let alice_deivce = Device::from(&alice); + let bob_device = Device::from(&bob); + alice.store.save_device(bob_device).await.unwrap(); + bob.store.save_device(alice_deivce).await.unwrap(); + + (alice, bob, otk) } #[tokio::test] @@ -1590,7 +1628,7 @@ mod test { #[tokio::test] async fn test_keys_query() { - let mut machine = get_prepared_machine().await; + let (mut machine, _) = get_prepared_machine().await; let response = keys_query_response(); let alice_id = UserId::try_from("@alice:example.org").unwrap(); let alice_device_id = "JLAFKJWSCS".to_owned(); @@ -1612,4 +1650,55 @@ mod test { assert_eq!(device.user_id(), &alice_id); assert_eq!(device.device_id(), &alice_device_id); } + + #[tokio::test] + async fn test_missing_sessions_calculation() { + let (mut machine, _) = get_machine_after_query().await; + + let alice = alice_id(); + let alice_device = alice_device_id(); + + let missing_sessions = machine + .get_missing_sessions([alice.clone()].iter()) + .await + .unwrap(); + + assert!(missing_sessions.contains_key(&alice)); + let user_sessions = missing_sessions.get(&alice).unwrap(); + assert!(user_sessions.contains_key(&alice_device)); + } + + #[tokio::test] + async fn test_key_claiming() { + let (mut alice_machine, bob_machine, one_time_keys) = get_machine_pair().await; + + let mut bob_keys = BTreeMap::new(); + + let one_time_key = one_time_keys.iter().nth(0).unwrap(); + let mut keys = BTreeMap::new(); + keys.insert(one_time_key.0.clone(), one_time_key.1.clone()); + bob_keys.insert(bob_machine.device_id.clone(), keys); + + let mut one_time_keys = BTreeMap::new(); + one_time_keys.insert(bob_machine.user_id.clone(), bob_keys); + + let response = keys::claim_keys::Response { + failures: BTreeMap::new(), + one_time_keys, + }; + + alice_machine + .receive_keys_claim_response(&response) + .await + .unwrap(); + + let session = alice_machine + .store + .get_sessions(bob_machine.account.identity_keys().curve25519()) + .await + .unwrap() + .unwrap(); + + assert!(!session.lock().await.is_empty()) + } }