// 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.

use std::{collections::BTreeMap, iter, ops::Not, sync::Arc, time::Duration};

use assert_matches2::{assert_let, assert_matches};
use futures_util::{pin_mut, FutureExt, StreamExt};
use itertools::Itertools;
use matrix_sdk_common::deserialized_responses::{
    UnableToDecryptInfo, UnableToDecryptReason, UnsignedDecryptionResult, UnsignedEventLocation,
    WithheldCode,
};
use matrix_sdk_test::{async_test, message_like_event_content, ruma_response_from_json, test_json};
use ruma::{
    api::client::{
        keys::{get_keys, upload_keys},
        sync::sync_events::DeviceLists,
    },
    device_id,
    events::{
        room::message::{
            AddMentions, MessageType, Relation, ReplyWithinThread, RoomMessageEventContent,
        },
        AnyMessageLikeEvent, AnyMessageLikeEventContent, AnyToDeviceEvent, MessageLikeEvent,
        OriginalMessageLikeEvent,
    },
    room_id,
    serde::Raw,
    uint, user_id, DeviceId, DeviceKeyAlgorithm, DeviceKeyId, MilliSecondsSinceUnixEpoch,
    OneTimeKeyAlgorithm, TransactionId, UserId,
};
use serde_json::json;
use vodozemac::{
    megolm::{GroupSession, SessionConfig},
    Ed25519PublicKey,
};

use super::CrossSigningBootstrapRequests;
use crate::{
    error::EventError,
    machine::{
        test_helpers::{
            get_machine_after_query_test_helper, get_machine_pair_with_session,
            get_machine_pair_with_setup_sessions_test_helper, get_prepared_machine_test_helper,
        },
        EncryptionSyncChanges, OlmMachine,
    },
    olm::{BackedUpRoomKey, ExportedRoomKey, SenderData, VerifyJson},
    session_manager::CollectStrategy,
    store::{BackupDecryptionKey, Changes, CryptoStore, MemoryStore},
    types::{
        events::{
            room::encrypted::{EncryptedToDeviceEvent, ToDeviceEncryptedEventContent},
            room_key_withheld::{MegolmV1AesSha2WithheldContent, RoomKeyWithheldContent},
            ToDeviceEvent,
        },
        requests::{AnyOutgoingRequest, ToDeviceRequest},
        DeviceKeys, SignedKey, SigningKeys,
    },
    utilities::json_convert,
    verification::tests::bob_id,
    Account, DecryptionSettings, DeviceData, EncryptionSettings, MegolmError, OlmError,
    RoomEventDecryptionResult, TrustRequirement,
};

mod decryption_verification_state;
mod interactive_verification;
mod megolm_sender_data;
mod olm_encryption;
mod room_settings;
mod send_encrypted_to_device;

fn alice_id() -> &'static UserId {
    user_id!("@alice:example.org")
}

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 json = &test_json::KEYS_UPLOAD;
    ruma_response_from_json(json)
}

fn keys_query_response() -> get_keys::v3::Response {
    let json = &test_json::KEYS_QUERY;
    ruma_response_from_json(json)
}

pub fn to_device_requests_to_content(
    requests: Vec<Arc<ToDeviceRequest>>,
) -> ToDeviceEncryptedEventContent {
    let to_device_request = &requests[0];

    to_device_request
        .messages
        .values()
        .next()
        .unwrap()
        .values()
        .next()
        .unwrap()
        .deserialize_as()
        .unwrap()
}

#[async_test]
async fn test_create_olm_machine() {
    let test_start_ts = MilliSecondsSinceUnixEpoch::now();
    let machine = OlmMachine::new(user_id(), alice_device_id()).await;

    let device_creation_time = machine.device_creation_time();
    assert!(device_creation_time <= MilliSecondsSinceUnixEpoch::now());
    assert!(device_creation_time >= test_start_ts);

    let cache = machine.store().cache().await.unwrap();
    let account = cache.account().await.unwrap();
    assert!(!account.shared());

    let own_device = machine
        .get_device(machine.user_id(), machine.device_id(), None)
        .await
        .unwrap()
        .expect("We should always have our own device in the store");

    assert!(own_device.is_locally_trusted(), "Our own device should always be locally trusted");
}

#[async_test]
async fn test_generate_one_time_keys() {
    let machine = OlmMachine::new(user_id(), alice_device_id()).await;

    machine
        .store()
        .with_transaction(|mut tr| async {
            let account = tr.account().await.unwrap();
            assert!(account.generate_one_time_keys_if_needed().is_some());
            Ok((tr, ()))
        })
        .await
        .unwrap();

    let mut response = keys_upload_response();

    machine.receive_keys_upload_response(&response).await.unwrap();

    machine
        .store()
        .with_transaction(|mut tr| async {
            let account = tr.account().await.unwrap();
            assert!(account.generate_one_time_keys_if_needed().is_some());
            Ok((tr, ()))
        })
        .await
        .unwrap();

    response.one_time_key_counts.insert(OneTimeKeyAlgorithm::SignedCurve25519, uint!(50));

    machine.receive_keys_upload_response(&response).await.unwrap();

    machine
        .store()
        .with_transaction(|mut tr| async {
            let account = tr.account().await.unwrap();
            assert!(account.generate_one_time_keys_if_needed().is_none());

            Ok((tr, ()))
        })
        .await
        .unwrap();
}

#[async_test]
async fn test_device_key_signing() {
    let machine = OlmMachine::new(user_id(), alice_device_id()).await;

    let (device_keys, identity_keys) = {
        let cache = machine.store().cache().await.unwrap();
        let account = cache.account().await.unwrap();
        let device_keys = account.device_keys();
        let identity_keys = account.identity_keys();
        (device_keys, identity_keys)
    };

    let ed25519_key = identity_keys.ed25519;

    let ret = ed25519_key.verify_json(
        machine.user_id(),
        &DeviceKeyId::from_parts(DeviceKeyAlgorithm::Ed25519, machine.device_id()),
        &device_keys,
    );
    ret.unwrap();
}

#[async_test]
async fn test_session_invalidation() {
    let machine = OlmMachine::new(user_id(), alice_device_id()).await;
    let room_id = room_id!("!test:example.org");

    machine.create_outbound_group_session_with_defaults_test_helper(room_id).await.unwrap();
    assert!(machine.inner.group_session_manager.get_outbound_group_session(room_id).is_some());

    machine.discard_room_key(room_id).await.unwrap();

    assert!(machine
        .inner
        .group_session_manager
        .get_outbound_group_session(room_id)
        .unwrap()
        .invalidated());
}

#[test]
fn test_invalid_signature() {
    let account = Account::with_device_id(user_id(), alice_device_id());

    let device_keys = account.device_keys();

    let key = Ed25519PublicKey::from_slice(&[0u8; 32]).unwrap();

    let ret = key.verify_json(
        account.user_id(),
        &DeviceKeyId::from_parts(DeviceKeyAlgorithm::Ed25519, account.device_id()),
        &device_keys,
    );
    ret.unwrap_err();
}

