refactor(keystore): MlsPendingMessage has no real primary key

keystore/Cargo.toml

@@ -42,6 +42,10 @@ log = { workspace = true }
 proteus-traits = { workspace = true, optional = true }
 itertools.workspace = true
 aes-gcm = "0.10"
+twox-hash = { version = "2.1.2", default-features = false, features = [
+  "alloc",
+  "xxhash3_128",
+] }
 
 [target.'cfg(target_os = "ios")'.dependencies]
 security-framework = "3.5"

keystore/src/connection/mod.rs

@@ -340,7 +340,7 @@ impl Database {
         conversation_id: &[u8],
     ) -> CryptoKeystoreResult<Vec<MlsPendingMessage>> {
         let mut conn = self.conn().await?;
-        let persisted_records = MlsPendingMessage::find_all_by_conversation_id(&mut conn, conversation_id).await?;
+        let persisted_records = MlsPendingMessage::find_all_matching(&mut conn, &conversation_id.into()).await?;
 
         let transaction_guard = self.transaction.lock().await;
         let Some(transaction) = transaction_guard.as_ref() else {

keystore/src/entities/mls.rs

@@ -1,4 +1,4 @@
-use zeroize::{Zeroize, ZeroizeOnDrop};
+use zeroize::Zeroize;
 
 use crate::{
     CryptoKeystoreResult, Sha256Hash,
@@ -73,74 +73,52 @@ pub struct PersistedMlsPendingGroup {
     pub custom_configuration: Vec<u8>,
 }
 
-/// [`MlsPendingMessage`]s have no distinct primary key;
-/// they must always be accessed via [`MlsPendingMessage::find_all_by_conversation_id`] and
-/// cleaned up with [`MlsPendingMessage::delete_by_conversation_id`]
-///
-/// However, we have to fake a primary key type in order to support
-/// `KeystoreTransaction::remove_pending_messages_by_conversation_id`. Additionally we need the same one in WASM, where
-/// it's necessary for item-level encryption.
+/// Typesafe reference to a conversation id.
 ///
-/// This implementation is fairly inefficient and hopefully temporary. But it at least implements the correct semantics.
-#[derive(ZeroizeOnDrop)]
-pub struct MlsPendingMessagePrimaryKey {
-    pub(crate) foreign_id: Vec<u8>,
-    message: Vec<u8>,
-}
+/// [`MlsPendingMessage`]s have no distinct primary key; they must always be accessed via
+/// collective accessors. This type makes that possible.
+#[derive(
+    Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, derive_more::AsRef, derive_more::Deref, derive_more::From,
+)]
+pub struct ConversationId<'a>(&'a [u8]);
 
-impl MlsPendingMessagePrimaryKey {
-    /// Construct a partial mls pending message primary key from only the conversation id.
-    ///
-    /// This does not in fact uniquely identify a single pending message--it should always uniquely
-    /// identify exactly 0 pending messages--but we have to have it so that we can search and delete
-    /// by conversation id within transactions.
-    pub(crate) fn from_conversation_id(conversation_id: impl AsRef<[u8]>) -> Self {
-        Self {
-            foreign_id: conversation_id.as_ref().to_owned(),
-            message: Vec::new(),
-        }
+impl<'a> KeyType for ConversationId<'a> {
+    fn bytes(&self) -> std::borrow::Cow<'_, [u8]> {
+        self.0.into()
     }
 }
 
+/// [`MlsPendingMessage`]s have no distinct primary key;
+/// they must always be accessed via the [`SearchableEntity`][crate::traits::SearchableEntity] and
+/// [`DeletableBySearchKey`][crate::traits::DeletableBySearchKey] traits.
+///
+/// However the keystore's support of internal transactions demands a primary key:
+/// ultimately that structure boils down to `Map<CollectionName, Map<PrimaryKey, Entity>>`, so anything other
+/// than a full primary key just breaks things.
+///
+/// We use `xxhash3` as a fast hash implementation, and take 128 bits of hash to ensure
+/// that the chance of a collision is effectively 0.
+pub struct MlsPendingMessagePrimaryKey(u128);
+
 impl From<&MlsPendingMessage> for MlsPendingMessagePrimaryKey {
     fn from(value: &MlsPendingMessage) -> Self {
-        Self {
-            foreign_id: value.foreign_id.clone(),
-            message: value.message.clone(),
-        }
+        let mut hasher = twox_hash::xxhash3_128::Hasher::new();
+        hasher.write(&value.foreign_id);
+        hasher.write(&value.message);
+        Self(hasher.finish_128())
     }
 }
 
 impl KeyType for MlsPendingMessagePrimaryKey {
     fn bytes(&self) -> std::borrow::Cow<'_, [u8]> {
-        // run-length encoding: 32 bits of size for each field, followed by the field
-        let fields = [&self.foreign_id, &self.message];
-        let mut key = Vec::with_capacity(
-            ((u32::BITS / u8::BITS) as usize * fields.len()) + self.foreign_id.len() + self.message.len(),
-        );
-        for field in fields {
-            key.extend((field.len() as u32).to_le_bytes());
-            key.extend(field.as_slice());
-        }
-        key.into()
+        self.0.to_be_bytes().as_slice().to_owned().into()
     }
 }
 
 impl OwnedKeyType for MlsPendingMessagePrimaryKey {
     fn from_bytes(bytes: &[u8]) -> Option<Self> {
-        // run-length decoding: 32 bits of size for each field, followed by the field
-        let (len, bytes) = bytes.split_at_checked(4)?;
-        let len = u32::from_le_bytes(len.try_into().ok()?);
-        let (foreign_id, bytes) = bytes.split_at_checked(len as _)?;
-
-        let (len, bytes) = bytes.split_at_checked(4)?;
-        let len = u32::from_le_bytes(len.try_into().ok()?);
-        let (message, bytes) = bytes.split_at_checked(len as _)?;
-
-        bytes.is_empty().then(|| Self {
-            foreign_id: foreign_id.to_owned(),
-            message: message.to_owned(),
-        })
+        let array = bytes.try_into().ok()?;
+        Some(Self(u128::from_be_bytes(array)))
     }
 }
 

