Hoist check for missing OTKs to SessionManager

`Account::create_outbound_session` no longer takes an entire list of keys;
rather it takes a single key and it is up to the caller to pick a key out of
the list.

This in turn means that `SessionCreationError` loses one of its reason codes.

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

@@ -242,13 +242,6 @@ pub enum SessionCreationError {
     )]
     OneTimeKeyNotSigned(OwnedUserId, OwnedDeviceId),
 
-    /// The signed one-time key is missing.
-    #[error(
-        "Tried to create a new Olm session for {0} {1}, but the signed \
-        one-time key is missing"
-    )]
-    OneTimeKeyMissing(OwnedUserId, OwnedDeviceId),
-
     /// The one-time key algorithm is unsupported.
     #[error(
         "Tried to create a new Olm session for {0} {1}, but the one-time \

crates/matrix-sdk-crypto/src/olm/account.rs

@@ -851,15 +851,8 @@ impl Account {
     pub fn create_outbound_session(
         &self,
         device: &ReadOnlyDevice,
-        key_map: &BTreeMap<OwnedDeviceKeyId, Raw<ruma::encryption::OneTimeKey>>,
+        one_time_key: &Raw<ruma::encryption::OneTimeKey>,
     ) -> Result<Session, SessionCreationError> {
-        let one_time_key = key_map.values().next().ok_or_else(|| {
-            SessionCreationError::OneTimeKeyMissing(
-                device.user_id().to_owned(),
-                device.device_id().into(),
-            )
-        })?;
-
         let one_time_key: SignedKey = match one_time_key.deserialize_as() {
             Ok(OneTimeKey::SignedKey(k)) => k,
             Ok(OneTimeKey::Key(_)) => {
@@ -956,11 +949,12 @@ impl Account {
         use ruma::events::dummy::ToDeviceDummyEventContent;
 
         other.generate_one_time_keys_helper(1);
-        let one_time = other.signed_one_time_keys();
+        let one_time_map = other.signed_one_time_keys();
+        let one_time = one_time_map.values().next().unwrap();
 
         let device = ReadOnlyDevice::from_account(other);
 
-        let mut our_session = self.create_outbound_session(&device, &one_time).unwrap();
+        let mut our_session = self.create_outbound_session(&device, one_time).unwrap();
 
         other.mark_keys_as_published();
 

crates/matrix-sdk-crypto/src/session_manager/sessions.rs

@@ -24,6 +24,7 @@ use ruma::{
     },
     assign,
     events::dummy::ToDeviceDummyEventContent,
+    serde::Raw,
     DeviceId, DeviceKeyAlgorithm, OwnedDeviceId, OwnedServerName, OwnedTransactionId, OwnedUserId,
     SecondsSinceUnixEpoch, ServerName, TransactionId, UserId,
 };
@@ -349,6 +350,37 @@ impl SessionManager {
         self.failures.extend(failed_servers);
         self.failures.remove(successful_servers);
 
+        // build a map of user_id -> device_id -> key for each device to we can start
+        // a session with.
+        let mut device_map: BTreeMap<
+            OwnedUserId,
+            BTreeMap<OwnedDeviceId, &Raw<ruma::encryption::OneTimeKey>>,
+        > = BTreeMap::new();
+
+        for (user_id, user_devices) in response.one_time_keys.iter() {
+            for (device_id, key_map) in user_devices {
+                match key_map.values().next() {
+                    Some(k) => {
+                        device_map.entry(user_id.clone()).or_default().insert(device_id.clone(), k);
+                    }
+                    None => {
+                        warn!(
+                            user_id = user_id.as_str(),
+                            device_id = device_id.as_str(),
+                            "Tried to create a new Olm session, but the signed one-time key is missing",
+                        );
+
+                        self.failed_devices
+                            .write()
+                            .unwrap()
+                            .entry(user_id.clone())
+                            .or_default()
+                            .insert(device_id.clone());
+                    }
+                };
+            }
+        }
+
         struct SessionInfo {
             session_id: String,
             algorithm: EventEncryptionAlgorithm,
@@ -370,8 +402,8 @@ impl SessionManager {
         let mut new_sessions: BTreeMap<&UserId, BTreeMap<&DeviceId, SessionInfo>> = BTreeMap::new();
 
         let mut store_transaction = self.store.transaction().await;
-        for (user_id, user_devices) in &response.one_time_keys {
-            for (device_id, key_map) in user_devices {
+        for (user_id, user_devices) in device_map.iter() {
+            for (device_id, one_time_key) in user_devices {
                 let device = match self.store.get_readonly_device(user_id, device_id).await {
                     Ok(Some(d)) => d,
                     Ok(None) => {
@@ -396,7 +428,7 @@ impl SessionManager {
                 };
 
                 let account = store_transaction.account().await?;
-                let session = match account.create_outbound_session(&device, key_map) {
+                let session = match account.create_outbound_session(&device, one_time_key) {
                     Ok(s) => s,
                     Err(e) => {
                         warn!(