#[test]
fn test_one_time_key_signing() {
    let mut account = Account::with_device_id(user_id(), alice_device_id());
    account.update_uploaded_key_count(49);
    account.generate_one_time_keys_if_needed();

    let mut one_time_keys = account.signed_one_time_keys();
    let ed25519_key = account.identity_keys().ed25519;

    let one_time_key: SignedKey = one_time_keys
        .values_mut()
        .next()
        .expect("One time keys should be generated")
        .deserialize_as()
        .unwrap();

    ed25519_key
        .verify_json(
            account.user_id(),
            &DeviceKeyId::from_parts(DeviceKeyAlgorithm::Ed25519, account.device_id()),
            &one_time_key,
        )
        .expect("One-time key has been signed successfully");
}

#[async_test]
async fn test_keys_for_upload() {
    let machine = OlmMachine::new(user_id(), alice_device_id()).await;

    let key_counts = BTreeMap::from([(OneTimeKeyAlgorithm::SignedCurve25519, 49u8.into())]);
    machine
        .receive_sync_changes(EncryptionSyncChanges {
            to_device_events: Vec::new(),
            changed_devices: &Default::default(),
            one_time_keys_counts: &key_counts,
            unused_fallback_keys: None,
            next_batch_token: None,
        })
        .await
        .expect("We should be able to update our one-time key counts");

    let (ed25519_key, mut request) = {
        let cache = machine.store().cache().await.unwrap();
        let account = cache.account().await.unwrap();
        let ed25519_key = account.identity_keys().ed25519;

        let request =
            machine.keys_for_upload(&account).await.expect("Can't prepare initial key upload");
        (ed25519_key, request)
    };

    let one_time_key: SignedKey = request
        .one_time_keys
        .values_mut()
        .next()
        .expect("One time keys should be generated")
        .deserialize_as()
        .unwrap();

    let ret = ed25519_key.verify_json(
        machine.user_id(),
        &DeviceKeyId::from_parts(DeviceKeyAlgorithm::Ed25519, machine.device_id()),
        &one_time_key,
    );
    ret.unwrap();

    let device_keys: DeviceKeys = request.device_keys.unwrap().deserialize_as().unwrap();

    let ret = ed25519_key.verify_json(
        machine.user_id(),
        &DeviceKeyId::from_parts(DeviceKeyAlgorithm::Ed25519, machine.device_id()),
        &device_keys,
    );
    ret.unwrap();

    let response = {
        let cache = machine.store().cache().await.unwrap();
        let account = cache.account().await.unwrap();

        let mut response = keys_upload_response();
        response.one_time_key_counts.insert(
            OneTimeKeyAlgorithm::SignedCurve25519,
            account.max_one_time_keys().try_into().unwrap(),
        );

        response
    };

    machine.receive_keys_upload_response(&response).await.unwrap();

    {
        let cache = machine.store().cache().await.unwrap();
        let account = cache.account().await.unwrap();
        let ret = machine.keys_for_upload(&account).await;
        assert!(ret.is_none());
    }
}

#[async_test]
async fn test_keys_query() {
    let (machine, _) = get_prepared_machine_test_helper(user_id(), false).await;
    let response = keys_query_response();
    let alice_id = user_id!("@alice:example.org");
    let alice_device_id: &DeviceId = device_id!("JLAFKJWSCS");

    let alice_devices = machine.store().get_user_devices(alice_id).await.unwrap();
    assert!(alice_devices.devices().peekable().peek().is_none());

    let req_id = TransactionId::new();
    machine.receive_keys_query_response(&req_id, &response).await.unwrap();

    let device = machine.store().get_device(alice_id, alice_device_id).await.unwrap().unwrap();
    assert_eq!(device.user_id(), alice_id);
    assert_eq!(device.device_id(), alice_device_id);
}

#[async_test]
async fn test_query_keys_for_users() {
    let (machine, _) = get_prepared_machine_test_helper(user_id(), false).await;
    let alice_id = user_id!("@alice:example.org");
    let (_, request) = machine.query_keys_for_users(vec![alice_id]);
    assert!(request.device_keys.contains_key(alice_id));
}

#[async_test]
async fn test_missing_sessions_calculation() {
    let (machine, _) = get_machine_after_query_test_helper().await;

    let alice = alice_id();
    let alice_device = alice_device_id();

    let (_, missing_sessions) =
        machine.get_missing_sessions(iter::once(alice)).await.unwrap().unwrap();

    assert!(missing_sessions.one_time_keys.contains_key(alice));
    let user_sessions = missing_sessions.one_time_keys.get(alice).unwrap();
    assert!(user_sessions.contains_key(alice_device));
}

#[async_test]
async fn test_room_key_sharing() {
    let (alice, bob) = get_machine_pair_with_session(alice_id(), user_id(), false).await;

    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),
    );
    let event = json_convert(&event).unwrap();

    let alice_session =
        alice.inner.group_session_manager.get_outbound_group_session(room_id).unwrap();

    let (decrypted, room_key_updates) = bob
        .receive_sync_changes(EncryptionSyncChanges {
            to_device_events: vec![event],
            changed_devices: &Default::default(),
            one_time_keys_counts: &Default::default(),
            unused_fallback_keys: None,
            next_batch_token: None,
        })
        .await
        .unwrap();

    let event = decrypted[0].deserialize().unwrap();

    if let AnyToDeviceEvent::RoomKey(event) = event {
        assert_eq!(&event.sender, alice.user_id());
        assert!(event.content.session_key.is_empty());
    } else {
        panic!("expected RoomKeyEvent found {event:?}");
    }

    let session = bob.store().get_inbound_group_session(room_id, alice_session.session_id()).await;

    assert!(session.unwrap().is_some());

    assert_eq!(room_key_updates.len(), 1);
    assert_eq!(room_key_updates[0].room_id, room_id);
    assert_eq!(room_key_updates[0].session_id, alice_session.session_id());
}

#[async_test]
async fn test_request_missing_secrets() {
    let (alice, _) = get_machine_pair_with_session(alice_id(), bob_id(), false).await;

    let should_query_secrets = alice.query_missing_secrets_from_other_sessions().await.unwrap();

    assert!(should_query_secrets);

    let outgoing_to_device = alice
        .outgoing_requests()
        .await
        .unwrap()
        .into_iter()
        .filter(|outgoing| match outgoing.request.as_ref() {
            AnyOutgoingRequest::ToDeviceRequest(request) => {
                request.event_type.to_string() == "m.secret.request"
            }
            _ => false,
        })
        .collect_vec();

    assert_eq!(outgoing_to_device.len(), 4);

    // The second time, as there are already in-flight requests, it should have no
    // effect.
    let should_query_secrets_now = alice.query_missing_secrets_from_other_sessions().await.unwrap();
    assert!(!should_query_secrets_now);
}