keystore/src/entities/platform/generic/mls/pending_message.rs

@@ -6,10 +6,12 @@ use rusqlite::{Row, params};
 use crate::{
     CryptoKeystoreResult,
     connection::{DatabaseConnection, KeystoreDatabaseConnection, TransactionWrapper},
-    entities::{MlsPendingMessage, count_helper, count_helper_tx, delete_helper, get_helper, load_all_helper},
+    entities::{
+        ConversationId, MlsPendingMessage, count_helper, count_helper_tx, delete_helper, get_helper, load_all_helper,
+    },
     traits::{
-        BorrowPrimaryKey, Entity, EntityBase, EntityDatabaseMutation, EntityDeleteBorrowed, EntityGetBorrowed, KeyType,
-        OwnedKeyType, PrimaryKey,
+        BorrowPrimaryKey, DeletableBySearchKey, Entity, EntityBase, EntityDatabaseMutation, EntityDeleteBorrowed,
+        EntityGetBorrowed, KeyType, OwnedKeyType, PrimaryKey, SearchableEntity,
     },
 };
 
@@ -19,27 +21,6 @@ impl MlsPendingMessage {
         let message = row.get("message")?;
         Ok(Self { foreign_id, message })
     }
-
-    pub async fn find_all_by_conversation_id(
-        conn: &mut <Self as EntityBase>::ConnectionType,
-        conversation_id: &[u8],
-    ) -> CryptoKeystoreResult<Vec<Self>> {
-        let conn = conn.conn().await;
-        let mut stmt = conn.prepare_cached("SELECT * FROM mls_pending_messages WHERE id = ?")?;
-        let values = stmt
-            .query_map([conversation_id], Self::from_row)?
-            .collect::<Result<_, _>>()?;
-        Ok(values)
-    }
-
-    pub async fn delete_by_conversation_id(
-        tx: &TransactionWrapper<'_>,
-        conversation_id: &[u8],
-    ) -> CryptoKeystoreResult<bool> {
-        // a slight misuse of this helper, but SQL doesn't care if we end up deleting N rows instead of 1
-        // with this query
-        delete_helper::<Self>(tx, "id", conversation_id).await
-    }
 }
 
 impl EntityBase for MlsPendingMessage {
@@ -85,3 +66,36 @@ impl<'a> EntityDatabaseMutation<'a> for MlsPendingMessage {
         panic!("cannot delete `MlsPendingMessage` by primary key as it has no distinct primary key")
     }
 }
