feat(crypto-bindings): Save/Load dehydrated pickle key

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

@@ -1,13 +1,15 @@
 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;
 use tokio::runtime::Handle;
-use zeroize::Zeroize;
 
 #[derive(Debug, thiserror::Error, uniffi::Error)]
 #[uniffi(flat_error)]
@@ -33,6 +35,9 @@ impl From<matrix_sdk_crypto::dehydrated_devices::DehydrationError> for Dehydrati
                 Self::MissingSigningKey(e)
             }
             matrix_sdk_crypto::dehydrated_devices::DehydrationError::Store(e) => Self::Store(e),
+            matrix_sdk_crypto::dehydrated_devices::DehydrationError::PickleKeyLength(l) => {
+                Self::PickleKeyLength(l)
+            }
         }
     }
 }
@@ -73,20 +78,18 @@ impl DehydratedDevices {
         let device_data: Raw<_> = serde_json::from_str(&device_data)?;
         let device_id: OwnedDeviceId = device_id.into();
 
-        let mut key = get_pickle_key(&pickle_key)?;
+        let key: DehydratedDeviceKey = pickle_key.try_into()?;
 
         let ret = RehydratedDevice {
             runtime: self.runtime.to_owned(),
             inner: ManuallyDrop::new(self.runtime.block_on(self.inner.rehydrate(
-                &key,
+                key,
                 &device_id,
                 device_data,
             ))?),
         }
         .into();
 
-        key.zeroize();
-
         Ok(ret)
     }
 }
@@ -140,12 +143,10 @@ impl DehydratedDevice {
         device_display_name: String,
         pickle_key: Vec<u8>,
     ) -> Result<UploadDehydratedDeviceRequest, DehydrationError> {
-        let mut key = get_pickle_key(&pickle_key)?;
+        let key: DehydratedDeviceKey = pickle_key.try_into()?;
 
         let request =
-            self.runtime.block_on(self.inner.keys_for_upload(device_display_name, &key))?;
-
-        key.zeroize();
+            self.runtime.block_on(self.inner.keys_for_upload(device_display_name, key))?;
 
         Ok(request.into())
     }
@@ -176,16 +177,3 @@ impl From<dehydrated_device::put_dehydrated_device::unstable::Request>
         Self { body }
     }
 }
-
-fn get_pickle_key(pickle_key: &[u8]) -> Result<Box<[u8; 32]>, 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);
-
-        Ok(key)
-    } else {
-        Err(DehydrationError::PickleKeyLength(pickle_key_length))
-    }
-}

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

@@ -822,6 +822,17 @@ impl TryFrom<matrix_sdk_crypto::store::BackupKeys> for BackupKeys {
     }
 }
 
+/// Dehydrated device key
+#[derive(uniffi::Record)]
+pub struct DehydratedDeviceKey {
+    pub(crate) inner: Vec<u8>,
+}
+
+impl From<matrix_sdk_crypto::store::DehydratedDeviceKey> for DehydratedDeviceKey {
+    fn from(pickle_key: matrix_sdk_crypto::store::DehydratedDeviceKey) -> Self {
+        DehydratedDeviceKey { inner: pickle_key.into() }
+    }
+}
 impl From<matrix_sdk_crypto::store::RoomKeyCounts> for RoomKeyCounts {
     fn from(count: matrix_sdk_crypto::store::RoomKeyCounts) -> Self {
         Self { total: count.total as i64, backed_up: count.backed_up as i64 }

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

@@ -56,10 +56,11 @@ use crate::{
     parse_user_id,
     responses::{response_from_string, OwnedResponse},
     BackupKeys, BackupRecoveryKey, BootstrapCrossSigningResult, CrossSigningKeyExport,
-    CrossSigningStatus, DecodeError, DecryptedEvent, Device, DeviceLists, EncryptionSettings,
-    EventEncryptionAlgorithm, KeyImportError, KeysImportResult, MegolmV1BackupKey,
-    ProgressListener, Request, RequestType, RequestVerificationResult, RoomKeyCounts, RoomSettings,
-    Sas, SignatureUploadRequest, StartSasResult, UserIdentity, Verification, VerificationRequest,
+    CrossSigningStatus, DecodeError, DecryptedEvent, DehydratedDeviceKey, Device, DeviceLists,
+    EncryptionSettings, EventEncryptionAlgorithm, KeyImportError, KeysImportResult,
+    MegolmV1BackupKey, ProgressListener, Request, RequestType, RequestVerificationResult,
+    RoomKeyCounts, RoomSettings, Sas, SignatureUploadRequest, StartSasResult, UserIdentity,
+    Verification, VerificationRequest,
 };
 
 /// The return value for the [`OlmMachine::receive_sync_changes()`] method.
@@ -1532,6 +1533,36 @@ impl OlmMachine {
         }
         .into()
     }
+
+    /// Get the cached dehydrated device pickle key if any.
+    ///
+    /// None if the key was not previously cached (via
+    /// [`Self::save_dehydrated_device_pickle_key`]).
+    ///
+    /// Should be used to periodically rotate the dehydrated device to avoid
+    /// OTK exhaustion and accumulation of to_device messages.
+    pub fn get_dehydrated_device_key(
+        &self,
+    ) -> Result<Option<DehydratedDeviceKey>, CryptoStoreError> {
+        Ok(self
+            .runtime
+            .block_on(self.inner.get_dehydrated_device_pickle_key())?
+            .map(DehydratedDeviceKey::from))
+    }
+
+    /// Store the dehydrated device pickle key in the crypto store.
+    ///
+    /// Use `None` to delete any previously saved pickle key.
+    ///
+    /// This is useful if the client wants to periodically rotate dehydrated
+    /// devices to avoid OTK exhaustion and accumulated to_device problems.
+    pub fn save_dehydrated_device_key(
+        &self,
+        pickle_key: Option<DehydratedDeviceKey>,
+    ) -> Result<(), CryptoStoreError> {
+        let sdk_pickle_key = pickle_key.map(|k| k.inner.try_into()).and_then(|res| res.ok());
+        Ok(self.runtime.block_on(self.inner.save_dehydrated_device_pickle_key(sdk_pickle_key))?)
+    }
 }
 
 impl OlmMachine {

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::{CryptoStoreWrapper, DehydratedDeviceKey, MemoryStore, RoomKeyInfo, Store},
     verification::VerificationMachine,
     Account, CryptoStoreError, EncryptionSyncChanges, OlmError, OlmMachine, SignatureError,
 };