#[async_test]
async fn test_request_missing_secrets_cross_signed() {
    let (alice, bob) = get_machine_pair_with_session(alice_id(), bob_id(), false).await;

    setup_cross_signing_for_machine_test_helper(&alice, &bob).await;

    let should_query_secrets = alice.query_missing_secrets_from_other_sessions().await.unwrap();

    assert!(should_query_secrets);

    let outgoing_to_device = alice
        .outgoing_requests()
        .await
        .unwrap()
        .into_iter()
        .filter(|outgoing| match outgoing.request.as_ref() {
            AnyOutgoingRequest::ToDeviceRequest(request) => {
                request.event_type.to_string() == "m.secret.request"
            }
            _ => false,
        })
        .collect_vec();
    assert_eq!(outgoing_to_device.len(), 1);

    // The second time, as there are already in-flight requests, it should have no
    // effect.
    let should_query_secrets_now = alice.query_missing_secrets_from_other_sessions().await.unwrap();
    assert!(!should_query_secrets_now);
}

#[async_test]
async fn test_megolm_encryption() {
    let (alice, bob) =
        get_machine_pair_with_setup_sessions_test_helper(alice_id(), user_id(), false).await;
    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),
    );

    let mut room_keys_received_stream = Box::pin(bob.store().room_keys_received_stream());

    let group_session = 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()
        .inbound_group_session
        .unwrap();
    bob.store().save_inbound_group_sessions(&[group_session.clone()]).await.unwrap();

    // when we decrypt the room key, the
    // inbound_group_session_streamroom_keys_received_stream should tell us
    // about it.
    let room_keys = room_keys_received_stream
        .next()
        .now_or_never()
        .flatten()
        .expect("We should have received an update of room key infos")
        .unwrap();
    assert_eq!(room_keys.len(), 1);
    assert_eq!(room_keys[0].session_id, group_session.session_id());

    let plaintext = "It is a secret to everybody";

    let content = RoomMessageEventContent::text_plain(plaintext);

    let encrypted_content = alice
        .encrypt_room_event(room_id, AnyMessageLikeEventContent::RoomMessage(content.clone()))
        .await
        .unwrap();

    let event = json!({
        "event_id": "$xxxxx:example.org",
        "origin_server_ts": MilliSecondsSinceUnixEpoch::now(),
        "sender": alice.user_id(),
        "type": "m.room.encrypted",
        "content": encrypted_content,
    });

    let event = json_convert(&event).unwrap();

    let decryption_settings =
        DecryptionSettings { sender_device_trust_requirement: TrustRequirement::Untrusted };

    let decryption_result =
        bob.try_decrypt_room_event(&event, room_id, &decryption_settings).await.unwrap();
    assert_let!(RoomEventDecryptionResult::Decrypted(decrypted_event) = decryption_result);
    let decrypted_event = decrypted_event.event.deserialize().unwrap();

    if let AnyMessageLikeEvent::RoomMessage(MessageLikeEvent::Original(
        OriginalMessageLikeEvent { sender, content, .. },
    )) = decrypted_event
    {
        assert_eq!(&sender, alice.user_id());
        if let MessageType::Text(c) = &content.msgtype {
            assert_eq!(&c.body, plaintext);
        } else {
            panic!("Decrypted event has a mismatched content");
        }
    } else {
        panic!("Decrypted room event has the wrong type");
    }

    // Just decrypting the event should *not* cause an update on the
    // inbound_group_session_stream.
    if let Some(igs) = room_keys_received_stream.next().now_or_never() {
        panic!("Session stream unexpectedly returned update: {igs:?}");
    }
}

#[async_test]
async fn test_withheld_unverified() {
    let (alice, bob) =
        get_machine_pair_with_setup_sessions_test_helper(alice_id(), user_id(), false).await;
    let room_id = room_id!("!test:example.org");

    let room_keys_withheld_received_stream = bob.store().room_keys_withheld_received_stream();
    pin_mut!(room_keys_withheld_received_stream);

    let encryption_settings = EncryptionSettings::default();
    let encryption_settings = EncryptionSettings {
        sharing_strategy: CollectStrategy::DeviceBasedStrategy {
            only_allow_trusted_devices: true,
            error_on_verified_user_problem: false,
        },
        ..encryption_settings
    };

    let to_device_requests = alice
        .share_room_key(room_id, iter::once(bob.user_id()), encryption_settings)
        .await
        .expect("Share room key should be ok");

    // Here there will be only one request, and it's for a m.room_key.withheld

    // Transform that into an event to feed it back to bob machine
    let wh_content = to_device_requests[0]
        .messages
        .values()
        .next()
        .unwrap()
        .values()
        .next()
        .unwrap()
        .deserialize_as::<RoomKeyWithheldContent>()
        .expect("Deserialize should work");

    let event = ToDeviceEvent::new(alice.user_id().to_owned(), wh_content);

    let event = json_convert(&event).unwrap();

    bob.receive_sync_changes(EncryptionSyncChanges {
        to_device_events: vec![event],
        changed_devices: &Default::default(),
        one_time_keys_counts: &Default::default(),
        unused_fallback_keys: None,
        next_batch_token: None,
    })
    .await
    .unwrap();

    // We should receive a notification on the room_keys_withheld_received_stream
    let withheld_received = room_keys_withheld_received_stream
        .next()
        .now_or_never()
        .flatten()
        .expect("We should have received a notification of room key being withheld");
    assert_eq!(withheld_received.len(), 1);

    assert_eq!(&withheld_received[0].room_id, room_id);
    assert_matches!(
        &withheld_received[0].withheld_event.content,
        RoomKeyWithheldContent::MegolmV1AesSha2(MegolmV1AesSha2WithheldContent::Unverified(
            unverified_withheld_content
        ))
    );
    assert_eq!(unverified_withheld_content.room_id, room_id);

    let plaintext = "You shouldn't be able to decrypt that message";

    let content = RoomMessageEventContent::text_plain(plaintext);

    let content = alice
        .encrypt_room_event(room_id, AnyMessageLikeEventContent::RoomMessage(content.clone()))
        .await
        .unwrap();

    let room_event = json!({
        "event_id": "$xxxxx:example.org",
        "origin_server_ts": MilliSecondsSinceUnixEpoch::now(),
        "sender": alice.user_id(),
        "type": "m.room.encrypted",
        "content": content,
    });
    let room_event = json_convert(&room_event).unwrap();

    let decryption_settings =
        DecryptionSettings { sender_device_trust_requirement: TrustRequirement::Untrusted };
    let decrypt_result = bob.decrypt_room_event(&room_event, room_id, &decryption_settings).await;

    assert_matches!(&decrypt_result, Err(MegolmError::MissingRoomKey(Some(_))));

    let err = decrypt_result.err().unwrap();
    assert_matches!(err, MegolmError::MissingRoomKey(Some(WithheldCode::Unverified)));

    // Also check `try_decrypt_room_event`.
    let decrypt_result =
        bob.try_decrypt_room_event(&room_event, room_id, &decryption_settings).await.unwrap();
    assert_let!(RoomEventDecryptionResult::UnableToDecrypt(utd_info) = decrypt_result);
    assert!(utd_info.session_id.is_some());
    assert_eq!(
        utd_info.reason,
        UnableToDecryptReason::MissingMegolmSession {
            withheld_code: Some(WithheldCode::Unverified)
        }
    );
}

