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

crates/matrix-sdk-crypto/CHANGELOG.md

@@ -6,6 +6,11 @@ All notable changes to this project will be documented in this file.
 
 ## [Unreleased] - ReleaseDate
 
+- Expose new API `OlmMachine::get_dehydrated_device_pickle_key` and `OlmMachine::save_dehydrated_device_pickle_key`
+  to store/load the dehydrated device pickle key. This allows client to automatically rotate the dehydrated device 
+  to avoid OTKs exhaustion and to_device accumulation.
+  ([#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/machine/mod.rs

@@ -72,9 +72,9 @@ use crate::{
     },
     session_manager::{GroupSessionManager, SessionManager},
     store::{
-        Changes, CryptoStoreWrapper, DeviceChanges, IdentityChanges, IntoCryptoStore, MemoryStore,
-        PendingChanges, Result as StoreResult, RoomKeyInfo, RoomSettings, SecretImportError, Store,
-        StoreCache, StoreTransaction,
+        Changes, CryptoStoreWrapper, DehydratedDeviceKey, DeviceChanges, IdentityChanges,
+        IntoCryptoStore, MemoryStore, PendingChanges, Result as StoreResult, RoomKeyInfo,
+        RoomSettings, SecretImportError, Store, StoreCache, StoreTransaction,
     },
     types::{
         events::{
@@ -2378,6 +2378,33 @@ impl OlmMachine {
         DehydratedDevices { inner: self.to_owned() }
     }
 
+    /// 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 async fn get_dehydrated_device_pickle_key(
+        &self,
+    ) -> StoreResult<Option<DehydratedDeviceKey>> {
+        self.inner.store.load_dehydrated_device_pickle_key().await
+    }
+
+    /// 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 async fn save_dehydrated_device_pickle_key(
+        &self,
+        dehydrated_device_pickle_key: Option<DehydratedDeviceKey>,
+    ) -> Result<(), CryptoStoreError> {
+        let changes = Changes { dehydrated_device_pickle_key, ..Default::default() };
+        self.inner.store.save_changes(changes).await
+    }
+
     /// Get the stored encryption settings for the given room, such as the
     /// encryption algorithm or whether to encrypt only for trusted devices.
     ///

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,31 @@ 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);
+
+                let changes = Changes { dehydrated_device_pickle_key: None, ..Default::default() };
+                store.save_changes(changes).await.unwrap();
+
+                let restored = store.load_backup_keys().await.unwrap();
+                assert!(restored.decryption_key.is_none(), "The pickle key should be deleted");
+            }
+
+
             #[async_test]
             async fn test_custom_value_saving() {
                 let (_, store) = get_loaded_store("custom_value_saving").await;

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

@@ -31,8 +31,8 @@ use vodozemac::Curve25519PublicKey;
 
 use super::{
     caches::{DeviceStore, GroupSessionStore},
-    Account, BackupKeys, Changes, CryptoStore, InboundGroupSession, PendingChanges, RoomKeyCounts,
-    RoomSettings, Session,
+    Account, BackupKeys, Changes, CryptoStore, DehydratedDeviceKey, InboundGroupSession,
+    PendingChanges, RoomKeyCounts, RoomSettings, Session,
 };
 use crate::{
     gossiping::{GossipRequest, GossippedSecret, SecretInfo},
@@ -93,6 +93,7 @@ pub struct MemoryStore {
     leases: StdRwLock<HashMap<String, (String, Instant)>>,
     secret_inbox: StdRwLock<HashMap<String, Vec<GossippedSecret>>>,
     backup_keys: RwLock<BackupKeys>,
+    dehydrated_device_pickle_key: RwLock<Option<DehydratedDeviceKey>>,
     next_batch_token: RwLock<Option<String>>,
     room_settings: StdRwLock<HashMap<OwnedRoomId, RoomSettings>>,
 }
@@ -116,6 +117,7 @@ impl Default for MemoryStore {
             custom_values: Default::default(),
             leases: Default::default(),
             backup_keys: Default::default(),
+            dehydrated_device_pickle_key: Default::default(),
             secret_inbox: Default::default(),
             next_batch_token: Default::default(),
             room_settings: Default::default(),
@@ -268,6 +270,11 @@ impl CryptoStore for MemoryStore {
             self.backup_keys.write().await.backup_version = Some(version);
         }
 
+        if let Some(pickle_key) = changes.dehydrated_device_pickle_key {
+            let mut lock = self.dehydrated_device_pickle_key.write().await;
+            *lock = Some(pickle_key);
+        }
+
         {
             let mut secret_inbox = self.secret_inbox.write().unwrap();
             for secret in changes.secrets {
@@ -486,6 +493,10 @@ impl CryptoStore for MemoryStore {
         Ok(self.backup_keys.read().await.to_owned())
     }
 
+    async fn load_dehydrated_device_pickle_key(&self) -> Result<Option<DehydratedDeviceKey>> {
+        Ok(self.dehydrated_device_pickle_key.read().await.to_owned())
+    }
+
     async fn get_outbound_group_session(
         &self,
         room_id: &RoomId,
@@ -1125,7 +1136,10 @@ mod integration_tests {
             InboundGroupSession, OlmMessageHash, OutboundGroupSession, PrivateCrossSigningIdentity,
             SenderDataType, StaticAccountData,
         },
-        store::{BackupKeys, Changes, CryptoStore, PendingChanges, RoomKeyCounts, RoomSettings},
+        store::{
+            BackupKeys, Changes, CryptoStore, DehydratedDeviceKey, PendingChanges, RoomKeyCounts,
+            RoomSettings,
+        },
         types::events::room_key_withheld::RoomKeyWithheldEvent,
         Account, DeviceData, GossipRequest, GossippedSecret, SecretInfo, Session, TrackedUser,
         UserIdentityData,
@@ -1288,6 +1302,12 @@ mod integration_tests {
             self.0.load_backup_keys().await
         }
 
+        async fn load_dehydrated_device_pickle_key(
+            &self,
+        ) -> Result<Option<DehydratedDeviceKey>, Self::Error> {
+            self.0.load_dehydrated_device_pickle_key().await
+        }
+
         async fn get_outbound_group_session(
             &self,
             room_id: &RoomId,

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

@@ -518,6 +518,7 @@ pub struct Changes {
     pub private_identity: Option<PrivateCrossSigningIdentity>,
     pub backup_version: Option<String>,
     pub backup_decryption_key: Option<BackupDecryptionKey>,
+    pub dehydrated_device_pickle_key: Option<DehydratedDeviceKey>,
     pub sessions: Vec<Session>,
     pub message_hashes: Vec<OlmMessageHash>,
     pub inbound_group_sessions: Vec<InboundGroupSession>,
@@ -550,6 +551,7 @@ impl Changes {
         self.private_identity.is_none()
             && self.backup_version.is_none()
             && self.backup_decryption_key.is_none()
+            && self.dehydrated_device_pickle_key.is_none()
             && self.sessions.is_empty()
             && self.message_hashes.is_empty()
             && self.inbound_group_sessions.is_empty()
@@ -749,6 +751,43 @@ impl Debug for BackupDecryptionKey {
     }
 }
 
+/// The pickle key used to safely store the dehydrated device pickle.
+///
+/// This input key material will be expanded using HKDF into an AES key, MAC
+/// key, and an initialization vector (IV).
+#[derive(Clone, Zeroize, ZeroizeOnDrop, Deserialize, Serialize)]
+#[serde(transparent)]
+pub struct DehydratedDeviceKey {
+    pub(crate) inner: Box<[u8; DehydratedDeviceKey::KEY_SIZE]>,
+}
+
+impl DehydratedDeviceKey {
+    /// The number of bytes the encryption key will hold.
+    pub const KEY_SIZE: usize = 32;
+
+    /// Generates a new random pickle key.
+    pub fn new() -> Result<Self, rand::Error> {
+        let mut rng = rand::thread_rng();
+
+        let mut key = Box::new([0u8; Self::KEY_SIZE]);
+        rand::Fill::try_fill(key.as_mut_slice(), &mut rng)?;
+
+        Ok(Self { inner: key })
+    }
+
+    /// Export the [`DehydratedDeviceKey`] as a base64 encoded string.
+    pub fn to_base64(&self) -> String {
+        base64_encode(self.inner.as_slice())
+    }
+}
+
+#[cfg(not(tarpaulin_include))]
+impl Debug for DehydratedDeviceKey {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_tuple("DehydratedDeviceKey").field(&"...").finish()
+    }
+}
+
 impl DeviceChanges {
     /// Merge the given `DeviceChanges` into this instance of `DeviceChanges`.
     pub fn extend(&mut self, other: DeviceChanges) {

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

@@ -22,7 +22,8 @@ use ruma::{
 use vodozemac::Curve25519PublicKey;
 
 use super::{
-    BackupKeys, Changes, CryptoStoreError, PendingChanges, Result, RoomKeyCounts, RoomSettings,
+    BackupKeys, Changes, CryptoStoreError, DehydratedDeviceKey, PendingChanges, Result,
+    RoomKeyCounts, RoomSettings,
 };
 #[cfg(doc)]
 use crate::olm::SenderData;
@@ -195,6 +196,11 @@ pub trait CryptoStore: AsyncTraitDeps {
     /// Get the backup keys we have stored.
     async fn load_backup_keys(&self) -> Result<BackupKeys, Self::Error>;
 
+    /// Get the dehydrated device pickle key we have stored.
+    async fn load_dehydrated_device_pickle_key(
+        &self,
+    ) -> Result<Option<DehydratedDeviceKey>, Self::Error>;
+
     /// Get the outbound group session we have stored that is used for the
     /// given room.
     async fn get_outbound_group_session(
@@ -465,6 +471,10 @@ impl<T: CryptoStore> CryptoStore for EraseCryptoStoreError<T> {
         self.0.load_backup_keys().await.map_err(Into::into)
     }
 
+    async fn load_dehydrated_device_pickle_key(&self) -> Result<Option<DehydratedDeviceKey>> {
+        self.0.load_dehydrated_device_pickle_key().await.map_err(Into::into)
+    }
+
     async fn get_outbound_group_session(
         &self,
         room_id: &RoomId,

crates/matrix-sdk-indexeddb/src/crypto_store/mod.rs

@@ -29,8 +29,8 @@ use matrix_sdk_crypto::{
         StaticAccountData,
     },
     store::{
-        BackupKeys, Changes, CryptoStore, CryptoStoreError, PendingChanges, RoomKeyCounts,
-        RoomSettings,
+        BackupKeys, Changes, CryptoStore, CryptoStoreError, DehydratedDeviceKey, PendingChanges,
+        RoomKeyCounts, RoomSettings,
     },
     types::events::room_key_withheld::RoomKeyWithheldEvent,
     vodozemac::base64_encode,
@@ -104,6 +104,9 @@ mod keys {
     /// with the client-side recovery key, which is actually an AES key for use
     /// with SSSS.
     pub const RECOVERY_KEY_V1: &str = "recovery_key_v1";
+
+    /// Indexeddb key for the dehydrated device pickle key.
+    pub const DEHYDRATION_PICKLE_KEY: &str = "dehydration_pickle_key";
 }
 
 /// An implementation of [CryptoStore] that uses [IndexedDB] for persistent
@@ -471,6 +474,7 @@ impl IndexeddbCryptoStore {
 
         let decryption_key_pickle = &changes.backup_decryption_key;
         let backup_version = &changes.backup_version;
+        let dehydration_pickle_key = &changes.dehydrated_device_pickle_key;
 
         let mut core = indexeddb_changes.get(keys::CORE);
         if let Some(next_batch) = &changes.next_batch_token {
@@ -487,6 +491,13 @@ impl IndexeddbCryptoStore {
             );
         }
 
+        if let Some(i) = &dehydration_pickle_key {
+            core.put(
+                JsValue::from_str(keys::DEHYDRATION_PICKLE_KEY),
+                self.serializer.serialize_value(i)?,
+            );
+        }
+
         if let Some(a) = &decryption_key_pickle {
             indexeddb_changes.get(keys::BACKUP_KEYS).put(
                 JsValue::from_str(keys::RECOVERY_KEY_V1),
@@ -1291,6 +1302,25 @@ impl_crypto_store! {
         Ok(key)
     }
 
+
+     async fn load_dehydrated_device_pickle_key(&self) -> Result<Option<DehydratedDeviceKey>> {
+       if let Some(pickle) = self
+            .inner
+            .transaction_on_one_with_mode(keys::CORE, IdbTransactionMode::Readonly)?
+            .object_store(keys::CORE)?
+            .get(&JsValue::from_str(keys::DEHYDRATION_PICKLE_KEY))?
+            .await?
+        {
+            let pickle: DehydratedDeviceKey = self.serializer.deserialize_value(pickle)?;
+
+            Ok(Some(pickle))
+        } else {
+            Ok(None)
+        }
+    }
+
+
+
     async fn get_withheld_info(
         &self,
         room_id: &RoomId,

crates/matrix-sdk-sqlite/src/crypto_store.rs

@@ -27,7 +27,10 @@ use matrix_sdk_crypto::{
         InboundGroupSession, OutboundGroupSession, PickledInboundGroupSession,
         PrivateCrossSigningIdentity, SenderDataType, Session, StaticAccountData,
     },
-    store::{BackupKeys, Changes, CryptoStore, PendingChanges, RoomKeyCounts, RoomSettings},
+    store::{
+        BackupKeys, Changes, CryptoStore, DehydratedDeviceKey, PendingChanges, RoomKeyCounts,
+        RoomSettings,
+    },
     types::events::room_key_withheld::RoomKeyWithheldEvent,
     Account, DeviceData, GossipRequest, GossippedSecret, SecretInfo, TrackedUser, UserIdentityData,
 };
@@ -189,6 +192,9 @@ impl SqliteCryptoStore {
 
 const DATABASE_VERSION: u8 = 9;
 
+/// key for the dehydrated device pickle key in the key/value table.
+const DEHYDRATED_DEVICE_PICKLE_KEY: &str = "dehydrated_device_pickle_key";
+
 /// Run migrations for the given version of the database.
 async fn run_migrations(conn: &SqliteAsyncConn, version: u8) -> Result<()> {
     if version == 0 {
@@ -846,6 +852,11 @@ impl CryptoStore for SqliteCryptoStore {
                     txn.set_kv("backup_version_v1", &serialized_backup_version)?;
                 }
 
+                if let Some(pickle_key) = &changes.dehydrated_device_pickle_key {
+                    let serialized_pickle_key = this.serialize_value(pickle_key)?;
+                    txn.set_kv(DEHYDRATED_DEVICE_PICKLE_KEY, &serialized_pickle_key)?;
+                }
+
                 for device in changes.devices.new.iter().chain(&changes.devices.changed) {
                     let user_id = this.encode_key("device", device.user_id().as_bytes());
                     let device_id = this.encode_key("device", device.device_id().as_bytes());
@@ -1091,6 +1102,15 @@ impl CryptoStore for SqliteCryptoStore {
         Ok(BackupKeys { backup_version, decryption_key })
     }
 
+    async fn load_dehydrated_device_pickle_key(&self) -> Result<Option<DehydratedDeviceKey>> {
+        let conn = self.acquire().await?;
+
+        conn.get_kv(DEHYDRATED_DEVICE_PICKLE_KEY)
+            .await?
+            .map(|value| self.deserialize_value(&value))
+            .transpose()
+    }
+
     async fn get_outbound_group_session(
         &self,
         room_id: &RoomId,