crypto: Expose the key-claiming methods in the clients.

src/async_client.rs

@@ -177,9 +177,7 @@ impl SyncSettings {
 }
 
 #[cfg(feature = "encryption")]
-use api::r0::keys::get_keys;
+use api::r0::keys::{claim_keys, get_keys, upload_keys, KeyAlgorithm};
-#[cfg(feature = "encryption")]
-use api::r0::keys::upload_keys;
 use api::r0::message::create_message_event;
 use api::r0::session::login;
 use api::r0::sync::sync_events;
@@ -640,6 +638,37 @@ impl AsyncClient {
         Ok(response)
     }
 
+    /// Claim one-time keys creating new Olm sessions.
+    ///
+    /// # Arguments
+    ///
+    /// * `users` - The list of user/device pairs that we should claim keys for.
+    ///
+    /// # Panics
+    ///
+    /// Panics if the client isn't logged in, or if no encryption keys need to
+    /// be uploaded.
+    #[cfg(feature = "encryption")]
+    #[cfg_attr(docsrs, doc(cfg(feature = "encryption")))]
+    #[instrument]
+    async fn claim_one_time_keys(
+        &self,
+        one_time_keys: HashMap<UserId, HashMap<DeviceId, KeyAlgorithm>>,
+    ) -> Result<claim_keys::Response> {
+        let request = claim_keys::Request {
+            timeout: None,
+            one_time_keys,
+        };
+
+        let response = self.send(request).await?;
+        self.base_client
+            .write()
+            .await
+            .receive_keys_claim_response(&response)
+            .await?;
+        Ok(response)
+    }
+
     /// Upload the E2E encryption keys.
     ///
     /// This uploads the long lived device keys as well as the required amount

src/base_client.rs

@@ -40,10 +40,12 @@ use tokio::sync::Mutex;
 use crate::crypto::{OlmMachine, OneTimeKeys};
 #[cfg(feature = "encryption")]
 use ruma_client_api::r0::keys::{
-    get_keys::Response as KeysQueryResponse, upload_keys::Response as KeysUploadResponse,
+    claim_keys::Response as KeysClaimResponse, get_keys::Response as KeysQueryResponse,
-    DeviceKeys,
+    upload_keys::Response as KeysUploadResponse, DeviceKeys, KeyAlgorithm,
 };
 use ruma_identifiers::RoomId;
+#[cfg(feature = "encryption")]
+use ruma_identifiers::{DeviceId, UserId as RumaUserId};
 
 pub type Token = String;
 pub type UserId = String;
@@ -394,6 +396,23 @@ impl Client {
         }
     }
 
+    /// Get a tuple of device and one-time keys that need to be uploaded.
+    ///
+    /// Returns an empty error if no keys need to be uploaded.
+    #[cfg(feature = "encryption")]
+    #[cfg_attr(docsrs, doc(cfg(feature = "encryption")))]
+    pub async fn get_missing_sessions(
+        &self,
+        users: impl Iterator<Item = &String>,
+    ) -> HashMap<RumaUserId, HashMap<DeviceId, KeyAlgorithm>> {
+        let mut olm = self.olm.lock().await;
+
+        match &mut *olm {
+            Some(o) => o.get_missing_sessions(users).await,
+            None => HashMap::new(),
+        }
+    }
+
     /// Get a tuple of device and one-time keys that need to be uploaded.
     ///
     /// Returns an empty error if no keys need to be uploaded.
@@ -443,6 +462,25 @@ impl Client {
         Ok(())
     }
 
+    /// Receive a successful keys claim response.
+    ///
+    /// # Arguments
+    ///
+    /// * `response` - The keys claim response of the request that the client
+    /// performed.
+    ///
+    /// # Panics
+    /// Panics if the client hasn't been logged in.
+    #[cfg(feature = "encryption")]
+    #[cfg_attr(docsrs, doc(cfg(feature = "encryption")))]
+    pub async fn receive_keys_claim_response(&self, response: &KeysClaimResponse) -> Result<()> {
+        let mut olm = self.olm.lock().await;
+
+        let o = olm.as_mut().expect("Client isn't logged in.");
+        o.receive_keys_claim_response(response).await?;
+        Ok(())
+    }
+
     /// Receive a successful keys query response.
     ///
     /// # Arguments

src/crypto/memory_stores.rs

@@ -116,6 +116,10 @@ impl UserDevices {
     pub fn keys(&self) -> impl Iterator<Item = &String> {
         self.entries.keys()
     }
+
+    pub fn devices(&self) -> impl Iterator<Item = &Device> {
+        self.entries.values()
+    }
 }
 
 impl DeviceStore {

src/crypto/store/memorystore.rs

@@ -97,4 +97,9 @@ impl CryptoStore for MemoryStore {
     async fn get_user_devices(&self, user_id: &str) -> Result<UserDevices> {
         Ok(self.devices.user_devices(user_id))
     }
+
+    async fn save_device(&self, device: Device) -> Result<()> {
+        self.devices.add(device);
+        Ok(())
+    }
 }

src/crypto/store/mod.rs

@@ -83,6 +83,7 @@ pub trait CryptoStore: Debug + Send + Sync {
     ) -> Result<Option<Arc<Mutex<InboundGroupSession>>>>;
     fn tracked_users(&self) -> &HashSet<String>;
     async fn add_user_for_tracking(&mut self, user: &str) -> Result<bool>;
+    async fn save_device(&self, device: Device) -> Result<()>;
     async fn get_device(&self, user_id: &str, device_id: &str) -> Result<Option<Device>>;
     async fn get_user_devices(&self, user_id: &str) -> Result<UserDevices>;
 }

src/crypto/store/sqlite.rs

@@ -417,6 +417,10 @@ impl CryptoStore for SqliteStore {
     async fn get_user_devices(&self, user_id: &str) -> Result<UserDevices> {
         todo!()
     }
+
+    async fn save_device(&self, device: Device) -> Result<()> {
+        todo!()
+    }
 }
 
 impl std::fmt::Debug for SqliteStore {