/// Test what happens when we feed an unencrypted event into the decryption
/// functions
#[async_test]
async fn test_decrypt_unencrypted_event() {
    let (bob, _) = get_prepared_machine_test_helper(user_id(), false).await;
    let room_id = room_id!("!test:example.org");

    let event = json!({
        "event_id": "$xxxxx:example.org",
        "origin_server_ts": MilliSecondsSinceUnixEpoch::now(),
        "sender": user_id(),
        // it's actually the lack of an `algorithm` that upsets it, rather than the event type.
        "type": "m.room.encrypted",
        "content":  RoomMessageEventContent::text_plain("plain"),
    });

    let event = json_convert(&event).unwrap();

    // decrypt_room_event should return an error
    let decryption_settings =
        DecryptionSettings { sender_device_trust_requirement: TrustRequirement::Untrusted };
    assert_matches!(
        bob.decrypt_room_event(&event, room_id, &decryption_settings).await,
        Err(MegolmError::JsonError(..))
    );

    // so should get_room_event_encryption_info
    assert_matches!(
        bob.get_room_event_encryption_info(&event, room_id).await,
        Err(MegolmError::JsonError(..))
    );
}

pub async fn setup_cross_signing_for_machine_test_helper(alice: &OlmMachine, bob: &OlmMachine) {
    let CrossSigningBootstrapRequests { upload_signing_keys_req: alice_upload_signing, .. } =
        alice.bootstrap_cross_signing(false).await.expect("Expect Alice x-signing key request");

    let CrossSigningBootstrapRequests { upload_signing_keys_req: bob_upload_signing, .. } =
        bob.bootstrap_cross_signing(false).await.expect("Expect Bob x-signing key request");

    let bob_device_keys = bob
        .get_device(bob.user_id(), bob.device_id(), None)
        .await
        .unwrap()
        .unwrap()
        .as_device_keys()
        .to_owned();

    let alice_device_keys = alice
        .get_device(alice.user_id(), alice.device_id(), None)
        .await
        .unwrap()
        .unwrap()
        .as_device_keys()
        .to_owned();

    // We only want to setup cross signing we don't actually sign the current
    // devices. so we ignore the new device signatures
    let json = json!({
        "device_keys": {
            bob.user_id() : { bob.device_id() : bob_device_keys},
            alice.user_id() : { alice.device_id():  alice_device_keys }
        },
        "failures": {},
        "master_keys": {
            bob.user_id() : bob_upload_signing.master_key.unwrap(),
            alice.user_id() : alice_upload_signing.master_key.unwrap()
        },
        "user_signing_keys": {
            bob.user_id() : bob_upload_signing.user_signing_key.unwrap(),
            alice.user_id() : alice_upload_signing.user_signing_key.unwrap()
        },
        "self_signing_keys": {
            bob.user_id() : bob_upload_signing.self_signing_key.unwrap(),
            alice.user_id() : alice_upload_signing.self_signing_key.unwrap()
        },
      }
    );

    let kq_response = ruma_response_from_json(&json);
    alice.receive_keys_query_response(&TransactionId::new(), &kq_response).await.unwrap();
    bob.receive_keys_query_response(&TransactionId::new(), &kq_response).await.unwrap();
}

async fn sign_alice_device_for_machine_test_helper(alice: &OlmMachine, bob: &OlmMachine) {
    let CrossSigningBootstrapRequests {
        upload_signing_keys_req: upload_signing,
        upload_signatures_req: upload_signature,
        ..
    } = alice.bootstrap_cross_signing(false).await.expect("Expect Alice x-signing key request");

    let mut device_keys = alice
        .get_device(alice.user_id(), alice.device_id(), None)
        .await
        .unwrap()
        .unwrap()
        .as_device_keys()
        .to_owned();

    let raw_extracted =
        upload_signature.signed_keys.get(alice.user_id()).unwrap().iter().next().unwrap().1.get();

    let new_signature: DeviceKeys = serde_json::from_str(raw_extracted).unwrap();

    let self_sign_key_id = upload_signing
        .self_signing_key
        .as_ref()
        .unwrap()
        .get_first_key_and_id()
        .unwrap()
        .0
        .to_owned();

    device_keys.signatures.add_signature(
        alice.user_id().to_owned(),
        self_sign_key_id.to_owned(),
        new_signature.signatures.get_signature(alice.user_id(), &self_sign_key_id).unwrap(),
    );

    let updated_keys_with_x_signing = json!({ device_keys.device_id.to_string(): device_keys });

    let json = json!({
        "device_keys": {
            alice.user_id() : updated_keys_with_x_signing
        },
        "failures": {},
        "master_keys": {
            alice.user_id() : upload_signing.master_key.unwrap(),
        },
        "user_signing_keys": {
            alice.user_id() : upload_signing.user_signing_key.unwrap(),
        },
        "self_signing_keys": {
            alice.user_id() : upload_signing.self_signing_key.unwrap(),
        },
      }
    );

    let kq_response = ruma_response_from_json(&json);
    alice.receive_keys_query_response(&TransactionId::new(), &kq_response).await.unwrap();
    bob.receive_keys_query_response(&TransactionId::new(), &kq_response).await.unwrap();
}

