From bfc7434f7e1ea24728dd2c236b5316cd49c59542 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Damir=20Jeli=C4=87?= <poljar@termina.org.uk>
Date: Tue, 20 Apr 2021 13:35:47 +0200
Subject: [PATCH] crypto: Move the outbound session filter logic into the group
 session cache

---
 matrix_sdk_crypto/src/key_request.rs          |  5 ++---
 .../src/session_manager/group_sessions.rs     | 21 +++++++++++++++++++
 2 files changed, 23 insertions(+), 3 deletions(-)

diff --git a/matrix_sdk_crypto/src/key_request.rs b/matrix_sdk_crypto/src/key_request.rs
index bc4df8c1..1eec2981 100644
--- a/matrix_sdk_crypto/src/key_request.rs
+++ b/matrix_sdk_crypto/src/key_request.rs
@@ -516,11 +516,10 @@ impl KeyRequestMachine {
     ) -> Result<Option<u32>, KeyshareDecision> {
         let outbound_session = self
             .outbound_group_sessions
-            .get_or_load(session.room_id())
+            .get_with_id(session.room_id(), session.session_id())
             .await
             .ok()
-            .flatten()
-            .filter(|o| session.session_id() == o.session_id());
+            .flatten();
 
         let own_device_check = || {
             if device.trust_state() {
diff --git a/matrix_sdk_crypto/src/session_manager/group_sessions.rs b/matrix_sdk_crypto/src/session_manager/group_sessions.rs
index a656d00e..71645714 100644
--- a/matrix_sdk_crypto/src/session_manager/group_sessions.rs
+++ b/matrix_sdk_crypto/src/session_manager/group_sessions.rs
@@ -62,6 +62,12 @@ impl GroupSessionCache {
         self.sessions.insert(session.room_id().to_owned(), session);
     }
 
+    /// Either get a session for the given room from the cache or load it from
+    /// the store.
+    ///
+    /// # Arguments
+    ///
+    /// * `room_id` - The id of the room this session is used for.
     pub async fn get_or_load(&self, room_id: &RoomId) -> StoreResult<Option<OutboundGroupSession>> {
         // Get the cached session, if there isn't one load one from the store
         // and put it in the cache.
@@ -89,6 +95,21 @@ impl GroupSessionCache {
     fn get(&self, room_id: &RoomId) -> Option<OutboundGroupSession> {
         self.sessions.get(room_id).map(|s| s.clone())
     }
+
+    /// Get or load the session for the given room with the given session id.
+    ///
+    /// This is the same as [get_or_load()](#method.get_or_load) but it will
+    /// filter out the session if it doesn't match the given session id.
+    pub async fn get_with_id(
+        &self,
+        room_id: &RoomId,
+        session_id: &str,
+    ) -> StoreResult<Option<OutboundGroupSession>> {
+        Ok(self
+            .get_or_load(room_id)
+            .await?
+            .filter(|o| session_id == o.session_id()))
+    }
 }
 
 #[derive(Debug, Clone)]