crypto: Make restored outbound sessions wait for requests if they have some

matrix_sdk_crypto/src/machine.rs

@@ -751,7 +751,9 @@ impl OlmMachine {
         self.key_request_machine
             .mark_outgoing_request_as_sent(request_id)
             .await?;
-        self.group_session_manager.mark_request_as_sent(request_id);
+        self.group_session_manager
+            .mark_request_as_sent(request_id)
+            .await?;
         self.session_manager
             .mark_outgoing_request_as_sent(request_id);
 

matrix_sdk_crypto/src/olm/group_sessions/outbound.rs

@@ -32,7 +32,7 @@ use std::{
     },
     time::Duration,
 };
-use tracing::debug;
+use tracing::{debug, error, trace};
 
 use matrix_sdk_common::{
     events::{
@@ -204,6 +204,11 @@ impl OutboundGroupSession {
     /// users/devices that received the session.
     pub fn mark_request_as_sent(&self, request_id: &Uuid) {
         if let Some((_, r)) = self.to_share_with_set.remove(request_id) {
+            trace!(
+                request_id = request_id.to_string().as_str(),
+                "Marking to-device request carrying a room key as sent"
+            );
+
             let user_pairs = r.0.messages.iter().map(|(u, v)| {
                 (
                     u.clone(),
@@ -233,6 +238,19 @@ impl OutboundGroupSession {
                 );
                 self.mark_as_shared();
             }
+        } else {
+            let request_ids: Vec<String> = self
+                .to_share_with_set
+                .iter()
+                .map(|e| e.key().to_string())
+                .collect();
+
+            error!(
+                all_request_ids = ?request_ids,
+                request_id = request_id.to_string().as_str(),
+                "Marking to-device request carrying a room key as sent but no \
+                    request found with the given id"
+            );
         }
     }
 
@@ -428,6 +446,11 @@ impl OutboundGroupSession {
             .collect()
     }
 
+    /// Get the list of request ids this session is waiting for to be sent out.
+    pub(crate) fn pending_request_ids(&self) -> Vec<Uuid> {
+        self.to_share_with_set.iter().map(|e| *e.key()).collect()
+    }
+
     /// Restore a Session from a previously pickled string.
     ///
     /// Returns the restored group session or a `OlmGroupSessionError` if there

matrix_sdk_crypto/src/session_manager/group_sessions.rs

@@ -28,12 +28,12 @@ use matrix_sdk_common::{
     uuid::Uuid,
 };
 use serde_json::Value;
-use tracing::{debug, info};
+use tracing::{debug, info, trace};
 
 use crate::{
     error::{EventError, MegolmResult, OlmResult},
     olm::{Account, InboundGroupSession, OutboundGroupSession, Session, ShareState},
-    store::{Changes, Store},
+    store::{Changes, Result as StoreResult, Store},
     Device, EncryptionSettings, OlmError, ToDeviceRequest,
 };
 
@@ -72,10 +72,21 @@ impl GroupSessionManager {
         }
     }
 
-    pub fn mark_request_as_sent(&self, request_id: &Uuid) {
+    pub async fn mark_request_as_sent(&self, request_id: &Uuid) -> StoreResult<()> {
         if let Some((_, s)) = self.outbound_sessions_being_shared.remove(request_id) {
             s.mark_request_as_sent(request_id);
+
+            let mut changes = Changes::default();
+            changes.outbound_group_sessions.push(s.clone());
+            self.store.save_changes(changes).await?;
+        } else {
+            trace!(
+                request_id = request_id.to_string().as_str(),
+                "Marking room key share request as sent but session found that owns the given id"
+            )
         }
+
+        Ok(())
     }
 
     /// Get an outbound group session for a room, if one exists.
@@ -144,6 +155,11 @@ impl GroupSessionManager {
         let outbound_session = if let Some(s) = self.outbound_group_sessions.get(room_id) {
             Some(s.clone())
         } else if let Some(s) = self.store.get_outbound_group_sessions(room_id).await? {
+            for request_id in s.pending_request_ids() {
+                self.outbound_sessions_being_shared
+                    .insert(request_id, s.clone());
+            }
+
             self.outbound_group_sessions
                 .insert(room_id.clone(), s.clone());