crypto: support for building key bundles

richvdh · richvdh · commit 2dea4d13aa76 · 2025-03-10T13:20:05.000Z
Add a method to CryptoStore which will construct a key bundle, ready for encrypting and sharing with invited users. Part of #4504
diff --git a/crates/matrix-sdk-crypto/src/store/mod.rs b/crates/matrix-sdk-crypto/src/store/mod.rs
@@ -53,7 +53,7 @@ use futures_util::StreamExt;
 use matrix_sdk_common::locks::RwLock as StdRwLock;
 use ruma::{
     encryption::KeyUsage, events::secret::request::SecretName, DeviceId, OwnedDeviceId,
-    OwnedRoomId, OwnedUserId, UserId,
+    OwnedRoomId, OwnedUserId, RoomId, UserId,
 };
 use serde::{de::DeserializeOwned, Deserialize, Serialize};
 use thiserror::Error;
@@ -94,10 +94,15 @@ pub mod integration_tests;
 use caches::{SequenceNumber, UsersForKeyQuery};
 pub(crate) use crypto_store_wrapper::CryptoStoreWrapper;
 pub use error::{CryptoStoreError, Result};
-use matrix_sdk_common::{store_locks::CrossProcessStoreLock, timeout::timeout};
+use matrix_sdk_common::{
+    deserialized_responses::WithheldCode, store_locks::CrossProcessStoreLock, timeout::timeout,
+};
 pub use memorystore::MemoryStore;
 pub use traits::{CryptoStore, DynCryptoStore, IntoCryptoStore};
 