#[async_test]
#[cfg(feature = "automatic-room-key-forwarding")]
async fn test_query_ratcheted_key() {
    let (alice, bob) =
        get_machine_pair_with_setup_sessions_test_helper(alice_id(), user_id(), false).await;
    let room_id = room_id!("!test:example.org");

    // Need a second bob session to check gossiping
    let bob_id = user_id();
    let bob_other_device = device_id!("OTHERBOB");
    let bob_other_machine = OlmMachine::new(bob_id, bob_other_device).await;
    let bob_other_device = DeviceData::from_machine_test_helper(&bob_other_machine).await.unwrap();
    bob.store().save_device_data(&[bob_other_device]).await.unwrap();
    bob.get_device(bob_id, device_id!("OTHERBOB"), None)
        .await
        .unwrap()
        .expect("should exist")
        .set_trust_state(crate::LocalTrust::Verified);

    alice.create_outbound_group_session_with_defaults_test_helper(room_id).await.unwrap();

    let plaintext = "It is a secret to everybody";

    let content = RoomMessageEventContent::text_plain(plaintext);

    let content = alice
        .encrypt_room_event(room_id, AnyMessageLikeEventContent::RoomMessage(content.clone()))
        .await
        .unwrap();

    let room_event = json!({
        "event_id": "$xxxxx:example.org",
        "origin_server_ts": MilliSecondsSinceUnixEpoch::now(),
        "sender": alice.user_id(),
        "type": "m.room.encrypted",
        "content": content,
    });

    // should share at index 1
    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),
    );

    let group_session = 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()
        .inbound_group_session;
    bob.store().save_inbound_group_sessions(&[group_session.unwrap()]).await.unwrap();

    let room_event = json_convert(&room_event).unwrap();

    let decryption_settings =
        DecryptionSettings { sender_device_trust_requirement: TrustRequirement::Untrusted };
    let decrypt_error =
        bob.decrypt_room_event(&room_event, room_id, &decryption_settings).await.unwrap_err();

    if let MegolmError::Decryption(vodo_error) = decrypt_error {
        if let vodozemac::megolm::DecryptionError::UnknownMessageIndex(_, _) = vodo_error {
            // check that key has been requested
            let outgoing_to_devices =
                bob.inner.key_request_machine.outgoing_to_device_requests().await.unwrap();
            assert_eq!(1, outgoing_to_devices.len());
        } else {
            panic!("Should be UnknownMessageIndex error ")
        }
    } else {
        panic!("Should have been unable to decrypt")
    }
}

#[async_test]
async fn test_room_key_over_megolm() {
    let (alice, bob) =
        get_machine_pair_with_setup_sessions_test_helper(alice_id(), user_id(), false).await;
    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 {
        sender: alice.user_id().to_owned(),
        content: to_device_requests_to_content(to_device_requests),
        other: Default::default(),
    };
    let event = json_convert(&event).unwrap();
    let changed_devices = DeviceLists::new();
    let key_counts: BTreeMap<_, _> = Default::default();

    let _ = bob
        .receive_sync_changes(EncryptionSyncChanges {
            to_device_events: vec![event],
            changed_devices: &changed_devices,
            one_time_keys_counts: &key_counts,
            unused_fallback_keys: None,
            next_batch_token: None,
        })
        .await
        .unwrap();

    let group_session = GroupSession::new(SessionConfig::version_1());
    let session_key = group_session.session_key();
    let session_id = group_session.session_id();

    let content = message_like_event_content!({
        "algorithm": "m.megolm.v1.aes-sha2",
        "room_id": room_id,
        "session_id": session_id,
        "session_key": session_key.to_base64(),
    });

    let encrypted_content =
        alice.encrypt_room_event_raw(room_id, "m.room_key", &content).await.unwrap();
    let event = json!({
        "sender": alice.user_id(),
        "content": encrypted_content,
        "type": "m.room.encrypted",
    });

    let event: EncryptedToDeviceEvent = serde_json::from_value(event).unwrap();

    let decrypt_result = 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;

    assert_matches!(decrypt_result, Err(OlmError::EventError(EventError::UnsupportedAlgorithm)));

    let event: Raw<AnyToDeviceEvent> = json_convert(&event).unwrap();

    bob.receive_sync_changes(EncryptionSyncChanges {
        to_device_events: vec![event],
        changed_devices: &changed_devices,
        one_time_keys_counts: &key_counts,
        unused_fallback_keys: None,
        next_batch_token: None,
    })
    .await
    .unwrap();

    let session = bob.store().get_inbound_group_session(room_id, &session_id).await;

    assert!(session.unwrap().is_none());
}

#[async_test]
async fn test_room_key_with_fake_identity_keys() {
    let room_id = room_id!("!test:localhost");
    let (alice, _) =
        get_machine_pair_with_setup_sessions_test_helper(alice_id(), user_id(), false).await;
    let device = DeviceData::from_machine_test_helper(&alice).await.unwrap();
    alice.store().save_device_data(&[device]).await.unwrap();

    let (outbound, mut inbound) = alice
        .store()
        .static_account()
        .create_group_session_pair(room_id, Default::default(), SenderData::unknown())
        .await
        .unwrap();

    let fake_key = Ed25519PublicKey::from_base64("ee3Ek+J2LkkPmjGPGLhMxiKnhiX//xcqaVL4RP6EypE")
        .unwrap()
        .into();
    let signing_keys = SigningKeys::from([(DeviceKeyAlgorithm::Ed25519, fake_key)]);
    inbound.creator_info.signing_keys = signing_keys.into();

    let content = message_like_event_content!({});
    let content = outbound.encrypt("m.dummy", &content).await;
    alice.store().save_inbound_group_sessions(&[inbound]).await.unwrap();

    let event = json!({
        "sender": alice.user_id(),
        "event_id": "$xxxxx:example.org",
        "origin_server_ts": MilliSecondsSinceUnixEpoch::now(),
        "type": "m.room.encrypted",
        "content": content,
    });
    let event = json_convert(&event).unwrap();

    let decryption_settings =
        DecryptionSettings { sender_device_trust_requirement: TrustRequirement::Untrusted };
    assert_matches!(
        alice.decrypt_room_event(&event, room_id, &decryption_settings).await,
        Err(MegolmError::MismatchedIdentityKeys { .. })
    );
}

#[async_test]
async fn test_importing_private_cross_signing_keys_verifies_the_public_identity() {
    async fn create_additional_machine(machine: &OlmMachine) -> OlmMachine {
        let second_machine = OlmMachine::new(machine.user_id(), "ADDITIONAL_MACHINE".into()).await;

        let identity = machine
            .get_identity(machine.user_id(), None)
            .await
            .unwrap()
            .expect("We should know about our own user identity if we bootstrapped it")
            .own()
            .unwrap();

        let mut changes = Changes::default();
        identity.mark_as_unverified();
        changes.identities.new.push(crate::UserIdentityData::Own(identity.inner));

        second_machine.store().save_changes(changes).await.unwrap();

        second_machine
    }

    let (alice, bob) =
        get_machine_pair_with_setup_sessions_test_helper(alice_id(), user_id(), false).await;
    setup_cross_signing_for_machine_test_helper(&alice, &bob).await;

    let second_alice = create_additional_machine(&alice).await;

    let export = alice
        .export_cross_signing_keys()
        .await
        .unwrap()
        .expect("We should be able to export our cross-signing keys");

    let identity = second_alice
        .get_identity(second_alice.user_id(), None)
        .await
        .unwrap()
        .expect("We should know about our own user identity")
        .own()
        .unwrap();

    assert!(!identity.is_verified(), "Initially our identity should not be verified");

    second_alice
        .import_cross_signing_keys(export)
        .await
        .expect("We should be able to import our cross-signing keys");

    let identity = second_alice
        .get_identity(second_alice.user_id(), None)
        .await
        .unwrap()
        .expect("We should know about our own user identity")
        .own()
        .unwrap();

    assert!(
        identity.is_verified(),
        "Our identity should be verified after we imported the private cross-signing keys"
    );

    let second_bob = create_additional_machine(&bob).await;

    let export = second_alice
        .export_cross_signing_keys()
        .await
        .unwrap()
        .expect("The machine should now be able to export cross-signing keys as well");

    second_bob.import_cross_signing_keys(export).await.expect_err(
        "Importing cross-signing keys that don't match our public identity should fail",
    );

    let identity = second_bob
        .get_identity(second_bob.user_id(), None)
        .await
        .unwrap()
        .expect("We should know about our own user identity")
        .own()
        .unwrap();

    assert!(
        !identity.is_verified(),
        "Our identity should not be verified when there's a mismatch in the cross-signing keys"
    );
}

