Merge pull request #4383 from matrix-org/valere/cache_dehydration_pic…

…kle_key

feat(crypto): Support storing the dehydrated device pickle key

bindings/matrix-sdk-crypto-ffi/src/dehydrated_devices.rs

@@ -1,8 +1,11 @@
 use std::{mem::ManuallyDrop, sync::Arc};
 
-use matrix_sdk_crypto::dehydrated_devices::{
-    DehydratedDevice as InnerDehydratedDevice, DehydratedDevices as InnerDehydratedDevices,
-    RehydratedDevice as InnerRehydratedDevice,
+use matrix_sdk_crypto::{
+    dehydrated_devices::{
+        DehydratedDevice as InnerDehydratedDevice, DehydratedDevices as InnerDehydratedDevices,
+        RehydratedDevice as InnerRehydratedDevice,
+    },
+    store::DehydratedDeviceKey,
 };
 use ruma::{api::client::dehydrated_device, events::AnyToDeviceEvent, serde::Raw, OwnedDeviceId};
 use serde_json::json;
@@ -177,13 +180,13 @@ impl From<dehydrated_device::put_dehydrated_device::unstable::Request>
     }
 }
 
-fn get_pickle_key(pickle_key: &[u8]) -> Result<Box<[u8; 32]>, DehydrationError> {
+fn get_pickle_key(pickle_key: &[u8]) -> Result<DehydratedDeviceKey, DehydrationError> {
     let pickle_key_length = pickle_key.len();
 
     if pickle_key_length == 32 {
-        let mut key = Box::new([0u8; 32]);
-        key.copy_from_slice(pickle_key);
-
+        let mut raw_bytes = [0u8; 32];
+        raw_bytes.copy_from_slice(pickle_key);
+        let key = DehydratedDeviceKey::from_bytes(&raw_bytes);
         Ok(key)
     } else {
         Err(DehydrationError::PickleKeyLength(pickle_key_length))

crates/matrix-sdk-crypto/CHANGELOG.md

@@ -6,6 +6,13 @@ All notable changes to this project will be documented in this file.
 
 ## [Unreleased] - ReleaseDate
 
+- Expose new API `DehydratedDevices::get_dehydrated_device_pickle_key`, `DehydratedDevices::save_dehydrated_device_pickle_key`
+  and `DehydratedDevices::delete_dehydrated_device_pickle_key` to store/load the dehydrated device pickle key.
+  This allows client to automatically rotate the dehydrated device to avoid one-time-keys exhaustion and to_device accumulation.
+  [**breaking**] `DehydratedDevices::keys_for_upload` and `DehydratedDevices::rehydrate` now use the `DehydratedDeviceKey`
+  as parameter instead of a raw byte array. Use `DehydratedDeviceKey::from_bytes` to migrate.
+  ([#4383](https://github.com/matrix-org/matrix-rust-sdk/pull/4383))
+
 - Added new `UtdCause` variants `WithheldForUnverifiedOrInsecureDevice` and `WithheldBySender`.
   These variants provide clearer categorization for expected Unable-To-Decrypt (UTD) errors 
   when the sender either did not wish to share or was unable to share the room_key.

crates/matrix-sdk-crypto/src/dehydrated_devices.rs

@@ -57,7 +57,7 @@ use tracing::{instrument, trace};
 use vodozemac::LibolmPickleError;
 
 use crate::{
-    store::{CryptoStoreWrapper, MemoryStore, RoomKeyInfo, Store},
+    store::{Changes, CryptoStoreWrapper, DehydratedDeviceKey, MemoryStore, RoomKeyInfo, Store},
     verification::VerificationMachine,
     Account, CryptoStoreError, EncryptionSyncChanges, OlmError, OlmMachine, SignatureError,
 };
@@ -132,15 +132,49 @@ impl DehydratedDevices {
     ///   private keys of the device.
     pub async fn rehydrate(
         &self,
-        pickle_key: &[u8; 32],
+        pickle_key: &DehydratedDeviceKey,
         device_id: &DeviceId,
         device_data: Raw<DehydratedDeviceData>,
     ) -> Result<RehydratedDevice, DehydrationError> {
-        let pickle_key = expand_pickle_key(pickle_key, device_id);
+        let pickle_key = expand_pickle_key(pickle_key.inner.as_ref(), device_id);
         let rehydrated = self.inner.rehydrate(&pickle_key, device_id, device_data).await?;
 
         Ok(RehydratedDevice { rehydrated, original: self.inner.to_owned() })
     }
+
+    /// Get the cached dehydrated device pickle key if any.
+    ///
+    /// None if the key was not previously cached (via
+    /// [`DehydratedDevices::save_dehydrated_device_pickle_key`]).
+    ///
+    /// Should be used to periodically rotate the dehydrated device to avoid
+    /// one-time keys exhaustion and accumulation of to_device messages.
+    pub async fn get_dehydrated_device_pickle_key(
+        &self,
+    ) -> Result<Option<DehydratedDeviceKey>, DehydrationError> {
+        Ok(self.inner.store().load_dehydrated_device_pickle_key().await?)
+    }
+
+    /// Store the dehydrated device pickle key in the crypto store.
+    ///
+    /// This is useful if the client wants to periodically rotate dehydrated
+    /// devices to avoid one-time keys exhaustion and accumulated to_device
+    /// problems.
+    pub async fn save_dehydrated_device_pickle_key(
+        &self,
+        dehydrated_device_pickle_key: &DehydratedDeviceKey,
+    ) -> Result<(), DehydrationError> {
+        let changes = Changes {
+            dehydrated_device_pickle_key: Some(dehydrated_device_pickle_key.clone()),
+            ..Default::default()
+        };
+        Ok(self.inner.store().save_changes(changes).await?)
+    }
+
+    /// Deletes the previously stored dehydrated device pickle key.
+    pub async fn delete_dehydrated_device_pickle_key(&self) -> Result<(), DehydrationError> {
+        Ok(self.inner.store().delete_dehydrated_device_pickle_key().await?)
+    }
 }
 
 /// A rehydraded device.
@@ -170,7 +204,7 @@ impl RehydratedDevice {
     ///
     /// ```no_run
     /// # use anyhow::Result;
-    /// # use matrix_sdk_crypto::OlmMachine;
+    /// # use matrix_sdk_crypto::{ OlmMachine, store::DehydratedDeviceKey };
     /// # use ruma::{api::client::dehydrated_device, DeviceId};
     /// # async fn example() -> Result<()> {
     /// # let machine: OlmMachine = unimplemented!();
@@ -184,9 +218,9 @@ impl RehydratedDevice {
     /// ) -> Result<dehydrated_device::get_events::unstable::Response> {
     ///     todo!("Download the to-device events of the dehydrated device");
     /// }
-    ///
-    /// // Don't use a zero key for real.
-    /// let pickle_key = [0u8; 32];
+    /// // Get the cached dehydrated key (got it after verification/recovery)
+    /// let pickle_key = machine
+    ///     .dehydrated_devices().get_dehydrated_device_pickle_key().await?.unwrap();
     ///
     /// // Fetch the dehydrated device from the server.
     /// let response = get_dehydrated_device().await?;
@@ -285,11 +319,13 @@ impl DehydratedDevice {
     /// # Examples
     ///
     /// ```no_run
-    /// # use matrix_sdk_crypto::OlmMachine;
-    /// # async fn example() -> anyhow::Result<()> {
+    /// # use matrix_sdk_crypto::OlmMachine;    /// #
+    /// use matrix_sdk_crypto::store::DehydratedDeviceKey;
+    ///
+    /// async fn example() -> anyhow::Result<()> {
     /// # let machine: OlmMachine = unimplemented!();
-    /// // Don't use a zero key for real.
-    /// let pickle_key = [0u8; 32];
+    /// // Create a new random key
+    /// let pickle_key = DehydratedDeviceKey::new()?;
     ///
     /// // Create the dehydrated device.
     /// let device = machine.dehydrated_devices().create().await?;
@@ -299,6 +335,9 @@ impl DehydratedDevice {
     ///     .keys_for_upload("Dehydrated device".to_owned(), &pickle_key)
     ///     .await?;
     ///
+    /// // Save the key if you want to later one rotate the dehydrated device
+    /// machine.dehydrated_devices().save_dehydrated_device_pickle_key(&pickle_key).await.unwrap();
+    ///
     /// // Send the request out using your HTTP client.
     /// // client.send(request).await?;
     /// # Ok(())
@@ -314,7 +353,7 @@ impl DehydratedDevice {
     pub async fn keys_for_upload(
         &self,
         initial_device_display_name: String,
-        pickle_key: &[u8; 32],
+        pickle_key: &DehydratedDeviceKey,
     ) -> Result<put_dehydrated_device::unstable::Request, DehydrationError> {
         let mut transaction = self.store.transaction().await;
 
@@ -330,7 +369,8 @@ impl DehydratedDevice {
 
         trace!("Creating an upload request for a dehydrated device");
 
-        let pickle_key = expand_pickle_key(pickle_key, &self.store.static_account().device_id);
+        let pickle_key =
+            expand_pickle_key(pickle_key.inner.as_ref(), &self.store.static_account().device_id);
         let device_id = self.store.static_account().device_id.clone();
         let device_data = account.dehydrate(&pickle_key);
         let initial_device_display_name = Some(initial_device_display_name);
@@ -393,12 +433,15 @@ mod tests {
             tests::to_device_requests_to_content,
         },
         olm::OutboundGroupSession,
+        store::DehydratedDeviceKey,
         types::{events::ToDeviceEvent, DeviceKeys as DeviceKeysType},
         utilities::json_convert,
         EncryptionSettings, OlmMachine,
     };
 
-    const PICKLE_KEY: &[u8; 32] = &[0u8; 32];
+    fn pickle_key() -> DehydratedDeviceKey {
+        DehydratedDeviceKey::from_bytes(&[0u8; 32])
+    }
 
     fn user_id() -> &'static UserId {
         user_id!("@alice:localhost")
@@ -467,7 +510,7 @@ mod tests {
         let dehydrated_device = olm_machine.dehydrated_devices().create().await.unwrap();
 
         let request = dehydrated_device
-            .keys_for_upload("Foo".to_owned(), PICKLE_KEY)
+            .keys_for_upload("Foo".to_owned(), &pickle_key())
             .await
             .expect("We should be able to create a request to upload a dehydrated device");
 
@@ -497,7 +540,7 @@ mod tests {
         let dehydrated_device = alice.dehydrated_devices().create().await.unwrap();
 
         let mut request = dehydrated_device
-            .keys_for_upload("Foo".to_owned(), PICKLE_KEY)
+            .keys_for_upload("Foo".to_owned(), &pickle_key())
             .await
             .expect("We should be able to create a request to upload a dehydrated device");
 
@@ -531,7 +574,7 @@ mod tests {
         // Rehydrate the device.
         let rehydrated = bob
             .dehydrated_devices()
-            .rehydrate(PICKLE_KEY, &request.device_id, request.device_data)
+            .rehydrate(&pickle_key(), &request.device_id, request.device_data)
             .await
             .expect("We should be able to rehydrate the device");
 
@@ -561,4 +604,43 @@ mod tests {
             "The session ids of the imported room key and the outbound group session should match"
         );
     }
+
+    #[async_test]
+    async fn test_dehydrated_device_pickle_key_cache() {
+        let alice = get_olm_machine().await;
+
+        let dehydrated_manager = alice.dehydrated_devices();
+
+        let stored_key = dehydrated_manager.get_dehydrated_device_pickle_key().await.unwrap();
+        assert!(stored_key.is_none());
+
+        let pickle_key = DehydratedDeviceKey::new().unwrap();
+
+        dehydrated_manager.save_dehydrated_device_pickle_key(&pickle_key).await.unwrap();
+
+        let stored_key =
+            dehydrated_manager.get_dehydrated_device_pickle_key().await.unwrap().unwrap();
+        assert_eq!(stored_key.to_base64(), pickle_key.to_base64());
+
+        let dehydrated_device = dehydrated_manager.create().await.unwrap();
+
+        let request = dehydrated_device
+            .keys_for_upload("Foo".to_owned(), &stored_key)
+            .await
+            .expect("We should be able to create a request to upload a dehydrated device");
+
+        // Rehydrate the device.
+        dehydrated_manager
+            .rehydrate(&stored_key, &request.device_id, request.device_data)
+            .await
+            .expect("We should be able to rehydrate the device");
+
+        dehydrated_manager
+            .delete_dehydrated_device_pickle_key()
+            .await
+            .expect("Should be able to delete the dehydrated device key");
+
+        let stored_key = dehydrated_manager.get_dehydrated_device_pickle_key().await.unwrap();
+        assert!(stored_key.is_none());
+    }
 }

crates/matrix-sdk-crypto/src/store/integration_tests.rs

@@ -51,7 +51,7 @@ macro_rules! cryptostore_integration_tests {
                     PrivateCrossSigningIdentity, SenderData, SenderDataType, Session
                 },
                 store::{
-                    BackupDecryptionKey, Changes, CryptoStore, DeviceChanges, GossipRequest,
+                    BackupDecryptionKey, Changes, CryptoStore, DehydratedDeviceKey, DeviceChanges, GossipRequest,
                     IdentityChanges, PendingChanges, RoomSettings,
                 },
                 testing::{get_device, get_other_identity, get_own_identity},
@@ -1217,6 +1217,53 @@ macro_rules! cryptostore_integration_tests {
                 assert!(restored.backup_version.is_some(), "The backup version should now be Some as well");
             }
 
+            #[async_test]
+            async fn test_dehydration_pickle_key_saving() {
+                let (_account, store) = get_loaded_store("dehydration_key_saving").await;
+
+                let restored = store.load_dehydrated_device_pickle_key().await.unwrap();
+                assert!(restored.is_none(), "Initially no pickle key should be present");
+
+                let dehydrated_device_pickle_key = Some(DehydratedDeviceKey::new().unwrap());
+                let exported_base64 = dehydrated_device_pickle_key.clone().unwrap().to_base64();
+
+                let changes = Changes { dehydrated_device_pickle_key, ..Default::default() };
+                store.save_changes(changes).await.unwrap();
+
+                let restored = store.load_dehydrated_device_pickle_key().await.unwrap();
+                assert!(restored.is_some(), "We should be able to restore a pickle key");
+                assert_eq!(restored.unwrap().to_base64(), exported_base64);
+
+                // If None, should not clear the existing saved key
+                let changes = Changes { dehydrated_device_pickle_key: None, ..Default::default() };
+                store.save_changes(changes).await.unwrap();
+
+                let restored = store.load_dehydrated_device_pickle_key().await.unwrap();
+                assert!(restored.is_some(), "We should be able to restore a pickle key");
+                assert_eq!(restored.unwrap().to_base64(), exported_base64);
+
+            }
+
+             #[async_test]
+            async fn test_delete_dehydration_pickle_key() {
+                let (_account, store) = get_loaded_store("dehydration_key_saving").await;
+
+                let dehydrated_device_pickle_key = DehydratedDeviceKey::new().unwrap();
+
+                let changes = Changes { dehydrated_device_pickle_key: Some(dehydrated_device_pickle_key), ..Default::default() };
+                store.save_changes(changes).await.unwrap();
+
+                let restored = store.load_dehydrated_device_pickle_key().await.unwrap();
+                assert!(restored.is_some(), "We should be able to restore a pickle key");
+
+                store.delete_dehydrated_device_pickle_key().await.unwrap();
+
+                let restored = store.load_dehydrated_device_pickle_key().await.unwrap();
+                assert!(restored.is_none(), "The previously saved key should be deleted");
+
+            }
+
+
             #[async_test]
             async fn test_custom_value_saving() {
                 let (_, store) = get_loaded_store("custom_value_saving").await;