Merge pull request matrix-org#4932 from matrix-org/rav/history_sharing/save_key_bundle_data

crypto: store received room key bundle data information

Add hooks to the memory store and sqlite store to stash the information about room key data.

Author: richvdh

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

@@ -75,11 +75,11 @@ use crate::{
     store::{
         Changes, CryptoStoreWrapper, DeviceChanges, IdentityChanges, IntoCryptoStore, MemoryStore,
         PendingChanges, Result as StoreResult, RoomKeyInfo, RoomSettings, SecretImportError, Store,
-        StoreCache, StoreTransaction,
+        StoreCache, StoreTransaction, StoredRoomKeyBundleData,
     },
     types::{
         events::{
-            olm_v1::{AnyDecryptedOlmEvent, DecryptedRoomKeyEvent},
+            olm_v1::{AnyDecryptedOlmEvent, DecryptedRoomKeyBundleEvent, DecryptedRoomKeyEvent},
             room::encrypted::{
                 EncryptedEvent, EncryptedToDeviceEvent, RoomEncryptedEventContent,
                 RoomEventEncryptionScheme, SupportedEventEncryptionSchemes,
@@ -939,6 +939,26 @@ impl OlmMachine {
         }
     }
 
+    #[instrument()]
+    async fn receive_room_key_bundle(
+        &self,
+        sender_key: Curve25519PublicKey,
+        event: &DecryptedRoomKeyBundleEvent,
+        changes: &mut Changes,
+    ) -> OlmResult<()> {
+        let Some(sender_device_keys) = &event.sender_device_keys else {
+            warn!("Received a room key bundle with no sender device keys: ignoring");
+            return Ok(());
+        };
+
+        changes.received_room_key_bundles.push(StoredRoomKeyBundleData {
+            sender_user: event.sender.clone(),
+            sender_data: SenderData::device_info(sender_device_keys.clone()),
+            bundle_data: event.content.clone(),
+        });
+        Ok(())
+    }
+
     fn add_withheld_info(&self, changes: &mut Changes, event: &RoomKeyWithheldEvent) {
         debug!(?event.content, "Processing `m.room_key.withheld` event");
 
@@ -1184,6 +1204,10 @@ impl OlmMachine {
             AnyDecryptedOlmEvent::Dummy(_) => {
                 debug!("Received an `m.dummy` event");
             }
+            AnyDecryptedOlmEvent::RoomKeyBundle(e) => {
+                debug!("Received a room key bundle event {:?}", e);
+                self.receive_room_key_bundle(decrypted.result.sender_key, e, changes).await?;
+            }
             AnyDecryptedOlmEvent::Custom(_) => {
                 warn!("Received an unexpected encrypted to-device event");
             }

crates/matrix-sdk-crypto/src/machine/tests/mod.rs

@@ -33,7 +33,7 @@ use ruma::{
             AddMentions, MessageType, Relation, ReplyWithinThread, RoomMessageEventContent,
         },
         AnyMessageLikeEvent, AnyMessageLikeEventContent, AnyToDeviceEvent, MessageLikeEvent,
-        OriginalMessageLikeEvent,
+        OriginalMessageLikeEvent, ToDeviceEventType,
     },
     room_id,
     serde::Raw,
@@ -115,6 +115,7 @@ pub fn to_device_requests_to_content(
     requests: Vec<Arc<ToDeviceRequest>>,
 ) -> ToDeviceEncryptedEventContent {
     let to_device_request = &requests[0];
+    assert_eq!(to_device_request.event_type, ToDeviceEventType::RoomEncrypted);
 
     to_device_request
         .messages

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

@@ -52,7 +52,7 @@ macro_rules! cryptostore_integration_tests {
                 },
                 store::{
                     BackupDecryptionKey, Changes, CryptoStore, DehydratedDeviceKey, DeviceChanges, GossipRequest,
-                    IdentityChanges, PendingChanges, RoomSettings,
+                    IdentityChanges, PendingChanges, RoomSettings, StoredRoomKeyBundleData,
                 },
                 testing::{get_device, get_other_identity, get_own_identity},
                 types::{
@@ -64,6 +64,7 @@ macro_rules! cryptostore_integration_tests {
                             CommonWithheldCodeContent, MegolmV1AesSha2WithheldContent,
                             RoomKeyWithheldContent,
                         },
+                        room_key_bundle::RoomKeyBundleContent,
                         secret_send::SecretSendContent,
                         ToDeviceEvent,
                     },
@@ -1276,6 +1277,57 @@ macro_rules! cryptostore_integration_tests {
                 assert_eq!(None, loaded_2);
             }
 
+            #[async_test]
+            #[ignore] // not yet implemented for all stores
+            async fn test_received_room_key_bundle() {
+                let store = get_store("received_room_key_bundle", None, true).await;
+                let test_room = room_id!("!room:example.org");
+
+                fn make_bundle_data(sender_user: &UserId, bundle_uri: &str) -> StoredRoomKeyBundleData {
+                    let jwk = ruma::events::room::JsonWebKeyInit {
+                        kty: "oct".to_owned(),
+                        key_ops: vec!["encrypt".to_owned(), "decrypt".to_owned()],
+                        alg: "A256CTR".to_owned(),
+                        k: ruma::serde::Base64::new(vec![0u8; 0]),
+                        ext: true,
+                    }.into();
+
+                    let file = ruma::events::room::EncryptedFileInit {
+                        url: ruma::OwnedMxcUri::from(bundle_uri),
+                        key: jwk,
+                        iv: ruma::serde::Base64::new(vec![0u8; 0]),
+                        hashes: Default::default(),
+                        v: "".to_owned(),
+                    }.into();
+
+                    StoredRoomKeyBundleData {
+                        sender_user: sender_user.to_owned(),
+                        sender_data: SenderData::unknown(),
+                        bundle_data: RoomKeyBundleContent {
+                            room_id: room_id!("!room:example.org").to_owned(),
+                            file,
+                        },
+                    }
+                }
+
+                // Add three entries
+                let changes = Changes {
+                    received_room_key_bundles: vec![
+                        make_bundle_data(user_id!("@alice:example.com"), "alice1"),
+                        make_bundle_data(user_id!("@bob:example.com"), "bob1"),
+                        make_bundle_data(user_id!("@alice:example.com"), "alice2"),
+                    ],
+                    ..Default::default()
+                };
+                store.save_changes(changes).await.unwrap();
+
+                // Check we get the right one
+                let bundle = store.get_received_room_key_bundle_data(
+                    test_room, user_id!("@alice:example.com")
+                ).await.unwrap().expect("Did not get any bundle data");
+                assert_eq!(bundle.bundle_data.file.url.to_string(), "alice2");
+            }
+
             fn session_info(session: &InboundGroupSession) -> (&RoomId, &str) {
                 (&session.room_id(), &session.session_id())
             }

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

@@ -33,6 +33,7 @@ use vodozemac::Curve25519PublicKey;
 use super::{
     caches::DeviceStore, Account, BackupKeys, Changes, CryptoStore, DehydratedDeviceKey,
     InboundGroupSession, PendingChanges, RoomKeyCounts, RoomSettings, Session,
+    StoredRoomKeyBundleData,
 };
 use crate::{
     gossiping::{GossipRequest, GossippedSecret, SecretInfo},
@@ -71,7 +72,7 @@ impl BackupVersion {
 }
 
 /// An in-memory only store that will forget all the E2EE key once it's dropped.
-#[derive(Debug)]
+#[derive(Default, Debug)]
 pub struct MemoryStore {
     static_account: Arc<StdRwLock<Option<StaticAccountData>>>,
 
@@ -102,36 +103,10 @@ pub struct MemoryStore {
     dehydrated_device_pickle_key: RwLock<Option<DehydratedDeviceKey>>,
     next_batch_token: RwLock<Option<String>>,
     room_settings: StdRwLock<HashMap<OwnedRoomId, RoomSettings>>,
-    save_changes_lock: Arc<Mutex<()>>,
-}
+    room_key_bundles:
+        StdRwLock<HashMap<OwnedRoomId, HashMap<OwnedUserId, StoredRoomKeyBundleData>>>,
 
-impl Default for MemoryStore {
-    fn default() -> Self {
-        MemoryStore {
-            static_account: Default::default(),
-            account: Default::default(),
-            sessions: Default::default(),
-            inbound_group_sessions: Default::default(),
-            inbound_group_sessions_backed_up_to: Default::default(),
-            outbound_group_sessions: Default::default(),
-            private_identity: Default::default(),
-            tracked_users: Default::default(),
-            olm_hashes: Default::default(),
-            devices: DeviceStore::new(),
-            identities: Default::default(),
-            outgoing_key_requests: Default::default(),
-            key_requests_by_info: Default::default(),
-            direct_withheld_info: Default::default(),
-            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(),
-            save_changes_lock: Default::default(),
-        }
-    }
+    save_changes_lock: Arc<Mutex<()>>,
 }
 
 impl MemoryStore {
@@ -348,6 +323,16 @@ impl CryptoStore for MemoryStore {
             settings.extend(changes.room_settings);
         }
 
+        if !changes.received_room_key_bundles.is_empty() {
+            let mut room_key_bundles = self.room_key_bundles.write();
+            for bundle in changes.received_room_key_bundles {
+                room_key_bundles
+                    .entry(bundle.bundle_data.room_id.clone())
+                    .or_default()
+                    .insert(bundle.sender_user.clone(), bundle);
+            }
+        }
+
         Ok(())
     }
 
@@ -719,6 +704,18 @@ impl CryptoStore for MemoryStore {
         Ok(self.room_settings.read().get(room_id).cloned())
     }
 
+    async fn get_received_room_key_bundle_data(
+        &self,
+        room_id: &RoomId,
+        user_id: &UserId,
+    ) -> Result<Option<StoredRoomKeyBundleData>> {
+        let guard = self.room_key_bundles.read();
+
+        let result = guard.get(room_id).and_then(|bundles| bundles.get(user_id).cloned());
+
+        Ok(result)
+    }
+
     async fn get_custom_value(&self, key: &str) -> Result<Option<Vec<u8>>> {
         Ok(self.custom_values.read().get(key).cloned())
     }
@@ -1249,7 +1246,7 @@ mod integration_tests {
         },
         store::{
             BackupKeys, Changes, CryptoStore, DehydratedDeviceKey, PendingChanges, RoomKeyCounts,
-            RoomSettings,
+            RoomSettings, StoredRoomKeyBundleData,
         },
         types::events::room_key_withheld::RoomKeyWithheldEvent,
         Account, DeviceData, GossipRequest, GossippedSecret, SecretInfo, Session, TrackedUser,
@@ -1511,6 +1508,14 @@ mod integration_tests {
             self.0.get_room_settings(room_id).await
         }
 
+        async fn get_received_room_key_bundle_data(
+            &self,
+            room_id: &RoomId,
+            user_id: &UserId,
+        ) -> crate::store::Result<Option<StoredRoomKeyBundleData>, Self::Error> {
+            self.0.get_received_room_key_bundle_data(room_id, user_id).await
+        }
+
         async fn get_custom_value(&self, key: &str) -> Result<Option<Vec<u8>>, Self::Error> {
             self.0.get_custom_value(key).await
         }

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

@@ -70,7 +70,7 @@ use crate::{
     identities::{user::UserIdentity, Device, DeviceData, UserDevices, UserIdentityData},
     olm::{
         Account, ExportedRoomKey, InboundGroupSession, OlmMessageHash, OutboundGroupSession,
-        PrivateCrossSigningIdentity, Session, StaticAccountData,
+        PrivateCrossSigningIdentity, SenderData, Session, StaticAccountData,
     },
     types::{
         events::room_key_withheld::RoomKeyWithheldEvent, BackupSecrets, CrossSigningSecrets,
@@ -101,7 +101,8 @@ pub use memorystore::MemoryStore;
 pub use traits::{CryptoStore, DynCryptoStore, IntoCryptoStore};
 
 use crate::types::{
-    events::room_key_withheld::RoomKeyWithheldContent, room_history::RoomKeyBundle,
+    events::{room_key_bundle::RoomKeyBundleContent, room_key_withheld::RoomKeyWithheldContent},
+    room_history::RoomKeyBundle,
 };
 pub use crate::{
     dehydrated_devices::DehydrationError,
@@ -541,6 +542,26 @@ pub struct Changes {
     pub room_settings: HashMap<OwnedRoomId, RoomSettings>,
     pub secrets: Vec<GossippedSecret>,
     pub next_batch_token: Option<String>,
+
+    /// Historical room key history bundles that we have received and should
+    /// store.
+    pub received_room_key_bundles: Vec<StoredRoomKeyBundleData>,
+}
+
+/// Information about an [MSC4268] room key bundle.
+///
+/// [MSC4268]: https://github.com/matrix-org/matrix-spec-proposals/pull/4268
+#[derive(Clone, Debug, Serialize, Deserialize)]
+pub struct StoredRoomKeyBundleData {
+    /// The user that sent us this data.
+    pub sender_user: OwnedUserId,
+
+    /// Information about the sender of this data and how much we trust that
+    /// information.
+    pub sender_data: SenderData,
+
+    /// The room key bundle data itself.
+    pub bundle_data: RoomKeyBundleContent,
 }
 
 /// A user for which we are tracking the list of devices.
@@ -573,6 +594,7 @@ impl Changes {
             && self.room_settings.is_empty()
             && self.secrets.is_empty()
             && self.next_batch_token.is_none()
+            && self.received_room_key_bundles.is_empty()
     }
 }
 

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

@@ -23,7 +23,7 @@ use vodozemac::Curve25519PublicKey;
 
 use super::{
     BackupKeys, Changes, CryptoStoreError, DehydratedDeviceKey, PendingChanges, Result,
-    RoomKeyCounts, RoomSettings,
+    RoomKeyCounts, RoomSettings, StoredRoomKeyBundleData,
 };
 #[cfg(doc)]
 use crate::olm::SenderData;
@@ -323,6 +323,14 @@ pub trait CryptoStore: AsyncTraitDeps {
         room_id: &RoomId,
     ) -> Result<Option<RoomSettings>, Self::Error>;
 
+    /// Get the details about the room key bundle data received from the given
+    /// user for the given room.
+    async fn get_received_room_key_bundle_data(
+        &self,
+        room_id: &RoomId,
+        user_id: &UserId,
+    ) -> Result<Option<StoredRoomKeyBundleData>, Self::Error>;
+
     /// Get arbitrary data from the store
     ///
     /// # Arguments
@@ -569,6 +577,14 @@ impl<T: CryptoStore> CryptoStore for EraseCryptoStoreError<T> {
         self.0.get_room_settings(room_id).await.map_err(Into::into)
     }
 
+    async fn get_received_room_key_bundle_data(
+        &self,
+        room_id: &RoomId,
+        user_id: &UserId,
+    ) -> Result<Option<StoredRoomKeyBundleData>> {
+        self.0.get_received_room_key_bundle_data(room_id, user_id).await.map_err(Into::into)
+    }
+
     async fn get_custom_value(&self, key: &str) -> Result<Option<Vec<u8>>, Self::Error> {
         self.0.get_custom_value(key).await.map_err(Into::into)
     }