#[async_test]
async fn test_wait_on_key_query_doesnt_block_store() {
    // Waiting for a key query shouldn't delay other write attempts to the store.
    // This test will end immediately if it works, and times out after a few seconds
    // if it failed.

    let machine = OlmMachine::new(bob_id(), bob_device_id()).await;

    // Mark Alice as a tracked user, so it gets into the groups of users for which
    // we need to query keys.
    machine.update_tracked_users([alice_id()]).await.unwrap();

    // Start a background task that will wait for the key query to finish silently
    // in the background.
    let machine_cloned = machine.clone();
    let wait = tokio::spawn(async move {
        let machine = machine_cloned;
        let user_devices =
            machine.get_user_devices(alice_id(), Some(Duration::from_secs(10))).await.unwrap();
        assert!(user_devices.devices().next().is_some());
    });

    // Let the background task work first.
    tokio::task::yield_now().await;

    // Create a key upload request and process it back immediately.
    let requests = machine.bootstrap_cross_signing(false).await.unwrap();

    let req = requests.upload_keys_req.expect("upload keys request should be there");
    let response = keys_upload_response();
    let mark_request_as_sent = machine.mark_request_as_sent(&req.request_id, &response);
    tokio::time::timeout(Duration::from_secs(5), mark_request_as_sent)
        .await
        .expect("no timeout")
        .expect("the underlying request has been marked as sent");

    // Answer the key query, so the background task completes immediately?
    let response = keys_query_response();
    let key_queries = machine.inner.identity_manager.users_for_key_query().await.unwrap();

    for (id, _) in key_queries {
        machine.mark_request_as_sent(&id, &response).await.unwrap();
    }

    // The waiting should successfully complete.
    wait.await.unwrap();
}

#[async_test]
async fn test_fix_incorrect_usage_of_backup_key_causing_decryption_errors() {
    let store = MemoryStore::new();

    let backup_decryption_key = BackupDecryptionKey::new().unwrap();

    store
        .save_changes(Changes {
            backup_decryption_key: Some(backup_decryption_key.clone()),
            backup_version: Some("1".to_owned()),
            ..Default::default()
        })
        .await
        .unwrap();

    // Some valid key data
    let data = json!({
       "algorithm": "m.megolm.v1.aes-sha2",
       "room_id": "!room:id",
       "sender_key": "FOvlmz18LLI3k/llCpqRoKT90+gFF8YhuL+v1YBXHlw",
       "session_id": "/2K+V777vipCxPZ0gpY9qcpz1DYaXwuMRIu0UEP0Wa0",
       "session_key": "AQAAAAAclzWVMeWBKH+B/WMowa3rb4ma3jEl6n5W4GCs9ue65CruzD3ihX+85pZ9hsV9Bf6fvhjp76WNRajoJYX0UIt7aosjmu0i+H+07hEQ0zqTKpVoSH0ykJ6stAMhdr6Q4uW5crBmdTTBIsqmoWsNJZKKoE2+ldYrZ1lrFeaJbjBIY/9ivle++74qQsT2dIKWPanKc9Q2Gl8LjESLtFBD9Fmt",
       "sender_claimed_keys": {
           "ed25519": "F4P7f1Z0RjbiZMgHk1xBCG3KC4/Ng9PmxLJ4hQ13sHA"
       },
       "forwarding_curve25519_key_chain": ["DBPC2zr6c9qimo9YRFK3RVr0Two/I6ODb9mbsToZN3Q", "bBc/qzZFOOKshMMT+i4gjS/gWPDoKfGmETs9yfw9430"]
    });

    let backed_up_room_key: BackedUpRoomKey = serde_json::from_value(data).unwrap();

    // Create the machine using `with_store` and without a call to enable_backup_v1,
    // like regenerate_olm would do
    let alice = OlmMachine::with_store(user_id(), alice_device_id(), store, None).await.unwrap();

    let exported_key = ExportedRoomKey::from_backed_up_room_key(
        room_id!("!room:id").to_owned(),
        "/2K+V777vipCxPZ0gpY9qcpz1DYaXwuMRIu0UEP0Wa0".into(),
        backed_up_room_key,
    );

    alice.store().import_exported_room_keys(vec![exported_key], |_, _| {}).await.unwrap();

    let (_, request) = alice.backup_machine().backup().await.unwrap().unwrap();

    let key_backup_data = request.rooms[&room_id!("!room:id").to_owned()]
        .sessions
        .get("/2K+V777vipCxPZ0gpY9qcpz1DYaXwuMRIu0UEP0Wa0")
        .unwrap()
        .deserialize()
        .unwrap();

    let ephemeral = key_backup_data.session_data.ephemeral.encode();
    let ciphertext = key_backup_data.session_data.ciphertext.encode();
    let mac = key_backup_data.session_data.mac.encode();

    // Prior to the fix for GHSA-9ggc-845v-gcgv, this would produce a
    // `Mac(MacError)`
    backup_decryption_key
        .decrypt_v1(&ephemeral, &mac, &ciphertext)
        .expect("The backed up key should be decrypted successfully");
}

#[async_test]
async fn test_olm_machine_with_custom_account() {
    let store = MemoryStore::new();
    let account = vodozemac::olm::Account::new();
    let curve_key = account.identity_keys().curve25519;

    let alice =
        OlmMachine::with_store(user_id(), alice_device_id(), store, Some(account)).await.unwrap();

    assert_eq!(
        alice.identity_keys().curve25519,
        curve_key,
        "The Olm machine should have used the Account we provided"
    );
}