@@ -69,6 +69,10 @@ pub enum DehydrationError {
     #[error(transparent)]
     Pickle(#[from] LibolmPickleError),
 
+    /// The pickle key has an invalid length
+    #[error("The pickle key has an invalid length, expected 32 bytes, got {0}")]
+    PickleKeyLength(usize),
+
     /// The dehydrated device could not be signed by our user identity,
     /// we're missing the self-signing key.
     #[error("The self-signing key is missing, can't create a dehydrated device")]
@@ -132,11 +136,11 @@ 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() })
@@ -314,7 +318,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 +334,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);
@@ -467,7 +472,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.into())
             .await
             .expect("We should be able to create a request to upload a dehydrated device");
 
@@ -497,7 +502,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.into())
             .await
             .expect("We should be able to create a request to upload a dehydrated device");
 
@@ -531,7 +536,7 @@ mod tests {
         // Rehydrate the device.
         let rehydrated = bob
             .dehydrated_devices()
-            .rehydrate(PICKLE_KEY, &request.device_id, request.device_data)
+            .rehydrate(PICKLE_KEY.into(), &request.device_id, request.device_data)
             .await
             .expect("We should be able to rehydrate the device");
 

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

@@ -96,6 +96,7 @@ use matrix_sdk_common::{store_locks::CrossProcessStoreLock, timeout::timeout};
 pub use memorystore::MemoryStore;
 pub use traits::{CryptoStore, DynCryptoStore, IntoCryptoStore};
 
+use crate::dehydrated_devices::DehydrationError;
 pub use crate::gossiping::{GossipRequest, SecretInfo};
 
 /// A wrapper for our CryptoStore trait object.
@@ -781,6 +782,33 @@ impl DehydratedDeviceKey {
     }
 }
 
+impl From<&[u8; 32]> for DehydratedDeviceKey {
+    fn from(value: &[u8; 32]) -> Self {
+        DehydratedDeviceKey { inner: Box::new(*value) }
+    }
+}
+
+impl From<DehydratedDeviceKey> for Vec<u8> {
+    fn from(key: DehydratedDeviceKey) -> Self {
+        key.inner.to_vec()
+    }
+}
+
+impl TryFrom<Vec<u8>> for DehydratedDeviceKey {
+    type Error = DehydrationError;
+
+    fn try_from(value: Vec<u8>) -> Result<Self, Self::Error> {
+        if value.len() == 32 {
+            // Convert the Vec<u8> to a Box<[u8; 32]>
+            let mut key = Box::new([0u8; 32]);
+            key.copy_from_slice(&value);
+            Ok(DehydratedDeviceKey { inner: key })
+        } else {
+            Err(DehydrationError::PickleKeyLength(value.len()))
+        }
+    }
+}
+
 #[cfg(not(tarpaulin_include))]
 impl Debug for DehydratedDeviceKey {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {