Reshare Megolm session after the other party unwedges (#3604)

For each recipient device, we keep an "Olm wedging index" that indicates (approximately) how many times they tried to unwedge the Olm session. When we share a Megolm session, we store the Olm wedging index. This allows us to tell whether the Olm session was unwedged since we shared the Megolm session, so that we know if we should re-share it.

We detect an attempt to unwedge the Olm session by checking if the Olm message that we decrypted is from a brand new Olm session. This is not entirely accurate, since this could happen, for example, if Alice and Bob create Olm sessions at the same time. This may result in some Megolm sessions being re-shared unnecessarily, but should not make a huge difference.

crates/matrix-sdk-crypto/src/gossiping/machine.rs

@@ -642,7 +642,9 @@ impl GossipMachine {
         // information is recorded there.
         } else if let Some(outbound) = outbound_session {
             match outbound.is_shared_with(&device.inner) {
-                ShareState::Shared(message_index) => Ok(Some(message_index)),
+                ShareState::Shared { message_index, olm_wedging_index: _ } => {
+                    Ok(Some(message_index))
+                }
                 ShareState::SharedButChangedSenderKey => Err(KeyForwardDecision::ChangedSenderKey),
                 ShareState::NotShared => Err(KeyForwardDecision::OutboundSessionNotShared),
             }

crates/matrix-sdk-crypto/src/identities/device.rs

@@ -43,7 +43,9 @@ use crate::{
     olm::{
         InboundGroupSession, OutboundGroupSession, Session, ShareInfo, SignedJsonObject, VerifyJson,
     },
-    store::{Changes, CryptoStoreWrapper, DeviceChanges, Result as StoreResult},
+    store::{
+        caches::SequenceNumber, Changes, CryptoStoreWrapper, DeviceChanges, Result as StoreResult,
+    },
     types::{
         events::{
             forwarded_room_key::ForwardedRoomKeyContent,
@@ -89,6 +91,10 @@ pub struct ReadOnlyDevice {
     /// Default to epoch for migration purpose.
     #[serde(default = "default_timestamp")]
     first_time_seen_ts: MilliSecondsSinceUnixEpoch,
+    /// The number of times the device has tried to unwedge Olm sessions with
+    /// us.
+    #[serde(default)]
+    pub(crate) olm_wedging_index: SequenceNumber,
 }
 
 fn default_timestamp() -> MilliSecondsSinceUnixEpoch {
@@ -580,6 +586,7 @@ impl ReadOnlyDevice {
             deleted: Arc::new(AtomicBool::new(false)),
             withheld_code_sent: Arc::new(AtomicBool::new(false)),
             first_time_seen_ts: MilliSecondsSinceUnixEpoch::now(),
+            olm_wedging_index: Default::default(),
         }
     }
 
@@ -833,7 +840,11 @@ impl ReadOnlyDevice {
 
         match self.encrypt(store, event_type, content).await {
             Ok((session, encrypted)) => Ok(MaybeEncryptedRoomKey::Encrypted {
-                share_info: ShareInfo::new_shared(session.sender_key().to_owned(), message_index),
+                share_info: ShareInfo::new_shared(
+                    session.sender_key().to_owned(),
+                    message_index,
+                    self.olm_wedging_index,
+                ),
                 used_session: session,
                 message: encrypted.cast(),
             }),
@@ -971,6 +982,7 @@ impl TryFrom<&DeviceKeys> for ReadOnlyDevice {
             trust_state: Arc::new(RwLock::new(LocalTrust::Unset)),
             withheld_code_sent: Arc::new(AtomicBool::new(false)),
             first_time_seen_ts: MilliSecondsSinceUnixEpoch::now(),
+            olm_wedging_index: Default::default(),
         };
 
         device.verify_device_keys(device_keys)?;

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

@@ -62,7 +62,7 @@ use crate::{
     error::{EventError, OlmResult, SessionCreationError},
     identities::ReadOnlyDevice,
     requests::UploadSigningKeysRequest,
-    store::{Changes, Store},
+    store::{Changes, DeviceChanges, Store},
     types::{
         events::{
             olm_v1::AnyDecryptedOlmEvent,
@@ -1327,12 +1327,24 @@ impl Account {
                 // we might try to create the same session again.
                 // TODO: separate the session cache from the storage so we only add
                 // it to the cache but don't store it.
-                store
-                    .save_changes(Changes {
-                        sessions: vec![result.session.clone()],
-                        ..Default::default()
-                    })
-                    .await?;
+                let mut changes =
+                    Changes { sessions: vec![result.session.clone()], ..Default::default() };
+
+                // Any new Olm session will bump the Olm wedging index for the
+                // sender's device, if we have their device, which will cause us
+                // to re-send existing Megolm sessions to them the next time we
+                // use the session.  If we don't have their device, this means
+                // that we haven't tried to send them any Megolm sessions yet,
+                // so we don't need to worry about it.
+                if let Some(device) = store.get_device_from_curve_key(sender, sender_key).await? {
+                    let mut read_only_device = device.inner;
+                    read_only_device.olm_wedging_index.increment();
+
+                    changes.devices =
+                        DeviceChanges { changed: vec![read_only_device], ..Default::default() };
+                }
+
+                store.save_changes(changes).await?;
 
                 Ok((SessionType::New(result.session), result.plaintext))
             }

crates/matrix-sdk-crypto/src/olm/group_sessions/outbound.rs

@@ -47,6 +47,7 @@ use super::SessionCreationError;
 use crate::types::events::room::encrypted::MegolmV2AesSha2Content;
 use crate::{
     session_manager::CollectStrategy,
+    store::caches::SequenceNumber,
     types::{
         events::{
             room::encrypted::{
@@ -67,10 +68,19 @@ const ROTATION_PERIOD: Duration = ONE_WEEK;
 const ROTATION_MESSAGES: u64 = 100;
 
 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
+/// Information about whether a session was shared with a device.
 pub(crate) enum ShareState {
+    /// The session was not shared with the device.
     NotShared,
+    /// The session was shared with the device with the given device ID, but
+    /// with a different curve25519 key.
     SharedButChangedSenderKey,
-    Shared(u32),
+    /// The session was shared with the device, at the given message index. The
+    /// `olm_wedging_index` is the value of the `olm_wedging_index` from the
+    /// `ReadOnlyDevice` at the time that we last shared the session with the
+    /// device, and indicates whether we need to re-share the session with the
+    /// device.
+    Shared { message_index: u32, olm_wedging_index: SequenceNumber },
 }
 
 /// Settings for an encrypted room.
@@ -169,8 +179,12 @@ pub enum ShareInfo {
 
 impl ShareInfo {
     /// Helper to create a SharedWith info
-    pub fn new_shared(sender_key: Curve25519PublicKey, message_index: u32) -> Self {
-        ShareInfo::Shared(SharedWith { sender_key, message_index })
+    pub fn new_shared(
+        sender_key: Curve25519PublicKey,
+        message_index: u32,
+        olm_wedging_index: SequenceNumber,
+    ) -> Self {
+        ShareInfo::Shared(SharedWith { sender_key, message_index, olm_wedging_index })
     }
 
     /// Helper to create a Withheld info
@@ -185,6 +199,9 @@ pub struct SharedWith {
     pub sender_key: Curve25519PublicKey,
     /// The message index that the device received.
     pub message_index: u32,
+    /// The Olm wedging index of the device at the time the session was shared.
+    #[serde(default)]
+    pub olm_wedging_index: SequenceNumber,
 }
 
 impl OutboundGroupSession {
@@ -527,7 +544,10 @@ impl OutboundGroupSession {
                 d.get(device.device_id()).map(|s| match s {
                     ShareInfo::Shared(s) => {
                         if device.curve25519_key() == Some(s.sender_key) {
-                            ShareState::Shared(s.message_index)
+                            ShareState::Shared {
+                                message_index: s.message_index,
+                                olm_wedging_index: s.olm_wedging_index,
+                            }
                         } else {
                             ShareState::SharedButChangedSenderKey
                         }
@@ -551,7 +571,10 @@ impl OutboundGroupSession {
                     Some(match info {
                         ShareInfo::Shared(info) => {
                             if device.curve25519_key() == Some(info.sender_key) {
-                                ShareState::Shared(info.message_index)
+                                ShareState::Shared {
+                                    message_index: info.message_index,
+                                    olm_wedging_index: info.olm_wedging_index,
+                                }
                             } else {
                                 ShareState::SharedButChangedSenderKey
                             }
@@ -598,12 +621,10 @@ impl OutboundGroupSession {
         sender_key: Curve25519PublicKey,
         index: u32,
     ) {
-        self.shared_with_set
-            .write()
-            .unwrap()
-            .entry(user_id.to_owned())
-            .or_default()
-            .insert(device_id.to_owned(), ShareInfo::new_shared(sender_key, index));
+        self.shared_with_set.write().unwrap().entry(user_id.to_owned()).or_default().insert(
+            device_id.to_owned(),
+            ShareInfo::new_shared(sender_key, index, Default::default()),
+        );
     }
 
     /// Mark the session as shared with the given user/device pair, starting
@@ -615,7 +636,8 @@ impl OutboundGroupSession {
         device_id: &DeviceId,
         sender_key: Curve25519PublicKey,
     ) {
-        let share_info = ShareInfo::new_shared(sender_key, self.message_index().await);
+        let share_info =
+            ShareInfo::new_shared(sender_key, self.message_index().await, Default::default());
         self.shared_with_set
             .write()
             .unwrap()