#[async_test]
async fn test_unsigned_decryption() {
    let (alice, bob) =
        get_machine_pair_with_setup_sessions_test_helper(alice_id(), user_id(), false).await;
    let room_id = room_id!("!test:example.org");

    // Share the room key for the first message.
    let to_device_requests = alice
        .share_room_key(room_id, iter::once(bob.user_id()), EncryptionSettings::default())
        .await
        .unwrap();
    let first_room_key_event = ToDeviceEvent::new(
        alice.user_id().to_owned(),
        to_device_requests_to_content(to_device_requests),
    );

    // Save the first room key.
    let group_session = bob
        .store()
        .with_transaction(|mut tr| async {
            let res = bob
                .decrypt_to_device_event(&mut tr, &first_room_key_event, &mut Changes::default())
                .await?;
            Ok((tr, res))
        })
        .await
        .unwrap()
        .inbound_group_session;
    bob.store().save_inbound_group_sessions(&[group_session.unwrap()]).await.unwrap();

    // Encrypt first message.
    let first_message_text = "This is the original message";
    let first_message_content = RoomMessageEventContent::text_plain(first_message_text);
    let first_message_encrypted_content =
        alice.encrypt_room_event(room_id, first_message_content).await.unwrap();

    let mut first_message_encrypted_event = json!({
        "event_id": "$message1",
        "origin_server_ts": MilliSecondsSinceUnixEpoch::now(),
        "sender": alice.user_id(),
        "type": "m.room.encrypted",
        "content": first_message_encrypted_content,
    });
    let raw_encrypted_event = json_convert(&first_message_encrypted_event).unwrap();

    // Bob has the room key, so first message should be decrypted successfully.
    let decryption_settings =
        DecryptionSettings { sender_device_trust_requirement: TrustRequirement::Untrusted };
    let raw_decrypted_event =
        bob.decrypt_room_event(&raw_encrypted_event, room_id, &decryption_settings).await.unwrap();

    let decrypted_event = raw_decrypted_event.event.deserialize().unwrap();
    assert_matches!(decrypted_event, AnyMessageLikeEvent::RoomMessage(first_message));

    let first_message = first_message.as_original().unwrap();
    assert_eq!(first_message.content.body(), first_message_text);
    assert!(first_message.unsigned.relations.is_empty());

    assert!(raw_decrypted_event.unsigned_encryption_info.is_none());

    // Get a new room key, but don't give it to Bob yet.
    alice.discard_room_key(room_id).await.unwrap();
    let to_device_requests = alice
        .share_room_key(room_id, iter::once(bob.user_id()), EncryptionSettings::default())
        .await
        .unwrap();
    let second_room_key_event = ToDeviceEvent::new(
        alice.user_id().to_owned(),
        to_device_requests_to_content(to_device_requests),
    );

    // Encrypt a second message, an edit.
    let second_message_text = "This is the ~~original~~ edited message";
    let second_message_content = RoomMessageEventContent::text_plain(second_message_text)
        .make_replacement(first_message, None);
    let second_message_encrypted_content =
        alice.encrypt_room_event(room_id, second_message_content).await.unwrap();

    let second_message_encrypted_event = json!({
        "event_id": "$message2",
        "origin_server_ts": MilliSecondsSinceUnixEpoch::now(),
        "sender": alice.user_id(),
        "type": "m.room.encrypted",
        "content": second_message_encrypted_content,
    });

    // Bundle the edit in the unsigned object of the first event.
    let relations = json!({
        "m.relations": {
            "m.replace": second_message_encrypted_event,
        },
    });
    first_message_encrypted_event.as_object_mut().unwrap().insert("unsigned".to_owned(), relations);
    let raw_encrypted_event = json_convert(&first_message_encrypted_event).unwrap();

    // Bob does not have the second room key, so second message should fail to
    // decrypt.
    let raw_decrypted_event =
        bob.decrypt_room_event(&raw_encrypted_event, room_id, &decryption_settings).await.unwrap();

    let decrypted_event = raw_decrypted_event.event.deserialize().unwrap();
    assert_matches!(decrypted_event, AnyMessageLikeEvent::RoomMessage(first_message));

    let first_message = first_message.as_original().unwrap();
    assert_eq!(first_message.content.body(), first_message_text);
    // Deserialization of the edit failed, but it was here.
    assert!(first_message.unsigned.relations.replace.is_none());
    assert!(first_message.unsigned.relations.has_replacement());

    let unsigned_encryption_info = raw_decrypted_event.unsigned_encryption_info.unwrap();
    assert_eq!(unsigned_encryption_info.len(), 1);
    let replace_encryption_result =
        unsigned_encryption_info.get(&UnsignedEventLocation::RelationsReplace).unwrap();
    assert_matches!(
        replace_encryption_result,
        UnsignedDecryptionResult::UnableToDecrypt(UnableToDecryptInfo {
            session_id: Some(second_room_key_session_id),
            reason: UnableToDecryptReason::MissingMegolmSession { withheld_code: None },
        })
    );

    // Give Bob the second room key.
    let group_session = bob
        .store()
        .with_transaction(|mut tr| async {
            let res = bob
                .decrypt_to_device_event(&mut tr, &second_room_key_event, &mut Changes::default())
                .await?;
            Ok((tr, res))
        })
        .await
        .unwrap()
        .inbound_group_session
        .unwrap();
    assert_eq!(group_session.session_id(), second_room_key_session_id);
    bob.store().save_inbound_group_sessions(&[group_session]).await.unwrap();

    // Second message should decrypt now.
    let raw_decrypted_event =
        bob.decrypt_room_event(&raw_encrypted_event, room_id, &decryption_settings).await.unwrap();

    let decrypted_event = raw_decrypted_event.event.deserialize().unwrap();
    assert_matches!(decrypted_event, AnyMessageLikeEvent::RoomMessage(first_message));

    let first_message = first_message.as_original().unwrap();
    assert_eq!(first_message.content.body(), first_message_text);
    let replace = first_message.unsigned.relations.replace.as_ref().unwrap();
    assert_matches!(&replace.content.relates_to, Some(Relation::Replacement(replace_content)));
    assert_eq!(replace_content.new_content.msgtype.body(), second_message_text);

    let unsigned_encryption_info = raw_decrypted_event.unsigned_encryption_info.unwrap();
    assert_eq!(unsigned_encryption_info.len(), 1);
    let replace_encryption_result =
        unsigned_encryption_info.get(&UnsignedEventLocation::RelationsReplace).unwrap();
    assert_matches!(replace_encryption_result, UnsignedDecryptionResult::Decrypted(_));

    // Get a new room key again, but don't give it to Bob yet.
    alice.discard_room_key(room_id).await.unwrap();
    let to_device_requests = alice
        .share_room_key(room_id, iter::once(bob.user_id()), EncryptionSettings::default())
        .await
        .unwrap();
    let third_room_key_event = ToDeviceEvent::new(
        alice.user_id().to_owned(),
        to_device_requests_to_content(to_device_requests),
    );

    // Encrypt a third message, a thread event.
    let third_message_text = "This a reply in a thread";
    let third_message_content = RoomMessageEventContent::text_plain(third_message_text)
        .make_for_thread(first_message, ReplyWithinThread::No, AddMentions::No);
    let third_message_encrypted_content =
        alice.encrypt_room_event(room_id, third_message_content).await.unwrap();

    let third_message_encrypted_event = json!({
        "event_id": "$message3",
        "origin_server_ts": MilliSecondsSinceUnixEpoch::now(),
        "sender": alice.user_id(),
        "type": "m.room.encrypted",
        "content": third_message_encrypted_content,
        "room_id": room_id,
    });

    // Bundle the edit in the unsigned object of the first event.
    let relations = json!({
        "m.relations": {
            "m.replace": second_message_encrypted_event,
            "m.thread": {
                "latest_event": third_message_encrypted_event,
                "count": 1,
                "current_user_participated": true,
            }
        },
    });
    first_message_encrypted_event.as_object_mut().unwrap().insert("unsigned".to_owned(), relations);
    let raw_encrypted_event = json_convert(&first_message_encrypted_event).unwrap();

    // Bob does not have the third room key, so third message should fail to
    // decrypt.
    let raw_decrypted_event =
        bob.decrypt_room_event(&raw_encrypted_event, room_id, &decryption_settings).await.unwrap();

    let decrypted_event = raw_decrypted_event.event.deserialize().unwrap();
    assert_matches!(decrypted_event, AnyMessageLikeEvent::RoomMessage(first_message));

    let first_message = first_message.as_original().unwrap();
    assert_eq!(first_message.content.body(), first_message_text);
    assert!(first_message.unsigned.relations.replace.is_some());
    // Deserialization of the thread event succeeded, but it is still encrypted.
    let thread = first_message.unsigned.relations.thread.as_ref().unwrap();
    assert_matches!(thread.latest_event.deserialize(), Ok(AnyMessageLikeEvent::RoomEncrypted(_)));

    let unsigned_encryption_info = raw_decrypted_event.unsigned_encryption_info.unwrap();
    assert_eq!(unsigned_encryption_info.len(), 2);
    let replace_encryption_result =
        unsigned_encryption_info.get(&UnsignedEventLocation::RelationsReplace).unwrap();
    assert_matches!(replace_encryption_result, UnsignedDecryptionResult::Decrypted(_));
    let thread_encryption_result =
        unsigned_encryption_info.get(&UnsignedEventLocation::RelationsThreadLatestEvent).unwrap();
    assert_matches!(
        thread_encryption_result,
        UnsignedDecryptionResult::UnableToDecrypt(UnableToDecryptInfo {
            session_id: Some(third_room_key_session_id),
            reason: UnableToDecryptReason::MissingMegolmSession { withheld_code: None },
        })
    );

    // Give Bob the third room key.
    let group_session = bob
        .store()
        .with_transaction(|mut tr| async {
            let res = bob
                .decrypt_to_device_event(&mut tr, &third_room_key_event, &mut Changes::default())
                .await?;
            Ok((tr, res))
        })
        .await
        .unwrap()
        .inbound_group_session
        .unwrap();
    assert_eq!(group_session.session_id(), third_room_key_session_id);
    bob.store().save_inbound_group_sessions(&[group_session]).await.unwrap();

    // Third message should decrypt now.
    let raw_decrypted_event =
        bob.decrypt_room_event(&raw_encrypted_event, room_id, &decryption_settings).await.unwrap();

    let decrypted_event = raw_decrypted_event.event.deserialize().unwrap();
    assert_matches!(decrypted_event, AnyMessageLikeEvent::RoomMessage(first_message));

    let first_message = first_message.as_original().unwrap();
    assert_eq!(first_message.content.body(), first_message_text);
    assert!(first_message.unsigned.relations.replace.is_some());
    let thread = &first_message.unsigned.relations.thread.as_ref().unwrap();
    assert_matches!(
        thread.latest_event.deserialize(),
        Ok(AnyMessageLikeEvent::RoomMessage(third_message))
    );
    let third_message = third_message.as_original().unwrap();
    assert_eq!(third_message.content.body(), third_message_text);

    let unsigned_encryption_info = raw_decrypted_event.unsigned_encryption_info.unwrap();
    assert_eq!(unsigned_encryption_info.len(), 2);
    let replace_encryption_result =
        unsigned_encryption_info.get(&UnsignedEventLocation::RelationsReplace).unwrap();
    assert_matches!(replace_encryption_result, UnsignedDecryptionResult::Decrypted(_));
    let thread_encryption_result =
        unsigned_encryption_info.get(&UnsignedEventLocation::RelationsThreadLatestEvent).unwrap();
    assert_matches!(thread_encryption_result, UnsignedDecryptionResult::Decrypted(_));
}