+
+#[async_trait]
+impl<'a> SearchableEntity<ConversationId<'a>> for MlsPendingMessage {
+    async fn find_all_matching(
+        conn: &mut Self::ConnectionType,
+        conversation_id: &ConversationId<'a>,
+    ) -> CryptoKeystoreResult<Vec<Self>> {
+        let conversation_id = *conversation_id.as_ref();
+        let conn = conn.conn().await;
+        let mut stmt = conn.prepare_cached("SELECT * FROM mls_pending_messages WHERE id = ?")?;
+        let values = stmt
+            .query_map([conversation_id], Self::from_row)?
+            .collect::<Result<_, _>>()?;
+        Ok(values)
+    }
+
+    fn matches(&self, conversation_id: &ConversationId<'a>) -> bool {
+        *conversation_id.as_ref() == self.foreign_id.as_slice()
+    }
+}
+
+#[async_trait]
+impl<'a> DeletableBySearchKey<'a, ConversationId<'a>> for MlsPendingMessage {
+    async fn delete_all_matching(
+        tx: &Self::Transaction,
+        conversation_id: &ConversationId<'a>,
+    ) -> CryptoKeystoreResult<()> {
+        // a slight misuse of this helper, but SQL doesn't care if we end up deleting N rows instead of 1
+        // with this query
+        delete_helper::<Self>(tx, "id", *conversation_id.as_ref()).await?;
+        Ok(())
+    }
+}

keystore/src/entities/platform/wasm/mls/pending_message.rs

@@ -6,31 +6,14 @@ use wasm_bindgen::JsValue;
 use crate::{
     CryptoKeystoreResult,
     connection::{KeystoreDatabaseConnection, TransactionWrapper},
-    entities::MlsPendingMessage,
+    entities::{ConversationId, MlsPendingMessage},
     traits::{
-        DecryptWithExplicitEncryptionKey as _, Decryptable, Decrypting, EncryptWithExplicitEncryptionKey as _,
-        Encrypting, EncryptionKey, Entity, EntityBase, EntityDatabaseMutation,
+        DecryptWithExplicitEncryptionKey as _, Decryptable, Decrypting, DeletableBySearchKey,
+        EncryptWithExplicitEncryptionKey as _, Encrypting, EncryptionKey, Entity, EntityBase, EntityDatabaseMutation,
+        SearchableEntity,
     },
 };
 
-impl MlsPendingMessage {
-    pub async fn find_all_by_conversation_id(
-        conn: &mut <Self as EntityBase>::ConnectionType,
-        conversation_id: &[u8],
-    ) -> crate::CryptoKeystoreResult<Vec<Self>> {
-        let storage = conn.storage();
-        let id = JsValue::from(Uint8Array::from(conversation_id));
-        storage.get_all_with_query(Some(id.into())).await
-    }
-
-    pub async fn delete_by_conversation_id(
-        tx: &<Self as EntityDatabaseMutation<'_>>::Transaction,
-        conversation_id: &[u8],
-    ) -> crate::CryptoKeystoreResult<bool> {
-        tx.delete::<Self>(conversation_id).await
-    }
-}
-
 impl EntityBase for MlsPendingMessage {
     type ConnectionType = KeystoreDatabaseConnection;
     type AutoGeneratedFields = ();
@@ -118,3 +101,31 @@ impl Decrypting<'static> for MlsPendingMessageDecrypt {
 impl Decryptable<'static> for MlsPendingMessage {
     type DecryptableFrom = MlsPendingMessageDecrypt;
 }
+
+#[async_trait(?Send)]
+impl<'a> SearchableEntity<ConversationId<'a>> for MlsPendingMessage {
+    async fn find_all_matching(
+        conn: &mut Self::ConnectionType,
+        conversation_id: &ConversationId<'a>,
+    ) -> CryptoKeystoreResult<Vec<Self>> {
+        let conversation_id = *conversation_id.as_ref();
+        let storage = conn.storage();
+        let id = JsValue::from(Uint8Array::from(conversation_id));
+        storage.get_all_with_query(Some(id.into())).await
+    }
+
+    fn matches(&self, conversation_id: &ConversationId<'a>) -> bool {
+        *conversation_id.as_ref() == self.foreign_id.as_slice()
+    }
+}
+
+#[async_trait(?Send)]
+impl<'a> DeletableBySearchKey<'a, ConversationId<'a>> for MlsPendingMessage {
+    async fn delete_all_matching(
+        tx: &Self::Transaction,
+        conversation_id: &ConversationId<'a>,
+    ) -> CryptoKeystoreResult<()> {
+        tx.delete::<Self>(*conversation_id.as_ref()).await?;
+        Ok(())
+    }
+}

keystore/src/traits/mod.rs

@@ -23,5 +23,5 @@ pub use item_encryption::{
 };
 pub use key_type::{KeyType, OwnedKeyType};
 pub use primary_key::{BorrowPrimaryKey, PrimaryKey};
-pub use searchable_entity::SearchableEntity;
+pub use searchable_entity::{DeletableBySearchKey, SearchableEntity};
 pub use unique_entity::{UniqueEntity, UniqueEntityExt, UniqueEntityImplementationHelper};

keystore/src/traits/searchable_entity.rs

