feat(crypto): Return the list of room keys we imported

Clients might want to retry decryption after they import room keys, for
this to be efficient they'll need to know which room keys got imported.

This patch extends the return value of the room key import result to
include the map of room keys that got imported.

Author: poljar

crates/matrix-sdk-crypto/src/file_encryption/key_export.rs

@@ -235,7 +235,10 @@ fn decrypt_helper(ciphertext: &str, passphrase: &str) -> Result<String, KeyExpor
 
 #[cfg(test)]
 mod test {
-    use std::io::Cursor;
+    use std::{
+        collections::{BTreeMap, BTreeSet},
+        io::Cursor,
+    };
 
     use indoc::indoc;
     use matrix_sdk_test::async_test;
@@ -313,7 +316,7 @@ mod test {
         assert_eq!(export, decrypted);
         assert_eq!(
             machine.import_keys(decrypted, false, |_, _| {}).await.unwrap(),
-            RoomKeyImportResult::new(0, 1)
+            RoomKeyImportResult::new(0, 1, BTreeMap::new())
         );
     }
 
@@ -325,29 +328,45 @@ mod test {
 
         let export = vec![session.export_at_index(10).await];
 
-        assert_eq!(
-            machine.import_keys(export.clone(), false, |_, _| {}).await?,
-            RoomKeyImportResult::new(1, 1)
+        let keys = RoomKeyImportResult::new(
+            1,
+            1,
+            BTreeMap::from([(
+                session.room_id().to_owned(),
+                BTreeMap::from([(
+                    session.sender_key().to_owned(),
+                    BTreeSet::from([session.session_id().to_owned()]),
+                )]),
+            )]),
         );
+
+        assert_eq!(machine.import_keys(export.clone(), false, |_, _| {}).await?, keys,);
         assert_eq!(
             machine.import_keys(export, false, |_, _| {}).await?,
-            RoomKeyImportResult::new(0, 1)
+            RoomKeyImportResult::new(0, 1, BTreeMap::new())
         );
 
         let better_export = vec![session.export().await];
 
-        assert_eq!(
-            machine.import_keys(better_export, false, |_, _| {}).await?,
-            RoomKeyImportResult::new(1, 1)
-        );
+        assert_eq!(machine.import_keys(better_export, false, |_, _| {}).await?, keys,);
 
         let another_session = machine.create_inbound_session(&room_id).await?;
         let export = vec![another_session.export_at_index(10).await];
-        assert_eq!(
-            machine.import_keys(export, false, |_, _| {}).await?,
-            RoomKeyImportResult::new(1, 1)
+
+        let keys = RoomKeyImportResult::new(
+            1,
+            1,
+            BTreeMap::from([(
+                another_session.room_id().to_owned(),
+                BTreeMap::from([(
+                    another_session.sender_key().to_owned(),
+                    BTreeSet::from([another_session.session_id().to_owned()]),
+                )]),
+            )]),
         );
 
+        assert_eq!(machine.import_keys(export, false, |_, _| {}).await?, keys,);
+
         Ok(())
     }
 

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

@@ -40,18 +40,31 @@ pub mod store;
 mod utilities;
 mod verification;
 
+use std::collections::{BTreeMap, BTreeSet};
+
+use ruma::RoomId;
+
 /// Return type for the room key importing.
 #[derive(Debug, Clone, PartialEq)]
 pub struct RoomKeyImportResult {
     /// The number of room keys that were imported.
     pub imported_count: usize,
     /// The total number of room keys that were found in the export.
     pub total_count: usize,
+    /// The map of keys that were imported.
+    ///
+    /// It's a map from room id to a map of the sender key to a set of session
+    /// ids.
+    pub keys: BTreeMap<RoomId, BTreeMap<String, BTreeSet<String>>>,
 }
 
 impl RoomKeyImportResult {
-    pub(crate) fn new(imported_count: usize, total_count: usize) -> Self {
-        Self { imported_count, total_count }
+    pub(crate) fn new(
+        imported_count: usize,
+        total_count: usize,
+        keys: BTreeMap<RoomId, BTreeMap<String, BTreeSet<String>>>,
+    ) -> Self {
+        Self { imported_count, total_count, keys }
     }
 }
 

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

@@ -15,7 +15,7 @@
 #[cfg(feature = "sled_cryptostore")]
 use std::path::Path;
 use std::{
-    collections::{BTreeMap, HashSet},
+    collections::{BTreeMap, BTreeSet, HashSet},
     mem,
     sync::Arc,
 };
@@ -1359,6 +1359,7 @@ impl OlmMachine {
         };
 
         let total_count = exported_keys.len();
+        let mut keys = BTreeMap::new();
 
         for (i, key) in exported_keys.into_iter().enumerate() {
             let session = InboundGroupSession::from_export(key)?;
@@ -1372,6 +1373,12 @@ impl OlmMachine {
                     session.mark_as_backed_up()
                 }
 
+                keys.entry(session.room_id().to_owned())
+                    .or_insert_with(BTreeMap::new)
+                    .entry(session.sender_key().to_owned())
+                    .or_insert_with(BTreeSet::new)
+                    .insert(session.session_id().to_owned());
+
                 sessions.push(session)
             }
 
@@ -1384,9 +1391,9 @@ impl OlmMachine {
 
         self.store.save_changes(changes).await?;
 
-        info!(total_count, imported_count, "Successfully imported room keys");
+        info!(total_count, imported_count, room_keys =? keys, "Successfully imported room keys");
 
-        Ok(RoomKeyImportResult::new(imported_count, total_count))
+        Ok(RoomKeyImportResult::new(imported_count, total_count, keys))
     }
 
     /// Export the keys that match the given predicate.