From 2797256ec9be1f81ce26ae44bd8022e08603832d Mon Sep 17 00:00:00 2001
From: Andy Balaam <andy.balaam@matrix.org>
Date: Fri, 2 Aug 2024 11:17:37 +0100
Subject: [PATCH] crypto: Create a SenderDataType enum

---
 .../src/olm/group_sessions/inbound.rs         | 11 ++++++++--
 .../src/olm/group_sessions/mod.rs             |  2 +-
 .../src/olm/group_sessions/sender_data.rs     | 22 +++++++++++++++++++
 crates/matrix-sdk-crypto/src/olm/mod.rs       |  2 +-
 4 files changed, 33 insertions(+), 4 deletions(-)

diff --git a/crates/matrix-sdk-crypto/src/olm/group_sessions/inbound.rs b/crates/matrix-sdk-crypto/src/olm/group_sessions/inbound.rs
index c46747a2c31..51b774fe28c 100644
--- a/crates/matrix-sdk-crypto/src/olm/group_sessions/inbound.rs
+++ b/crates/matrix-sdk-crypto/src/olm/group_sessions/inbound.rs
@@ -36,8 +36,8 @@ use vodozemac::{
 };
 
 use super::{
-    BackedUpRoomKey, ExportedRoomKey, OutboundGroupSession, SenderData, SessionCreationError,
-    SessionKey,
+    BackedUpRoomKey, ExportedRoomKey, OutboundGroupSession, SenderData, SenderDataType,
+    SessionCreationError, SessionKey,
 };
 use crate::{
     error::{EventError, MegolmResult},
@@ -477,6 +477,13 @@ impl InboundGroupSession {
     pub(crate) fn mark_as_imported(&mut self) {
         self.imported = true;
     }
+
+    /// Return the [`SenderDataType`] of our [`SenderData`]. This is used during
+    /// serialization, to allow us to store the type in a separate queryable
+    /// column/property.
+    pub fn sender_data_type(&self) -> SenderDataType {
+        self.sender_data.to_type()
+    }
 }
 
 #[cfg(not(tarpaulin_include))]
diff --git a/crates/matrix-sdk-crypto/src/olm/group_sessions/mod.rs b/crates/matrix-sdk-crypto/src/olm/group_sessions/mod.rs
index 33da6296414..dce57b15a94 100644
--- a/crates/matrix-sdk-crypto/src/olm/group_sessions/mod.rs
+++ b/crates/matrix-sdk-crypto/src/olm/group_sessions/mod.rs
@@ -25,7 +25,7 @@ pub(crate) use outbound::ShareState;
 pub use outbound::{
     EncryptionSettings, OutboundGroupSession, PickledOutboundGroupSession, ShareInfo,
 };
-pub use sender_data::SenderData;
+pub use sender_data::{SenderData, SenderDataType};
 pub(crate) use sender_data_finder::SenderDataFinder;
 use thiserror::Error;
 pub use vodozemac::megolm::{ExportedSessionKey, SessionKey};
diff --git a/crates/matrix-sdk-crypto/src/olm/group_sessions/sender_data.rs b/crates/matrix-sdk-crypto/src/olm/group_sessions/sender_data.rs
index c59814b0d9b..d41b4e56242 100644
--- a/crates/matrix-sdk-crypto/src/olm/group_sessions/sender_data.rs
+++ b/crates/matrix-sdk-crypto/src/olm/group_sessions/sender_data.rs
@@ -155,6 +155,15 @@ impl SenderData {
             SenderData::SenderKnown { master_key_verified: true, .. } => 3,
         }
     }
+
+    /// Return our type: `UnknownDevice`, `DeviceInfo`, or `SenderKnown`.
+    pub fn to_type(&self) -> SenderDataType {
+        match self {
+            Self::UnknownDevice { .. } => SenderDataType::UnknownDevice,
+            Self::DeviceInfo { .. } => SenderDataType::DeviceInfo,
+            Self::SenderKnown { .. } => SenderDataType::SenderKnown,
+        }
+    }
 }
 
 /// Used when deserialising and the sender_data property is missing.
@@ -169,6 +178,19 @@ impl Default for SenderData {
     }
 }
 
+/// Used when serializing [`crate::olm::group_sessions::InboundGroupSession`]s.
+/// We want just the type of the session's [`SenderData`] to be queryable, so we
+/// store the type as a separate column/property in the database.
+#[derive(Clone, Copy, Debug, PartialEq, Deserialize, Serialize)]
+pub enum SenderDataType {
+    /// The [`SenderData`] is of type `UnknownDevice`.
+    UnknownDevice = 1,
+    /// The [`SenderData`] is of type `DeviceInfo`.
+    DeviceInfo = 2,
+    /// The [`SenderData`] is of type `SenderKnown`.
+    SenderKnown = 3,
+}
+
 #[cfg(test)]
 mod tests {
     use std::{cmp::Ordering, collections::BTreeMap};
diff --git a/crates/matrix-sdk-crypto/src/olm/mod.rs b/crates/matrix-sdk-crypto/src/olm/mod.rs
index 01bf6f2a1ff..c8fa21307b0 100644
--- a/crates/matrix-sdk-crypto/src/olm/mod.rs
+++ b/crates/matrix-sdk-crypto/src/olm/mod.rs
@@ -28,7 +28,7 @@ pub(crate) use account::{OlmDecryptionInfo, SessionType};
 pub use group_sessions::{
     BackedUpRoomKey, EncryptionSettings, ExportedRoomKey, InboundGroupSession,
     OutboundGroupSession, PickledInboundGroupSession, PickledOutboundGroupSession, SenderData,
-    SessionCreationError, SessionExportError, SessionKey, ShareInfo,
+    SenderDataType, SessionCreationError, SessionExportError, SessionKey, ShareInfo,
 };
 pub(crate) use group_sessions::{SenderDataFinder, ShareState};
 pub use session::{PickledSession, Session};