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] 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()) + } }