crypto: update senderdata integration tests

Extend the integration tests for megolm sender data to check that we update
existing inbound group sessions when we get a `/keys/query` response.

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

@@ -17,17 +17,24 @@
 
 use std::collections::BTreeMap;
 
+use as_variant::as_variant;
 use matrix_sdk_test::{ruma_response_from_json, test_json};
 use ruma::{
-    api::client::keys::{claim_keys, get_keys, upload_keys},
+    api::client::keys::{
+        claim_keys, get_keys, get_keys::v3::Response as KeysQueryResponse, upload_keys,
+    },
     device_id,
     encryption::OneTimeKey,
     events::dummy::ToDeviceDummyEventContent,
     serde::Raw,
     user_id, DeviceId, OwnedDeviceKeyId, TransactionId, UserId,
 };
+use serde_json::json;
 
-use crate::{store::Changes, types::events::ToDeviceEvent, DeviceData, OlmMachine};
+use crate::{
+    store::Changes, types::events::ToDeviceEvent, CrossSigningBootstrapRequests, DeviceData,
+    OlmMachine, OutgoingRequests,
+};
 
 /// These keys need to be periodically uploaded to the server.
 type OneTimeKeys = BTreeMap<OwnedDeviceKeyId, Raw<OneTimeKey>>;
@@ -182,3 +189,36 @@ pub async fn create_session(
     let response = claim_keys::v3::Response::new(one_time_keys);
     machine.inner.session_manager.create_sessions(&response).await.unwrap();
 }
+
+/// Given a set of requests returned by `bootstrap_cross_signing` for one user,
+/// return a `/keys/query` response which might be returned to another user/
+pub fn bootstrap_requests_to_keys_query_response(
+    bootstrap_requests: CrossSigningBootstrapRequests,
+) -> KeysQueryResponse {
+    let mut kq_response = json!({});
+
+    // If we have a master key, add that to the response
+    bootstrap_requests.upload_signing_keys_req.master_key.map(|key| {
+        let user_id = key.user_id.clone();
+        kq_response["master_keys"] = json!({user_id: key});
+    });
+
+    // If we have a self-signing key, add that
+    bootstrap_requests.upload_signing_keys_req.self_signing_key.map(|key| {
+        let user_id = key.user_id.clone();
+        kq_response["self_signing_keys"] = json!({user_id: key});
+    });
+
+    // And if we have a device, add that
+    bootstrap_requests
+        .upload_keys_req
+        .and_then(|req| as_variant!(req.request.as_ref(), OutgoingRequests::KeysUpload).cloned())
+        .and_then(|keys_upload_request| keys_upload_request.device_keys)
+        .map(|dk| {
+            let user_id: String = dk.get_field("user_id").unwrap().unwrap();
+            let device_id: String = dk.get_field("device_id").unwrap().unwrap();
+            kq_response["device_keys"] = json!({user_id: { device_id: dk }});
+        });
+
+    matrix_sdk_test::ruma_response_from_json(&kq_response)
+}

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

@@ -26,13 +26,16 @@ use serde_json::json;
 
 use crate::{
     machine::{
-        test_helpers::get_machine_pair_with_setup_sessions_test_helper,
+        test_helpers::{
+            bootstrap_requests_to_keys_query_response,
+            get_machine_pair_with_setup_sessions_test_helper,
+        },
         tests::to_device_requests_to_content,
     },
     olm::{InboundGroupSession, SenderData},
     store::RoomKeyInfo,
     types::events::{room::encrypted::ToDeviceEncryptedEventContent, EventType, ToDeviceEvent},
-    EncryptionSettings, EncryptionSyncChanges, OlmMachine, Session,
+    DeviceData, EncryptionSettings, EncryptionSyncChanges, OlmMachine, Session,
 };
 
 /// Test the behaviour when a megolm session is received from an unknown device,
@@ -105,6 +108,105 @@ async fn test_receive_megolm_session_from_known_device() {
     );
 }
 
