crypto: Save the olm message hash.

matrix_sdk_crypto/src/machine.rs

@@ -834,6 +834,7 @@ impl OlmMachine {
                     };
 
                     changes.sessions.push(decrypted.session);
+                    changes.message_hashes.push(decrypted.message_hash);
 
                     if let Some(group_session) = decrypted.inbound_group_session {
                         changes.inbound_group_sessions.push(group_session);

matrix_sdk_crypto/src/olm/account.rs

@@ -72,15 +72,27 @@ pub struct Account {
     pub(crate) store: Store,
 }
 
+#[derive(Debug, Clone)]
 pub struct OlmDecryptionInfo {
     pub session: Session,
-    pub message_hash: String,
+    pub message_hash: OlmMessageHash,
     pub event: Raw<AnyToDeviceEvent>,
     pub signing_key: String,
     pub sender_key: String,
     pub inbound_group_session: Option<InboundGroupSession>,
 }
 
+/// A hash of a succesfully decrypted Olm message.
+///
+/// Can be used to check if a message has been replayed to us.
+#[derive(Debug, Clone)]
+pub struct OlmMessageHash {
+    /// The curve25519 key of the sender that sent us the Olm message.
+    pub sender_key: String,
+    /// The hash of the message.
+    pub hash: String,
+}
+
 impl Deref for Account {
     type Target = ReadOnlyAccount;
 
@@ -119,7 +131,7 @@ impl Account {
                 .chain(&[message_type])
                 .chain(&ciphertext.body);
 
-            let message_hash = encode(sha.finalize().as_slice());
+            let hash = encode(sha.finalize().as_slice());
 
             // Create a OlmMessage from the ciphertext and the type.
             let message =
@@ -135,7 +147,10 @@ impl Account {
 
             Ok(OlmDecryptionInfo {
                 session,
-                message_hash,
+                message_hash: OlmMessageHash {
+                    hash,
+                    sender_key: content.sender_key.clone(),
+                },
                 event,
                 signing_key,
                 sender_key: content.sender_key.clone(),

matrix_sdk_crypto/src/olm/mod.rs

@@ -24,7 +24,7 @@ mod signing;
 mod utility;
 
 pub(crate) use account::{Account, OlmDecryptionInfo};
-pub use account::{AccountPickle, PickledAccount, ReadOnlyAccount};
+pub use account::{AccountPickle, OlmMessageHash, PickledAccount, ReadOnlyAccount};
 pub use group_sessions::{
     EncryptionSettings, ExportedRoomKey, InboundGroupSession, InboundGroupSessionPickle,
     PickledInboundGroupSession,

matrix_sdk_crypto/src/store/memorystore.rs

@@ -40,6 +40,7 @@ pub struct MemoryStore {
     inbound_group_sessions: GroupSessionStore,
     tracked_users: Arc<DashSet<UserId>>,
     users_for_key_query: Arc<DashSet<UserId>>,
+    olm_hashes: Arc<DashMap<String, DashSet<String>>>,
     devices: DeviceStore,
     identities: Arc<DashMap<UserId, UserIdentities>>,
     values: Arc<DashMap<String, String>>,
@@ -52,6 +53,7 @@ impl Default for MemoryStore {
             inbound_group_sessions: GroupSessionStore::new(),
             tracked_users: Arc::new(DashSet::new()),
             users_for_key_query: Arc::new(DashSet::new()),
+            olm_hashes: Arc::new(DashMap::new()),
             devices: DeviceStore::new(),
             identities: Arc::new(DashMap::new()),
             values: Arc::new(DashMap::new()),
@@ -120,6 +122,13 @@ impl CryptoStore for MemoryStore {
                 .insert(identity.user_id().to_owned(), identity.clone());
         }
 
+        for hash in changes.message_hashes {
+            self.olm_hashes
+                .entry(hash.sender_key.to_owned())
+                .or_insert_with(DashSet::new)
+                .insert(hash.hash.clone());
+        }
+
         Ok(())
     }
 
@@ -218,14 +227,22 @@ impl CryptoStore for MemoryStore {
     async fn load_identity(&self) -> Result<Option<PrivateCrossSigningIdentity>> {
         Ok(None)
     }
+
+    async fn is_message_known(&self, message_hash: &crate::olm::OlmMessageHash) -> Result<bool> {
+        Ok(self
+            .olm_hashes
+            .entry(message_hash.sender_key.to_owned())
+            .or_insert_with(DashSet::new)
+            .contains(&message_hash.hash))
+    }
 }
 
 #[cfg(test)]
 mod test {
     use crate::{
         identities::device::test::get_device,
-        olm::{test::get_account_and_session, InboundGroupSession},
-        store::{memorystore::MemoryStore, CryptoStore},
+        olm::{test::get_account_and_session, InboundGroupSession, OlmMessageHash},
+        store::{memorystore::MemoryStore, Changes, CryptoStore},
     };
     use matrix_sdk_common::identifiers::room_id;
 
@@ -329,4 +346,21 @@ mod test {
 
         assert!(store.is_user_tracked(device.user_id()));
     }
+
+    #[tokio::test]
+    async fn test_message_hash() {
+        let store = MemoryStore::new();
+
+        let hash = OlmMessageHash {
+            sender_key: "test_sender".to_owned(),
+            hash: "test_hash".to_owned(),
+        };
+
+        let mut changes = Changes::default();
+        changes.message_hashes.push(hash.clone());
+
+        assert!(!store.is_message_known(&hash).await.unwrap());
+        store.save_changes(changes).await.unwrap();
+        assert!(store.is_message_known(&hash).await.unwrap());
+    }
 }

matrix_sdk_crypto/src/store/mod.rs

@@ -82,7 +82,9 @@ use matrix_sdk_common_macros::send_sync;
 use crate::{
     error::SessionUnpicklingError,
     identities::{Device, ReadOnlyDevice, UserDevices, UserIdentities},
-    olm::{InboundGroupSession, PrivateCrossSigningIdentity, ReadOnlyAccount, Session},
+    olm::{
+        InboundGroupSession, OlmMessageHash, PrivateCrossSigningIdentity, ReadOnlyAccount, Session,
+    },
     verification::VerificationMachine,
 };
 
@@ -108,6 +110,7 @@ pub(crate) struct Store {
 pub struct Changes {
     pub account: Option<ReadOnlyAccount>,
     pub sessions: Vec<Session>,
+    pub message_hashes: Vec<OlmMessageHash>,
     pub inbound_group_sessions: Vec<InboundGroupSession>,
     pub identities: IdentityChanges,
     pub devices: DeviceChanges,
@@ -444,4 +447,7 @@ pub trait CryptoStore: Debug {
 
     /// Load a serializeable object from the store.
     async fn get_value(&self, key: &str) -> Result<Option<String>>;
+
+    /// Check if a hash for an Olm message stored in the database.
+    async fn is_message_known(&self, message_hash: &OlmMessageHash) -> Result<bool>;
 }

matrix_sdk_crypto/src/store/sqlite.rs

@@ -42,8 +42,9 @@ use crate::{
     identities::{LocalTrust, OwnUserIdentity, ReadOnlyDevice, UserIdentities, UserIdentity},
     olm::{
         AccountPickle, IdentityKeys, InboundGroupSession, InboundGroupSessionPickle,
-        PickledAccount, PickledCrossSigningIdentity, PickledInboundGroupSession, PickledSession,
-        PicklingMode, PrivateCrossSigningIdentity, ReadOnlyAccount, Session, SessionPickle,
+        OlmMessageHash, PickledAccount, PickledCrossSigningIdentity, PickledInboundGroupSession,
+        PickledSession, PicklingMode, PrivateCrossSigningIdentity, ReadOnlyAccount, Session,
+        SessionPickle,
     },
 };
 
@@ -491,6 +492,24 @@ impl SqliteStore {
             )
             .await?;
 
+        connection
+            .execute(
+                r#"
+            CREATE TABLE IF NOT EXISTS olm_hashes (
+                "id" INTEGER NOT NULL PRIMARY KEY,
+                "account_id" INTEGER NOT NULL,
+                "sender_key" TEXT NOT NULL,
+                "hash" TEXT NOT NULL,
+                FOREIGN KEY ("account_id") REFERENCES "accounts" ("id")
+                    ON DELETE CASCADE
+                UNIQUE(account_id,sender_key,hash)
+            );
+
+            CREATE INDEX IF NOT EXISTS "olm_hashes_index" ON "olm_hashes" ("account_id");
+        "#,
+            )
+            .await?;
+
         Ok(())
     }
 
@@ -1466,6 +1485,25 @@ impl SqliteStore {
         Ok(())
     }
 
+    async fn save_olm_hashses(
+        &self,
+        connection: &mut SqliteConnection,
+        hashes: &[OlmMessageHash],
+    ) -> Result<()> {
+        let account_id = self.account_id().ok_or(CryptoStoreError::AccountUnset)?;
+
+        for hash in hashes {
+            query("REPLACE INTO olm_hashes (account_id, sender_key, hash) VALUES (?1, ?2, ?3)")
+                .bind(account_id)
+                .bind(&hash.sender_key)
+                .bind(&hash.hash)
+                .execute(&mut *connection)
+                .await?;
+        }
+
+        Ok(())
+    }
+
     async fn save_user_helper(
         &self,
         mut connection: &mut SqliteConnection,
@@ -1681,6 +1719,9 @@ impl CryptoStore for SqliteStore {
         self.save_user_identities(&mut transaction, &changes.identities.changed)
             .await?;
 
+        self.save_olm_hashses(&mut transaction, &changes.message_hashes)
+            .await?;
+
         transaction.commit().await?;
 
         Ok(())
@@ -1796,6 +1837,22 @@ impl CryptoStore for SqliteStore {
 
         Ok(row.map(|r| r.0))
     }
+
+    async fn is_message_known(&self, message_hash: &OlmMessageHash) -> Result<bool> {
+        let account_id = self.account_id().ok_or(CryptoStoreError::AccountUnset)?;
+        let mut connection = self.connection.lock().await;
+
+        let row: Option<(String,)> = query_as(
+            "SELECT hash FROM olm_hashes WHERE account_id = ? and sender_key = ? and hash = ?",
+        )
+        .bind(account_id)
+        .bind(&message_hash.sender_key)
+        .bind(&message_hash.hash)
+        .fetch_optional(&mut *connection)
+        .await?;
+
+        Ok(row.is_some())
+    }
 }
 
 #[cfg(not(tarpaulin_include))]
@@ -1817,8 +1874,8 @@ mod test {
             user::test::{get_other_identity, get_own_identity},
         },
         olm::{
-            GroupSessionKey, InboundGroupSession, PrivateCrossSigningIdentity, ReadOnlyAccount,
-            Session,
+            GroupSessionKey, InboundGroupSession, OlmMessageHash, PrivateCrossSigningIdentity,
+            ReadOnlyAccount, Session,
         },
         store::{Changes, DeviceChanges, IdentityChanges},
     };
@@ -2371,4 +2428,21 @@ mod test {
         store.remove_value(&key).await.unwrap();
         assert!(store.get_value(&key).await.unwrap().is_none());
     }
+
+    #[tokio::test(threaded_scheduler)]
+    async fn olm_hash_saving() {
+        let (_, store, _dir) = get_loaded_store().await;
+
+        let hash = OlmMessageHash {
+            sender_key: "test_sender".to_owned(),
+            hash: "test_hash".to_owned(),
+        };
+
+        let mut changes = Changes::default();
+        changes.message_hashes.push(hash.clone());
+
+        assert!(!store.is_message_known(&hash).await.unwrap());
+        store.save_changes(changes).await.unwrap();
+        assert!(store.is_message_known(&hash).await.unwrap());
+    }
 }