#[async_test]
async fn test_mark_all_tracked_users_as_dirty() {
    let store = MemoryStore::new();
    let account = vodozemac::olm::Account::new();

    // Put some tracked users
    let damir = user_id!("@damir:localhost");
    let ben = user_id!("@ben:localhost");
    let ivan = user_id!("@ivan:localhost");

    // Mark them as not dirty.
    store.save_tracked_users(&[(damir, false), (ben, false), (ivan, false)]).await.unwrap();

    // Let's imagine the migration has been done: this is useful so that tracked
    // users are not marked as dirty when creating the `OlmMachine`.
    store
        .set_custom_value(OlmMachine::key_for_has_migrated_verification_latch(), vec![0])
        .await
        .unwrap();

    let alice =
        OlmMachine::with_store(user_id(), alice_device_id(), store, Some(account)).await.unwrap();

    // All users are marked as not dirty.
    alice.store().load_tracked_users().await.unwrap().iter().for_each(|tracked_user| {
        assert!(tracked_user.dirty.not());
    });

    // Now, mark all tracked users as dirty.
    alice.mark_all_tracked_users_as_dirty().await.unwrap();

    // All users are now marked as dirty.
    alice.store().load_tracked_users().await.unwrap().iter().for_each(|tracked_user| {
        assert!(tracked_user.dirty);
    });
}

#[async_test]
async fn test_verified_latch_migration() {
    let store = MemoryStore::new();
    let account = vodozemac::olm::Account::new();

    // put some tracked users
    let bob_id = user_id!("@bob:localhost");
    let carol_id = user_id!("@carol:localhost");

    // Mark them as not dirty
    let to_track_not_dirty = vec![(bob_id, false), (carol_id, false)];
    store.save_tracked_users(&to_track_not_dirty).await.unwrap();

    let alice =
        OlmMachine::with_store(user_id(), alice_device_id(), store, Some(account)).await.unwrap();

    let alice_store = alice.store();

    // A migration should have occurred and all users should be marked as dirty
    alice_store.load_tracked_users().await.unwrap().iter().for_each(|tu| {
        assert!(tu.dirty);
    });

    // Ensure it does so only once
    alice_store.save_tracked_users(&to_track_not_dirty).await.unwrap();

    OlmMachine::migration_post_verified_latch_support(alice_store, alice.identity_manager())
        .await
        .unwrap();

    // Migration already done, so user should not be marked as dirty
    alice_store.load_tracked_users().await.unwrap().iter().for_each(|tu| {
        assert!(!tu.dirty);
    });
}