@@ -2,7 +2,7 @@ use async_trait::async_trait;
 
 use crate::{
     CryptoKeystoreResult,
-    traits::{Entity, KeyType},
+    traits::{Entity, EntityDatabaseMutation, KeyType},
 };
 
 /// Entities implementing `SearchableEntity` have a distinct search key which
@@ -36,3 +36,25 @@ pub trait SearchableEntity<SearchKey: KeyType>: Entity {
     /// on the entity.
     fn matches(&self, search_key: &SearchKey) -> bool;
 }
+
+/// Entities implementing `DeletableBySearchKey` can be deleted by that search key.
+///
+/// This is a way at the type-system level to implement `WHERE`-clause deletion.
+///
+/// This trait can potentially be implemented multiple times per entity, in case there
+/// are a variety of interesting searches.
+///
+/// While the trait design does not require it, implementations should take advantage of
+/// database features such as indices to ensure that deletion by a search key is efficient.
+#[cfg_attr(target_family = "wasm", async_trait(?Send))]
+#[cfg_attr(not(target_family = "wasm"), async_trait)]
+pub trait DeletableBySearchKey<'a, SearchKey: KeyType>:
+    SearchableEntity<SearchKey> + EntityDatabaseMutation<'a>
+{
+    /// Delete all entities matching the search key.
+    ///
+    /// The specific meaning of "matching" the search key will depend on the entity
+    /// in question, it should always have the same meaning for this implementation
+    /// and the equivalent [`SearchableEntity`] implementation.
+    async fn delete_all_matching(tx: &Self::Transaction, search_key: &SearchKey) -> CryptoKeystoreResult<()>;
+}

keystore/src/transaction/dynamic_dispatch/entity_id.rs

@@ -13,7 +13,7 @@ use crate::{
         StoredBufferedCommit, StoredCredential, StoredE2eiEnrollment, StoredEncryptionKeyPair,
         StoredEpochEncryptionKeypair, StoredHpkePrivateKey, StoredKeypackage, StoredPskBundle,
     },
-    traits::{BorrowPrimaryKey, Entity, EntityDatabaseMutation, KeyType, OwnedKeyType as _},
+    traits::{BorrowPrimaryKey, DeletableBySearchKey as _, Entity, EntityDatabaseMutation, KeyType, OwnedKeyType as _},
     transaction::dynamic_dispatch::EntityType,
 };
 
@@ -39,7 +39,7 @@ impl EntityId {
             .ok_or(CryptoKeystoreError::InvalidPrimaryKeyBytes(self.typ.collection_name()))
     }
 
-    fn from_key<E>(primary_key: Cow<'_, [u8]>) -> Option<Self>
+    pub(crate) fn from_key<E>(primary_key: Cow<'_, [u8]>) -> Option<Self>
     where
         E: Entity,
     {
@@ -98,10 +98,9 @@ impl EntityId {
             EntityType::PersistedMlsPendingGroup => {
                 PersistedMlsPendingGroup::delete(tx, &self.primary_key::<PersistedMlsPendingGroup>()?).await
             }
-            EntityType::MlsPendingMessage => {
-                let primary_key = self.primary_key::<MlsPendingMessage>()?;
-                MlsPendingMessage::delete_by_conversation_id(tx, &primary_key.foreign_id).await
-            }
+            EntityType::MlsPendingMessage => MlsPendingMessage::delete_all_matching(tx, &self.id.as_slice().into())
+                .await
+                .map(|_| false),
             EntityType::StoredE2eiEnrollment => {
                 StoredE2eiEnrollment::delete(tx, &self.primary_key::<StoredE2eiEnrollment>()?).await
             }

keystore/src/transaction/mod.rs

@@ -10,7 +10,7 @@ use itertools::Itertools;
 use crate::{
     CryptoKeystoreError, CryptoKeystoreResult,
     connection::{Database, KeystoreDatabaseConnection},
-    entities::{MlsPendingMessage, MlsPendingMessagePrimaryKey, PersistedMlsGroup},
+    entities::{MlsPendingMessage, PersistedMlsGroup},
     traits::{
         BorrowPrimaryKey, Entity, EntityBase as _, EntityDatabaseMutation, EntityDeleteBorrowed, KeyType,
         SearchableEntity,
@@ -165,10 +165,8 @@ impl KeystoreTransaction {
 
         let mut deleted_set = self.deleted.write().await;
         deleted_set.insert(
-            EntityId::from_primary_key::<MlsPendingMessage>(&MlsPendingMessagePrimaryKey::from_conversation_id(
-                conversation_id,
-            ))
-            .expect("mls pending messages are proper entities which can be parsed"),
+            EntityId::from_key::<MlsPendingMessage>(conversation_id.into())
+                .expect("mls pending messages are proper entities which can be parsed"),
         );
     }