+use crate::types::{
+    events::room_key_withheld::RoomKeyWithheldContent, room_history::RoomKeyBundle,
+};
 pub use crate::{
     dehydrated_devices::DehydrationError,
     gossiping::{GossipRequest, SecretInfo},
@@ -1987,6 +1992,36 @@ impl Store {
         Ok(futures_util::stream::iter(sessions.into_iter().filter(predicate))
             .then(|session| async move { session.export().await }))
     }
+
+    /// Assemble a room key bundle for sharing encrypted history, as per
+    /// MSC4268.
+    pub async fn build_room_key_bundle(
+        &self,
+        room_id: &RoomId,
+    ) -> std::result::Result<RoomKeyBundle, CryptoStoreError> {
+        // TODO: make this WAY more efficient. We should only fetch sessions for the
+        //   correct room.
+        let mut sessions = self.get_inbound_group_sessions().await?;
+        sessions.retain(|session| session.room_id == room_id);
+
+        let mut bundle = RoomKeyBundle::default();
+        for session in sessions {
+            if session.shared_history() {
+                bundle.room_keys.push(session.export().await.into());
+            } else {
+                bundle.withheld.push(RoomKeyWithheldContent::new(
+                    session.algorithm().to_owned(),
+                    WithheldCode::Unauthorised,
+                    session.room_id().to_owned(),
+                    session.session_id().to_owned(),
+                    session.sender_key().to_owned(),
+                    self.device_id().to_owned(),
+                ));
+            }
+        }
+
+        Ok(bundle)
+    }
 }
 
 impl Deref for Store {
@@ -2021,12 +2056,17 @@ mod tests {
     use std::pin::pin;
 
     use futures_util::StreamExt;
+    use insta::assert_json_snapshot;
     use matrix_sdk_test::async_test;
-    use ruma::{room_id, user_id};
+    use ruma::{device_id, room_id, user_id, DeviceId, RoomId, UserId};
+    use vodozemac::megolm::SessionKey;
 
     use crate::{
-        machine::test_helpers::get_machine_pair, store::DehydratedDeviceKey,
+        machine::test_helpers::get_machine_pair,
+        olm::{InboundGroupSession, SenderData},
+        store::{DehydratedDeviceKey, MemoryStore},
         types::EventEncryptionAlgorithm,
+        OlmMachine,
     };
 
     #[async_test]
@@ -2190,4 +2230,148 @@ mod tests {
 
         assert!(pickle_key.is_err());
     }
+
+    #[async_test]
+    async fn test_build_room_key_bundle() {
+        // Given: Alice has sent a number of room keys to Bob, including some in the
+        // wrong room, and some that are not marked as shared...
+        let alice = olm_machine_from_account_pickle(
+            user_id!("@a:s.co"),
+            device_id!("ALICE"),
+            olm_pickle_data1(),
+        )
+        .await;
+
+        let bob = olm_machine_from_account_pickle(
+            user_id!("@b:s.co"),
+            device_id!("BOB"),
+            olm_pickle_data2(),
+        )
+        .await;
+
+        let room1_id = room_id!("!room1:localhost");
+        let room2_id = room_id!("!room2:localhost");
+
+        /* base64-encoded Megolm session keys, created with:
+
+           println!(
+               "{}",
+               vodozemac::megolm::GroupSession::new(Default::default()).session_key().to_base64()
+           );
+        */
+
+        let session_key1 = "AgAAAAC2XHVzsMBKs4QCRElJ92CJKyGtknCSC8HY7cQ7UYwndMKLQAejXLh5UA0l6s736mgctcUMNvELScUWrObdflrHo+vth/gWreXOaCnaSxmyjjKErQwyIYTkUfqbHy40RJfEesLwnN23on9XAkch/iy8R2+Jz7B8zfG01f2Ow2SxPQFnAndcO1ZSD2GmXgedy6n4B20MWI1jGP2wiexOWbFSya8DO/VxC9m5+/mF+WwYqdpKn9g4Y05Yw4uz7cdjTc3rXm7xK+8E7hI//5QD1nHPvuKYbjjM9u2JSL+Bzp61Cw";
+        let session_key2 = "AgAAAAC1BXreFTUQQSBGekTEuYxhdytRKyv4JgDGcG+VOBYdPNGgs807SdibCGJky4lJ3I+7ZDGHoUzZPZP/4ogGu4kxni0PWdtWuN7+5zsuamgoFF/BkaGeUUGv6kgIkx8pyPpM5SASTUEP9bN2loDSpUPYwfiIqz74DgC4WQ4435sTBctYvKz8n+TDJwdLXpyT6zKljuqADAioud+s/iqx9LYn9HpbBfezZcvbg67GtE113pLrvde3IcPI5s6dNHK2onGO2B2eoaobcen18bbEDnlUGPeIivArLya7Da6us14jBQ";
+        let session_key3 = "AgAAAAAM9KFsliaUUhGSXgwOzM5UemjkNH4n8NHgvC/y8hhw13zTF+ooGD4uIYEXYX630oNvQm/EvgZo+dkoc0re+vsqsx4sQeNODdSjcBsWOa0oDF+irQn9oYoLUDPI1IBtY1rX+FV99Zm/xnG7uFOX7aTVlko2GSdejy1w9mfobmfxu5aUc04A9zaKJP1pOthZvRAlhpymGYHgsDtWPrrjyc/yypMflE4kIUEEEtu1kT6mrAmcl615XYRAHYK9G2+fZsGvokwzbkl4nulGwcZMpQEoM0nD2o3GWgX81HW3nGfKBg";
+        let session_key4 = "AgAAAAA4Kkesxq2h4v9PLD6Sm3Smxspz1PXTqytQPCMQMkkrHNmzV2bHlJ+6/Al9cu8vh1Oj69AK0WUAeJOJuaiskEeg/PI3P03+UYLeC379RzgqwSHdBgdQ41G2vD6zpgmE/8vYToe+qpCZACtPOswZxyqxHH+T/Iq0nv13JmlFGIeA6fEPfr5Y28B49viG74Fs9rxV9EH5PfjbuPM/p+Sz5obShuaBPKQBX1jT913nEXPoIJ06exNZGr0285nw/LgVvNlmWmbqNnbzO2cNZjQWA+xZYz5FSfyCxwqEBbEdUCuRCQ";
+        let session1 = create_inbound_group_session_with_visibility(
+            &alice,
+            &room1_id,
+            &SessionKey::from_base64(session_key1).unwrap(),
+            true,
+        )
+        .await;
+        let session2 = create_inbound_group_session_with_visibility(
+            &alice,
+            &room1_id,
+            &SessionKey::from_base64(session_key2).unwrap(),
+            true,
+        )
+        .await;
+
+        let unshared_session = create_inbound_group_session_with_visibility(
+            &alice,
+            &room1_id,
+            &SessionKey::from_base64(session_key3).unwrap(),
+            false,
+        )
+        .await;
+        let room2_session = create_inbound_group_session_with_visibility(
+            &alice,
+            &room2_id,
+            &SessionKey::from_base64(session_key4).unwrap(),
+            true,
+        )
+        .await;
+
+        bob.store()
+            .save_inbound_group_sessions(&[
+                session1.clone(),
+                session2.clone(),
+                unshared_session,
+                room2_session,
+            ])
+            .await
+            .unwrap();
+
+        // When I build the bundle
+        let mut bundle = bob.store().build_room_key_bundle(&room1_id).await.unwrap();
+
+        // Then the bundle matches the snapshot.
+        // We sort the sessions in the bundle, so that the snapshot is stable.
+        bundle.room_keys.sort_by_key(|session| session.session_id.clone());
+        assert_json_snapshot!(bundle);
+    }
+
+    /// Create a test [`OlmMachine`], backed by an in-memory store, based on the
+    /// given pickle data.
+    async fn olm_machine_from_account_pickle(
+        user_id: &UserId,
+        device_id: &DeviceId,
+        pickle: vodozemac::olm::AccountPickle,
+    ) -> OlmMachine {
+        OlmMachine::with_store(user_id, device_id, MemoryStore::new(), Some(pickle.into()))
+            .await
+            .unwrap()
+    }
+
+    /// A hardcoded set of device keys, suitable for creating a test olm machine
+    /// with.
+    fn olm_pickle_data1() -> vodozemac::olm::AccountPickle {
+        let alice_account_pickle = serde_json::json!({
+            "signing_key": {"Normal": b"alicesigningkey12345678901234567"},
+            "diffie_hellman_key": b"alicediffiehelmankey123456789012",
+            "one_time_keys":{"next_key_id":0,"public_keys":{},"private_keys":{}},
+            "fallback_keys":{"key_id":0,"fallback_key":null,"previous_fallback_key":null}
+        });
+
+        serde_json::from_value(alice_account_pickle).unwrap()
+    }
+
+    /// Another set hardcoded set of device keys, suitable for creating a test
+    /// olm machine with.
+    fn olm_pickle_data2() -> vodozemac::olm::AccountPickle {
+        let alice_account_pickle = serde_json::json!({
+            "signing_key": {"Normal": b"alicesigningkey12345678901234567"},
+            "diffie_hellman_key": b"alicediffiehelmankey123456789012",
+            "one_time_keys":{"next_key_id":0,"public_keys":{},"private_keys":{}},
+            "fallback_keys":{"key_id":0,"fallback_key":null,"previous_fallback_key":null}
+        });
+
+        serde_json::from_value(alice_account_pickle).unwrap()
+    }
+
+    async fn create_inbound_group_session_with_visibility(
+        olm_machine: &OlmMachine,
+        room_id: &RoomId,
+        session_key: &SessionKey,
+        shared_history: bool,
+    ) -> InboundGroupSession {
+        let static_account = olm_machine.store().static_account();
+        let sender_key = static_account.identity_keys.curve25519;
+        let signing_key = static_account.identity_keys.ed25519;
+
+        let session = InboundGroupSession::new(
+            sender_key,
+            signing_key,
+            room_id,
+            session_key,
+            SenderData::unknown(),
+            EventEncryptionAlgorithm::MegolmV1AesSha2,
+            None,
+            shared_history,
+        )
+        .unwrap();
+        session
+    }
 }
diff --git a/crates/matrix-sdk-crypto/src/store/snapshots/matrix_sdk_crypto__store__tests__build_room_key_bundle.snap b/crates/matrix-sdk-crypto/src/store/snapshots/matrix_sdk_crypto__store__tests__build_room_key_bundle.snap
@@ -0,0 +1,39 @@
+---
+source: crates/matrix-sdk-crypto/src/store/mod.rs
+expression: bundle
+---
+{
+  "room_keys": [
+    {
+      "algorithm": "m.megolm.v1.aes-sha2",
+      "room_id": "!room1:localhost",
+      "sender_key": "oW3g1uEcvDKEPNpyIfbAj2zz4U+qgO2jIWHgmNxAlyQ",
+      "session_id": "AWcCd1w7VlIPYaZeB53LqfgHbQxYjWMY/bCJ7E5ZsVI",
+      "session_key": "AQAAAAC2XHVzsMBKs4QCRElJ92CJKyGtknCSC8HY7cQ7UYwndMKLQAejXLh5UA0l6s736mgctcUMNvELScUWrObdflrHo+vth/gWreXOaCnaSxmyjjKErQwyIYTkUfqbHy40RJfEesLwnN23on9XAkch/iy8R2+Jz7B8zfG01f2Ow2SxPQFnAndcO1ZSD2GmXgedy6n4B20MWI1jGP2wiexOWbFS",
+      "sender_claimed_keys": {
+        "ed25519": "o45ljUqAb374+5wsaAcX2XBrSZQt0Sr2XnxnJvaBWEc"
+      }
+    },
+    {
+      "algorithm": "m.megolm.v1.aes-sha2",
+      "room_id": "!room1:localhost",
+      "sender_key": "oW3g1uEcvDKEPNpyIfbAj2zz4U+qgO2jIWHgmNxAlyQ",
+      "session_id": "y1i8rPyf5MMnB0tenJPrMqWO6oAMCKi536z+KrH0tic",
+      "session_key": "AQAAAAC1BXreFTUQQSBGekTEuYxhdytRKyv4JgDGcG+VOBYdPNGgs807SdibCGJky4lJ3I+7ZDGHoUzZPZP/4ogGu4kxni0PWdtWuN7+5zsuamgoFF/BkaGeUUGv6kgIkx8pyPpM5SASTUEP9bN2loDSpUPYwfiIqz74DgC4WQ4435sTBctYvKz8n+TDJwdLXpyT6zKljuqADAioud+s/iqx9LYn",
+      "sender_claimed_keys": {
+        "ed25519": "o45ljUqAb374+5wsaAcX2XBrSZQt0Sr2XnxnJvaBWEc"
+      }
+    }
+  ],
+  "withheld": [
+    {
+      "algorithm": "m.megolm.v1.aes-sha2",
+      "reason": "You are not authorised to read the message.",
+      "code": "m.unauthorised",
+      "from_device": "BOB",
+      "room_id": "!room1:localhost",
+      "sender_key": "oW3g1uEcvDKEPNpyIfbAj2zz4U+qgO2jIWHgmNxAlyQ",
+      "session_id": "lpRzTgD3Nook/Wk62Fm9ECWGnKYZgeCwO1Y+uuPJz/I"
+    }
+  ]
+}
diff --git a/crates/matrix-sdk-crypto/src/types/mod.rs b/crates/matrix-sdk-crypto/src/types/mod.rs
@@ -49,6 +49,7 @@ pub mod events;
 mod one_time_keys;
 pub mod qr_login;
 pub mod requests;
+pub mod room_history;
 
 pub use self::{backup::*, cross_signing::*, device_keys::*, one_time_keys::*};
 use crate::store::BackupDecryptionKey;
diff --git a/crates/matrix-sdk-crypto/src/types/room_history.rs b/crates/matrix-sdk-crypto/src/types/room_history.rs
@@ -0,0 +1,102 @@
+/*
+Copyright 2025 The Matrix.org Foundation C.I.C.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+//! Types for sharing encrypted room history, per MSC4268
+
+use std::fmt::Debug;
+
+use ruma::{DeviceKeyAlgorithm, OwnedRoomId};
+use serde::{Deserialize, Serialize};
+use vodozemac::{megolm::ExportedSessionKey, Curve25519PublicKey};
+
+#[cfg(doc)]
+use crate::olm::InboundGroupSession;
+use crate::{
+    olm::ExportedRoomKey,
+    types::{
+        deserialize_curve_key, events::room_key_withheld::RoomKeyWithheldContent,
+        serialize_curve_key, EventEncryptionAlgorithm, SigningKeys,
+    },
+};
+
+/// A bundle of room keys, for sharing encrypted room history.
+#[derive(Deserialize, Serialize, Debug, Default)]
+pub struct RoomKeyBundle {
+    /// Keys that we are sharing with the recipient.
+    pub room_keys: Vec<SharedRoomKey>,
+
+    /// Keys that we are *not* sharing with the recipient.
+    pub withheld: Vec<RoomKeyWithheldContent>,
+}
+
+/// An [`InboundGroupSession`] for sharing as part of the key bundle.
+#[derive(Deserialize, Serialize)]
+pub struct SharedRoomKey {
+    /// The encryption algorithm that the session uses.
+    pub algorithm: EventEncryptionAlgorithm,
+
+    /// The room where the session is used.
+    pub room_id: OwnedRoomId,
+
+    /// The Curve25519 key of the device which initiated the session originally.
+    #[serde(deserialize_with = "deserialize_curve_key", serialize_with = "serialize_curve_key")]
+    pub sender_key: Curve25519PublicKey,
+
+    /// The ID of the session that the key is for.
+    pub session_id: String,
+
+    /// The key for the session.
+    pub session_key: ExportedSessionKey,
+
+    /// The Ed25519 key of the device which initiated the session originally.
+    #[serde(default)]
+    pub sender_claimed_keys: SigningKeys<DeviceKeyAlgorithm>,
+}
+
+impl Debug for SharedRoomKey {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("SharedRoomKey")
+            .field("algorithm", &self.algorithm)
+            .field("room_id", &self.room_id)
+            .field("sender_key", &self.sender_key)
+            .field("session_id", &self.session_id)
+            .field("sender_claimed_keys", &self.sender_claimed_keys)
+            .finish_non_exhaustive()
+    }
+}
+
+impl From<ExportedRoomKey> for SharedRoomKey {
+    fn from(exported_room_key: ExportedRoomKey) -> Self {
+        let ExportedRoomKey {
+            algorithm,
+            room_id,
+            sender_key,
+            session_id,
+            session_key,
+            sender_claimed_keys,
+            shared_history: _,
+            forwarding_curve25519_key_chain: _,
+        } = exported_room_key;
+        SharedRoomKey {
+            algorithm,
+            room_id,
+            sender_key,
+            session_id,
+            session_key,
+            sender_claimed_keys,
+        }
+    }
+}
