diff --git a/crates/matrix-sdk-crypto/src/dehydrated_devices.rs b/crates/matrix-sdk-crypto/src/dehydrated_devices.rs index 61e3265a5..c1ced5201 100644 --- a/crates/matrix-sdk-crypto/src/dehydrated_devices.rs +++ b/crates/matrix-sdk-crypto/src/dehydrated_devices.rs @@ -388,8 +388,9 @@ mod tests { }; use crate::{ - machine::tests::{ - create_session, get_prepared_machine_test_helper, to_device_requests_to_content, + machine::{ + test_helpers::get_prepared_machine_test_helper, + tests::{create_session, to_device_requests_to_content}, }, olm::OutboundGroupSession, types::{events::ToDeviceEvent, DeviceKeys as DeviceKeysType}, diff --git a/crates/matrix-sdk-crypto/src/file_encryption/key_export.rs b/crates/matrix-sdk-crypto/src/file_encryption/key_export.rs index 820b9c436..01deb50a3 100644 --- a/crates/matrix-sdk-crypto/src/file_encryption/key_export.rs +++ b/crates/matrix-sdk-crypto/src/file_encryption/key_export.rs @@ -243,7 +243,8 @@ mod tests { encrypt_room_key_export, }; use crate::{ - error::OlmResult, machine::tests::get_prepared_machine_test_helper, RoomKeyImportResult, + error::OlmResult, machine::test_helpers::get_prepared_machine_test_helper, + RoomKeyImportResult, }; const PASSPHRASE: &str = "1234"; diff --git a/crates/matrix-sdk-crypto/src/gossiping/machine.rs b/crates/matrix-sdk-crypto/src/gossiping/machine.rs index d13c3dedf..c01903c1f 100644 --- a/crates/matrix-sdk-crypto/src/gossiping/machine.rs +++ b/crates/matrix-sdk-crypto/src/gossiping/machine.rs @@ -1937,7 +1937,8 @@ mod tests { use tokio_stream::StreamExt; use crate::{ - machine::tests::get_machine_pair_with_setup_sessions_test_helper, EncryptionSyncChanges, + machine::test_helpers::get_machine_pair_with_setup_sessions_test_helper, + EncryptionSyncChanges, }; let alice_id = user_id!("@alice:localhost"); diff --git a/crates/matrix-sdk-crypto/src/machine/mod.rs b/crates/matrix-sdk-crypto/src/machine/mod.rs index 3a2c5a5a1..643ee5056 100644 --- a/crates/matrix-sdk-crypto/src/machine/mod.rs +++ b/crates/matrix-sdk-crypto/src/machine/mod.rs @@ -2443,7 +2443,14 @@ pub(crate) mod tests { use super::{testing::response_from_file, CrossSigningBootstrapRequests}; use crate::{ error::{EventError, SetRoomSettingsError}, - machine::{EncryptionSyncChanges, OlmMachine}, + machine::{ + test_helpers::{ + get_machine_after_query_test_helper, get_machine_pair, + get_machine_pair_with_session, get_machine_pair_with_setup_sessions_test_helper, + get_prepared_machine_test_helper, + }, + EncryptionSyncChanges, OlmMachine, + }, olm::{ BackedUpRoomKey, ExportedRoomKey, InboundGroupSession, OutboundGroupSession, SenderData, VerifyJson, @@ -2469,9 +2476,6 @@ pub(crate) mod tests { OlmError, OutgoingRequests, ToDeviceRequest, UserIdentities, }; - /// These keys need to be periodically uploaded to the server. - type OneTimeKeys = BTreeMap>; - fn alice_id() -> &'static UserId { user_id!("@alice:example.org") } @@ -2517,119 +2521,6 @@ pub(crate) mod tests { .unwrap() } - pub(crate) async fn get_prepared_machine_test_helper( - user_id: &UserId, - use_fallback_key: bool, - ) -> (OlmMachine, OneTimeKeys) { - let machine = OlmMachine::new(user_id, bob_device_id()).await; - - let request = machine - .store() - .with_transaction(|mut tr| async { - let account = tr.account().await.unwrap(); - account.generate_fallback_key_if_needed(); - account.update_uploaded_key_count(0); - account.generate_one_time_keys_if_needed(); - let request = machine - .keys_for_upload(account) - .await - .expect("Can't prepare initial key upload"); - Ok((tr, request)) - }) - .await - .unwrap(); - - let response = keys_upload_response(); - machine.receive_keys_upload_response(&response).await.unwrap(); - - let keys = if use_fallback_key { request.fallback_keys } else { request.one_time_keys }; - - (machine, keys) - } - - async fn get_machine_after_query_test_helper() -> (OlmMachine, OneTimeKeys) { - let (machine, otk) = get_prepared_machine_test_helper(user_id(), false).await; - let response = keys_query_response(); - let req_id = TransactionId::new(); - - machine.receive_keys_query_response(&req_id, &response).await.unwrap(); - - (machine, otk) - } - - pub async fn get_machine_pair( - alice: &UserId, - bob: &UserId, - use_fallback_key: bool, - ) -> (OlmMachine, OlmMachine, OneTimeKeys) { - let (bob, otk) = get_prepared_machine_test_helper(bob, use_fallback_key).await; - - let alice_device = alice_device_id(); - let alice = OlmMachine::new(alice, alice_device).await; - - let alice_device = DeviceData::from_machine_test_helper(&alice).await.unwrap(); - let bob_device = DeviceData::from_machine_test_helper(&bob).await.unwrap(); - alice.store().save_device_data(&[bob_device]).await.unwrap(); - bob.store().save_device_data(&[alice_device]).await.unwrap(); - - (alice, bob, otk) - } - - async fn get_machine_pair_with_session( - alice: &UserId, - bob: &UserId, - use_fallback_key: bool, - ) -> (OlmMachine, OlmMachine) { - let (alice, bob, mut one_time_keys) = get_machine_pair(alice, bob, use_fallback_key).await; - - let (device_key_id, one_time_key) = one_time_keys.pop_first().unwrap(); - - let one_time_keys = BTreeMap::from([( - bob.user_id().to_owned(), - BTreeMap::from([( - bob.device_id().to_owned(), - BTreeMap::from([(device_key_id, one_time_key)]), - )]), - )]); - - let response = claim_keys::v3::Response::new(one_time_keys); - alice.inner.session_manager.create_sessions(&response).await.unwrap(); - - (alice, bob) - } - - pub(crate) async fn get_machine_pair_with_setup_sessions_test_helper( - alice: &UserId, - bob: &UserId, - use_fallback_key: bool, - ) -> (OlmMachine, OlmMachine) { - let (alice, bob) = get_machine_pair_with_session(alice, bob, use_fallback_key).await; - - let bob_device = - alice.get_device(bob.user_id(), bob.device_id(), None).await.unwrap().unwrap(); - - let (session, content) = - bob_device.encrypt("m.dummy", ToDeviceDummyEventContent::new()).await.unwrap(); - alice.store().save_sessions(&[session]).await.unwrap(); - - let event = - ToDeviceEvent::new(alice.user_id().to_owned(), content.deserialize_as().unwrap()); - - let decrypted = bob - .store() - .with_transaction(|mut tr| async { - let res = - bob.decrypt_to_device_event(&mut tr, &event, &mut Changes::default()).await?; - Ok((tr, res)) - }) - .await - .unwrap(); - - bob.store().save_sessions(&[decrypted.session.session()]).await.unwrap(); - - (alice, bob) - } - #[async_test] async fn test_create_olm_machine() { let test_start_ts = MilliSecondsSinceUnixEpoch::now(); @@ -5010,3 +4901,6 @@ pub(crate) mod tests { assert_matches!(thread_encryption_result, UnsignedDecryptionResult::Decrypted(_)); } } + +#[cfg(test)] +pub(crate) mod test_helpers; diff --git a/crates/matrix-sdk-crypto/src/machine/test_helpers.rs b/crates/matrix-sdk-crypto/src/machine/test_helpers.rs new file mode 100644 index 000000000..3fb69bd5d --- /dev/null +++ b/crates/matrix-sdk-crypto/src/machine/test_helpers.rs @@ -0,0 +1,171 @@ +// Copyright 2020 The Matrix.org Foundation C.I.C. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! A set of helper functions for creating [`OlmMachine`]s, and pairs of +//! interconnected machines. + +use std::collections::BTreeMap; + +use matrix_sdk_test::test_json; +use ruma::{ + api::{ + client::keys::{claim_keys, get_keys, upload_keys}, + IncomingResponse, + }, + device_id, + encryption::OneTimeKey, + events::dummy::ToDeviceDummyEventContent, + serde::Raw, + user_id, DeviceId, OwnedDeviceKeyId, TransactionId, UserId, +}; + +use crate::{ + machine::testing::response_from_file, store::Changes, types::events::ToDeviceEvent, DeviceData, + OlmMachine, +}; + +/// These keys need to be periodically uploaded to the server. +type OneTimeKeys = BTreeMap>; + +fn alice_device_id() -> &'static DeviceId { + device_id!("JLAFKJWSCS") +} + +fn bob_device_id() -> &'static DeviceId { + device_id!("NTHHPZDPRN") +} + +fn user_id() -> &'static UserId { + user_id!("@bob:example.com") +} + +fn keys_upload_response() -> upload_keys::v3::Response { + let data = response_from_file(&test_json::KEYS_UPLOAD); + upload_keys::v3::Response::try_from_http_response(data) + .expect("Can't parse the `/keys/upload` response") +} + +fn keys_query_response() -> get_keys::v3::Response { + let data = response_from_file(&test_json::KEYS_QUERY); + get_keys::v3::Response::try_from_http_response(data) + .expect("Can't parse the `/keys/upload` response") +} + +pub async fn get_prepared_machine_test_helper( + user_id: &UserId, + use_fallback_key: bool, +) -> (OlmMachine, OneTimeKeys) { + let machine = OlmMachine::new(user_id, bob_device_id()).await; + + let request = machine + .store() + .with_transaction(|mut tr| async { + let account = tr.account().await.unwrap(); + account.generate_fallback_key_if_needed(); + account.update_uploaded_key_count(0); + account.generate_one_time_keys_if_needed(); + let request = + machine.keys_for_upload(account).await.expect("Can't prepare initial key upload"); + Ok((tr, request)) + }) + .await + .unwrap(); + + let response = keys_upload_response(); + machine.receive_keys_upload_response(&response).await.unwrap(); + + let keys = if use_fallback_key { request.fallback_keys } else { request.one_time_keys }; + + (machine, keys) +} + +pub async fn get_machine_after_query_test_helper() -> (OlmMachine, OneTimeKeys) { + let (machine, otk) = get_prepared_machine_test_helper(user_id(), false).await; + let response = keys_query_response(); + let req_id = TransactionId::new(); + + machine.receive_keys_query_response(&req_id, &response).await.unwrap(); + + (machine, otk) +} + +pub async fn get_machine_pair( + alice: &UserId, + bob: &UserId, + use_fallback_key: bool, +) -> (OlmMachine, OlmMachine, OneTimeKeys) { + let (bob, otk) = get_prepared_machine_test_helper(bob, use_fallback_key).await; + + let alice_device = alice_device_id(); + let alice = OlmMachine::new(alice, alice_device).await; + + let alice_device = DeviceData::from_machine_test_helper(&alice).await.unwrap(); + let bob_device = DeviceData::from_machine_test_helper(&bob).await.unwrap(); + alice.store().save_device_data(&[bob_device]).await.unwrap(); + bob.store().save_device_data(&[alice_device]).await.unwrap(); + + (alice, bob, otk) +} + +pub async fn get_machine_pair_with_session( + alice: &UserId, + bob: &UserId, + use_fallback_key: bool, +) -> (OlmMachine, OlmMachine) { + let (alice, bob, mut one_time_keys) = get_machine_pair(alice, bob, use_fallback_key).await; + + let (device_key_id, one_time_key) = one_time_keys.pop_first().unwrap(); + + let one_time_keys = BTreeMap::from([( + bob.user_id().to_owned(), + BTreeMap::from([( + bob.device_id().to_owned(), + BTreeMap::from([(device_key_id, one_time_key)]), + )]), + )]); + + let response = claim_keys::v3::Response::new(one_time_keys); + alice.inner.session_manager.create_sessions(&response).await.unwrap(); + + (alice, bob) +} + +pub async fn get_machine_pair_with_setup_sessions_test_helper( + alice: &UserId, + bob: &UserId, + use_fallback_key: bool, +) -> (OlmMachine, OlmMachine) { + let (alice, bob) = get_machine_pair_with_session(alice, bob, use_fallback_key).await; + + let bob_device = alice.get_device(bob.user_id(), bob.device_id(), None).await.unwrap().unwrap(); + + let (session, content) = + bob_device.encrypt("m.dummy", ToDeviceDummyEventContent::new()).await.unwrap(); + alice.store().save_sessions(&[session]).await.unwrap(); + + let event = ToDeviceEvent::new(alice.user_id().to_owned(), content.deserialize_as().unwrap()); + + let decrypted = bob + .store() + .with_transaction(|mut tr| async { + let res = bob.decrypt_to_device_event(&mut tr, &event, &mut Changes::default()).await?; + Ok((tr, res)) + }) + .await + .unwrap(); + + bob.store().save_sessions(&[decrypted.session.session()]).await.unwrap(); + + (alice, bob) +} diff --git a/crates/matrix-sdk-crypto/src/store/crypto_store_wrapper.rs b/crates/matrix-sdk-crypto/src/store/crypto_store_wrapper.rs index 96ea8e237..f09cedfcc 100644 --- a/crates/matrix-sdk-crypto/src/store/crypto_store_wrapper.rs +++ b/crates/matrix-sdk-crypto/src/store/crypto_store_wrapper.rs @@ -289,7 +289,7 @@ mod test { use ruma::user_id; use super::*; - use crate::machine::tests::get_machine_pair_with_setup_sessions_test_helper; + use crate::machine::test_helpers::get_machine_pair_with_setup_sessions_test_helper; #[async_test] async fn cache_cleared_after_device_update() { diff --git a/crates/matrix-sdk-crypto/src/store/mod.rs b/crates/matrix-sdk-crypto/src/store/mod.rs index 85712c3c0..5a752f066 100644 --- a/crates/matrix-sdk-crypto/src/store/mod.rs +++ b/crates/matrix-sdk-crypto/src/store/mod.rs @@ -1941,7 +1941,7 @@ mod tests { use matrix_sdk_test::async_test; use ruma::{room_id, user_id}; - use crate::{machine::tests::get_machine_pair, types::EventEncryptionAlgorithm}; + use crate::{machine::test_helpers::get_machine_pair, types::EventEncryptionAlgorithm}; #[async_test] async fn import_room_keys_notifies_stream() {