1use std::sync::Arc;
45
46use ruma::{
47 api::client::dehydrated_device::{put_dehydrated_device, DehydratedDeviceData},
48 assign,
49 events::AnyToDeviceEvent,
50 serde::Raw,
51 DeviceId,
52};
53use thiserror::Error;
54use tracing::{instrument, trace};
55use vodozemac::{DehydratedDeviceError, LibolmPickleError};
56
57use crate::{
58 store::{
59 types::{Changes, DehydratedDeviceKey, RoomKeyInfo},
60 CryptoStoreWrapper, MemoryStore, Store,
61 },
62 verification::VerificationMachine,
63 Account, CryptoStoreError, DecryptionSettings, EncryptionSyncChanges, OlmError, OlmMachine,
64 SignatureError,
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(
146 &self,
147 pickle_key: &DehydratedDeviceKey,
148 device_id: &DeviceId,
149 device_data: Raw<DehydratedDeviceData>,
150 ) -> Result<RehydratedDevice, DehydrationError> {
151 let rehydrated =
152 self.inner.rehydrate(pickle_key.inner.as_ref(), device_id, device_data).await?;
153
154 Ok(RehydratedDevice { rehydrated, original: self.inner.to_owned() })
155 }
156
157 pub async fn get_dehydrated_device_pickle_key(
165 &self,
166 ) -> Result<Option<DehydratedDeviceKey>, DehydrationError> {
167 Ok(self.inner.store().load_dehydrated_device_pickle_key().await?)
168 }
169
170 pub async fn save_dehydrated_device_pickle_key(
176 &self,
177 dehydrated_device_pickle_key: &DehydratedDeviceKey,
178 ) -> Result<(), DehydrationError> {
179 let changes = Changes {
180 dehydrated_device_pickle_key: Some(dehydrated_device_pickle_key.clone()),
181 ..Default::default()
182 };
183 Ok(self.inner.store().save_changes(changes).await?)
184 }
185
186 pub async fn delete_dehydrated_device_pickle_key(&self) -> Result<(), DehydrationError> {
188 Ok(self.inner.store().delete_dehydrated_device_pickle_key().await?)
189 }
190}
191
192#[derive(Debug)]
197pub struct RehydratedDevice {
198 rehydrated: OlmMachine,
199 original: OlmMachine,
200}
201
202impl RehydratedDevice {
203 #[instrument(
272 skip_all,
273 fields(
274 user_id = ?self.original.user_id(),
275 rehydrated_device_id = ?self.rehydrated.device_id(),
276 original_device_id = ?self.original.device_id()
277 )
278 )]
279 pub async fn receive_events(
280 &self,
281 events: Vec<Raw<AnyToDeviceEvent>>,
282 decryption_settings: &DecryptionSettings,
283 ) -> Result<Vec<RoomKeyInfo>, OlmError> {
284 trace!("Receiving events for a rehydrated Device");
285
286 let sync_changes = EncryptionSyncChanges {
287 to_device_events: events,
288 next_batch_token: None,
289 one_time_keys_counts: &Default::default(),
290 changed_devices: &Default::default(),
291 unused_fallback_keys: None,
292 };
293
294 let mut rehydrated_transaction = self.rehydrated.store().transaction().await;
297
298 let (_, changes) = self
299 .rehydrated
300 .preprocess_sync_changes(&mut rehydrated_transaction, sync_changes, decryption_settings)
301 .await?;
302
303 let room_keys = &changes.inbound_group_sessions;
305 let updates = room_keys.iter().map(Into::into).collect();
306
307 trace!(room_key_count = room_keys.len(), "Collected room keys from the rehydrated device");
308
309 self.original.store().save_inbound_group_sessions(room_keys).await?;
310
311 rehydrated_transaction.commit().await?;
312 self.rehydrated.store().save_changes(changes).await?;
313
314 Ok(updates)
315 }
316}
317
318#[derive(Debug)]
323pub struct DehydratedDevice {
324 store: Store,
325}
326
327impl DehydratedDevice {
328 #[instrument(
365 skip_all, fields(
366 user_id = ?self.store.static_account().user_id,
367 device_id = ?self.store.static_account().device_id,
368 identity_keys = ?self.store.static_account().identity_keys,
369 )
370 )]
371 pub async fn keys_for_upload(
372 &self,
373 initial_device_display_name: String,
374 pickle_key: &DehydratedDeviceKey,
375 ) -> Result<put_dehydrated_device::unstable::Request, DehydrationError> {
376 let mut transaction = self.store.transaction().await;
377
378 let account = transaction.account().await?;
379 account.generate_fallback_key_if_needed();
380
381 let (device_keys, one_time_keys, fallback_keys) = account.keys_for_upload();
382
383 let mut device_keys = device_keys
384 .expect("We should always try to upload device keys for a dehydrated device.");
385
386 self.store.private_identity().lock().await.sign_device_keys(&mut device_keys).await?;
387
388 trace!("Creating an upload request for a dehydrated device");
389
390 let device_id = self.store.static_account().device_id.clone();
391 let device_data = account.dehydrate(pickle_key.inner.as_ref());
392 let initial_device_display_name = Some(initial_device_display_name);
393
394 transaction.commit().await?;
395
396 Ok(
397 assign!(put_dehydrated_device::unstable::Request::new(device_id, device_data, device_keys.to_raw()), {
398 one_time_keys, fallback_keys, initial_device_display_name
399 }),
400 )
401 }
402}
403
404#[cfg(test)]
405mod tests {
406 use std::{collections::BTreeMap, iter};
407
408 use js_option::JsOption;
409 use matrix_sdk_test::async_test;
410 use ruma::{
411 api::client::{
412 dehydrated_device::put_dehydrated_device,
413 keys::get_keys::v3::Response as KeysQueryResponse,
414 },
415 assign,
416 encryption::DeviceKeys,
417 events::AnyToDeviceEvent,
418 room_id,
419 serde::Raw,
420 user_id, DeviceId, RoomId, TransactionId, UserId,
421 };
422
423 use crate::{
424 dehydrated_devices::DehydratedDevice,
425 machine::{
426 test_helpers::{create_session, get_prepared_machine_test_helper},
427 tests::to_device_requests_to_content,
428 },
429 olm::OutboundGroupSession,
430 store::types::DehydratedDeviceKey,
431 types::{events::ToDeviceEvent, DeviceKeys as DeviceKeysType},
432 utilities::json_convert,
433 DecryptionSettings, EncryptionSettings, OlmMachine, TrustRequirement,
434 };
435
436 fn pickle_key() -> DehydratedDeviceKey {
437 DehydratedDeviceKey::from_bytes(&[0u8; 32])
438 }
439
440 fn user_id() -> &'static UserId {
441 user_id!("@alice:localhost")
442 }
443
444 async fn get_olm_machine() -> OlmMachine {
445 let (olm_machine, _) = get_prepared_machine_test_helper(user_id(), false).await;
446 olm_machine.bootstrap_cross_signing(false).await.unwrap();
447
448 olm_machine
449 }
450
451 async fn receive_device_keys(
454 olm_machine: &OlmMachine,
455 user_id: &UserId,
456 device_id: &DeviceId,
457 device_keys: Raw<DeviceKeys>,
458 ) {
459 let device_keys = BTreeMap::from([(device_id.to_owned(), device_keys)]);
460
461 let keys_query_response = assign!(
462 KeysQueryResponse::new(), {
463 device_keys: BTreeMap::from([(user_id.to_owned(), device_keys)]),
464 }
465 );
466
467 olm_machine
468 .mark_request_as_sent(&TransactionId::new(), &keys_query_response)
469 .await
470 .unwrap();
471 }
472
473 async fn send_room_key(
474 machine: &OlmMachine,
475 room_id: &RoomId,
476 recipient: &UserId,
477 ) -> (Raw<AnyToDeviceEvent>, OutboundGroupSession) {
478 let to_device_requests = machine
479 .share_room_key(room_id, iter::once(recipient), EncryptionSettings::default())
480 .await
481 .unwrap();
482
483 let event = ToDeviceEvent::new(
484 user_id().to_owned(),
485 to_device_requests_to_content(to_device_requests),
486 );
487
488 let session =
489 machine.inner.group_session_manager.get_outbound_group_session(room_id).expect(
490 "An outbound group session should have been created when the room key was shared",
491 );
492
493 (
494 json_convert(&event)
495 .expect("We should be able to convert the to-device event into it's Raw variatn"),
496 session,
497 )
498 }
499
500 #[async_test]
501 async fn test_dehydrated_device_creation() {
502 let olm_machine = get_olm_machine().await;
503
504 let dehydrated_device = olm_machine.dehydrated_devices().create().await.unwrap();
505
506 let request = dehydrated_device
507 .keys_for_upload("Foo".to_owned(), &pickle_key())
508 .await
509 .expect("We should be able to create a request to upload a dehydrated device");
510
511 assert!(
512 !request.one_time_keys.is_empty(),
513 "The dehydrated device creation request should contain some one-time keys"
514 );
515
516 assert!(
517 !request.fallback_keys.is_empty(),
518 "The dehydrated device creation request should contain some fallback keys"
519 );
520
521 let device_keys: DeviceKeysType = request.device_keys.deserialize_as().unwrap();
522 assert_eq!(
523 device_keys.dehydrated,
524 JsOption::Some(true),
525 "The device keys of the dehydrated device should be marked as dehydrated."
526 );
527 }
528
529 #[async_test]
530 async fn test_dehydrated_device_rehydration() {
531 let room_id = room_id!("!test:example.org");
532 let alice = get_olm_machine().await;
533
534 let dehydrated_device = alice.dehydrated_devices().create().await.unwrap();
535
536 let mut request = dehydrated_device
537 .keys_for_upload("Foo".to_owned(), &pickle_key())
538 .await
539 .expect("We should be able to create a request to upload a dehydrated device");
540
541 let (key_id, one_time_key) = request
542 .one_time_keys
543 .pop_first()
544 .expect("The dehydrated device creation request should contain a one-time key");
545
546 receive_device_keys(&alice, user_id(), &request.device_id, request.device_keys).await;
548 create_session(&alice, user_id(), &request.device_id, key_id, one_time_key).await;
550
551 let (event, group_session) = send_room_key(&alice, room_id, user_id()).await;
553
554 let bob = get_olm_machine().await;
556
557 let room_key = bob
558 .store()
559 .get_inbound_group_session(room_id, group_session.session_id())
560 .await
561 .unwrap();
562
563 assert!(
564 room_key.is_none(),
565 "We should not have access to the room key that was only sent to the dehydrated device"
566 );
567
568 let rehydrated = bob
570 .dehydrated_devices()
571 .rehydrate(&pickle_key(), &request.device_id, request.device_data)
572 .await
573 .expect("We should be able to rehydrate the device");
574
575 assert_eq!(rehydrated.rehydrated.device_id(), request.device_id);
576 assert_eq!(rehydrated.original.device_id(), alice.device_id());
577
578 let decryption_settings =
579 DecryptionSettings { sender_device_trust_requirement: TrustRequirement::Untrusted };
580
581 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().unwrap();
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;
671
672 let room_key = bob
673 .store()
674 .get_inbound_group_session(room_id, group_session.session_id())
675 .await
676 .unwrap();
677
678 assert!(
679 room_key.is_none(),
680 "We should not have access to the room key that was only sent to the dehydrated device"
681 );
682
683 let rehydrated = bob
685 .dehydrated_devices()
686 .rehydrate(&pickle_key(), &device_id, request.device_data)
687 .await
688 .expect("We should be able to rehydrate the device");
689
690 assert_eq!(rehydrated.rehydrated.device_id(), &device_id);
691 assert_eq!(rehydrated.original.device_id(), alice.device_id());
692
693 let decryption_settings =
694 DecryptionSettings { sender_device_trust_requirement: TrustRequirement::Untrusted };
695
696 let ret = rehydrated
698 .receive_events(vec![event], &decryption_settings)
699 .await
700 .expect("We should be able to push to-device events into the rehydrated device");
701
702 assert_eq!(ret.len(), 1, "The rehydrated device should have imported a room key");
703
704 let room_key = bob
707 .store()
708 .get_inbound_group_session(room_id, group_session.session_id())
709 .await
710 .unwrap()
711 .expect("We should now have access to the room key, since the rehydrated device imported it for us");
712
713 assert_eq!(
714 room_key.session_id(),
715 group_session.session_id(),
716 "The session ids of the imported room key and the outbound group session should match"
717 );
718 }
719
720 async fn legacy_dehydrated_device_keys_for_upload(
724 dehydrated_device: &DehydratedDevice,
725 pickle_key: &DehydratedDeviceKey,
726 ) -> put_dehydrated_device::unstable::Request {
727 let mut transaction = dehydrated_device.store.transaction().await;
728 let account = transaction.account().await.unwrap();
729 account.generate_fallback_key_if_needed();
730
731 let (device_keys, one_time_keys, fallback_keys) = account.keys_for_upload();
732 let mut device_keys = device_keys.unwrap();
733 dehydrated_device
734 .store
735 .private_identity()
736 .lock()
737 .await
738 .sign_device_keys(&mut device_keys)
739 .await
740 .expect("Should be able to cross-sign a device");
741
742 let device_id = account.device_id().to_owned();
743 let device_data = account.legacy_dehydrate(pickle_key.inner.as_ref());
744 transaction.commit().await.unwrap();
745
746 assign!(put_dehydrated_device::unstable::Request::new(device_id, device_data, device_keys.to_raw()), {
747 one_time_keys, fallback_keys
748 })
749 }
750}