From 9617d9aac91cfc41792ee27802400bd6db580f27 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Damir=20Jeli=C4=87?= <poljar@termina.org.uk>
Date: Wed, 9 Sep 2020 16:10:16 +0200
Subject: [PATCH] crypto: Test the import/export of group sessions.

---
 .../src/olm/group_sessions/inbound.rs         | 43 ++++++++++++++++++-
 matrix_sdk_crypto/src/olm/mod.rs              | 21 ++++++++-
 2 files changed, 62 insertions(+), 2 deletions(-)

diff --git a/matrix_sdk_crypto/src/olm/group_sessions/inbound.rs b/matrix_sdk_crypto/src/olm/group_sessions/inbound.rs
index ad9edf77..c0a295f5 100644
--- a/matrix_sdk_crypto/src/olm/group_sessions/inbound.rs
+++ b/matrix_sdk_crypto/src/olm/group_sessions/inbound.rs
@@ -95,6 +95,20 @@ impl InboundGroupSession {
         })
     }
 
+    /// Create a InboundGroupSession from an exported version of the group
+    /// session.
+    ///
+    /// Most notably this can be called with an `ExportedRoomKey` from a
+    /// previous [`export()`] call.
+    ///
+    ///
+    /// [`export()`]: #method.export
+    pub fn from_export(
+        exported_session: impl Into<ExportedRoomKey>,
+    ) -> Result<Self, OlmGroupSessionError> {
+        Self::try_from(exported_session.into())
+    }
+
     /// Store the group session as a base64 encoded string.
     ///
     /// # Arguments
@@ -114,7 +128,10 @@ impl InboundGroupSession {
         }
     }
 
-    /// Export this session.
+    /// Export this session at the first known message index.
+    ///
+    /// If only a limited part of this session should be exported use
+    /// [`export_at_index()`](#method.export_at_index).
     pub async fn export(&self) -> ExportedRoomKey {
         self.export_at_index(self.first_known_index().await)
             .await
@@ -295,3 +312,27 @@ impl InboundGroupSessionPickle {
         &self.0
     }
 }
+
+impl TryFrom<ExportedRoomKey> for InboundGroupSession {
+    type Error = OlmGroupSessionError;
+
+    fn try_from(key: ExportedRoomKey) -> Result<Self, Self::Error> {
+        let session = OlmInboundGroupSession::import(&key.session_key.0)?;
+
+        let forwarding_chains = if key.forwarding_curve25519_key_chain.is_empty() {
+            None
+        } else {
+            Some(key.forwarding_curve25519_key_chain)
+        };
+
+        Ok(InboundGroupSession {
+            inner: Arc::new(Mutex::new(session)),
+            session_id: Arc::new(key.session_id),
+            sender_key: Arc::new(key.sender_key),
+            signing_key: Arc::new(key.sender_claimed_keys),
+            room_id: Arc::new(key.room_id),
+            forwarding_chains: Arc::new(Mutex::new(forwarding_chains)),
+            imported: Arc::new(true),
+        })
+    }
+}
diff --git a/matrix_sdk_crypto/src/olm/mod.rs b/matrix_sdk_crypto/src/olm/mod.rs
index 05e6b874..e88b4959 100644
--- a/matrix_sdk_crypto/src/olm/mod.rs
+++ b/matrix_sdk_crypto/src/olm/mod.rs
@@ -38,10 +38,11 @@ pub(crate) mod test {
     use crate::olm::{Account, InboundGroupSession, Session};
     use matrix_sdk_common::{
         api::r0::keys::SignedKey,
+        events::forwarded_room_key::ForwardedRoomKeyEventContent,
         identifiers::{room_id, user_id, DeviceId, UserId},
     };
     use olm_rs::session::OlmMessage;
-    use std::collections::BTreeMap;
+    use std::{collections::BTreeMap, convert::TryInto};
 
     fn alice_id() -> UserId {
         user_id!("@alice:example.org")
@@ -222,4 +223,22 @@ pub(crate) mod test {
             inbound.decrypt_helper(ciphertext).await.unwrap().0
         );
     }
+
+    #[tokio::test]
+    async fn group_session_export() {
+        let alice = Account::new(&alice_id(), &alice_device_id());
+        let room_id = room_id!("!test:localhost");
+
+        let (_, inbound) = alice
+            .create_group_session_pair(&room_id, Default::default())
+            .await
+            .unwrap();
+
+        let export = inbound.export().await;
+        let export: ForwardedRoomKeyEventContent = export.try_into().unwrap();
+
+        let imported = InboundGroupSession::from_export(export).unwrap();
+
+        assert_eq!(inbound.session_id(), imported.session_id());
+    }
 }