1use std::sync::Arc;
45
46use ruma::{
47 DeviceId,
48 api::client::dehydrated_device::{DehydratedDeviceData, put_dehydrated_device},
49 assign,
50 events::AnyToDeviceEvent,
51 serde::Raw,
52};
53use thiserror::Error;
54use tracing::{instrument, trace};
55use vodozemac::{DehydratedDeviceError, LibolmPickleError};
56
57use crate::{
58 Account, CryptoStoreError, DecryptionSettings, EncryptionSyncChanges, OlmError, OlmMachine,
59 SignatureError,
60 store::{
61 CryptoStoreWrapper, MemoryStore, Store,
62 types::{Changes, DehydratedDeviceKey, RoomKeyInfo},
63 },
64 verification::VerificationMachine,
65};
66
67#[derive(Debug, Error)]
69pub enum DehydrationError {
70 #[error(transparent)]
72 LegacyPickle(#[from] LibolmPickleError),
73
74 #[error(transparent)]
76 Pickle(#[from] DehydratedDeviceError),
77
78 #[error("The pickle key has an invalid length, expected 32 bytes, got {0}")]
80 PickleKeyLength(usize),
81
82 #[error("The self-signing key is missing, can't create a dehydrated device")]
85 MissingSigningKey(#[from] SignatureError),
86
87 #[error(transparent)]
89 Json(#[from] serde_json::Error),
90
91 #[error(transparent)]
93 Store(#[from] CryptoStoreError),
94}
95
96#[derive(Debug)]
98pub struct DehydratedDevices {
99 pub(crate) inner: OlmMachine,
100}
101
102impl DehydratedDevices {
103 pub async fn create(&self) -> Result<DehydratedDevice, DehydrationError> {
105 let user_id = self.inner.user_id();
106 let user_identity = self.inner.store().private_identity();
107
108 let account = Account::new_dehydrated(user_id);
109 let store =
110 Arc::new(CryptoStoreWrapper::new(user_id, account.device_id(), MemoryStore::new()));
111
112 let verification_machine = VerificationMachine::new(
113 account.static_data().clone(),
114 user_identity.clone(),
115 store.clone(),
116 );
117
118 let store =
119 Store::new(account.static_data().clone(), user_identity, store, verification_machine);
120 store
121 .save_pending_changes(crate::store::types::PendingChanges { account: Some(account) })
122 .await?;
123
124 Ok(DehydratedDevice { store })
125 }
126
127 pub async fn rehydrate(
144 &self,
145 pickle_key: &DehydratedDeviceKey,
146 device_id: &DeviceId,
147 device_data: Raw<DehydratedDeviceData>,
148 ) -> Result<RehydratedDevice, DehydrationError> {
149 let rehydrated =
150 self.inner.rehydrate(pickle_key.inner.as_ref(), device_id, device_data).await?;
151
152 Ok(RehydratedDevice { rehydrated, original: self.inner.to_owned() })
153 }
154
155 pub async fn get_dehydrated_device_pickle_key(
163 &self,
164 ) -> Result<Option<DehydratedDeviceKey>, DehydrationError> {
165 Ok(self.inner.store().load_dehydrated_device_pickle_key().await?)
166 }
167
168 pub async fn save_dehydrated_device_pickle_key(
174 &self,
175 dehydrated_device_pickle_key: &DehydratedDeviceKey,
176 ) -> Result<(), DehydrationError> {
177 let changes = Changes {
178 dehydrated_device_pickle_key: Some(dehydrated_device_pickle_key.clone()),
179 ..Default::default()
180 };
181 Ok(self.inner.store().save_changes(changes).await?)
182 }
183
184 pub async fn delete_dehydrated_device_pickle_key(&self) -> Result<(), DehydrationError> {
186 Ok(self.inner.store().delete_dehydrated_device_pickle_key().await?)
187 }
188}
189
190#[derive(Debug)]
195pub struct RehydratedDevice {
196 rehydrated: OlmMachine,
197 original: OlmMachine,
198}
199
200impl RehydratedDevice {
201 #[instrument(
269 skip_all,
270 fields(
271 user_id = ?self.original.user_id(),
272 rehydrated_device_id = ?self.rehydrated.device_id(),
273 original_device_id = ?self.original.device_id()
274 )
275 )]
276 pub async fn receive_events(
277 &self,
278 events: Vec<Raw<AnyToDeviceEvent>>,
279 decryption_settings: &DecryptionSettings,
280 ) -> Result<Vec<RoomKeyInfo>, OlmError> {
281 trace!("Receiving events for a rehydrated Device");
282
283 let sync_changes = EncryptionSyncChanges {
284 to_device_events: events,
285 next_batch_token: None,
286 one_time_keys_counts: &Default::default(),
287 changed_devices: &Default::default(),
288 unused_fallback_keys: None,
289 };
290
291 let mut rehydrated_transaction = self.rehydrated.store().transaction().await;
294
295 let (_, changes) = self
296 .rehydrated
297 .preprocess_sync_changes(&mut rehydrated_transaction, sync_changes, decryption_settings)
298 .await?;
299
300 let room_keys = &changes.inbound_group_sessions;
302 let updates = room_keys.iter().map(Into::into).collect();
303
304 trace!(room_key_count = room_keys.len(), "Collected room keys from the rehydrated device");
305
306 self.original.store().save_inbound_group_sessions(room_keys).await?;
307
308 rehydrated_transaction.commit().await?;
309 self.rehydrated.store().save_changes(changes).await?;
310
311 Ok(updates)
312 }
313}
314
315#[derive(Debug)]
320pub struct DehydratedDevice {
321 store: Store,
322}
323
324impl DehydratedDevice {
325 #[instrument(
362 skip_all, fields(
363 user_id = ?self.store.static_account().user_id,
364 device_id = ?self.store.static_account().device_id,
365 identity_keys = ?self.store.static_account().identity_keys,
366 )
367 )]
368 pub async fn keys_for_upload(
369 &self,
370 initial_device_display_name: String,
371 pickle_key: &DehydratedDeviceKey,
372 ) -> Result<put_dehydrated_device::unstable::Request, DehydrationError> {
373 let mut transaction = self.store.transaction().await;
374
375 let account = transaction.account().await?;
376 account.generate_fallback_key_if_needed();
377
378 let (device_keys, one_time_keys, fallback_keys) = account.keys_for_upload();
379
380 let mut device_keys = device_keys
381 .expect("We should always try to upload device keys for a dehydrated device.");
382
383 self.store.private_identity().lock().await.sign_device_keys(&mut device_keys).await?;
384
385 trace!("Creating an upload request for a dehydrated device");
386
387 let device_id = self.store.static_account().device_id.clone();
388 let device_data = account.dehydrate(pickle_key.inner.as_ref());
389 let initial_device_display_name = Some(initial_device_display_name);
390
391 transaction.commit().await?;
392
393 Ok(
394 assign!(put_dehydrated_device::unstable::Request::new(device_id, device_data, device_keys.to_raw()), {
395 one_time_keys, fallback_keys, initial_device_display_name
396 }),
397 )
398 }
399}
400
401#[cfg(test)]
402mod tests {
403 use std::{collections::BTreeMap, iter};
404
405 use js_option::JsOption;
406 use matrix_sdk_test::async_test;
407 use ruma::{
408 DeviceId, RoomId, TransactionId, UserId,
409 api::client::{
410 dehydrated_device::put_dehydrated_device,
411 keys::get_keys::v3::Response as KeysQueryResponse,
412 },
413 assign,
414 encryption::DeviceKeys,
415 events::AnyToDeviceEvent,
416 room_id,
417 serde::Raw,
418 user_id,
419 };
420
421 use crate::{
422 DecryptionSettings, EncryptionSettings, OlmMachine, TrustRequirement,
423 dehydrated_devices::DehydratedDevice,
424 machine::{
425 test_helpers::{create_session, get_prepared_machine_test_helper},
426 tests::to_device_requests_to_content,
427 },
428 olm::OutboundGroupSession,
429 store::types::DehydratedDeviceKey,
430 types::{DeviceKeys as DeviceKeysType, events::ToDeviceEvent},
431 utilities::json_convert,
432 };
433
434 fn pickle_key() -> DehydratedDeviceKey {
435 DehydratedDeviceKey::from_bytes(&[0u8; 32])
436 }
437
438 fn user_id() -> &'static UserId {
439 user_id!("@alice:localhost")
440 }
441
442 async fn get_olm_machine() -> OlmMachine {
443 let (olm_machine, _) = get_prepared_machine_test_helper(user_id(), false).await;
444 olm_machine.bootstrap_cross_signing(false).await.unwrap();
445
446 olm_machine
447 }
448
449 async fn receive_device_keys(
452 olm_machine: &OlmMachine,
453 user_id: &UserId,
454 device_id: &DeviceId,
455 device_keys: Raw<DeviceKeys>,
456 ) {
457 let device_keys = BTreeMap::from([(device_id.to_owned(), device_keys)]);
458
459 let keys_query_response = assign!(
460 KeysQueryResponse::new(), {
461 device_keys: BTreeMap::from([(user_id.to_owned(), device_keys)]),
462 }
463 );
464
465 olm_machine
466 .mark_request_as_sent(&TransactionId::new(), &keys_query_response)
467 .await
468 .unwrap();
469 }
470
471 async fn send_room_key(
472 machine: &OlmMachine,
473 room_id: &RoomId,
474 recipient: &UserId,
475 ) -> (Raw<AnyToDeviceEvent>, OutboundGroupSession) {
476 let to_device_requests = machine
477 .share_room_key(room_id, iter::once(recipient), EncryptionSettings::default())
478 .await
479 .unwrap();
480
481 let event = ToDeviceEvent::new(
482 user_id().to_owned(),
483 to_device_requests_to_content(to_device_requests),
484 );
485
486 let session =
487 machine.inner.group_session_manager.get_outbound_group_session(room_id).expect(
488 "An outbound group session should have been created when the room key was shared",
489 );
490
491 (
492 json_convert(&event)
493 .expect("We should be able to convert the to-device event into it's Raw variatn"),
494 session,
495 )
496 }
497
498 #[async_test]
499 async fn test_dehydrated_device_creation() {
500 let olm_machine = get_olm_machine().await;
501
502 let dehydrated_device = olm_machine.dehydrated_devices().create().await.unwrap();
503
504 let request = dehydrated_device
505 .keys_for_upload("Foo".to_owned(), &pickle_key())
506 .await
507 .expect("We should be able to create a request to upload a dehydrated device");
508
509 assert!(
510 !request.one_time_keys.is_empty(),
511 "The dehydrated device creation request should contain some one-time keys"
512 );
513
514 assert!(
515 !request.fallback_keys.is_empty(),
516 "The dehydrated device creation request should contain some fallback keys"
517 );
518
519 let device_keys: DeviceKeysType = request.device_keys.deserialize_as().unwrap();
520 assert_eq!(
521 device_keys.dehydrated,
522 JsOption::Some(true),
523 "The device keys of the dehydrated device should be marked as dehydrated."
524 );
525 }
526
527 #[async_test]
528 async fn test_dehydrated_device_rehydration() {
529 let room_id = room_id!("!test:example.org");
530 let alice = get_olm_machine().await;
531
532 let dehydrated_device = alice.dehydrated_devices().create().await.unwrap();
533
534 let mut request = dehydrated_device
535 .keys_for_upload("Foo".to_owned(), &pickle_key())
536 .await
537 .expect("We should be able to create a request to upload a dehydrated device");
538
539 let (key_id, one_time_key) = request
540 .one_time_keys
541 .pop_first()
542 .expect("The dehydrated device creation request should contain a one-time key");
543
544 receive_device_keys(&alice, user_id(), &request.device_id, request.device_keys).await;
546 create_session(&alice, user_id(), &request.device_id, key_id, one_time_key).await;
548
549 let (event, group_session) = send_room_key(&alice, room_id, user_id()).await;
551
552 let bob = get_olm_machine().await;
555
556 let room_key = bob
557 .store()
558 .get_inbound_group_session(room_id, group_session.session_id())
559 .await
560 .unwrap();
561
562 assert!(
563 room_key.is_none(),
564 "We should not have access to the room key that was only sent to the dehydrated device"
565 );
566
567 let rehydrated = bob
569 .dehydrated_devices()
570 .rehydrate(&pickle_key(), &request.device_id, request.device_data)
571 .await
572 .expect("We should be able to rehydrate the device");
573
574 assert_eq!(rehydrated.rehydrated.device_id(), request.device_id);
575 assert_eq!(rehydrated.original.device_id(), alice.device_id());
576
577 let decryption_settings =
578 DecryptionSettings { sender_device_trust_requirement: TrustRequirement::Untrusted };
579
580 let ret = rehydrated
583 .receive_events(vec![event], &decryption_settings)
584 .await
585 .expect("We should be able to push to-device events into the rehydrated device");
586
587 assert_eq!(ret.len(), 1, "The rehydrated device should have imported a room key");
588
589 let room_key = bob
592 .store()
593 .get_inbound_group_session(room_id, group_session.session_id())
594 .await
595 .unwrap()
596 .expect("We should now have access to the room key, since the rehydrated device imported it for us");
597
598 assert_eq!(
599 room_key.session_id(),
600 group_session.session_id(),
601 "The session ids of the imported room key and the outbound group session should match"
602 );
603 }
604
605 #[async_test]
606 async fn test_dehydrated_device_pickle_key_cache() {
607 let alice = get_olm_machine().await;
608
609 let dehydrated_manager = alice.dehydrated_devices();
610
611 let stored_key = dehydrated_manager.get_dehydrated_device_pickle_key().await.unwrap();
612 assert!(stored_key.is_none());
613
614 let pickle_key = DehydratedDeviceKey::new();
615
616 dehydrated_manager.save_dehydrated_device_pickle_key(&pickle_key).await.unwrap();
617
618 let stored_key =
619 dehydrated_manager.get_dehydrated_device_pickle_key().await.unwrap().unwrap();
620 assert_eq!(stored_key.to_base64(), pickle_key.to_base64());
621
622 let dehydrated_device = dehydrated_manager.create().await.unwrap();
623
624 let request = dehydrated_device
625 .keys_for_upload("Foo".to_owned(), &stored_key)
626 .await
627 .expect("We should be able to create a request to upload a dehydrated device");
628
629 dehydrated_manager
631 .rehydrate(&stored_key, &request.device_id, request.device_data)
632 .await
633 .expect("We should be able to rehydrate the device");
634
635 dehydrated_manager
636 .delete_dehydrated_device_pickle_key()
637 .await
638 .expect("Should be able to delete the dehydrated device key");
639
640 let stored_key = dehydrated_manager.get_dehydrated_device_pickle_key().await.unwrap();
641 assert!(stored_key.is_none());
642 }
643
644 #[async_test]
646 async fn test_legacy_dehydrated_device_rehydration() {
647 let room_id = room_id!("!test:example.org");
648 let alice = get_olm_machine().await;
649
650 let dehydrated_device = alice.dehydrated_devices().create().await.unwrap();
651 let mut request =
652 legacy_dehydrated_device_keys_for_upload(&dehydrated_device, &pickle_key()).await;
653
654 let (key_id, one_time_key) = request
655 .one_time_keys
656 .pop_first()
657 .expect("The dehydrated device creation request should contain a one-time key");
658
659 let device_id = request.device_id;
660
661 receive_device_keys(&alice, user_id(), &device_id, request.device_keys).await;
663 create_session(&alice, user_id(), &device_id, key_id, one_time_key).await;
665
666 let (event, group_session) = send_room_key(&alice, room_id, user_id()).await;
668
669 let bob = get_olm_machine().await;
672
673 let room_key = bob
674 .store()
675 .get_inbound_group_session(room_id, group_session.session_id())
676 .await
677 .unwrap();
678
679 assert!(
680 room_key.is_none(),
681 "We should not have access to the room key that was only sent to the dehydrated device"
682 );
683
684 let rehydrated = bob
686 .dehydrated_devices()
687 .rehydrate(&pickle_key(), &device_id, request.device_data)
688 .await
689 .expect("We should be able to rehydrate the device");
690
691 assert_eq!(rehydrated.rehydrated.device_id(), &device_id);
692 assert_eq!(rehydrated.original.device_id(), alice.device_id());
693
694 let decryption_settings =
695 DecryptionSettings { sender_device_trust_requirement: TrustRequirement::Untrusted };
696
697 let ret = rehydrated
700 .receive_events(vec![event], &decryption_settings)
701 .await
702 .expect("We should be able to push to-device events into the rehydrated device");
703
704 assert_eq!(ret.len(), 1, "The rehydrated device should have imported a room key");
705
706 let room_key = bob
709 .store()
710 .get_inbound_group_session(room_id, group_session.session_id())
711 .await
712 .unwrap()
713 .expect("We should now have access to the room key, since the rehydrated device imported it for us");
714
715 assert_eq!(
716 room_key.session_id(),
717 group_session.session_id(),
718 "The session ids of the imported room key and the outbound group session should match"
719 );
720 }
721
722 async fn legacy_dehydrated_device_keys_for_upload(
726 dehydrated_device: &DehydratedDevice,
727 pickle_key: &DehydratedDeviceKey,
728 ) -> put_dehydrated_device::unstable::Request {
729 let mut transaction = dehydrated_device.store.transaction().await;
730 let account = transaction.account().await.unwrap();
731 account.generate_fallback_key_if_needed();
732
733 let (device_keys, one_time_keys, fallback_keys) = account.keys_for_upload();
734 let mut device_keys = device_keys.unwrap();
735 dehydrated_device
736 .store
737 .private_identity()
738 .lock()
739 .await
740 .sign_device_keys(&mut device_keys)
741 .await
742 .expect("Should be able to cross-sign a device");
743
744 let device_id = account.device_id().to_owned();
745 let device_data = account.legacy_dehydrate(pickle_key.inner.as_ref());
746 transaction.commit().await.unwrap();
747
748 assign!(put_dehydrated_device::unstable::Request::new(device_id, device_data, device_keys.to_raw()), {
749 one_time_keys, fallback_keys
750 })
751 }
752}