Skip to main content

matrix_sdk_crypto/store/
integration_tests.rs

1/// A macro which will run the CryptoStore integration test suite.
2///
3/// You need to provide a `async fn get_store() -> StoreResult<impl StateStore>`
4/// providing a fresh store on the same level you invoke the macro.
5///
6/// ## Usage example
7///
8/// ```no_run
9/// # use matrix_sdk_crypto::store::{
10/// #    MemoryStore as MyCryptoStore,
11/// # };
12///
13/// #[cfg(test)]
14/// mod tests {
15///     use super::MyCryptoStore;
16///
17///     async fn get_store(
18///         name: &str,
19///         passphrase: Option<&str>,
20///         clear_data: bool,
21///     ) -> MyCryptoStore {
22///         let store = MyCryptoStore::new();
23///         if clear_data {
24///             store.clear();
25///         }
26///         store
27///     }
28///
29///     cryptostore_integration_tests!();
30/// }
31/// ```
32#[allow(unused_macros)]
33#[macro_export]
34macro_rules! cryptostore_integration_tests {
35    () => {
36        mod cryptostore_integration_tests {
37            use std::collections::{BTreeMap, HashMap, HashSet};
38            use std::ops::Deref;
39            use std::time::Duration;
40
41            use assert_matches::assert_matches;
42            use matrix_sdk_test::async_test;
43            use ruma::{
44                device_id, events::secret::request::SecretName, room_id, serde::Raw, owned_room_id,
45                to_device::DeviceIdOrAllDevices, user_id, DeviceId, RoomId, TransactionId, UserId,
46            };
47            use serde_json::value::to_raw_value;
48            use serde_json::json;
49            use matrix_sdk_common::deserialized_responses::WithheldCode;
50            use $crate::{
51                olm::{
52                    Account, Curve25519PublicKey, InboundGroupSession, OlmMessageHash,
53                    PrivateCrossSigningIdentity, SenderData, SenderDataType, Session
54                },
55                store::{
56                    types::{
57                        BackupDecryptionKey, Changes, DehydratedDeviceKey, DeviceChanges,
58                        IdentityChanges, PendingChanges, StoredRoomKeyBundleData, RoomKeyWithheldEntry,
59                        RoomSettings
60                    },
61                    CryptoStore, GossipRequest,
62                },
63                testing::{get_device, get_other_identity, get_own_identity},
64                types::{
65                    events::{
66                        dummy::DummyEventContent,
67                        olm_v1::{DecryptedSecretSendEvent, OlmV1Keys},
68                        room_key_request::MegolmV1AesSha2Content,
69                        room_key_withheld::{
70                            CommonWithheldCodeContent, MegolmV1AesSha2WithheldContent,
71                            RoomKeyWithheldContent,
72                        },
73                        room_key_bundle::RoomKeyBundleContent,
74                        secret_send::SecretSendContent,
75                        ToDeviceEvent,
76                    },
77                    requests::ToDeviceRequest,
78                    DeviceKeys,
79                    EventEncryptionAlgorithm,
80                },
81                vodozemac::megolm::{GroupSession, SessionConfig}, DeviceData, GossippedSecret, LocalTrust,  SecretInfo,
82                TrackedUser,
83            };
84            #[cfg(feature = "experimental-push-secrets")]
85            use $crate::types::events::secret_push::SecretPushContent;
86
87            use super::get_store;
88
89            fn alice_id() -> &'static UserId {
90                user_id!("@alice:example.org")
91            }
92
93            fn alice_device_id() -> &'static DeviceId {
94                device_id!("ALICEDEVICE")
95            }
96
97            fn bob_id() -> &'static UserId {
98                user_id!("@bob:example.org")
99            }
100
101            fn bob_device_id() -> &'static DeviceId {
102                device_id!("BOBDEVICE")
103            }
104
105            pub async fn get_loaded_store(name: &str) -> (Account, impl CryptoStore + use<>) {
106                let store = get_store(name, None, true).await;
107                let account = get_account();
108
109                store.save_pending_changes(PendingChanges { account: Some(account.deep_clone()), }).await.expect("Can't save account");
110
111                (account, store)
112            }
113
114            fn get_account() -> Account {
115                Account::with_device_id(alice_id(), alice_device_id())
116            }
117
118            pub(crate) async fn get_account_and_session() -> (Account, Session) {
119                let alice = Account::with_device_id(alice_id(), alice_device_id());
120                let mut bob = Account::with_device_id(bob_id(), bob_device_id());
121
122                bob.generate_one_time_keys(1);
123                let one_time_key = *bob.one_time_keys().values().next().unwrap();
124                let sender_key = bob.identity_keys().curve25519;
125                let session = alice.create_outbound_session_helper(
126                    Default::default(),
127                    sender_key,
128                    one_time_key,
129                    false,
130                    alice.device_keys(),
131                ).unwrap();
132
133                (alice, session)
134            }
135
136            #[async_test]
137            async fn test_save_account_via_generic_save() {
138                let store = get_store("save_account_via_generic", None, true).await;
139                assert!(store.get_static_account().is_none());
140                assert!(store.load_account().await.unwrap().is_none());
141                let account = get_account();
142
143                store
144                    .save_pending_changes(PendingChanges { account: Some(account) })
145                    .await
146                    .expect("Can't save account");
147                assert!(store.get_static_account().is_some());
148            }
149
150            #[async_test]
151            async fn test_save_account() {
152                let store = get_store("save_account", None, true).await;
153                assert!(store.get_static_account().is_none());
154                assert!(store.load_account().await.unwrap().is_none());
155                let account = get_account();
156
157                store
158                    .save_pending_changes(PendingChanges { account: Some(account) })
159                    .await
160                    .expect("Can't save account");
161                assert!(store.get_static_account().is_some());
162            }
163
164            #[async_test]
165            async fn test_load_account() {
166                let store = get_store("load_account", None, true).await;
167                let account = get_account();
168
169                store
170                    .save_pending_changes(PendingChanges { account: Some(account.deep_clone()) })
171                    .await
172                    .expect("Can't save account");
173
174                let loaded_account = store.load_account().await.expect("Can't load account");
175                let loaded_account = loaded_account.unwrap();
176
177                assert_eq!(account, loaded_account);
178            }
179
180            #[async_test]
181            async fn test_load_account_with_passphrase() {
182                let passphrase = Some("secret_passphrase");
183                let store = get_store("load_account_with_passphrase", passphrase, true).await;
184                let account = get_account();
185
186                store
187                    .save_pending_changes(PendingChanges { account: Some(account.deep_clone()) })
188                    .await
189                    .expect("Can't save account");
190
191                let loaded_account = store.load_account().await.expect("Can't load account");
192                let loaded_account = loaded_account.unwrap();
193
194                assert_eq!(account, loaded_account);
195            }
196
197            #[async_test]
198            async fn test_save_and_share_account() {
199                let store = get_store("save_and_share_account", None, true).await;
200                let mut account = get_account();
201
202                store
203                    .save_pending_changes(PendingChanges { account: Some(account.deep_clone()) })
204                    .await
205                    .expect("Can't save account");
206
207                account.mark_as_shared();
208                account.update_uploaded_key_count(50);
209
210                store
211                    .save_pending_changes(PendingChanges { account: Some(account.deep_clone()) })
212                    .await
213                    .expect("Can't save account");
214
215                let loaded_account = store.load_account().await.expect("Can't load account");
216                let loaded_account = loaded_account.unwrap();
217
218                assert_eq!(account, loaded_account);
219                assert_eq!(account.uploaded_key_count(), loaded_account.uploaded_key_count());
220            }
221
222            #[async_test]
223            async fn test_load_sessions() {
224                let store = get_store("load_sessions", None, true).await;
225                let (account, session) = get_account_and_session().await;
226                store
227                    .save_pending_changes(PendingChanges { account: Some(account.deep_clone()) })
228                    .await
229                    .expect("Can't save account");
230
231                let changes = Changes {
232                    sessions: vec![session.clone()],
233                    devices: DeviceChanges { new: vec![DeviceData::from_account(&account)], ..Default::default() },
234                    ..Default::default()
235                };
236
237                store.save_changes(changes).await.unwrap();
238
239                let sessions = store
240                    .get_sessions(&session.sender_key.to_base64())
241                    .await
242                    .expect("Can't load sessions")
243                    .unwrap();
244                let loaded_session = sessions.get(0).cloned().expect("We should find the session in the store.");
245
246                assert_eq!(&session, &loaded_session, "The loaded session should be the same one we put into the store.");
247            }
248
249            #[async_test]
250            async fn test_add_and_save_session() {
251                let store_name = "add_and_save_session";
252
253                // Given we created a session and saved it in the store
254                let (session_id, account, sender_key) = {
255                    let store = get_store(store_name, None, true).await;
256                    let (account, session) = get_account_and_session().await;
257                    let sender_key = session.sender_key.to_base64();
258                    let session_id = session.session_id().to_owned();
259
260                    store
261                        .save_pending_changes(PendingChanges {
262                            account: Some(account.deep_clone()),
263                        })
264                        .await
265                        .expect("Can't save account");
266                    store
267                        .save_changes(Changes {
268                            devices: DeviceChanges {
269                                new: vec![DeviceData::from_account(&account)],
270                                ..Default::default()
271                            },
272                            ..Default::default()
273                        })
274                        .await
275                        .unwrap();
276
277                    let changes = Changes { sessions: vec![session.clone()], ..Default::default() };
278                    store.save_changes(changes).await.unwrap();
279
280                    let sessions = store.get_sessions(&sender_key).await.unwrap().unwrap();
281                    let session = &sessions[0];
282
283                    assert_eq!(session_id, session.session_id());
284
285                    (session_id, account, sender_key)
286                };
287
288                // When we reload the store
289                let store = get_store(store_name, None, false).await;
290
291                // Then the same account and session info was reloaded
292                let loaded_account = store.load_account().await.unwrap().unwrap();
293                assert_eq!(account, loaded_account);
294
295                let sessions = store.get_sessions(&sender_key).await.unwrap().unwrap();
296                let session = &sessions[0];
297
298                assert_eq!(session_id, session.session_id());
299            }
300
301            #[async_test]
302            async fn test_load_outbound_group_session() {
303                let dir = "load_outbound_group_session";
304                let room_id = room_id!("!test:localhost");
305
306                // Given we saved an outbound group session
307                {
308                    let (account, store) = get_loaded_store(dir.clone()).await;
309                    assert!(
310                        store.get_outbound_group_session(&room_id).await.unwrap().is_none(),
311                        "Initially there should be no outbound group session"
312                    );
313
314                    let (session, _) =
315                        account.create_group_session_pair_with_defaults(&room_id).await;
316
317                    let user_id = user_id!("@example:localhost");
318                    let request = ToDeviceRequest::new(
319                        user_id,
320                        DeviceIdOrAllDevices::AllDevices,
321                        "m.dummy",
322                        Raw::from_json(to_raw_value(&DummyEventContent::new()).unwrap()),
323                    );
324
325                    session.add_request(TransactionId::new(), request.into(), Default::default());
326
327                    let changes = Changes {
328                        outbound_group_sessions: vec![session.clone()],
329                        ..Default::default()
330                    };
331
332                    store.save_changes(changes).await.expect("Can't save group session");
333                    assert!(
334                        store.get_outbound_group_session(&room_id).await.unwrap().is_some(),
335                        "Sanity: after we've saved one, there should be an outbound_group_session"
336                    );
337                }
338
339                // When we reload the account
340                let store = get_store(dir, None, false).await;
341                store.load_account().await.unwrap();
342
343                // Then the saved session is restored
344                assert!(
345                    store.get_outbound_group_session(&room_id).await.unwrap().is_some(),
346                    "The outbound_group_session should have been loaded"
347                );
348            }
349
350            /// Test that we can import an inbound group session via [`CryptoStore::save_changes`]
351            #[async_test]
352            async fn test_save_changes_save_inbound_group_session() {
353                let (account, store) = get_loaded_store("save_inbound_group_session").await;
354
355                let room_id = &room_id!("!test:localhost");
356                let (_, session) = account.create_group_session_pair_with_defaults(room_id).await;
357
358                let changes =
359                    Changes { inbound_group_sessions: vec![session], ..Default::default() };
360
361                store.save_changes(changes).await.expect("Can't save group session");
362            }
363
364            /// Test that we can import a backed-up group session via
365            /// [`CryptoStore::save_inbound_group_sessions`]
366            #[async_test]
367            async fn test_save_inbound_group_session_from_backup() {
368                let (account, store) =
369                    get_loaded_store("save_inbound_group_session_from_backup").await;
370
371                let room_id = &room_id!("!test:localhost");
372                let (_, session) = account.create_group_session_pair_with_defaults(room_id).await;
373
374                session.mark_as_backed_up();
375                store
376                    .save_inbound_group_sessions(vec![session.clone()], Some(&"bkpver1"))
377                    .await
378                    .expect("could not save sessions");
379
380                let loaded_session = store
381                    .get_inbound_group_session(&session.room_id, session.session_id())
382                    .await
383                    .expect("error when loading session")
384                    .expect("session not found in store");
385                assert_eq!(session, loaded_session);
386                assert_eq!(store.get_inbound_group_sessions().await.unwrap().len(), 1);
387                assert_eq!(store.inbound_group_session_counts(None).await.unwrap().total, 1);
388
389                // It should *not* be returned by a request for backup for the same backup version
390                let to_back_up = store.inbound_group_sessions_for_backup("bkpver1", 1).await.unwrap();
391                assert_eq!(to_back_up.len(), 0, "backup was returned by backup query");
392                assert_eq!(
393                    store.inbound_group_session_counts(Some(&"bkpver1")).await.unwrap().backed_up, 1,
394                    "backed_up count",
395                );
396            }
397
398            /// Test that the behaviour of a key imported from an _old_ backup
399            /// is correct
400            ///
401            /// This currently only works on the MemoryStore, so is ignored. The
402            /// other stores are waiting for more work on
403            /// https://github.com/element-hq/element-web/issues/26892.
404            #[ignore]
405            #[async_test]
406            async fn test_save_inbound_group_session_from_old_backup() {
407                let (account, store) =
408                    get_loaded_store("save_inbound_group_session_from_old_backup").await;
409
410                let room_id = &room_id!("!test:localhost");
411                let (_, session) = account.create_group_session_pair_with_defaults(room_id).await;
412
413                session.mark_as_backed_up();
414                store
415                    .save_inbound_group_sessions(vec![session.clone()], Some(&"bkpver1"))
416                    .await
417                    .expect("could not save sessions");
418
419                // The session should be returned by a request for backup from a different backup version.
420                let to_back_up = store.inbound_group_sessions_for_backup("bkpver2", 1).await.unwrap();
421                assert_eq!(to_back_up, vec![session]);
422                assert_eq!(
423                    store.inbound_group_session_counts(Some(&"bkpver2")).await.unwrap().backed_up, 0,
424                    "backed_up count for backup version 2",
425                );
426            }
427
428            /// Test that we can import a not-backed-up group session via
429            /// [`CryptoStore::save_inbound_group_sessions`]
430            #[async_test]
431            async fn test_save_inbound_group_session_from_import() {
432                let (account, store) =
433                    get_loaded_store("save_inbound_group_session_from_import").await;
434
435                let room_id = &room_id!("!test:localhost");
436                let (_, session) = account.create_group_session_pair_with_defaults(room_id).await;
437
438                store
439                    .save_inbound_group_sessions(vec![session.clone()], None)
440                    .await
441                    .expect("could not save sessions");
442
443                let loaded_session = store
444                    .get_inbound_group_session(&session.room_id, session.session_id())
445                    .await
446                    .expect("error when loading session")
447                    .expect("session not found in store");
448                assert_eq!(session, loaded_session);
449                assert_eq!(store.get_inbound_group_sessions().await.unwrap().len(), 1);
450                assert_eq!(store.inbound_group_session_counts(None).await.unwrap().total, 1);
451                assert_eq!(store.inbound_group_session_counts(None).await.unwrap().backed_up, 0);
452
453                // It should be returned by a request for backup
454                let to_back_up = store.inbound_group_sessions_for_backup("bkpver1", 1).await.unwrap();
455                assert_eq!(to_back_up, vec![session]);
456            }
457
458            #[async_test]
459            async fn test_mark_inbound_group_sessions_as_backed_up() {
460                // Given a store exists with multiple unbacked-up sessions
461                let (account, store) =
462                    get_loaded_store("mark_inbound_group_sessions_as_backed_up").await;
463                let room_id = &room_id!("!test:localhost");
464                let mut sessions: Vec<InboundGroupSession> = Vec::with_capacity(10);
465                for _i in 0..10 {
466                    sessions.push(account.create_group_session_pair_with_defaults(room_id).await.1);
467                }
468                let changes = Changes { inbound_group_sessions: sessions.clone(), ..Default::default() };
469                store.save_changes(changes).await.expect("Can't save group session");
470                assert_eq!(store.inbound_group_sessions_for_backup("bkpver", 100).await.unwrap().len(), 10);
471
472                // When I mark some as backed up
473                store.mark_inbound_group_sessions_as_backed_up("bkpver", &[
474                    session_info(&sessions[1]),
475                    session_info(&sessions[3]),
476                    session_info(&sessions[5]),
477                    session_info(&sessions[7]),
478                    session_info(&sessions[9]),
479                ]).await.expect("Failed to mark sessions as backed up");
480
481                // And ask which still need backing up
482                let to_back_up = store.inbound_group_sessions_for_backup("bkpver", 10).await.unwrap();
483                let needs_backing_up = |i: usize| to_back_up.iter().any(|s| s.session_id() == sessions[i].session_id());
484
485                // Then the sessions we said were backed up no longer need backing up
486                assert!(!needs_backing_up(1));
487                assert!(!needs_backing_up(3));
488                assert!(!needs_backing_up(5));
489                assert!(!needs_backing_up(7));
490                assert!(!needs_backing_up(9));
491
492                // And the sessions we didn't mention still need backing up
493                assert!(needs_backing_up(0));
494                assert!(needs_backing_up(2));
495                assert!(needs_backing_up(4));
496                assert!(needs_backing_up(6));
497                assert!(needs_backing_up(8));
498                assert_eq!(to_back_up.len(), 5);
499            }
500
501            #[async_test]
502            async fn test_reset_inbound_group_session_for_backup() {
503                // Given a store exists where all sessions are backed up to backup_1
504                let (account, store) =
505                    get_loaded_store("reset_inbound_group_session_for_backup").await;
506                let room_id = &room_id!("!test:localhost");
507                let mut sessions: Vec<InboundGroupSession> = Vec::with_capacity(10);
508                for _ in 0..10 {
509                    sessions.push(account.create_group_session_pair_with_defaults(room_id).await.1);
510                }
511                let changes = Changes { inbound_group_sessions: sessions.clone(), ..Default::default() };
512                store.save_changes(changes).await.expect("Can't save group session");
513                assert_eq!(store.inbound_group_sessions_for_backup("backup_1", 100).await.unwrap().len(), 10);
514                store.mark_inbound_group_sessions_as_backed_up(
515                    "backup_1",
516                    &(0..10).map(|i| session_info(&sessions[i])).collect::<Vec<_>>(),
517                ).await.expect("Failed to mark sessions as backed up");
518
519                // Sanity: none need backing up to the same backup
520                {
521                    let to_back_up_old = store.inbound_group_sessions_for_backup("backup_1", 10).await.unwrap();
522                    assert_eq!(to_back_up_old.len(), 0);
523                }
524
525                // Some stores ignore backup_version and just reset when you
526                // tell them to. Tell them here.
527                store.reset_backup_state().await.expect("reset failed");
528
529                // When we ask what needs backing up to a different backup version
530                let to_back_up = store.inbound_group_sessions_for_backup("backup_02", 10).await.unwrap();
531
532                // Then the answer is everything
533                let needs_backing_up = |i: usize| to_back_up.iter().any(|s| s.session_id() == sessions[i].session_id());
534                assert!(needs_backing_up(0));
535                assert!(needs_backing_up(1));
536                assert!(needs_backing_up(8));
537                assert!(needs_backing_up(9));
538                assert_eq!(to_back_up.len(), 10);
539            }
540
541            #[async_test]
542            async fn test_load_inbound_group_session() {
543                let dir = "load_inbound_group_session";
544                let (account, store) = get_loaded_store(dir).await;
545                assert_eq!(store.get_inbound_group_sessions().await.unwrap().len(), 0);
546
547                let room_id = &room_id!("!test:localhost");
548                let (_, session) = account.create_group_session_pair_with_defaults(room_id).await;
549
550                let export = session.export().await;
551
552                let session = InboundGroupSession::from_export(&export).unwrap();
553
554                let changes =
555                    Changes { inbound_group_sessions: vec![session.clone()], ..Default::default() };
556
557                store.save_changes(changes).await.expect("Can't save group session");
558
559                drop(store);
560
561                let store = get_store(dir, None, false).await;
562
563                store.load_account().await.unwrap();
564
565                let loaded_session = store
566                    .get_inbound_group_session(&session.room_id, session.session_id())
567                    .await
568                    .unwrap()
569                    .unwrap();
570                assert_eq!(session, loaded_session);
571                loaded_session.export().await;
572
573                assert_eq!(store.get_inbound_group_sessions().await.unwrap().len(), 1);
574                assert_eq!(store.inbound_group_session_counts(None).await.unwrap().total, 1);
575            }
576
577            #[async_test]
578            async fn test_get_inbound_group_sessions_by_room_id_empty() {
579                let dir = "get_inbound_group_session_by_room_id_empty";
580                let (_, store) = get_loaded_store(dir).await;
581                assert_eq!(store.get_inbound_group_sessions().await.unwrap().len(), 0);
582
583                let room_id = &room_id!("!testing:localhost");
584                assert_eq!(store.get_inbound_group_sessions_by_room_id(room_id).await.unwrap().len(), 0);
585            }
586
587            #[async_test]
588            async fn test_get_inbound_group_sessions_by_room_id() {
589                let dir = "get_inbound_group_session_by_room_id";
590                let (account, store) = get_loaded_store(dir).await;
591                assert_eq!(store.get_inbound_group_sessions().await.unwrap().len(), 0);
592
593                let room_id = &room_id!("!testing:localhost");
594                let (_, session_1) = account.create_group_session_pair_with_defaults(room_id).await;
595                let (_, session_2) = account.create_group_session_pair_with_defaults(room_id).await;
596
597                let second_room_id = &room_id!("!other_room_testing:localhost");
598                let (_, session_3) = account.create_group_session_pair_with_defaults(second_room_id).await;
599
600                let mut sessions = vec![
601                    session_1,
602                    session_2,
603                    session_3
604                ];
605
606                let changes = Changes {
607                    inbound_group_sessions: sessions.clone(),
608                    ..Default::default()
609                };
610                store.save_changes(changes).await.expect("Can't save group session");
611
612                drop(store);
613
614                // The last session is in a different room, so should not be
615                // returned by get_inbound_group_sessions_by_room_id. Remove it
616                // from the list.
617                sessions.pop();
618
619                let store = get_store(dir, None, false).await;
620                // Make sure all the sessions are in the store
621                assert_eq!(store.get_inbound_group_sessions().await.unwrap().len(), 3);
622
623                store.load_account().await.unwrap();
624
625                let loaded_sessions = store
626                    .get_inbound_group_sessions_by_room_id(room_id)
627                    .await
628                    .unwrap();
629
630                assert_eq!(loaded_sessions.len(), 2);
631                assert_session_lists_eq(sessions, loaded_sessions, "room by id sessions");
632            }
633
634            #[async_test]
635            async fn test_fetch_inbound_group_sessions_for_device() {
636                // Given a store exists, containing inbound group sessions from different devices
637                let (account, store) =
638                    get_loaded_store("fetch_inbound_group_sessions_for_device").await;
639
640                let dev1 = Curve25519PublicKey::from_base64(
641                    "wjLpTLRqbqBzLs63aYaEv2Boi6cFEbbM/sSRQ2oAKk4"
642                ).unwrap();
643                let dev2 = Curve25519PublicKey::from_base64(
644                    "LTpv2DGMhggPAXO02+7f68CNEp6A40F0Yl8B094Y8gc"
645                ).unwrap();
646
647                let dev_1_unknown_a = create_session(&account, &dev1, SenderDataType::UnknownDevice).await;
648                let dev_1_unknown_b = create_session(&account, &dev1, SenderDataType::UnknownDevice).await;
649
650                let dev_1_keys_a = create_session(&account, &dev1, SenderDataType::DeviceInfo).await;
651                let dev_1_keys_b = create_session(&account, &dev1, SenderDataType::DeviceInfo).await;
652                let dev_1_keys_c = create_session(&account, &dev1, SenderDataType::DeviceInfo).await;
653                let dev_1_keys_d = create_session(&account, &dev1, SenderDataType::DeviceInfo).await;
654
655                let dev_2_unknown = create_session(
656                    &account, &dev2, SenderDataType::UnknownDevice).await;
657
658                let dev_2_keys = create_session(
659                    &account, &dev2, SenderDataType::DeviceInfo).await;
660
661                let sessions = vec![
662                    dev_1_unknown_a.clone(),
663                    dev_1_unknown_b.clone(),
664                    dev_1_keys_a.clone(),
665                    dev_1_keys_b.clone(),
666                    dev_1_keys_c.clone(),
667                    dev_1_keys_d.clone(),
668                    dev_2_unknown.clone(),
669                    dev_2_keys.clone(),
670                ];
671
672                let changes = Changes {
673                    inbound_group_sessions: sessions,
674                    ..Default::default()
675                };
676                store.save_changes(changes).await.expect("Can't save group session");
677
678                // When we fetch the list of sessions for device 1, unknown
679                let sessions_1_u = store.get_inbound_group_sessions_for_device_batch(
680                    dev1,
681                    SenderDataType::UnknownDevice,
682                    None,
683                    10
684                ).await.expect("Failed to get sessions for dev1");
685
686                // Then the expected sessions are returned
687                assert_session_lists_eq(sessions_1_u, [dev_1_unknown_a, dev_1_unknown_b], "device 1 sessions");
688
689                // And when we ask for the list of sessions for device 2, with device keys
690                let sessions_2_d = store
691                    .get_inbound_group_sessions_for_device_batch(dev2, SenderDataType::DeviceInfo, None, 10)
692                    .await
693                    .expect("Failed to get sessions for dev2");
694
695                // Then the matching session is returned
696                assert_eq!(sessions_2_d, vec![dev_2_keys], "device 2 sessions");
697
698                // And we can fetch device 1, keys in batches. We call the batch
699                // function repeatedly, to ensure it terminates correctly.
700                let mut sessions_1_k = Vec::new();
701                let mut previous_last_session_id: Option<String> = None;
702                loop {
703                    let mut sessions_1_k_batch = store.get_inbound_group_sessions_for_device_batch(
704                        dev1,
705                        SenderDataType::DeviceInfo,
706                        previous_last_session_id,
707                        2
708                    ).await.expect("Failed to get batch 1");
709
710                    // If there are no results in the batch, we have reached the end of the results.
711                    let Some(last_session) = sessions_1_k_batch.last() else {
712                        break;
713                    };
714
715                    // Check that there are exactly two results in the batch
716                    assert_eq!(sessions_1_k_batch.len(), 2);
717
718                    previous_last_session_id = Some(last_session.session_id().to_owned());
719
720                    // Modify one of the results, to check that that doesn't break iteration
721                    let mut last_session = last_session.clone();
722                    last_session.sender_data = SenderData::unknown();
723                    store.save_inbound_group_sessions(vec![last_session], None).await.unwrap();
724
725                    sessions_1_k.append(&mut sessions_1_k_batch);
726                }
727
728                assert_session_lists_eq(
729                    sessions_1_k,
730                    [dev_1_keys_a, dev_1_keys_b, dev_1_keys_c, dev_1_keys_d],
731                    "device 1 batched results"
732                );
733            }
734
735            /// Assert that two lists of sessions are the same, modulo ordering.
736            ///
737            /// There is no requirement for
738            /// `get_inbound_group_sessions_for_device_batch` to return the
739            /// results in a specific order. This helper ensures that the two
740            /// lists of inbound group sessions are equivalent, without worrying
741            /// about the ordering.
742            fn assert_session_lists_eq<I, J>(actual: I, expected: J, message: &str)
743                where I: IntoIterator<Item = InboundGroupSession>, J: IntoIterator<Item = InboundGroupSession>
744            {
745                let sorter = |a: &InboundGroupSession, b: &InboundGroupSession| Ord::cmp(a.session_id(), b.session_id());
746
747                let mut actual = Vec::from_iter(actual);
748                actual.sort_unstable_by(sorter);
749                let mut expected = Vec::from_iter(expected);
750                expected.sort_unstable_by(sorter);
751                assert_eq!(actual, expected, "{}", message);
752            }
753
754            #[async_test]
755            async fn test_tracked_users() {
756                let dir = "test_tracked_users";
757                let (_account, store) = get_loaded_store(dir.clone()).await;
758
759                let alice = user_id!("@alice:example.org");
760                let bob = user_id!("@bob:example.org");
761                let candy = user_id!("@candy:example.org");
762
763                let loaded = store.load_tracked_users().await.unwrap();
764                assert!(loaded.is_empty(), "Initially there are no tracked users");
765
766                let users = vec![(alice, true), (bob, false)];
767                store.save_tracked_users(&users).await.unwrap();
768
769                let check_loaded_users = |loaded: Vec<TrackedUser>| {
770                    let loaded: HashMap<_, _> =
771                        loaded.into_iter().map(|u| (u.user_id.to_owned(), u)).collect();
772
773                    let loaded_alice =
774                        loaded.get(alice).expect("Alice should be in the store as a tracked user");
775                    let loaded_bob =
776                        loaded.get(bob).expect("Bob should be in the store as as tracked user");
777
778                    assert!(!loaded.contains_key(candy), "Candy shouldn't be part of the store");
779                    assert_eq!(loaded.len(), 2, "Candy shouldn't be part of the store");
780
781                    assert!(loaded_alice.dirty, "Alice should be considered to be dirty");
782                    assert!(!loaded_bob.dirty, "Bob should not be considered to be dirty");
783                };
784
785                let loaded = store.load_tracked_users().await.unwrap();
786                check_loaded_users(loaded);
787
788                drop(store);
789
790                let name = dir.clone();let store = get_store(name, None, false).await;
791                let loaded = store.load_tracked_users().await.unwrap();
792                check_loaded_users(loaded);
793            }
794
795            #[async_test]
796            async fn test_device_saving() {
797                let dir = "device_saving";
798                let (_account, store) = get_loaded_store(dir.clone()).await;
799
800                let alice_device_1 = DeviceData::from_account(&Account::with_device_id(
801                    "@alice:localhost".try_into().unwrap(),
802                    "FIRSTDEVICE".into(),
803                ));
804
805                let alice_device_2 = DeviceData::from_account(&Account::with_device_id(
806                    "@alice:localhost".try_into().unwrap(),
807                    "SECONDDEVICE".into(),
808                ));
809
810                let json = json!({
811                    "algorithms": ["m.olm.v1.curve25519-aes-sha2", "m.megolm.v1.aes-sha2"],
812                    "user_id": "@bob:localhost",
813                    "device_id": "BOBDEVICE",
814                    "extra_property": "somevalue",
815                    "keys": {
816                        "curve25519:BOBDEVICE": "n0zs7qnaPLLf/OTL+dDLcI5kaPexbUeQ8jLQ2q6sO0E",
817                        "ed25519:BOBDEVICE": "RrKiu4+5EHRBWY6Qj6OtQGC0txpmEeanOz2irEZ/IN4",
818                    },
819                    "signatures": {
820                        "@bob:localhost": {
821                            "ed25519:BOBDEVICE": "9NjPewVHfB7Ah32mJ+CBx64mVoiQ8gbh+/2pc9WfAgut/H0Kqd/bbpgJq9Pn518szaXcGqEq0DxDP6CABBX8CQ",
822                        },
823                    },
824                });
825
826                let bob_device_1_keys: DeviceKeys = serde_json::from_value(json).unwrap();
827                let bob_device_1 = DeviceData::new(bob_device_1_keys, LocalTrust::Unset);
828
829                let changes = Changes {
830                    devices: DeviceChanges {
831                        new: vec![alice_device_1.clone(), alice_device_2.clone(), bob_device_1.clone()],
832                        ..Default::default()
833                    },
834                    ..Default::default()
835                };
836
837                store.save_changes(changes).await.unwrap();
838
839                drop(store);
840
841                let store = get_store(dir, None, false).await;
842
843                store.load_account().await.unwrap();
844
845                let loaded_device = store
846                    .get_device(alice_device_1.user_id(), alice_device_1.device_id())
847                    .await
848                    .unwrap()
849                    .unwrap();
850
851                assert_eq!(alice_device_1, loaded_device);
852
853                for algorithm in loaded_device.algorithms() {
854                    assert!(alice_device_1.algorithms().contains(algorithm));
855                }
856                assert_eq!(alice_device_1.algorithms().len(), loaded_device.algorithms().len());
857                assert_eq!(alice_device_1.keys(), loaded_device.keys());
858
859                let user_devices = store.get_user_devices(alice_device_1.user_id()).await.unwrap();
860                assert_eq!(user_devices.len(), 2);
861
862                let bob_device = store
863                    .get_device(bob_device_1.user_id(), bob_device_1.device_id())
864                    .await
865                    .unwrap();
866
867                let bob_device_json = serde_json::to_value(bob_device).unwrap();
868                assert_eq!(bob_device_json["device_keys"]["extra_property"], json!("somevalue"));
869            }
870
871            #[async_test]
872            async fn test_device_deleting() {
873                let dir = "device_deleting";
874                let (_account, store) = get_loaded_store(dir.clone()).await;
875                let device = get_device();
876
877                let changes = Changes {
878                    devices: DeviceChanges { changed: vec![device.clone()], ..Default::default() },
879                    ..Default::default()
880                };
881
882                store.save_changes(changes).await.unwrap();
883
884                let changes = Changes {
885                    devices: DeviceChanges { deleted: vec![device.clone()], ..Default::default() },
886                    ..Default::default()
887                };
888
889                store.save_changes(changes).await.unwrap();
890                drop(store);
891
892                let store = get_store(dir, None, false).await;
893
894                store.load_account().await.unwrap();
895
896                let loaded_device =
897                    store.get_device(device.user_id(), device.device_id()).await.unwrap();
898
899                assert!(loaded_device.is_none());
900            }
901
902            #[async_test]
903            async fn test_user_saving() {
904                let dir = "user_saving";
905
906                let user_id = user_id!("@example:localhost");
907                let device_id: &DeviceId = "WSKKLTJZCL".into();
908
909                let store = get_store(dir, None, true).await;
910
911                let account = Account::with_device_id(&user_id, device_id);
912
913                store.save_pending_changes(PendingChanges { account: Some(account), })
914                    .await
915                    .expect("Can't save account");
916
917                let own_identity = get_own_identity();
918
919                let changes = Changes {
920                    identities: IdentityChanges {
921                        changed: vec![own_identity.clone().into()],
922                        ..Default::default()
923                    },
924                    ..Default::default()
925                };
926
927                store.save_changes(changes).await.expect("Can't save identity");
928
929                drop(store);
930
931                let store = get_store(dir, None, false).await;
932
933                store.load_account().await.unwrap();
934
935                let loaded_user =
936                    store.get_user_identity(own_identity.user_id()).await.unwrap().unwrap();
937
938                assert_eq!(loaded_user.master_key(), own_identity.master_key());
939                assert_eq!(loaded_user.self_signing_key(), own_identity.self_signing_key());
940                assert_eq!(loaded_user.own().unwrap().clone(), own_identity.clone());
941
942                let other_identity = get_other_identity();
943
944                let changes = Changes {
945                    identities: IdentityChanges {
946                        changed: vec![other_identity.clone().into()],
947                        ..Default::default()
948                    },
949                    ..Default::default()
950                };
951
952                store.save_changes(changes).await.unwrap();
953
954                let loaded_user =
955                    store.get_user_identity(other_identity.user_id()).await.unwrap().unwrap();
956
957                assert_eq!(loaded_user.master_key(), other_identity.master_key());
958                assert_eq!(loaded_user.self_signing_key(), other_identity.self_signing_key());
959                assert_eq!(loaded_user.user_id(), other_identity.user_id());
960                assert_eq!(loaded_user.other().unwrap().clone(), other_identity);
961
962                own_identity.mark_as_verified();
963
964                let changes = Changes {
965                    identities: IdentityChanges {
966                        changed: vec![own_identity.into()],
967                        ..Default::default()
968                    },
969                    ..Default::default()
970                };
971
972                store.save_changes(changes).await.unwrap();
973                let loaded_user = store.get_user_identity(&user_id).await.unwrap().unwrap();
974                assert!(loaded_user.own().unwrap().is_verified())
975            }
976
977            #[async_test]
978            async fn test_private_identity_saving() {
979                let (_, store) = get_loaded_store("private_identity_saving").await;
980                assert!(store.load_identity().await.unwrap().is_none());
981                let identity = PrivateCrossSigningIdentity::new(alice_id().to_owned());
982
983                let changes =
984                    Changes { private_identity: Some(identity.clone()), ..Default::default() };
985
986                store.save_changes(changes).await.unwrap();
987                let loaded_identity = store.load_identity().await.unwrap().unwrap();
988                assert_eq!(identity.user_id(), loaded_identity.user_id());
989            }
990
991            #[async_test]
992            async fn test_olm_hash_saving() {
993                let (_, store) = get_loaded_store("olm_hash_saving").await;
994
995                let hash = OlmMessageHash {
996                    sender_key: "test_sender".to_owned(),
997                    hash: "test_hash".to_owned(),
998                };
999
1000                let mut changes = Changes::default();
1001                changes.message_hashes.push(hash.clone());
1002
1003                assert!(!store.is_message_known(&hash).await.unwrap());
1004                store.save_changes(changes).await.unwrap();
1005                assert!(store.is_message_known(&hash).await.unwrap());
1006            }
1007
1008            #[async_test]
1009            async fn test_key_request_saving() {
1010                let (account, store) = get_loaded_store("key_request_saving").await;
1011                let sender_key =
1012                    Curve25519PublicKey::from_base64("Nn0L2hkcCMFKqynTjyGsJbth7QrVmX3lbrksMkrGOAw")
1013                        .unwrap();
1014
1015                let id = TransactionId::new();
1016                let info: SecretInfo = MegolmV1AesSha2Content {
1017                    room_id: owned_room_id!("!test:localhost"),
1018                    sender_key: Some(sender_key),
1019                    session_id: "test_session_id".to_owned(),
1020                }
1021                .into();
1022
1023                let request = GossipRequest {
1024                    request_recipient: account.user_id().to_owned(),
1025                    request_id: id.clone(),
1026                    info: info.clone(),
1027                    sent_out: false,
1028                };
1029
1030                assert!(store.get_outgoing_secret_requests(&id).await.unwrap().is_none());
1031
1032                let mut changes = Changes::default();
1033                changes.key_requests.push(request.clone());
1034                store.save_changes(changes).await.unwrap();
1035
1036                let request = Some(request);
1037
1038                let stored_request = store.get_outgoing_secret_requests(&id).await.unwrap();
1039                assert_eq!(request, stored_request);
1040
1041                let stored_request = store.get_secret_request_by_info(&info).await.unwrap();
1042                assert_eq!(request, stored_request);
1043                assert!(!store.get_unsent_secret_requests().await.unwrap().is_empty());
1044
1045                let request = GossipRequest {
1046                    request_recipient: account.user_id().to_owned(),
1047                    request_id: id.clone(),
1048                    info: info.clone(),
1049                    sent_out: true,
1050                };
1051
1052                let mut changes = Changes::default();
1053                changes.key_requests.push(request.clone());
1054                store.save_changes(changes).await.unwrap();
1055
1056                assert!(store.get_unsent_secret_requests().await.unwrap().is_empty());
1057                let stored_request = store.get_outgoing_secret_requests(&id).await.unwrap();
1058                assert_eq!(Some(request), stored_request);
1059
1060                store.delete_outgoing_secret_requests(&id).await.unwrap();
1061
1062                let stored_request = store.get_outgoing_secret_requests(&id).await.unwrap();
1063                assert_eq!(None, stored_request);
1064
1065                let stored_request = store.get_secret_request_by_info(&info).await.unwrap();
1066                assert_eq!(None, stored_request);
1067                assert!(store.get_unsent_secret_requests().await.unwrap().is_empty());
1068            }
1069
1070            /// Test that if we try to store multiple gossip requests for the
1071            /// same secret, the latter request will replace the former.
1072            #[async_test]
1073            async fn test_secret_request_replacement() {
1074                let (account, store) = get_loaded_store("key_request_saving").await;
1075
1076                let id1 = TransactionId::new();
1077                let request1 = GossipRequest {
1078                    request_recipient: account.user_id().to_owned(),
1079                    request_id: id1.clone(),
1080                    info: SecretName::RecoveryKey.into(),
1081                    sent_out: false,
1082                };
1083                let msk_id = TransactionId::new();
1084                let msk_request = GossipRequest {
1085                    request_recipient: account.user_id().to_owned(),
1086                    request_id: msk_id.clone(),
1087                    info: SecretName::CrossSigningMasterKey.into(),
1088                    sent_out: false,
1089                };
1090
1091
1092                assert!(store.get_outgoing_secret_requests(&id1).await.unwrap().is_none());
1093
1094                let mut changes = Changes::default();
1095                changes.key_requests.push(request1.clone());
1096                changes.key_requests.push(msk_request.clone());
1097                store.save_changes(changes).await.unwrap();
1098
1099                assert!(store.get_outgoing_secret_requests(&id1).await.unwrap().is_some());
1100                assert!(store.get_outgoing_secret_requests(&msk_id).await.unwrap().is_some());
1101
1102                let id2 = TransactionId::new();
1103                let request2 = GossipRequest {
1104                    request_recipient: account.user_id().to_owned(),
1105                    request_id: id2.clone(),
1106                    info: SecretName::RecoveryKey.into(),
1107                    sent_out: false,
1108                };
1109
1110                let mut changes = Changes::default();
1111                changes.key_requests.push(request2.clone());
1112                store.save_changes(changes).await.unwrap();
1113
1114                // The first request for the recovery key should be replaced by
1115                // the second request, but request for the MSK should remain.
1116                assert!(store.get_outgoing_secret_requests(&id1).await.unwrap().is_none());
1117                assert!(store.get_outgoing_secret_requests(&id2).await.unwrap().is_some());
1118                assert!(store.get_outgoing_secret_requests(&msk_id).await.unwrap().is_some());
1119            }
1120
1121            #[async_test]
1122            async fn test_gossipped_secret_saving() {
1123                let (account, store) = get_loaded_store("gossipped_secret_saving").await;
1124
1125                let secret = "It is a secret to everybody";
1126
1127                let id = TransactionId::new();
1128                let info: SecretInfo = MegolmV1AesSha2Content {
1129                    room_id: owned_room_id!("!test:localhost"),
1130                    sender_key: Some(account.identity_keys().curve25519),
1131                    session_id: "test_session_id".to_owned(),
1132                }
1133                .into();
1134
1135                let gossip_request = GossipRequest {
1136                    request_recipient: account.user_id().to_owned(),
1137                    request_id: id.clone(),
1138                    info: info.clone(),
1139                    sent_out: true,
1140                };
1141
1142                let mut event = DecryptedSecretSendEvent {
1143                    sender: account.user_id().to_owned(),
1144                    recipient: account.user_id().to_owned(),
1145                    keys: OlmV1Keys {
1146                        ed25519: account.identity_keys().ed25519,
1147                    },
1148                    recipient_keys: OlmV1Keys {
1149                        ed25519: account.identity_keys().ed25519,
1150                    },
1151                    sender_device_keys: None,
1152                    content: SecretSendContent::new(id.to_owned(), secret.to_owned()),
1153                };
1154
1155                let value = GossippedSecret {
1156                    secret_name: SecretName::RecoveryKey,
1157                    gossip_request: gossip_request.to_owned(),
1158                    event: event.to_owned(),
1159                };
1160
1161                assert!(
1162                    store.get_secrets_from_inbox(&SecretName::RecoveryKey).await.unwrap().is_empty(),
1163                    "No secret should initially be found in the store"
1164                );
1165
1166                let mut changes = Changes::default();
1167                changes.secrets.push(value.into());
1168                store.save_changes(changes).await.unwrap();
1169
1170                let restored = store.get_secrets_from_inbox(&SecretName::RecoveryKey).await.unwrap();
1171                let first_secret = restored.first().expect("We should have restored a secret now");
1172                assert_eq!(first_secret.deref(), secret);
1173                assert_eq!(restored.len(), 1, "We should only have one secret stored for now");
1174
1175                let secret2 = "It is another secret to everybody";
1176                event.content.request_id = TransactionId::new();
1177                event.content.secret = secret2.to_string();
1178                let another_secret = GossippedSecret {
1179                    secret_name: SecretName::RecoveryKey,
1180                    gossip_request,
1181                    event,
1182                };
1183
1184                let mut changes = Changes::default();
1185                changes.secrets.push(another_secret.clone().into());
1186                store.save_changes(changes).await.unwrap();
1187
1188                let restored = store.get_secrets_from_inbox(&SecretName::RecoveryKey).await.unwrap();
1189                assert_eq!(restored.len(), 2, "We should only have two secrets stored");
1190
1191                let restored = store.get_secrets_from_inbox(&SecretName::CrossSigningMasterKey).await.unwrap();
1192                assert!(restored.is_empty(), "We should not have secrets of a different type stored");
1193
1194                // if we push a secret with the same name and secret, it should
1195                // succeed
1196                let mut changes = Changes::default();
1197                changes.secrets.push(another_secret.into());
1198                store.save_changes(changes).await.unwrap();
1199
1200                let restored = store.get_secrets_from_inbox(&SecretName::RecoveryKey).await.unwrap();
1201                // the store may store the secrets separately, or combine them
1202                // if they have the same name and secret
1203                assert!(restored.len() == 2 || restored.len() == 3, "We should only have two or three secrets stored");
1204
1205                store.delete_secrets_from_inbox(&SecretName::RecoveryKey).await.unwrap();
1206
1207                let restored = store.get_secrets_from_inbox(&SecretName::RecoveryKey).await.unwrap();
1208                assert!(restored.is_empty(), "We should not have any secrets after we have deleted them");
1209            }
1210
1211            #[async_test]
1212            #[cfg(feature = "experimental-push-secrets")]
1213            async fn test_pushed_secret_saving() {
1214                let (_account, store) = get_loaded_store("pushed_secret_saving").await;
1215
1216                let secret = "It is a secret to everybody";
1217
1218                let value = SecretPushContent::new(
1219                    SecretName::RecoveryKey,
1220                    secret.to_owned(),
1221                );
1222
1223                assert!(
1224                    store.get_secrets_from_inbox(&SecretName::RecoveryKey).await.unwrap().is_empty(),
1225                    "No secret should initially be found in the store"
1226                );
1227
1228                let mut changes = Changes::default();
1229                changes.secrets.push(value.clone().into());
1230                store.save_changes(changes).await.unwrap();
1231
1232                let restored = store.get_secrets_from_inbox(&SecretName::RecoveryKey).await.unwrap();
1233                let first_secret = restored.first().expect("We should have restored a secret now");
1234                assert_eq!(first_secret.deref(), secret);
1235                assert_eq!(restored.len(), 1, "We should only have one secret stored for now");
1236
1237                let secret2 = "It is another secret to everybody";
1238                let another_secret = SecretPushContent::new(
1239                    SecretName::RecoveryKey,
1240                    secret2.to_owned(),
1241                );
1242
1243                let mut changes = Changes::default();
1244                changes.secrets.push(another_secret.into());
1245                store.save_changes(changes).await.unwrap();
1246
1247                let restored = store.get_secrets_from_inbox(&SecretName::RecoveryKey).await.unwrap();
1248                assert_eq!(restored.len(), 2, "We should only have two secrets stored");
1249
1250                let restored = store.get_secrets_from_inbox(&SecretName::CrossSigningMasterKey).await.unwrap();
1251                assert!(restored.is_empty(), "We should not have secrets of a different type stored");
1252
1253                store.delete_secrets_from_inbox(&SecretName::RecoveryKey).await.unwrap();
1254
1255                let restored = store.get_secrets_from_inbox(&SecretName::RecoveryKey).await.unwrap();
1256                assert!(restored.is_empty(), "We should not have any secrets after we have deleted them");
1257            }
1258
1259            #[async_test]
1260            async fn test_withheld_info_storage() {
1261                let (account, store) = get_loaded_store("withheld_info_storage").await;
1262
1263                let user_id = account.user_id().to_owned();
1264                let room_id = room_id!("!DwLygpkclUAfQNnfva:example.com");
1265                let session_id_1 = "GBnDxGP9i3IkPsz3/ihNr6P7qjIXxSRVWZ1MYmSn09w";
1266                let session_id_2 = "IDLtnNCH2kIr3xIf1B7JFkGpQmTjyMca2jww+X6zeOE";
1267
1268                {
1269                    let mut info_list: BTreeMap<_, BTreeMap<_, RoomKeyWithheldEntry>> = BTreeMap::new();
1270
1271                    let content = RoomKeyWithheldContent::MegolmV1AesSha2(
1272                        MegolmV1AesSha2WithheldContent::Unverified(
1273                            CommonWithheldCodeContent::new(
1274                                room_id.to_owned(),
1275                                session_id_1.into(),
1276                                Curve25519PublicKey::from_base64(
1277                                    "9n7mdWKOjr9c4NTlG6zV8dbFtNK79q9vZADoh7nMUwA",
1278                                )
1279                                .unwrap(),
1280                                "DEVICEID".into(),
1281                            )
1282                            .into(),
1283                        ),
1284                    );
1285                    let event = ToDeviceEvent::new(user_id.to_owned(), content);
1286                    info_list
1287                        .entry(room_id.to_owned())
1288                        .or_default()
1289                        .insert(session_id_1.to_owned(), event.into());
1290
1291                    let content = RoomKeyWithheldContent::MegolmV1AesSha2(
1292                        MegolmV1AesSha2WithheldContent::BlackListed(
1293                            CommonWithheldCodeContent::new(
1294                                room_id.to_owned(),
1295                                session_id_2.into(),
1296                                Curve25519PublicKey::from_base64(
1297                                    "9n7mdWKOjr9c4NTlG6zV8dbFtNK79q9vZADoh7nMUwA",
1298                                )
1299                                .unwrap(),
1300                                "DEVICEID".into(),
1301                            )
1302                            .into(),
1303                        ),
1304                    );
1305                    let event = ToDeviceEvent::new(user_id.to_owned(), content);
1306                    info_list
1307                        .entry(room_id.to_owned())
1308                        .or_default()
1309                        .insert(session_id_2.to_owned(), event.into());
1310
1311                    let changes = Changes { withheld_session_info: info_list, ..Default::default() };
1312                    store.save_changes(changes).await.unwrap();
1313                }
1314
1315                // Test `get_withheld_info`
1316
1317                let is_withheld = store.get_withheld_info(room_id, session_id_1).await.unwrap();
1318
1319                assert_matches!(
1320                    is_withheld, Some(event)
1321                    if event.content.algorithm() == EventEncryptionAlgorithm::MegolmV1AesSha2 &&
1322                    event.content.withheld_code() == WithheldCode::Unverified
1323                );
1324
1325                let is_withheld = store.get_withheld_info(room_id, session_id_2).await.unwrap();
1326
1327                assert_matches!(
1328                    is_withheld, Some(event)
1329                    if event.content.algorithm() == EventEncryptionAlgorithm::MegolmV1AesSha2 &&
1330                    event.content.withheld_code() == WithheldCode::Blacklisted
1331                );
1332
1333                let other_room_id = room_id!("!nQRyiRFuyUhXeaQfiR:example.com");
1334
1335                let is_withheld =
1336                    store.get_withheld_info(other_room_id, session_id_2).await.unwrap();
1337
1338                assert!(is_withheld.is_none());
1339
1340                // Test `get_withheld_sessions_by_room_id`
1341                let withhelds = store.get_withheld_sessions_by_room_id(room_id).await.expect("Error getting withheld sessions by room ID");
1342                assert_eq!(withhelds.len(), 2);
1343                let withheld1 = withhelds.iter().find(|entry| entry.content.megolm_session_id() == Some(session_id_1)).expect("Did not find session 1 in withhelds list");
1344                assert_eq!(withheld1.content.withheld_code(), WithheldCode::Unverified)
1345            }
1346
1347            #[async_test]
1348            async fn test_room_settings_saving() {
1349                let (_, store) = get_loaded_store("room_settings_saving").await;
1350
1351                let room_1 = room_id!("!test_1:localhost");
1352                let settings_1 = RoomSettings {
1353                    algorithm: EventEncryptionAlgorithm::MegolmV1AesSha2,
1354                    #[cfg(feature = "experimental-encrypted-state-events")]
1355                    encrypt_state_events: false,
1356                    only_allow_trusted_devices: true,
1357                    session_rotation_period: Some(Duration::from_secs(10)),
1358                    session_rotation_period_messages: Some(123),
1359                };
1360
1361                let room_2 = room_id!("!test_2:localhost");
1362                let settings_2 = RoomSettings {
1363                    algorithm: EventEncryptionAlgorithm::OlmV1Curve25519AesSha2,
1364                    only_allow_trusted_devices: false,
1365                    ..Default::default()
1366                };
1367
1368                let room_3 = room_id!("!test_3:localhost");
1369
1370                let changes = Changes {
1371                    room_settings: HashMap::from([
1372                        (room_1.into(), settings_1.clone()),
1373                        (room_2.into(), settings_2.clone()),
1374                    ]),
1375                    ..Default::default()
1376                };
1377
1378                store.save_changes(changes).await.unwrap();
1379
1380                let loaded_settings_1 = store.get_room_settings(room_1).await.unwrap();
1381                assert_eq!(Some(settings_1), loaded_settings_1);
1382
1383                let loaded_settings_2 = store.get_room_settings(room_2).await.unwrap();
1384                assert_eq!(Some(settings_2), loaded_settings_2);
1385
1386                let loaded_settings_3 = store.get_room_settings(room_3).await.unwrap();
1387                assert_eq!(None, loaded_settings_3);
1388            }
1389
1390            #[async_test]
1391            async fn test_backup_keys_saving() {
1392                let (_account, store) = get_loaded_store("backup_keys_saving").await;
1393
1394                let restored = store.load_backup_keys().await.unwrap();
1395                assert!(restored.decryption_key.is_none(), "Initially no backup decryption key should be present");
1396
1397                let backup_decryption_key = Some(BackupDecryptionKey::new());
1398
1399                let changes = Changes { backup_decryption_key, ..Default::default() };
1400                store.save_changes(changes).await.unwrap();
1401
1402                let restored = store.load_backup_keys().await.unwrap();
1403                assert!(restored.decryption_key.is_some(), "We should be able to restore a backup decryption key");
1404                assert!(restored.backup_version.is_none(), "The backup version should still be None");
1405
1406                let changes = Changes { backup_version: Some("some_version".to_owned()), ..Default::default() };
1407                store.save_changes(changes).await.unwrap();
1408
1409                let restored = store.load_backup_keys().await.unwrap();
1410                assert!(restored.decryption_key.is_some(), "The backup decryption key should still be known");
1411                assert!(restored.backup_version.is_some(), "The backup version should now be Some as well");
1412            }
1413
1414            #[async_test]
1415            async fn test_dehydration_pickle_key_saving() {
1416                let (_account, store) = get_loaded_store("dehydration_pickle_key_saving").await;
1417
1418                let restored = store.load_dehydrated_device_pickle_key().await.unwrap();
1419                assert!(restored.is_none(), "Initially no pickle key should be present");
1420
1421                let dehydrated_device_pickle_key = Some(DehydratedDeviceKey::new());
1422                let exported_base64 = dehydrated_device_pickle_key.clone().unwrap().to_base64();
1423
1424                let changes = Changes { dehydrated_device_pickle_key, ..Default::default() };
1425                store.save_changes(changes).await.unwrap();
1426
1427                let restored = store.load_dehydrated_device_pickle_key().await.unwrap();
1428                assert!(restored.is_some(), "We should be able to restore a pickle key");
1429                assert_eq!(restored.unwrap().to_base64(), exported_base64);
1430
1431                // If None, should not clear the existing saved key
1432                let changes = Changes { dehydrated_device_pickle_key: None, ..Default::default() };
1433                store.save_changes(changes).await.unwrap();
1434
1435                let restored = store.load_dehydrated_device_pickle_key().await.unwrap();
1436                assert!(restored.is_some(), "We should be able to restore a pickle key");
1437                assert_eq!(restored.unwrap().to_base64(), exported_base64);
1438
1439            }
1440
1441             #[async_test]
1442            async fn test_delete_dehydration_pickle_key() {
1443                let (_account, store) = get_loaded_store("delete_dehydration_pickle_key").await;
1444
1445                let dehydrated_device_pickle_key = DehydratedDeviceKey::new();
1446
1447                let changes = Changes { dehydrated_device_pickle_key: Some(dehydrated_device_pickle_key), ..Default::default() };
1448                store.save_changes(changes).await.unwrap();
1449
1450                let restored = store.load_dehydrated_device_pickle_key().await.unwrap();
1451                assert!(restored.is_some(), "We should be able to restore a pickle key");
1452
1453                store.delete_dehydrated_device_pickle_key().await.unwrap();
1454
1455                let restored = store.load_dehydrated_device_pickle_key().await.unwrap();
1456                assert!(restored.is_none(), "The previously saved key should be deleted");
1457
1458            }
1459
1460
1461            #[async_test]
1462            async fn test_custom_value_saving() {
1463                let (_, store) = get_loaded_store("custom_value_saving").await;
1464                store.set_custom_value("A", "Hello".as_bytes().to_vec()).await.unwrap();
1465
1466                let loaded_1 = store.get_custom_value("A").await.unwrap();
1467                assert_eq!(Some("Hello".as_bytes().to_vec()), loaded_1);
1468
1469                let loaded_2 = store.get_custom_value("B").await.unwrap();
1470                assert_eq!(None, loaded_2);
1471            }
1472
1473            #[async_test]
1474            async fn test_received_room_key_bundle() {
1475                let store = get_store("received_room_key_bundle", None, true).await;
1476                let test_room = room_id!("!room:example.org");
1477
1478                fn make_bundle_data(sender_user: &UserId, bundle_uri: &str) -> StoredRoomKeyBundleData {
1479                    let info = ruma::events::room::V2EncryptedFileInfo::encode([0; 32], [0;16]).into();
1480
1481                    let file = ruma::events::room::EncryptedFile::new(
1482                        ruma::OwnedMxcUri::from(bundle_uri),
1483                        info,
1484                        Default::default()
1485                    );
1486
1487                    StoredRoomKeyBundleData {
1488                        sender_user: sender_user.to_owned(),
1489                        sender_key: Curve25519PublicKey::from_bytes([0u8; 32]),
1490                        sender_data: SenderData::unknown(),
1491                        bundle_data: RoomKeyBundleContent {
1492                            room_id: owned_room_id!("!room:example.org"),
1493                            file,
1494                        },
1495                    }
1496                }
1497
1498                // Add three entries
1499                let changes = Changes {
1500                    received_room_key_bundles: vec![
1501                        make_bundle_data(user_id!("@alice:example.com"), "alice1"),
1502                        make_bundle_data(user_id!("@bob:example.com"), "bob1"),
1503                        make_bundle_data(user_id!("@alice:example.com"), "alice2"),
1504                    ],
1505                    ..Default::default()
1506                };
1507                store.save_changes(changes).await.unwrap();
1508
1509                // Check we get the right one
1510                let bundle = store.get_received_room_key_bundle_data(
1511                    test_room, user_id!("@alice:example.com")
1512                ).await.unwrap().expect("Did not get any bundle data");
1513                assert_eq!(bundle.bundle_data.file.url.to_string(), "alice2");
1514            }
1515
1516            #[async_test]
1517            async fn test_room_pending_key_bundle() {
1518                use $crate::store::types::RoomPendingKeyBundleDetails;
1519                let store = get_store("room_pending_key_bundle", None, true).await;
1520                let test_room = room_id!("!room:example.org");
1521                let test_user = user_id!("@user:example.com");
1522                let timestamp = ruma::MilliSecondsSinceUnixEpoch::now();
1523
1524                // Empty to start with
1525                assert!(store.get_pending_key_bundle_details_for_room(test_room).await.unwrap().is_none());
1526
1527                // Now add an entry, and check it comes back correctly
1528                store.save_changes(Changes {
1529                    rooms_pending_key_bundle: HashMap::from([(
1530                        test_room.to_owned(),
1531                        Some(RoomPendingKeyBundleDetails {
1532                            room_id: test_room.to_owned(),
1533                            invite_accepted_at: timestamp,
1534                            inviter: test_user.to_owned(),
1535                        }),
1536                    )]),
1537                    ..Default::default()
1538                }).await.unwrap();
1539
1540                let details = store.get_pending_key_bundle_details_for_room(test_room).await.unwrap();
1541                assert_matches!(details, Some(details) => {
1542                    assert_eq!(details.room_id, test_room);
1543                    assert_eq!(details.inviter, test_user);
1544                    assert_eq!(details.invite_accepted_at, timestamp);
1545                });
1546
1547                let all_rooms = store.get_all_rooms_pending_key_bundles().await.unwrap();
1548                assert_eq!(all_rooms.len(), 1);
1549                assert_eq!(all_rooms[0].room_id, test_room);
1550                assert_eq!(all_rooms[0].inviter, test_user);
1551                assert_eq!(all_rooms[0].invite_accepted_at, timestamp);
1552
1553                // Clear the entry, and check it is blank again
1554                store.save_changes(Changes {
1555                    rooms_pending_key_bundle: HashMap::from([(test_room.to_owned(), None)]),
1556                    ..Default::default()
1557                }).await.unwrap();
1558                assert!(
1559                    store.get_pending_key_bundle_details_for_room(test_room).await.unwrap().is_none(),
1560                    "Pending key bundle details were present after being cleared"
1561                );
1562            }
1563
1564            #[async_test]
1565            async fn test_set_has_downloaded_all_room_keys() {
1566                let store = get_store("room_key_backups_fully_downloaded", None, true).await;
1567                let test_room = room_id!("!room:example.org");
1568
1569                let changes = Changes {
1570                    room_key_backups_fully_downloaded: HashSet::from_iter([test_room.to_owned()]),
1571                    ..Default::default()
1572                };
1573                store
1574                    .save_changes(changes)
1575                    .await
1576                    .expect("We should be able to save the changes to the store");
1577
1578                assert!(
1579                    store
1580                    .has_downloaded_all_room_keys(test_room)
1581                    .await
1582                    .expect("We should be able to check if we have downloaded room keys")
1583                );
1584            }
1585
1586            fn session_info(session: &InboundGroupSession) -> (&RoomId, &str) {
1587                (&session.room_id(), &session.session_id())
1588            }
1589
1590            async fn create_session(
1591                account: &Account,
1592                device_curve_key: &Curve25519PublicKey,
1593                sender_data_type: SenderDataType,
1594            ) -> InboundGroupSession {
1595                let sender_data = match sender_data_type {
1596                    SenderDataType::UnknownDevice => {
1597                        SenderData::UnknownDevice { legacy_session: false, owner_check_failed: false }
1598                    }
1599                    SenderDataType::DeviceInfo => SenderData::DeviceInfo {
1600                        device_keys: account.device_keys().clone(),
1601                        legacy_session: false,
1602                    },
1603                    SenderDataType::VerificationViolation => panic!("VerificationViolation not supported"),
1604                    SenderDataType::SenderUnverified=> panic!("SenderUnverified not supported"),
1605                    SenderDataType::SenderVerified => panic!("SenderVerified not supported"),
1606                };
1607
1608                let session_key = GroupSession::new(SessionConfig::default()).session_key();
1609
1610                InboundGroupSession::new(
1611                    device_curve_key.clone(),
1612                    account.device_keys().ed25519_key().unwrap(),
1613                    room_id!("!r:s.co"),
1614                    &session_key,
1615                    sender_data,
1616                    None,
1617                    EventEncryptionAlgorithm::MegolmV1AesSha2,
1618                    None,
1619                    false,
1620                )
1621                .unwrap()
1622            }
1623        }
1624    };
1625}
1626
1627#[allow(unused_macros)]
1628#[macro_export]
1629macro_rules! cryptostore_integration_tests_time {
1630    () => {
1631        mod cryptostore_integration_tests_time {
1632            use std::time::Duration;
1633
1634            use matrix_sdk_test::async_test;
1635            use $crate::store::CryptoStore as _;
1636
1637            use super::cryptostore_integration_tests::*;
1638
1639            #[async_test]
1640            async fn test_lease_locks() {
1641                let (_account, store) = get_loaded_store("lease_locks").await;
1642
1643                let acquired0 = store.try_take_leased_lock(0, "key", "alice").await.unwrap();
1644                assert_eq!(acquired0, Some(1)); // first generation
1645
1646                // Should extend the lease automatically (same holder).
1647                let acquired2 = store.try_take_leased_lock(300, "key", "alice").await.unwrap();
1648                assert_eq!(acquired2, Some(1)); // same lock generation
1649
1650                // Should extend the lease automatically (same holder + time is
1651                // ok).
1652                let acquired3 = store.try_take_leased_lock(300, "key", "alice").await.unwrap();
1653                assert_eq!(acquired3, Some(1)); // same lock generation
1654
1655                // Another attempt at taking the lock should fail, because it's
1656                // taken.
1657                let acquired4 = store.try_take_leased_lock(300, "key", "bob").await.unwrap();
1658                assert!(acquired4.is_none()); // not acquired
1659
1660                // Even if we insist.
1661                let acquired5 = store.try_take_leased_lock(300, "key", "bob").await.unwrap();
1662                assert!(acquired5.is_none());
1663
1664                // That's a nice test we got here, go take a little nap.
1665                tokio::time::sleep(Duration::from_millis(50)).await;
1666
1667                // Still too early.
1668                let acquired55 = store.try_take_leased_lock(300, "key", "bob").await.unwrap();
1669                assert!(acquired55.is_none()); // not acquired
1670
1671                // Ok you can take another nap then.
1672                tokio::time::sleep(Duration::from_millis(250)).await;
1673
1674                // At some point, we do get the lock.
1675                let acquired6 = store.try_take_leased_lock(0, "key", "bob").await.unwrap();
1676                assert_eq!(acquired6, Some(2)); // new lock generation!
1677
1678                tokio::time::sleep(Duration::from_millis(1)).await;
1679
1680                // The other gets it almost immediately too.
1681                let acquired7 = store.try_take_leased_lock(0, "key", "alice").await.unwrap();
1682                assert_eq!(acquired7, Some(3)); // new lock generation!
1683
1684                tokio::time::sleep(Duration::from_millis(1)).await;
1685
1686                // But when we take a longer lease…
1687                let acquired8 = store.try_take_leased_lock(300, "key", "bob").await.unwrap();
1688                assert_eq!(acquired8, Some(4)); // new lock generation!
1689
1690                // It blocks the other user.
1691                let acquired9 = store.try_take_leased_lock(300, "key", "alice").await.unwrap();
1692                assert!(acquired9.is_none()); // not acquired
1693
1694                // We can hold onto our lease.
1695                let acquired10 = store.try_take_leased_lock(300, "key", "bob").await.unwrap();
1696                assert_eq!(acquired10, Some(4)); // same lock generation
1697            }
1698        }
1699    };
1700}