+/// If we have a megolm session from an unknown device, test what happens when
+/// we get a /keys/query response that includes that device.
+#[async_test]
+async fn test_update_unknown_device_senderdata_on_keys_query() {
+    // GIVEN we have a megolm session from an unknown device
+
+    let (alice, bob) = get_machine_pair().await;
+    let mut bob_room_keys_received_stream = Box::pin(bob.store().room_keys_received_stream());
+
+    // `get_machine_pair_with_setup_sessions_test_helper` tells Bob about Alice's
+    // device keys, so to run this test, we need to make him forget them.
+    forget_devices_for_user(&bob, alice.user_id()).await;
+
+    // Alice starts a megolm session and shares the key with Bob, *without* sending
+    // the sender data.
+    let room_id = room_id!("!test:example.org");
+    let event = create_and_share_session_without_sender_data(&alice, &bob, room_id).await;
+
+    // Bob receives the to-device message
+    receive_to_device_event(&bob, &event).await;
+
+    // and now Bob should know about the session.
+    let room_key_info = get_room_key_received_update(&mut bob_room_keys_received_stream);
+    let session = get_inbound_group_session_or_assert(&bob, &room_key_info).await;
+
+    // Double-check that it is, in fact, an unknown device session.
+    assert_matches!(session.sender_data, SenderData::UnknownDevice { .. });
+
+    // WHEN Bob gets a /keys/query response for Alice, that includes the
+    // sending device...
+
+    let alice_device = DeviceData::from_machine_test_helper(&alice).await.unwrap();
+    let kq_response = json!({
+        "device_keys": { alice.user_id() : { alice.device_id():  alice_device.as_device_keys()}}
+    });
+    bob.receive_keys_query_response(
+        &TransactionId::new(),
+        &matrix_sdk_test::ruma_response_from_json(&kq_response),
+    )
+    .await
+    .unwrap();
+
+    // THEN Bob should have received an update about the session, and it should now
+    // be `SenderData::DeviceInfo`
+    let room_key_info = get_room_key_received_update(&mut bob_room_keys_received_stream);
+    let session = get_inbound_group_session_or_assert(&bob, &room_key_info).await;
+
+    assert_matches!(
+        session.sender_data,
+        SenderData::DeviceInfo {legacy_session, ..} => {
+            assert_eq!(legacy_session, true); // TODO: change when https://github.com/matrix-org/matrix-rust-sdk/pull/3785 lands
+        }
+    );
+}
+
+/// If we have a megolm session from an unsigned device, test what happens when
+/// we get a /keys/query response that includes that device.
+#[async_test]
+async fn test_update_device_info_senderdata_on_keys_query() {
+    // GIVEN we have a megolm session from an unsigned device
+
+    let (alice, bob) = get_machine_pair().await;
+    let mut bob_room_keys_received_stream = Box::pin(bob.store().room_keys_received_stream());
+
+    // Alice starts a megolm session and shares the key with Bob
+    let room_id = room_id!("!test:example.org");
+
+    let to_device_requests = alice
+        .share_room_key(room_id, iter::once(bob.user_id()), EncryptionSettings::default())
+        .await
+        .unwrap();
+    let event = ToDeviceEvent::new(
+        alice.user_id().to_owned(),
+        to_device_requests_to_content(to_device_requests),
+    );
+    // Bob receives the to-device message
+    receive_to_device_event(&bob, &event).await;
+
+    // and now Bob should know about the session.
+    let room_key_info = get_room_key_received_update(&mut bob_room_keys_received_stream);
+    let session = get_inbound_group_session_or_assert(&bob, &room_key_info).await;
+
+    // Double-check that it is, in fact, an unverified device session.
+    assert_matches!(session.sender_data, SenderData::DeviceInfo { .. });
+
+    // WHEN Bob receives a /keys/query response for Alice that includes a verifiable
+    // signature for her device
+    let bootstrap_requests = alice.bootstrap_cross_signing(false).await.unwrap();
+    let kq_response = bootstrap_requests_to_keys_query_response(bootstrap_requests);
+    bob.receive_keys_query_response(&TransactionId::new(), &kq_response).await.unwrap();
+
+    // THEN Bob should have received an update about the session, and it should now
+    // be `SenderData::SenderUnverified`
+    let room_key_info = get_room_key_received_update(&mut bob_room_keys_received_stream);
+    let session = get_inbound_group_session_or_assert(&bob, &room_key_info).await;
+
+    assert_matches!(session.sender_data, SenderData::SenderUnverified(_));
+}
+
 /// Convenience wrapper for [`get_machine_pair_with_setup_sessions_test_helper`]
 /// using standard user ids.
 async fn get_machine_pair() -> (OlmMachine, OlmMachine) {