From b1e0aca34ced476904f8a0a583e375382c18f414 Mon Sep 17 00:00:00 2001
From: Richard van der Hoff <richard@matrix.org>
Date: Fri, 9 Aug 2024 15:09:05 +0100
Subject: [PATCH] wip

---
 crates/matrix-sdk-crypto/src/machine/mod.rs   |  2 +-
 .../src/machine/tests/megolm_sender_data.rs   | 57 ++++++++++++++++++-
 2 files changed, 57 insertions(+), 2 deletions(-)

diff --git a/crates/matrix-sdk-crypto/src/machine/mod.rs b/crates/matrix-sdk-crypto/src/machine/mod.rs
index cd0cf50b205..15a361363a9 100644
--- a/crates/matrix-sdk-crypto/src/machine/mod.rs
+++ b/crates/matrix-sdk-crypto/src/machine/mod.rs
@@ -870,7 +870,7 @@ impl OlmMachine {
                 session.sender_data = sender_data;
 
                 if self.store().compare_group_session(&session).await? == SessionOrdering::Better {
-                    info!("Received a new megolm room key");
+                    info!(?session, "Received a new megolm room key");
 
                     Ok(Some(session))
                 } else {
diff --git a/crates/matrix-sdk-crypto/src/machine/tests/megolm_sender_data.rs b/crates/matrix-sdk-crypto/src/machine/tests/megolm_sender_data.rs
index 3e4b09c67f3..60ac677ee01 100644
--- a/crates/matrix-sdk-crypto/src/machine/tests/megolm_sender_data.rs
+++ b/crates/matrix-sdk-crypto/src/machine/tests/megolm_sender_data.rs
@@ -32,7 +32,7 @@ use crate::{
     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 +105,61 @@ 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
+        }
+    );
+}
+
 /// Convenience wrapper for [`get_machine_pair_with_setup_sessions_test_helper`]
 /// using standard user ids.
 async fn get_machine_pair() -> (OlmMachine, OlmMachine) {