From 85f22d5bb6976a720d52b1353519f654f6615acf Mon Sep 17 00:00:00 2001
From: Peter Goodspeed-Niklaus <peter.goodspeed-niklaus@wire.com>
Date: Wed, 17 Dec 2025 17:40:43 +0100
Subject: [PATCH 01/48] feat(keystore): introduce `trait OwnedKeyType`

We already knew that there is a real difference between key types
which are owned and those which are borrowed. We already knew that
entity keys need to be of the owned kind.

Adding this trait allows us to formalize that notion, and add
a utility to construct an owned variant from some arbitrary bytes.
---
 keystore/src/hash.rs            | 11 ++++++-
 keystore/src/traits/key_type.rs | 53 ++++++++++++++++++++-------------
 keystore/src/traits/mod.rs      |  2 +-
 3 files changed, 44 insertions(+), 22 deletions(-)

diff --git a/keystore/src/hash.rs b/keystore/src/hash.rs
index 33d8f3be47..9e7e8f1d37 100644
--- a/keystore/src/hash.rs
+++ b/keystore/src/hash.rs
@@ -2,7 +2,10 @@ use std::fmt;
 
 use sha2::{Digest, Sha256};
 
-use crate::{CryptoKeystoreResult, traits::KeyType};
+use crate::{
+    CryptoKeystoreResult,
+    traits::{KeyType, OwnedKeyType},
+};
 
 /// Used to calculate ID hashes for some MlsEntities' SQLite tables (not used on wasm).
 /// We only use sha256 on platforms where we use SQLite.
@@ -69,6 +72,12 @@ impl KeyType for Sha256Hash {
     }
 }
 
+impl OwnedKeyType for Sha256Hash {
+    fn from_bytes(bytes: &[u8]) -> Option<Self> {
+        bytes.try_into().ok().map(Self)
+    }
+}
+
 #[cfg(not(target_family = "wasm"))]
 impl rusqlite::ToSql for Sha256Hash {
     fn to_sql(&self) -> rusqlite::Result<rusqlite::types::ToSqlOutput<'_>> {
diff --git a/keystore/src/traits/key_type.rs b/keystore/src/traits/key_type.rs
index 53b7f5cb5f..d00d875669 100644
--- a/keystore/src/traits/key_type.rs
+++ b/keystore/src/traits/key_type.rs
@@ -4,7 +4,7 @@ use std::borrow::Cow;
 ///
 /// This might be a primary key, in which case the key uniquely identifies either 0 or 1 entries in the database.
 /// Or it might be a search key, in which case the key could match any number of entries.
-pub trait KeyType: Send + Sync {
+pub trait KeyType: Send + Sync + Sized {
     /// Get a unique binary representation of this key.
     ///
     /// For simple keys it can just be the borrowed form of the key itself,
@@ -12,35 +12,52 @@ pub trait KeyType: Send + Sync {
     fn bytes(&self) -> Cow<'_, [u8]>;
 }
 
-// useful for unique entities; non-allocating
-impl KeyType for () {
-    fn bytes(&self) -> Cow<'_, [u8]> {
-        Vec::new().into()
-    }
+/// An owned key type can be converted to from arbitrary bytes.
+pub trait OwnedKeyType: 'static + KeyType {
+    /// Parse some bytes into an instance of this type.
+    ///
+    /// We're just going with `Option` instead of `CryptoKeystoreResult` for now because
+    /// the hopeful assumption is that this is going to be a rare occurrence that doesn't
+    /// need much explanation.
+    fn from_bytes(bytes: &[u8]) -> Option<Self>;
 }
 
 macro_rules! impl_keytype {
-    ($t:ty, |$self:ident| $impl:expr) => {
+    ($t:ty, |$self:ident| $bytes:expr) => {
         impl KeyType for $t {
             fn bytes(&$self) -> Cow<'_, [u8]> {
-                $impl.into()
+                $bytes.into()
             }
         }
     };
+    ($t:ty, |$self:ident| $bytes:expr, |$bytes_id:ident| $from_bytes:expr) => {
+        impl_keytype!($t, |$self| $bytes);
+
+        impl OwnedKeyType for $t {
+            fn from_bytes($bytes_id: &[u8]) -> Option<Self> {
+                $from_bytes
+            }
+        }
+    };
+
+
 }
 
+// useful for unique entities; non-allocating
+impl_keytype!((), |self| Vec::new(), |bytes| bytes.is_empty().then_some(()));
 impl_keytype!(&[u8], |self| *self);
-impl_keytype!(Vec<u8>, |self| self.as_slice());
+impl_keytype!(Vec<u8>, |self| self.as_slice(), |bytes| Some(bytes.into()));
 impl_keytype!(&str, |self| self.as_bytes());
-impl_keytype!(String, |self| self.as_bytes());
+impl_keytype!(String, |self| self.as_bytes(), |bytes| str::from_utf8(bytes)
+    .ok()
+    .map(ToOwned::to_owned));
 
 macro_rules! impl_keytype_for_integer {
     ($t:ty) => {
-        impl KeyType for $t {
-            fn bytes(&self) -> Cow<'_, [u8]> {
-                Vec::from(self.to_le_bytes()).into()
-            }
-        }
+        impl_keytype!($t, |self| Vec::from(self.to_le_bytes()), |bytes| {
+            let array = bytes.try_into().ok()?;
+            Some(<$t>::from_le_bytes(array))
+        });
     };
 }
 
@@ -56,8 +73,4 @@ impl_keytype_for_integer!(i64);
 impl_keytype_for_integer!(i128);
 
 /// Some unique entities use a single byte as a key type
-impl KeyType for [u8; 1] {
-    fn bytes(&self) -> Cow<'_, [u8]> {
-        self.into()
-    }
-}
+impl_keytype!([u8; 1], |self| self, |bytes| bytes.try_into().ok());
diff --git a/keystore/src/traits/mod.rs b/keystore/src/traits/mod.rs
index 5068dba59e..ae9d54a922 100644
--- a/keystore/src/traits/mod.rs
+++ b/keystore/src/traits/mod.rs
@@ -16,5 +16,5 @@ pub use entity_base::EntityBase;
 pub use entity_database_mutation::{EntityDatabaseMutation, EntityDeleteBorrowed};
 pub use fetch_from_database::FetchFromDatabase;
 pub use item_encryption::{DecryptData, Decryptable, Decrypting, EncryptData, Encrypting};
-pub use key_type::KeyType;
+pub use key_type::{KeyType, OwnedKeyType};
 pub use unique_entity::{UniqueEntity, UniqueEntityExt, UniqueEntityImplementationHelper};

From 44bbf959c11c3bfacd90bd04b005e184ff8f0fb0 Mon Sep 17 00:00:00 2001
From: Peter Goodspeed-Niklaus <peter.goodspeed-niklaus@wire.com>
Date: Wed, 17 Dec 2025 17:44:58 +0100
Subject: [PATCH 02/48] refactor(keystore): formally require that primary key
 be `OwnedKeyType`

---
 keystore/src/traits/entity.rs | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/keystore/src/traits/entity.rs b/keystore/src/traits/entity.rs
index c1a30a58b7..61aa5ee9da 100644
--- a/keystore/src/traits/entity.rs
+++ b/keystore/src/traits/entity.rs
@@ -4,7 +4,7 @@ use async_trait::async_trait;
 
 use crate::{
     CryptoKeystoreResult,
-    traits::{EntityBase, KeyType},
+    traits::{EntityBase, KeyType, OwnedKeyType},
 };
 
 /// Something which can be stored in our database.
@@ -16,7 +16,7 @@ pub trait Entity: EntityBase {
     /// Each distinct `PrimaryKey` uniquely identifies either 0 or 1 instance.
     ///
     /// This constraint should be enforced at the DB level.
-    type PrimaryKey: KeyType;
+    type PrimaryKey: OwnedKeyType;
 
     /// Get this entity's primary key.
     ///

From 7c6b1e84dd1804da146a95c15abdfb7d6079ca97 Mon Sep 17 00:00:00 2001
From: Peter Goodspeed-Niklaus <peter.goodspeed-niklaus@wire.com>
Date: Thu, 18 Dec 2025 16:39:44 +0100
Subject: [PATCH 03/48] feat(keystore): complete the `FetchFromDatabase` API

Just looking to give it complete coverage of the kinds of things that
the new entity traits offer.
---
 keystore/src/traits/fetch_from_database.rs | 19 ++++++++++++++++---
 1 file changed, 16 insertions(+), 3 deletions(-)

diff --git a/keystore/src/traits/fetch_from_database.rs b/keystore/src/traits/fetch_from_database.rs
index 360070550c..b2401c9544 100644
--- a/keystore/src/traits/fetch_from_database.rs
+++ b/keystore/src/traits/fetch_from_database.rs
@@ -3,7 +3,7 @@ use async_trait::async_trait;
 use crate::{
     CryptoKeystoreResult,
     connection::KeystoreDatabaseConnection,
-    traits::{Entity, UniqueEntity},
+    traits::{BorrowPrimaryKey, Entity, UniqueEntity, UniqueEntityExt},
 };
 
 /// Interface to fetch from the database either from the connection directly or through a
@@ -29,8 +29,21 @@ pub trait FetchFromDatabase: Send + Sync {
     where
         E: Entity<ConnectionType = KeystoreDatabaseConnection>;
 
+    /// Get an instance of `E` from the database by the borrowed form of its primary key.
+    async fn get_borrowed<E>(
+        &self,
+        id: &<E as BorrowPrimaryKey>::BorrowedPrimaryKey,
+    ) -> CryptoKeystoreResult<Option<E>>
+    where
+        E: BorrowPrimaryKey;
+
     /// Get the requested unique entity from the database.
-    async fn get_unique<U>(&self) -> CryptoKeystoreResult<Option<U>>
+    async fn get_unique<'a, U>(&self) -> CryptoKeystoreResult<Option<U>>
+    where
+        U: UniqueEntityExt<'a>;
+
+    /// Determine whether a unique entity is present in the database.
+    async fn exists<'a, U>(&self) -> CryptoKeystoreResult<bool>
     where
-        U: UniqueEntity<ConnectionType = KeystoreDatabaseConnection>;
+        U: UniqueEntityExt<'a>;
 }

From 9f9ce0a590e75ff5afb8f7c1bdd7234028350e61 Mon Sep 17 00:00:00 2001
From: Peter Goodspeed-Niklaus <peter.goodspeed-niklaus@wire.com>
Date: Wed, 17 Dec 2025 17:46:36 +0100
Subject: [PATCH 04/48] refactor(keystore): rework the dynamic dispatch module

- break up the impls into separate files for easier navigation
- separate the concept of an entity type and an entity id
- borrow the actual entity instead of owning it
- make `execute_save` and `execute_delete` methods not free functions
- restrict visibility: nothing external should really need this dynamic
  dispatch stuff
---
 keystore/src/error.rs                         |   2 +
 keystore/src/transaction/dynamic_dispatch.rs  | 211 ------------------
 .../transaction/dynamic_dispatch/entity.rs    |  89 ++++++++
 .../transaction/dynamic_dispatch/entity_id.rs |  91 ++++++++
 .../dynamic_dispatch/entity_type.rs           |  95 ++++++++
 .../src/transaction/dynamic_dispatch/mod.rs   |   7 +
 6 files changed, 284 insertions(+), 211 deletions(-)
 delete mode 100644 keystore/src/transaction/dynamic_dispatch.rs
 create mode 100644 keystore/src/transaction/dynamic_dispatch/entity.rs
 create mode 100644 keystore/src/transaction/dynamic_dispatch/entity_id.rs
 create mode 100644 keystore/src/transaction/dynamic_dispatch/entity_type.rs
 create mode 100644 keystore/src/transaction/dynamic_dispatch/mod.rs

diff --git a/keystore/src/error.rs b/keystore/src/error.rs
index cc3c55be2e..8ff9477c1e 100644
--- a/keystore/src/error.rs
+++ b/keystore/src/error.rs
@@ -152,6 +152,8 @@ pub enum CryptoKeystoreError {
     MigrationNotSupported(u32),
     #[error("The migration failed: {0}")]
     MigrationFailed(String),
+    #[error("the provided bytes could not be interpreted as the primary key of {0}")]
+    InvalidPrimaryKeyBytes(&'static str),
 }
 
 #[cfg(target_family = "wasm")]
diff --git a/keystore/src/transaction/dynamic_dispatch.rs b/keystore/src/transaction/dynamic_dispatch.rs
deleted file mode 100644
index df993237ef..0000000000
--- a/keystore/src/transaction/dynamic_dispatch.rs
+++ /dev/null
@@ -1,211 +0,0 @@
-//! This module exists merely because the `Entity` trait is not object safe.
-//! See <https://doc.rust-lang.org/reference/items/traits.html#object-safety.>.
-
-#[cfg(target_family = "wasm")]
-use crate::entities::E2eiRefreshToken;
-#[cfg(feature = "proteus-keystore")]
-use crate::entities::{ProteusIdentity, ProteusPrekey, ProteusSession};
-use crate::{
-    CryptoKeystoreError, CryptoKeystoreResult,
-    connection::TransactionWrapper,
-    entities::{
-        ConsumerData, E2eiAcmeCA, E2eiCrl, E2eiIntermediateCert, EntityBase, EntityTransactionExt, MlsPendingMessage,
-        PersistedMlsGroup, PersistedMlsPendingGroup, StoredBufferedCommit, StoredCredential, StoredE2eiEnrollment,
-        StoredEncryptionKeyPair, StoredEpochEncryptionKeypair, StoredHpkePrivateKey, StoredKeypackage, StoredPskBundle,
-        StringEntityId, UniqueEntity,
-    },
-};
-
-#[derive(Debug)]
-pub enum Entity {
-    ConsumerData(ConsumerData),
-    HpkePrivateKey(StoredHpkePrivateKey),
-    StoredKeypackage(StoredKeypackage),
-    PskBundle(StoredPskBundle),
-    EncryptionKeyPair(StoredEncryptionKeyPair),
-    StoredEpochEncryptionKeypair(StoredEpochEncryptionKeypair),
-    StoredCredential(StoredCredential),
-    StoredBufferedCommit(StoredBufferedCommit),
-    PersistedMlsGroup(PersistedMlsGroup),
-    PersistedMlsPendingGroup(PersistedMlsPendingGroup),
-    MlsPendingMessage(MlsPendingMessage),
-    StoredE2eiEnrollment(StoredE2eiEnrollment),
-    #[cfg(target_family = "wasm")]
-    E2eiRefreshToken(E2eiRefreshToken),
-    E2eiAcmeCA(E2eiAcmeCA),
-    E2eiIntermediateCert(E2eiIntermediateCert),
-    E2eiCrl(E2eiCrl),
-    #[cfg(feature = "proteus-keystore")]
-    ProteusIdentity(ProteusIdentity),
-    #[cfg(feature = "proteus-keystore")]
-    ProteusPrekey(ProteusPrekey),
-    #[cfg(feature = "proteus-keystore")]
-    ProteusSession(ProteusSession),
-}
-
-#[derive(Debug, Clone, PartialEq)]
-pub enum EntityId {
-    HpkePrivateKey(Vec<u8>),
-    KeyPackage(Vec<u8>),
-    PskBundle(Vec<u8>),
-    EncryptionKeyPair(Vec<u8>),
-    EpochEncryptionKeyPair(Vec<u8>),
-    StoredCredential(Vec<u8>),
-    StoredBufferedCommit(Vec<u8>),
-    PersistedMlsGroup(Vec<u8>),
-    PersistedMlsPendingGroup(Vec<u8>),
-    MlsPendingMessage(Vec<u8>),
-    StoredE2eiEnrollment(Vec<u8>),
-    #[cfg(target_family = "wasm")]
-    E2eiRefreshToken(Vec<u8>),
-    E2eiAcmeCA(Vec<u8>),
-    E2eiIntermediateCert(Vec<u8>),
-    E2eiCrl(Vec<u8>),
-    #[cfg(feature = "proteus-keystore")]
-    ProteusIdentity(Vec<u8>),
-    #[cfg(feature = "proteus-keystore")]
-    ProteusPrekey(Vec<u8>),
-    #[cfg(feature = "proteus-keystore")]
-    ProteusSession(Vec<u8>),
-}
-
-impl EntityId {
-    fn as_id(&self) -> StringEntityId<'_> {
-        match self {
-            EntityId::HpkePrivateKey(vec) => vec.as_slice().into(),
-            EntityId::KeyPackage(vec) => vec.as_slice().into(),
-            EntityId::PskBundle(vec) => vec.as_slice().into(),
-            EntityId::EncryptionKeyPair(vec) => vec.as_slice().into(),
-            EntityId::EpochEncryptionKeyPair(vec) => vec.as_slice().into(),
-            EntityId::StoredCredential(vec) => vec.as_slice().into(),
-            EntityId::StoredBufferedCommit(vec) => vec.as_slice().into(),
-            EntityId::PersistedMlsGroup(vec) => vec.as_slice().into(),
-            EntityId::PersistedMlsPendingGroup(vec) => vec.as_slice().into(),
-            EntityId::MlsPendingMessage(vec) => vec.as_slice().into(),
-            EntityId::StoredE2eiEnrollment(vec) => vec.as_slice().into(),
-            #[cfg(target_family = "wasm")]
-            EntityId::E2eiRefreshToken(vec) => vec.as_slice().into(),
-            EntityId::E2eiAcmeCA(vec) => vec.as_slice().into(),
-            EntityId::E2eiIntermediateCert(vec) => vec.as_slice().into(),
-            EntityId::E2eiCrl(vec) => vec.as_slice().into(),
-            #[cfg(feature = "proteus-keystore")]
-            EntityId::ProteusIdentity(vec) => vec.as_slice().into(),
-            #[cfg(feature = "proteus-keystore")]
-            EntityId::ProteusSession(id) => id.as_slice().into(),
-            #[cfg(feature = "proteus-keystore")]
-            EntityId::ProteusPrekey(vec) => vec.as_slice().into(),
-        }
-    }
-
-    pub(crate) fn from_collection_name(entity_id: &'static str, id: &[u8]) -> CryptoKeystoreResult<Self> {
-        match entity_id {
-            StoredHpkePrivateKey::COLLECTION_NAME => Ok(Self::HpkePrivateKey(id.into())),
-            StoredKeypackage::COLLECTION_NAME => Ok(Self::KeyPackage(id.into())),
-            StoredPskBundle::COLLECTION_NAME => Ok(Self::PskBundle(id.into())),
-            StoredEncryptionKeyPair::COLLECTION_NAME => Ok(Self::EncryptionKeyPair(id.into())),
-            StoredEpochEncryptionKeypair::COLLECTION_NAME => Ok(Self::EpochEncryptionKeyPair(id.into())),
-            StoredBufferedCommit::COLLECTION_NAME => Ok(Self::StoredBufferedCommit(id.into())),
-            PersistedMlsGroup::COLLECTION_NAME => Ok(Self::PersistedMlsGroup(id.into())),
-            PersistedMlsPendingGroup::COLLECTION_NAME => Ok(Self::PersistedMlsPendingGroup(id.into())),
-            StoredCredential::COLLECTION_NAME => Ok(Self::StoredCredential(id.into())),
-            MlsPendingMessage::COLLECTION_NAME => Ok(Self::MlsPendingMessage(id.into())),
-            StoredE2eiEnrollment::COLLECTION_NAME => Ok(Self::StoredE2eiEnrollment(id.into())),
-            E2eiCrl::COLLECTION_NAME => Ok(Self::E2eiCrl(id.into())),
-            E2eiAcmeCA::COLLECTION_NAME => Ok(Self::E2eiAcmeCA(id.into())),
-            #[cfg(target_family = "wasm")]
-            E2eiRefreshToken::COLLECTION_NAME => Ok(Self::E2eiRefreshToken(id.into())),
-            E2eiIntermediateCert::COLLECTION_NAME => Ok(Self::E2eiIntermediateCert(id.into())),
-            #[cfg(feature = "proteus-keystore")]
-            ProteusIdentity::COLLECTION_NAME => Ok(Self::ProteusIdentity(id.into())),
-            #[cfg(feature = "proteus-keystore")]
-            ProteusPrekey::COLLECTION_NAME => Ok(Self::ProteusPrekey(id.into())),
-            #[cfg(feature = "proteus-keystore")]
-            ProteusSession::COLLECTION_NAME => Ok(Self::ProteusSession(id.into())),
-            _ => Err(CryptoKeystoreError::NotImplemented),
-        }
-    }
-
-    pub(crate) fn collection_name(&self) -> &'static str {
-        match self {
-            EntityId::KeyPackage(_) => StoredKeypackage::COLLECTION_NAME,
-            EntityId::PskBundle(_) => StoredPskBundle::COLLECTION_NAME,
-            EntityId::EncryptionKeyPair(_) => StoredEncryptionKeyPair::COLLECTION_NAME,
-            EntityId::EpochEncryptionKeyPair(_) => StoredEpochEncryptionKeypair::COLLECTION_NAME,
-            EntityId::StoredCredential(_) => StoredCredential::COLLECTION_NAME,
-            EntityId::StoredBufferedCommit(_) => StoredBufferedCommit::COLLECTION_NAME,
-            EntityId::PersistedMlsGroup(_) => PersistedMlsGroup::COLLECTION_NAME,
-            EntityId::PersistedMlsPendingGroup(_) => PersistedMlsPendingGroup::COLLECTION_NAME,
-            EntityId::MlsPendingMessage(_) => MlsPendingMessage::COLLECTION_NAME,
-            EntityId::StoredE2eiEnrollment(_) => StoredE2eiEnrollment::COLLECTION_NAME,
-            #[cfg(target_family = "wasm")]
-            EntityId::E2eiRefreshToken(_) => E2eiRefreshToken::COLLECTION_NAME,
-            EntityId::E2eiAcmeCA(_) => E2eiAcmeCA::COLLECTION_NAME,
-            EntityId::E2eiIntermediateCert(_) => E2eiIntermediateCert::COLLECTION_NAME,
-            EntityId::E2eiCrl(_) => E2eiCrl::COLLECTION_NAME,
-            #[cfg(feature = "proteus-keystore")]
-            EntityId::ProteusIdentity(_) => ProteusIdentity::COLLECTION_NAME,
-            #[cfg(feature = "proteus-keystore")]
-            EntityId::ProteusPrekey(_) => ProteusPrekey::COLLECTION_NAME,
-            #[cfg(feature = "proteus-keystore")]
-            EntityId::ProteusSession(_) => ProteusSession::COLLECTION_NAME,
-            EntityId::HpkePrivateKey(_) => StoredHpkePrivateKey::COLLECTION_NAME,
-        }
-    }
-}
-
-pub async fn execute_save(tx: &TransactionWrapper<'_>, entity: &Entity) -> CryptoKeystoreResult<()> {
-    match entity {
-        Entity::ConsumerData(consumer_data) => consumer_data.replace(tx).await,
-        Entity::HpkePrivateKey(mls_hpke_private_key) => mls_hpke_private_key.save(tx).await,
-        Entity::StoredKeypackage(mls_key_package) => mls_key_package.save(tx).await,
-        Entity::PskBundle(mls_psk_bundle) => mls_psk_bundle.save(tx).await,
-        Entity::EncryptionKeyPair(mls_encryption_key_pair) => mls_encryption_key_pair.save(tx).await,
-        Entity::StoredEpochEncryptionKeypair(mls_epoch_encryption_key_pair) => {
-            mls_epoch_encryption_key_pair.save(tx).await
-        }
-        Entity::StoredCredential(mls_credential) => mls_credential.save(tx).await,
-        Entity::StoredBufferedCommit(mls_pending_commit) => mls_pending_commit.save(tx).await,
-        Entity::PersistedMlsGroup(persisted_mls_group) => persisted_mls_group.save(tx).await,
-        Entity::PersistedMlsPendingGroup(persisted_mls_pending_group) => persisted_mls_pending_group.save(tx).await,
-        Entity::MlsPendingMessage(mls_pending_message) => mls_pending_message.save(tx).await,
-        Entity::StoredE2eiEnrollment(e2ei_enrollment) => e2ei_enrollment.save(tx).await,
-        #[cfg(target_family = "wasm")]
-        Entity::E2eiRefreshToken(e2ei_refresh_token) => e2ei_refresh_token.replace(tx).await,
-        Entity::E2eiAcmeCA(e2ei_acme_ca) => e2ei_acme_ca.replace(tx).await,
-        Entity::E2eiIntermediateCert(e2ei_intermediate_cert) => e2ei_intermediate_cert.save(tx).await,
-        Entity::E2eiCrl(e2ei_crl) => e2ei_crl.save(tx).await,
-        #[cfg(feature = "proteus-keystore")]
-        Entity::ProteusSession(record) => record.save(tx).await,
-        #[cfg(feature = "proteus-keystore")]
-        Entity::ProteusIdentity(record) => record.save(tx).await,
-        #[cfg(feature = "proteus-keystore")]
-        Entity::ProteusPrekey(record) => record.save(tx).await,
-    }
-}
-
-pub async fn execute_delete(tx: &TransactionWrapper<'_>, entity_id: &EntityId) -> CryptoKeystoreResult<()> {
-    match entity_id {
-        id @ EntityId::HpkePrivateKey(_) => StoredHpkePrivateKey::delete(tx, id.as_id()).await,
-        id @ EntityId::KeyPackage(_) => StoredKeypackage::delete(tx, id.as_id()).await,
-        id @ EntityId::PskBundle(_) => StoredPskBundle::delete(tx, id.as_id()).await,
-        id @ EntityId::EncryptionKeyPair(_) => StoredEncryptionKeyPair::delete(tx, id.as_id()).await,
-        id @ EntityId::EpochEncryptionKeyPair(_) => StoredEpochEncryptionKeypair::delete(tx, id.as_id()).await,
-        id @ EntityId::StoredCredential(_) => StoredCredential::delete(tx, id.as_id()).await,
-        id @ EntityId::StoredBufferedCommit(_) => StoredBufferedCommit::delete(tx, id.as_id()).await,
-        id @ EntityId::PersistedMlsGroup(_) => PersistedMlsGroup::delete(tx, id.as_id()).await,
-        id @ EntityId::PersistedMlsPendingGroup(_) => PersistedMlsPendingGroup::delete(tx, id.as_id()).await,
-        id @ EntityId::MlsPendingMessage(_) => MlsPendingMessage::delete(tx, id.as_id()).await,
-        id @ EntityId::StoredE2eiEnrollment(_) => StoredE2eiEnrollment::delete(tx, id.as_id()).await,
-        #[cfg(target_family = "wasm")]
-        id @ EntityId::E2eiRefreshToken(_) => E2eiRefreshToken::delete(tx, id.as_id()).await,
-        id @ EntityId::E2eiAcmeCA(_) => E2eiAcmeCA::delete(tx, id.as_id()).await,
-        id @ EntityId::E2eiIntermediateCert(_) => E2eiIntermediateCert::delete(tx, id.as_id()).await,
-        id @ EntityId::E2eiCrl(_) => E2eiCrl::delete(tx, id.as_id()).await,
-        #[cfg(feature = "proteus-keystore")]
-        id @ EntityId::ProteusSession(_) => ProteusSession::delete(tx, id.as_id()).await,
-        #[cfg(feature = "proteus-keystore")]
-        id @ EntityId::ProteusIdentity(_) => ProteusIdentity::delete(tx, id.as_id()).await,
-        #[cfg(feature = "proteus-keystore")]
-        id @ EntityId::ProteusPrekey(_) => ProteusPrekey::delete(tx, id.as_id()).await,
-    }
-}
diff --git a/keystore/src/transaction/dynamic_dispatch/entity.rs b/keystore/src/transaction/dynamic_dispatch/entity.rs
new file mode 100644
index 0000000000..9daa2f1319
--- /dev/null
+++ b/keystore/src/transaction/dynamic_dispatch/entity.rs
@@ -0,0 +1,89 @@
+#[cfg(target_family = "wasm")]
+use crate::entities::E2eiRefreshToken;
+#[cfg(feature = "proteus-keystore")]
+use crate::entities::{ProteusIdentity, ProteusPrekey, ProteusSession};
+use crate::{
+    CryptoKeystoreResult,
+    connection::TransactionWrapper,
+    entities::{
+        ConsumerData, E2eiAcmeCA, E2eiCrl, E2eiIntermediateCert, MlsPendingMessage, PersistedMlsGroup,
+        PersistedMlsPendingGroup, StoredBufferedCommit, StoredCredential, StoredE2eiEnrollment,
+        StoredEncryptionKeyPair, StoredEpochEncryptionKeypair, StoredHpkePrivateKey, StoredKeypackage, StoredPskBundle,
+    },
+    traits::{EntityDatabaseMutation as _, UniqueEntityExt as _},
+};
+
+#[derive(Debug)]
+pub(crate) enum Entity<'a> {
+    ConsumerData(&'a ConsumerData),
+    HpkePrivateKey(&'a StoredHpkePrivateKey),
+    StoredKeypackage(&'a StoredKeypackage),
+    PskBundle(&'a StoredPskBundle),
+    EncryptionKeyPair(&'a StoredEncryptionKeyPair),
+    StoredEpochEncryptionKeypair(&'a StoredEpochEncryptionKeypair),
+    StoredCredential(&'a StoredCredential),
+    StoredBufferedCommit(&'a StoredBufferedCommit),
+    PersistedMlsGroup(&'a PersistedMlsGroup),
+    PersistedMlsPendingGroup(&'a PersistedMlsPendingGroup),
+    MlsPendingMessage(&'a MlsPendingMessage),
+    StoredE2eiEnrollment(&'a StoredE2eiEnrollment),
+    #[cfg(target_family = "wasm")]
+    E2eiRefreshToken(&'a E2eiRefreshToken),
+    E2eiAcmeCA(&'a E2eiAcmeCA),
+    E2eiIntermediateCert(&'a E2eiIntermediateCert),
+    E2eiCrl(&'a E2eiCrl),
+    #[cfg(feature = "proteus-keystore")]
+    ProteusIdentity(&'a ProteusIdentity),
+    #[cfg(feature = "proteus-keystore")]
+    ProteusPrekey(&'a ProteusPrekey),
+    #[cfg(feature = "proteus-keystore")]
+    ProteusSession(&'a ProteusSession),
+}
+
+impl<'a, E> From<&'a E> for Entity<'a>
+where
+    E: crate::traits::Entity,
+{
+    #[inline]
+    fn from(value: &'a E) -> Self {
+        value.to_transaction_entity()
+    }
+}
+
+impl Entity<'_> {
+    pub(crate) async fn execute_save(&self, tx: &TransactionWrapper<'_>) -> CryptoKeystoreResult<()> {
+        match self {
+            Entity::ConsumerData(consumer_data) => {
+                consumer_data.set_and_replace(tx).await;
+                Ok(())
+            }
+            Entity::HpkePrivateKey(mls_hpke_private_key) => mls_hpke_private_key.save(tx).await,
+            Entity::StoredKeypackage(mls_key_package) => mls_key_package.save(tx).await,
+            Entity::PskBundle(mls_psk_bundle) => mls_psk_bundle.save(tx).await,
+            Entity::EncryptionKeyPair(mls_encryption_key_pair) => mls_encryption_key_pair.save(tx).await,
+            Entity::StoredEpochEncryptionKeypair(mls_epoch_encryption_key_pair) => {
+                mls_epoch_encryption_key_pair.save(tx).await
+            }
+            Entity::StoredCredential(mls_credential) => mls_credential.save(tx).await,
+            Entity::StoredBufferedCommit(mls_pending_commit) => mls_pending_commit.save(tx).await,
+            Entity::PersistedMlsGroup(persisted_mls_group) => persisted_mls_group.save(tx).await,
+            Entity::PersistedMlsPendingGroup(persisted_mls_pending_group) => persisted_mls_pending_group.save(tx).await,
+            Entity::MlsPendingMessage(mls_pending_message) => mls_pending_message.save(tx).await,
+            Entity::StoredE2eiEnrollment(e2ei_enrollment) => e2ei_enrollment.save(tx).await,
+            #[cfg(target_family = "wasm")]
+            Entity::E2eiRefreshToken(e2ei_refresh_token) => e2ei_refresh_token.replace(tx).await,
+            Entity::E2eiAcmeCA(e2ei_acme_ca) => {
+                e2ei_acme_ca.set_and_replace(tx).await;
+                Ok(())
+            }
+            Entity::E2eiIntermediateCert(e2ei_intermediate_cert) => e2ei_intermediate_cert.save(tx).await,
+            Entity::E2eiCrl(e2ei_crl) => e2ei_crl.save(tx).await,
+            #[cfg(feature = "proteus-keystore")]
+            Entity::ProteusSession(record) => record.save(tx).await,
+            #[cfg(feature = "proteus-keystore")]
+            Entity::ProteusIdentity(record) => record.save(tx).await,
+            #[cfg(feature = "proteus-keystore")]
+            Entity::ProteusPrekey(record) => record.save(tx).await,
+        }
+    }
+}
diff --git a/keystore/src/transaction/dynamic_dispatch/entity_id.rs b/keystore/src/transaction/dynamic_dispatch/entity_id.rs
new file mode 100644
index 0000000000..1d38313485
--- /dev/null
+++ b/keystore/src/transaction/dynamic_dispatch/entity_id.rs
@@ -0,0 +1,91 @@
+#[cfg(target_family = "wasm")]
+use crate::entities::E2eiRefreshToken;
+#[cfg(feature = "proteus-keystore")]
+use crate::entities::{ProteusIdentity, ProteusPrekey, ProteusSession};
+use crate::{
+    CryptoKeystoreError, CryptoKeystoreResult,
+    connection::TransactionWrapper,
+    entities::{
+        E2eiAcmeCA, E2eiCrl, E2eiIntermediateCert, MlsPendingMessage, PersistedMlsGroup, PersistedMlsPendingGroup,
+        StoredBufferedCommit, StoredCredential, StoredE2eiEnrollment, StoredEncryptionKeyPair,
+        StoredEpochEncryptionKeypair, StoredHpkePrivateKey, StoredKeypackage, StoredPskBundle,
+    },
+    traits::{Entity, EntityDatabaseMutation, KeyType as _, OwnedKeyType as _},
+    transaction::dynamic_dispatch::EntityType,
+};
+
+#[derive(Debug, Clone, PartialEq, Eq, Hash)]
+pub(crate) struct EntityId {
+    typ: EntityType,
+    id: Vec<u8>,
+}
+
+impl EntityId {
+    fn primary_key<E>(&self) -> CryptoKeystoreResult<E::PrimaryKey>
+    where
+        E: Entity,
+    {
+        <E as Entity>::PrimaryKey::from_bytes(&self.id)
+            .ok_or(CryptoKeystoreError::InvalidPrimaryKeyBytes(self.typ.collection_name()))
+    }
+
+    pub(crate) fn from_entity<E>(entity: &E) -> Self
+    where
+        E: Entity,
+    {
+        // assumption: nobody outside this crate will ever implement `Entity` on a foreign type
+        let typ =
+            EntityType::from_collection_name(E::COLLECTION_NAME).expect("all entities have a valid collection name");
+        let id = entity.primary_key().bytes().into_owned();
+        Self { typ, id }
+    }
+
+    pub(crate) async fn execute_delete(&self, tx: &TransactionWrapper<'_>) -> CryptoKeystoreResult<bool> {
+        match self.typ {
+            EntityType::HpkePrivateKey => {
+                StoredHpkePrivateKey::delete(tx, &self.primary_key::<StoredHpkePrivateKey>()?).await
+            }
+            EntityType::KeyPackage => StoredKeypackage::delete(tx, &self.primary_key::<StoredKeypackage>()?).await,
+            EntityType::PskBundle => StoredPskBundle::delete(tx, &self.primary_key::<StoredPskBundle>()?).await,
+            EntityType::EncryptionKeyPair => {
+                StoredEncryptionKeyPair::delete(tx, &self.primary_key::<StoredEncryptionKeyPair>()?).await
+            }
+            EntityType::EpochEncryptionKeyPair => {
+                StoredEpochEncryptionKeypair::delete(tx, &self.primary_key::<StoredEpochEncryptionKeypair>()?).await
+            }
+            EntityType::StoredCredential => {
+                StoredCredential::delete(tx, &self.primary_key::<StoredCredential>()?).await
+            }
+            EntityType::StoredBufferedCommit => {
+                StoredBufferedCommit::delete(tx, &self.primary_key::<StoredBufferedCommit>()?).await
+            }
+            EntityType::PersistedMlsGroup => {
+                PersistedMlsGroup::delete(tx, &self.primary_key::<PersistedMlsGroup>()?).await
+            }
+            EntityType::PersistedMlsPendingGroup => {
+                PersistedMlsPendingGroup::delete(tx, &self.primary_key::<PersistedMlsPendingGroup>()?).await
+            }
+            EntityType::MlsPendingMessage => {
+                MlsPendingMessage::delete(tx, &self.primary_key::<MlsPendingMessage>()?).await
+            }
+            EntityType::StoredE2eiEnrollment => {
+                StoredE2eiEnrollment::delete(tx, &self.primary_key::<StoredE2eiEnrollment>()?).await
+            }
+            #[cfg(target_family = "wasm")]
+            EntityType::E2eiRefreshToken => {
+                E2eiRefreshToken::delete(tx, &self.primary_key::<E2eiRefreshToken>()?).await
+            }
+            EntityType::E2eiAcmeCA => E2eiAcmeCA::delete(tx, &self.primary_key::<E2eiAcmeCA>()?).await,
+            EntityType::E2eiIntermediateCert => {
+                E2eiIntermediateCert::delete(tx, &self.primary_key::<E2eiIntermediateCert>()?).await
+            }
+            EntityType::E2eiCrl => E2eiCrl::delete(tx, &self.primary_key::<E2eiCrl>()?).await,
+            #[cfg(feature = "proteus-keystore")]
+            EntityType::ProteusSession => ProteusSession::delete(tx, &self.primary_key::<ProteusSession>()?).await,
+            #[cfg(feature = "proteus-keystore")]
+            EntityType::ProteusIdentity => ProteusIdentity::delete(tx, &self.primary_key::<ProteusIdentity>()?).await,
+            #[cfg(feature = "proteus-keystore")]
+            EntityType::ProteusPrekey => ProteusPrekey::delete(tx, &self.primary_key::<ProteusPrekey>()?).await,
+        }
+    }
+}
diff --git a/keystore/src/transaction/dynamic_dispatch/entity_type.rs b/keystore/src/transaction/dynamic_dispatch/entity_type.rs
new file mode 100644
index 0000000000..af224af14d
--- /dev/null
+++ b/keystore/src/transaction/dynamic_dispatch/entity_type.rs
@@ -0,0 +1,95 @@
+#[cfg(target_family = "wasm")]
+use crate::entities::E2eiRefreshToken;
+#[cfg(feature = "proteus-keystore")]
+use crate::entities::{ProteusIdentity, ProteusPrekey, ProteusSession};
+use crate::{
+    entities::{
+        E2eiAcmeCA, E2eiCrl, E2eiIntermediateCert, MlsPendingMessage, PersistedMlsGroup, PersistedMlsPendingGroup,
+        StoredBufferedCommit, StoredCredential, StoredE2eiEnrollment, StoredEncryptionKeyPair,
+        StoredEpochEncryptionKeypair, StoredHpkePrivateKey, StoredKeypackage, StoredPskBundle,
+    },
+    traits::EntityBase as _,
+};
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
+pub(crate) enum EntityType {
+    HpkePrivateKey,
+    KeyPackage,
+    PskBundle,
+    EncryptionKeyPair,
+    EpochEncryptionKeyPair,
+    StoredCredential,
+    StoredBufferedCommit,
+    PersistedMlsGroup,
+    PersistedMlsPendingGroup,
+    MlsPendingMessage,
+    StoredE2eiEnrollment,
+    #[cfg(target_family = "wasm")]
+    E2eiRefreshToken,
+    E2eiAcmeCA,
+    E2eiIntermediateCert,
+    E2eiCrl,
+    #[cfg(feature = "proteus-keystore")]
+    ProteusIdentity,
+    #[cfg(feature = "proteus-keystore")]
+    ProteusPrekey,
+    #[cfg(feature = "proteus-keystore")]
+    ProteusSession,
+}
+
+impl EntityType {
+    pub(crate) fn from_collection_name(collection_name: &'static str) -> Option<Self> {
+        match collection_name {
+            StoredHpkePrivateKey::COLLECTION_NAME => Some(Self::HpkePrivateKey),
+            StoredKeypackage::COLLECTION_NAME => Some(Self::KeyPackage),
+            StoredPskBundle::COLLECTION_NAME => Some(Self::PskBundle),
+            StoredEncryptionKeyPair::COLLECTION_NAME => Some(Self::EncryptionKeyPair),
+            StoredEpochEncryptionKeypair::COLLECTION_NAME => Some(Self::EpochEncryptionKeyPair),
+            StoredBufferedCommit::COLLECTION_NAME => Some(Self::StoredBufferedCommit),
+            PersistedMlsGroup::COLLECTION_NAME => Some(Self::PersistedMlsGroup),
+            PersistedMlsPendingGroup::COLLECTION_NAME => Some(Self::PersistedMlsPendingGroup),
+            StoredCredential::COLLECTION_NAME => Some(Self::StoredCredential),
+            MlsPendingMessage::COLLECTION_NAME => Some(Self::MlsPendingMessage),
+            StoredE2eiEnrollment::COLLECTION_NAME => Some(Self::StoredE2eiEnrollment),
+            E2eiCrl::COLLECTION_NAME => Some(Self::E2eiCrl),
+            E2eiAcmeCA::COLLECTION_NAME => Some(Self::E2eiAcmeCA),
+            #[cfg(target_family = "wasm")]
+            E2eiRefreshToken::COLLECTION_NAME => Some(Self::E2eiRefreshToken),
+            E2eiIntermediateCert::COLLECTION_NAME => Some(Self::E2eiIntermediateCert),
+            #[cfg(feature = "proteus-keystore")]
+            ProteusIdentity::COLLECTION_NAME => Some(Self::ProteusIdentity),
+            #[cfg(feature = "proteus-keystore")]
+            ProteusPrekey::COLLECTION_NAME => Some(Self::ProteusPrekey),
+            #[cfg(feature = "proteus-keystore")]
+            ProteusSession::COLLECTION_NAME => Some(Self::ProteusSession),
+            _ => None,
+        }
+    }
+
+    pub(crate) fn collection_name(&self) -> &'static str {
+        match self {
+            Self::KeyPackage => StoredKeypackage::COLLECTION_NAME,
+            Self::PskBundle => StoredPskBundle::COLLECTION_NAME,
+            Self::EncryptionKeyPair => StoredEncryptionKeyPair::COLLECTION_NAME,
+            Self::EpochEncryptionKeyPair => StoredEpochEncryptionKeypair::COLLECTION_NAME,
+            Self::StoredCredential => StoredCredential::COLLECTION_NAME,
+            Self::StoredBufferedCommit => StoredBufferedCommit::COLLECTION_NAME,
+            Self::PersistedMlsGroup => PersistedMlsGroup::COLLECTION_NAME,
+            Self::PersistedMlsPendingGroup => PersistedMlsPendingGroup::COLLECTION_NAME,
+            Self::MlsPendingMessage => MlsPendingMessage::COLLECTION_NAME,
+            Self::StoredE2eiEnrollment => StoredE2eiEnrollment::COLLECTION_NAME,
+            #[cfg(target_family = "wasm")]
+            Self::E2eiRefreshToken => E2eiRefreshToken::COLLECTION_NAME,
+            Self::E2eiAcmeCA => E2eiAcmeCA::COLLECTION_NAME,
+            Self::E2eiIntermediateCert => E2eiIntermediateCert::COLLECTION_NAME,
+            Self::E2eiCrl => E2eiCrl::COLLECTION_NAME,
+            #[cfg(feature = "proteus-keystore")]
+            Self::ProteusIdentity => ProteusIdentity::COLLECTION_NAME,
+            #[cfg(feature = "proteus-keystore")]
+            Self::ProteusPrekey => ProteusPrekey::COLLECTION_NAME,
+            #[cfg(feature = "proteus-keystore")]
+            Self::ProteusSession => ProteusSession::COLLECTION_NAME,
+            Self::HpkePrivateKey => StoredHpkePrivateKey::COLLECTION_NAME,
+        }
+    }
+}
diff --git a/keystore/src/transaction/dynamic_dispatch/mod.rs b/keystore/src/transaction/dynamic_dispatch/mod.rs
new file mode 100644
index 0000000000..514a90c358
--- /dev/null
+++ b/keystore/src/transaction/dynamic_dispatch/mod.rs
@@ -0,0 +1,7 @@
+mod entity;
+mod entity_id;
+mod entity_type;
+
+pub(crate) use entity::Entity;
+pub(crate) use entity_id::EntityId;
+pub(crate) use entity_type::EntityType;

From 2975abbb90e343d462e1ef187ceedd4430663718 Mon Sep 17 00:00:00 2001
From: Peter Goodspeed-Niklaus <peter.goodspeed-niklaus@wire.com>
Date: Thu, 18 Dec 2025 16:33:17 +0100
Subject: [PATCH 05/48] refactor(keystore): hold onto entity items in dynamic
 dispatch

We used to do some crazy thing involving unencrypted serialization
for some reason. It's much simpler this way: when an entity turns
itself into a dynamic entity, we just keep it. We do Arc-wrap it,
because there are times we want to iterate over a bunch of owned
instances and don't want the expense of cloning these potentially-
large items every time. But from there the API is straightforward.

While it was cool we could make the borrowing approach work for
dynamic entities, in practice it turned out that in general when
we hand an entity to the database to save, we really want it to
take over ownership, which means this pattern is better.
---
 keystore/src/traits/entity_base.rs            | 11 +++
 .../transaction/dynamic_dispatch/entity.rs    | 89 ++++++++++++-------
 2 files changed, 68 insertions(+), 32 deletions(-)

diff --git a/keystore/src/traits/entity_base.rs b/keystore/src/traits/entity_base.rs
index c11cb3bf75..df7015d6ed 100644
--- a/keystore/src/traits/entity_base.rs
+++ b/keystore/src/traits/entity_base.rs
@@ -1,3 +1,5 @@
+use std::{any::Any, sync::Arc};
+
 use crate::connection::DatabaseConnection;
 
 /// A supertrait that all entities must implement. This handles multiplexing over the two different database backends.
@@ -16,5 +18,14 @@ pub trait EntityBase: 'static + Sized {
         as_dyn_any.downcast_ref()
     }
 
+    fn downcast_arc<T>(self: Arc<Self>) -> Option<Arc<T>>
+    where
+        Self: Send + Sync,
+        T: EntityBase + Send + Sync,
+    {
+        let as_dyn_any = self as Arc<dyn Any + Send + Sync>;
+        as_dyn_any.downcast().ok()
+    }
+
     fn to_transaction_entity(self) -> crate::transaction::dynamic_dispatch::Entity;
 }
diff --git a/keystore/src/transaction/dynamic_dispatch/entity.rs b/keystore/src/transaction/dynamic_dispatch/entity.rs
index 9daa2f1319..d048ba6d90 100644
--- a/keystore/src/transaction/dynamic_dispatch/entity.rs
+++ b/keystore/src/transaction/dynamic_dispatch/entity.rs
@@ -1,3 +1,5 @@
+use std::sync::Arc;
+
 #[cfg(target_family = "wasm")]
 use crate::entities::E2eiRefreshToken;
 #[cfg(feature = "proteus-keystore")]
@@ -10,51 +12,74 @@ use crate::{
         PersistedMlsPendingGroup, StoredBufferedCommit, StoredCredential, StoredE2eiEnrollment,
         StoredEncryptionKeyPair, StoredEpochEncryptionKeypair, StoredHpkePrivateKey, StoredKeypackage, StoredPskBundle,
     },
-    traits::{EntityDatabaseMutation as _, UniqueEntityExt as _},
+    traits::{EntityBase, EntityDatabaseMutation as _, UniqueEntityExt as _},
 };
 
 #[derive(Debug)]
-pub(crate) enum Entity<'a> {
-    ConsumerData(&'a ConsumerData),
-    HpkePrivateKey(&'a StoredHpkePrivateKey),
-    StoredKeypackage(&'a StoredKeypackage),
-    PskBundle(&'a StoredPskBundle),
-    EncryptionKeyPair(&'a StoredEncryptionKeyPair),
-    StoredEpochEncryptionKeypair(&'a StoredEpochEncryptionKeypair),
-    StoredCredential(&'a StoredCredential),
-    StoredBufferedCommit(&'a StoredBufferedCommit),
-    PersistedMlsGroup(&'a PersistedMlsGroup),
-    PersistedMlsPendingGroup(&'a PersistedMlsPendingGroup),
-    MlsPendingMessage(&'a MlsPendingMessage),
-    StoredE2eiEnrollment(&'a StoredE2eiEnrollment),
+pub enum Entity {
+    ConsumerData(Arc<ConsumerData>),
+    HpkePrivateKey(Arc<StoredHpkePrivateKey>),
+    StoredKeypackage(Arc<StoredKeypackage>),
+    PskBundle(Arc<StoredPskBundle>),
+    EncryptionKeyPair(Arc<StoredEncryptionKeyPair>),
+    StoredEpochEncryptionKeypair(Arc<StoredEpochEncryptionKeypair>),
+    StoredCredential(Arc<StoredCredential>),
+    StoredBufferedCommit(Arc<StoredBufferedCommit>),
+    PersistedMlsGroup(Arc<PersistedMlsGroup>),
+    PersistedMlsPendingGroup(Arc<PersistedMlsPendingGroup>),
+    MlsPendingMessage(Arc<MlsPendingMessage>),
+    StoredE2eiEnrollment(Arc<StoredE2eiEnrollment>),
     #[cfg(target_family = "wasm")]
-    E2eiRefreshToken(&'a E2eiRefreshToken),
-    E2eiAcmeCA(&'a E2eiAcmeCA),
-    E2eiIntermediateCert(&'a E2eiIntermediateCert),
-    E2eiCrl(&'a E2eiCrl),
+    E2eiRefreshToken(Arc<E2eiRefreshToken>),
+    E2eiAcmeCA(Arc<E2eiAcmeCA>),
+    E2eiIntermediateCert(Arc<E2eiIntermediateCert>),
+    E2eiCrl(Arc<E2eiCrl>),
     #[cfg(feature = "proteus-keystore")]
-    ProteusIdentity(&'a ProteusIdentity),
+    ProteusIdentity(Arc<ProteusIdentity>),
     #[cfg(feature = "proteus-keystore")]
-    ProteusPrekey(&'a ProteusPrekey),
+    ProteusPrekey(Arc<ProteusPrekey>),
     #[cfg(feature = "proteus-keystore")]
-    ProteusSession(&'a ProteusSession),
+    ProteusSession(Arc<ProteusSession>),
 }
 
-impl<'a, E> From<&'a E> for Entity<'a>
-where
-    E: crate::traits::Entity,
-{
-    #[inline]
-    fn from(value: &'a E) -> Self {
-        value.to_transaction_entity()
+impl Entity {
+    /// Downcast this entity to an instance of the requested type.
+    ///
+    /// This increments the smart pointer counter instead of cloning the potentially large item instance.
+    pub(crate) fn downcast<E>(&self) -> Option<Arc<E>>
+    where
+        E: EntityBase + Send + Sync,
+    {
+        match self {
+            Entity::ConsumerData(consumer_data) => consumer_data.clone().downcast_arc(),
+            Entity::HpkePrivateKey(stored_hpke_private_key) => stored_hpke_private_key.clone().downcast_arc(),
+            Entity::StoredKeypackage(stored_keypackage) => stored_keypackage.clone().downcast_arc(),
+            Entity::PskBundle(stored_psk_bundle) => stored_psk_bundle.clone().downcast_arc(),
+            Entity::EncryptionKeyPair(stored_encryption_key_pair) => stored_encryption_key_pair.clone().downcast_arc(),
+            Entity::StoredEpochEncryptionKeypair(stored_epoch_encryption_keypair) => {
+                stored_epoch_encryption_keypair.clone().downcast_arc()
+            }
+            Entity::StoredCredential(stored_credential) => stored_credential.clone().downcast_arc(),
+            Entity::StoredBufferedCommit(stored_buffered_commit) => stored_buffered_commit.clone().downcast_arc(),
+            Entity::PersistedMlsGroup(persisted_mls_group) => persisted_mls_group.clone().downcast_arc(),
+            Entity::PersistedMlsPendingGroup(persisted_mls_pending_group) => {
+                persisted_mls_pending_group.clone().downcast_arc()
+            }
+            Entity::MlsPendingMessage(mls_pending_message) => mls_pending_message.clone().downcast_arc(),
+            Entity::StoredE2eiEnrollment(stored_e2ei_enrollment) => stored_e2ei_enrollment.clone().downcast_arc(),
+            Entity::E2eiAcmeCA(e2ei_acme_ca) => e2ei_acme_ca.clone().downcast_arc(),
+            Entity::E2eiIntermediateCert(e2ei_intermediate_cert) => e2ei_intermediate_cert.clone().downcast_arc(),
+            Entity::E2eiCrl(e2ei_crl) => e2ei_crl.clone().downcast_arc(),
+            Entity::ProteusIdentity(proteus_identity) => proteus_identity.clone().downcast_arc(),
+            Entity::ProteusPrekey(proteus_prekey) => proteus_prekey.clone().downcast_arc(),
+            Entity::ProteusSession(proteus_session) => proteus_session.clone().downcast_arc(),
+        }
     }
-}
 
-impl Entity<'_> {
     pub(crate) async fn execute_save(&self, tx: &TransactionWrapper<'_>) -> CryptoKeystoreResult<()> {
         match self {
             Entity::ConsumerData(consumer_data) => {
-                consumer_data.set_and_replace(tx).await;
+                consumer_data.set_and_replace(tx).await?;
                 Ok(())
             }
             Entity::HpkePrivateKey(mls_hpke_private_key) => mls_hpke_private_key.save(tx).await,
@@ -73,7 +98,7 @@ impl Entity<'_> {
             #[cfg(target_family = "wasm")]
             Entity::E2eiRefreshToken(e2ei_refresh_token) => e2ei_refresh_token.replace(tx).await,
             Entity::E2eiAcmeCA(e2ei_acme_ca) => {
-                e2ei_acme_ca.set_and_replace(tx).await;
+                e2ei_acme_ca.set_and_replace(tx).await?;
                 Ok(())
             }
             Entity::E2eiIntermediateCert(e2ei_intermediate_cert) => e2ei_intermediate_cert.save(tx).await,

From 291336dcff581c89533cc24dccd635c91695bd2f Mon Sep 17 00:00:00 2001
From: Peter Goodspeed-Niklaus <peter.goodspeed-niklaus@wire.com>
Date: Thu, 18 Dec 2025 17:46:57 +0100
Subject: [PATCH 06/48] refactor(keystore): simplify parameter type for
 borrowing trait methods

The traits didn't previously allow for multiple distinct implementations
of borrowing types sharing an owned type (I think) anyway, and we
already had an associated type naming a specific borrowed primary
key, so it was pointless overkill to accept a complicated generic
parameter type in these methods. Better and simpler to just use the
associated type.
---
 keystore/src/traits/entity.rs                   |  8 +++++---
 keystore/src/traits/entity_database_mutation.rs | 11 +++++------
 2 files changed, 10 insertions(+), 9 deletions(-)

diff --git a/keystore/src/traits/entity.rs b/keystore/src/traits/entity.rs
index 61aa5ee9da..ee817f35c3 100644
--- a/keystore/src/traits/entity.rs
+++ b/keystore/src/traits/entity.rs
@@ -61,8 +61,10 @@ pub trait BorrowPrimaryKey: Entity {
     ///
     /// The type signature here is somewhat complicated, but it breaks down simply: if our primary key is something
     /// like `Vec<u8>`, we want to be able to use this method even if what we have on hand is `&[u8]`.
-    async fn get_borrowed<Q>(conn: &mut Self::ConnectionType, key: &Q) -> CryptoKeystoreResult<Option<Self>>
+    async fn get_borrowed(
+        conn: &mut Self::ConnectionType,
+        key: &Self::BorrowedPrimaryKey,
+    ) -> CryptoKeystoreResult<Option<Self>>
     where
-        Self::PrimaryKey: Borrow<Q>,
-        Q: KeyType;
+        for<'pk> &'pk Self::BorrowedPrimaryKey: KeyType;
 }
diff --git a/keystore/src/traits/entity_database_mutation.rs b/keystore/src/traits/entity_database_mutation.rs
index 22f3820b3e..e094f466af 100644
--- a/keystore/src/traits/entity_database_mutation.rs
+++ b/keystore/src/traits/entity_database_mutation.rs
@@ -4,7 +4,7 @@ use async_trait::async_trait;
 
 use crate::{
     CryptoKeystoreResult,
-    traits::{Entity, KeyType},
+    traits::{BorrowPrimaryKey, Entity, KeyType},
 };
 
 /// Extend an [`Entity`] with db-mutating operations which can be performed when provided with a transaction.
@@ -45,16 +45,15 @@ pub trait EntityDatabaseMutation<'a>: Entity<ConnectionType = crate::connection:
 /// Extend an [`Entity`] with db-mutating operations which can be performed when provided with a transaction.
 #[cfg_attr(target_family = "wasm", async_trait(?Send))]
 #[cfg_attr(not(target_family = "wasm"), async_trait)]
-pub trait EntityDeleteBorrowed<'a>: EntityDatabaseMutation<'a> {
+pub trait EntityDeleteBorrowed<'a>: EntityDatabaseMutation<'a> + BorrowPrimaryKey {
     /// Delete an entity by a borrowed form of its primary key.
     ///
     /// The type signature here is somewhat complicated, but it breaks down simply: if our primary key is something
     /// like `Vec<u8>`, we want to be able to use this method even if what we have on hand is `&[u8]`.
-    async fn delete_borrowed<Q>(
+    async fn delete_borrowed(
         tx: &<Self as EntityDatabaseMutation<'a>>::Transaction,
-        id: &Q,
+        id: &<Self as BorrowPrimaryKey>::BorrowedPrimaryKey,
     ) -> CryptoKeystoreResult<bool>
     where
-        Self::PrimaryKey: Borrow<Q>,
-        Q: KeyType;
+        for<'pk> &'pk Self::BorrowedPrimaryKey: KeyType;
 }

From 807e7908691eae8582ecb77a529eca587f642f94 Mon Sep 17 00:00:00 2001
From: Peter Goodspeed-Niklaus <peter.goodspeed-niklaus@wire.com>
Date: Thu, 18 Dec 2025 16:38:46 +0100
Subject: [PATCH 07/48] chore(keystore): improve trait item docs

---
 keystore/src/traits/entity_base.rs              | 6 ++++++
 keystore/src/traits/entity_database_mutation.rs | 3 +++
 2 files changed, 9 insertions(+)

diff --git a/keystore/src/traits/entity_base.rs b/keystore/src/traits/entity_base.rs
index df7015d6ed..64413b5bff 100644
--- a/keystore/src/traits/entity_base.rs
+++ b/keystore/src/traits/entity_base.rs
@@ -7,6 +7,12 @@ use crate::connection::DatabaseConnection;
 /// This trait should be removed once the persistence layers are unified. See WPB-16241.
 pub trait EntityBase: 'static + Sized {
     type ConnectionType: for<'a> DatabaseConnection<'a>;
+
+    /// Entities which implement `EntityDatabaseMutation` have a `pre_save` method which might generate or
+    /// update some fields of the item. The canonical example is an `updated_at` field.
+    ///
+    /// This type must contain a copy of each modification to the item, so that the caller of a `.save(entity)`
+    /// function can know what has changed and what the new values are.
     type AutoGeneratedFields: Default;
 
     /// Beware: if you change the value of this constant on any WASM entity, you'll need to do a data migration
diff --git a/keystore/src/traits/entity_database_mutation.rs b/keystore/src/traits/entity_database_mutation.rs
index e094f466af..d8c599135c 100644
--- a/keystore/src/traits/entity_database_mutation.rs
+++ b/keystore/src/traits/entity_database_mutation.rs
@@ -17,6 +17,9 @@ pub trait EntityDatabaseMutation<'a>: Entity<ConnectionType = crate::connection:
     ///
     /// This will be called by the transaction's implementation of `save_mut` and should not be called
     /// as part of the `save` implementation.
+    ///
+    /// Returns, in a format specific to the entity, all the data which has been generated or updated
+    /// within this method.
     async fn pre_save(&mut self) -> CryptoKeystoreResult<Self::AutoGeneratedFields> {
         Ok(Default::default())
     }

From 7c2fed6051fa5edd28275f48c3c88be6631e2738 Mon Sep 17 00:00:00 2001
From: Peter Goodspeed-Niklaus <peter.goodspeed-niklaus@wire.com>
Date: Fri, 19 Dec 2025 11:17:28 +0100
Subject: [PATCH 08/48] chore(keystore): adjust old `UniqueEntity` trait to
 still compile

---
 keystore/src/entities/mls.rs | 9 ++++++---
 1 file changed, 6 insertions(+), 3 deletions(-)

diff --git a/keystore/src/entities/mls.rs b/keystore/src/entities/mls.rs
index 857118686d..a909fe3bde 100644
--- a/keystore/src/entities/mls.rs
+++ b/keystore/src/entities/mls.rs
@@ -312,7 +312,7 @@ pub trait UniqueEntity: EntityBase<ConnectionType = crate::connection::KeystoreD
         }
     }
 
-    async fn find_all(conn: &mut Self::ConnectionType, _params: EntityFindParams) -> CryptoKeystoreResult<Vec<Self>> {
+    async fn find_all(conn: &mut Self::ConnectionType) -> CryptoKeystoreResult<Vec<Self>> {
         match Self::find_unique(conn).await {
             Ok(record) => Ok(vec![record]),
             Err(CryptoKeystoreError::NotFound(..)) => Ok(vec![]),
@@ -366,7 +366,10 @@ pub trait UniqueEntity: EntityBase<ConnectionType = crate::connection::KeystoreD
 
 #[cfg_attr(target_family = "wasm", async_trait::async_trait(?Send))]
 #[cfg_attr(not(target_family = "wasm"), async_trait::async_trait)]
-impl<T: UniqueEntity + Send + Sync> EntityTransactionExt for T {
+impl<T> crate::entities::EntityTransactionExt for T
+where
+    T: crate::entities::Entity + UniqueEntity + Send + Sync,
+{
     #[cfg(not(target_family = "wasm"))]
     async fn save(&self, tx: &TransactionWrapper<'_>) -> CryptoKeystoreResult<()> {
         self.replace(tx).await
@@ -380,7 +383,7 @@ impl<T: UniqueEntity + Send + Sync> EntityTransactionExt for T {
     #[cfg(not(target_family = "wasm"))]
     async fn delete_fail_on_missing_id(
         _: &TransactionWrapper<'_>,
-        _id: StringEntityId<'_>,
+        _id: crate::entities::StringEntityId<'_>,
     ) -> CryptoKeystoreResult<()> {
         Err(CryptoKeystoreError::NotImplemented)
     }

From 9a3893ccb4d26235d73ae15ef4e71e45cd3c7fa6 Mon Sep 17 00:00:00 2001
From: Peter Goodspeed-Niklaus <peter.goodspeed-niklaus@wire.com>
Date: Fri, 19 Dec 2025 11:50:02 +0100
Subject: [PATCH 09/48] chore(keystore): update old blanket `impl Entity for
 UniqueEntity`

---
 keystore/src/entities/mod.rs | 8 +++++---
 1 file changed, 5 insertions(+), 3 deletions(-)

diff --git a/keystore/src/entities/mod.rs b/keystore/src/entities/mod.rs
index fe82483be4..b934ff11cb 100644
--- a/keystore/src/entities/mod.rs
+++ b/keystore/src/entities/mod.rs
@@ -324,13 +324,15 @@ cfg_if::cfg_if! {
         }
 
         #[async_trait::async_trait]
-        impl<T: UniqueEntity> Entity for T {
+        impl<T> Entity for T
+        where T : UniqueEntity + crate::entities::EntityBase<ConnectionType = crate::connection::KeystoreDatabaseConnection>
+        {
             fn id_raw(&self) -> &[u8] {
                 &[Self::ID as u8]
             }
 
-            async fn find_all(conn: &mut Self::ConnectionType, params: EntityFindParams) -> CryptoKeystoreResult<Vec<Self>> {
-                <Self as UniqueEntity>::find_all(conn, params).await
+            async fn find_all(conn: &mut Self::ConnectionType, _params: EntityFindParams) -> CryptoKeystoreResult<Vec<Self>> {
+                <Self as UniqueEntity>::find_all(conn).await
             }
 
             async fn find_one(conn: &mut Self::ConnectionType, _id: &StringEntityId) -> CryptoKeystoreResult<Option<Self>> {

From 3aa9d5cca917ca9f85893285f9f2fdab1fa8180d Mon Sep 17 00:00:00 2001
From: Peter Goodspeed-Niklaus <peter.goodspeed-niklaus@wire.com>
Date: Thu, 18 Dec 2025 16:48:27 +0100
Subject: [PATCH 10/48] chore(keystore): impl `BorrowPrimaryKey` for
 `MlsPendingMessage`

Pending messages have no real primary key and are always/only accessed
and removed in bulk, which turns out to be _really annoying_ with regard
to the new entity traits, which assumes that some primary key does exist.

For now we're faking the foreign id as the primary key so that encryption
works at all, and also for dynamic dispatch. But I'd really love to
rework this at some point in the future.
---
 .../platform/generic/mls/pending_message.rs   | 28 +++++++++++++++++--
 .../platform/wasm/mls/pending_message.rs      | 21 +++++++++++++-
 2 files changed, 45 insertions(+), 4 deletions(-)

diff --git a/keystore/src/entities/platform/generic/mls/pending_message.rs b/keystore/src/entities/platform/generic/mls/pending_message.rs
index ed832f27e4..94e7f80dc5 100644
--- a/keystore/src/entities/platform/generic/mls/pending_message.rs
+++ b/keystore/src/entities/platform/generic/mls/pending_message.rs
@@ -216,13 +216,18 @@ impl NewEntityBase for MlsPendingMessage {
     }
 }
 
-#[async_trait]
 /// Pending messages 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 it here to support `KeystoreTransaction::remove_pending_messages_by_conversation_id`.
+/// This is temporary! Post WPB-20844, we should remove that whole API and also reset the primary key type here to `()`.
+#[async_trait]
 impl NewEntity for MlsPendingMessage {
-    type PrimaryKey = ();
-    fn primary_key(&self) -> Self::PrimaryKey {}
+    type PrimaryKey = Vec<u8>;
+    fn primary_key(&self) -> Self::PrimaryKey {
+        self.foreign_id.clone()
+    }
 
     async fn get(conn: &mut Self::ConnectionType, key: &Self::PrimaryKey) -> CryptoKeystoreResult<Option<Self>> {
         panic!("cannot get `MlsPendingMessage` by primary key as it has no distinct primary key")
@@ -237,6 +242,23 @@ impl NewEntity for MlsPendingMessage {
     }
 }
 
+#[async_trait]
+impl BorrowPrimaryKey for MlsPendingMessage {
+    type BorrowedPrimaryKey = [u8];
+
+    fn borrow_primary_key(&self) -> &Self::BorrowedPrimaryKey {
+        &self.foreign_id
+    }
+
+    async fn get_borrowed<Q>(conn: &mut Self::ConnectionType, key: &Q) -> CryptoKeystoreResult<Option<Self>>
+    where
+        Self::PrimaryKey: Borrow<Q>,
+        Q: KeyType,
+    {
+        panic!("cannot get `MlsPendingMessage` by primary key as it has no distinct primary key")
+    }
+}
+
 #[async_trait]
 impl<'a> EntityDatabaseMutation<'a> for MlsPendingMessage {
     type Transaction = TransactionWrapper<'a>;
diff --git a/keystore/src/entities/platform/wasm/mls/pending_message.rs b/keystore/src/entities/platform/wasm/mls/pending_message.rs
index ef1ee2f5a7..4af774e43d 100644
--- a/keystore/src/entities/platform/wasm/mls/pending_message.rs
+++ b/keystore/src/entities/platform/wasm/mls/pending_message.rs
@@ -116,13 +116,13 @@ impl NewEntityBase for MlsPendingMessage {
     }
 }
 
-#[async_trait(?Send)]
 /// Pending messages 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 it as a byte vector in this impl in order for encryption and decryption
 /// to work.
+#[async_trait(?Send)]
 impl NewEntity for MlsPendingMessage {
     type PrimaryKey = Vec<u8>;
 
@@ -143,6 +143,25 @@ impl NewEntity for MlsPendingMessage {
     }
 }
 
+/// This implementation is purely to support `KeystoreTransaction::remove_pending_messages_by_conversation_id`;
+/// after WPB-20844, we should remove that whole API and also this implementation.
+#[async_trait(?Send)]
+impl BorrowPrimaryKey for MlsPendingMessage {
+    type BorrowedPrimaryKey = [u8];
+
+    fn borrow_primary_key(&self) -> &Self::BorrowedPrimaryKey {
+        &self.foreign_id
+    }
+
+    async fn get_borrowed<Q>(conn: &mut Self::ConnectionType, key: &Q) -> CryptoKeystoreResult<Option<Self>>
+    where
+        Self::PrimaryKey: Borrow<Q>,
+        Q: KeyType,
+    {
+        panic!("cannot get `MlsPendingMessage` by primary key as it has no distinct primary key")
+    }
+}
+
 #[async_trait(?Send)]
 impl<'a> EntityDatabaseMutation<'a> for MlsPendingMessage {
     type Transaction = TransactionWrapper<'a>;

From 86849c300d3141f09ca89a102004d4ed9cb468ad Mon Sep 17 00:00:00 2001
From: Peter Goodspeed-Niklaus <peter.goodspeed-niklaus@wire.com>
Date: Thu, 18 Dec 2025 16:51:34 +0100
Subject: [PATCH 11/48] chore(keystore): adjust impls of
 `to_transaction_entity`

---
 crypto-macros/src/entity_derive/derive_impl.rs                | 2 +-
 crypto-macros/src/entity_derive_new/derive_impl.rs            | 2 +-
 keystore/src/entities/platform/generic/general.rs             | 4 ++--
 keystore/src/entities/platform/generic/mls/credential.rs      | 4 ++--
 keystore/src/entities/platform/generic/mls/e2ei_acme_ca.rs    | 4 ++--
 .../src/entities/platform/generic/mls/encryption_keypair.rs   | 4 ++--
 .../src/entities/platform/generic/mls/hpke_private_key.rs     | 4 ++--
 keystore/src/entities/platform/generic/mls/pending_group.rs   | 4 ++--
 keystore/src/entities/platform/generic/mls/pending_message.rs | 4 ++--
 keystore/src/entities/platform/generic/mls/psk_bundle.rs      | 4 ++--
 keystore/src/entities/platform/generic/proteus/identity.rs    | 4 ++--
 keystore/src/entities/platform/generic/proteus/prekey.rs      | 4 ++--
 keystore/src/entities/platform/wasm/general.rs                | 4 ++--
 keystore/src/entities/platform/wasm/mls/credential.rs         | 4 ++--
 keystore/src/entities/platform/wasm/mls/e2ei_acme_ca.rs       | 4 ++--
 keystore/src/entities/platform/wasm/mls/encryption_keypair.rs | 4 ++--
 keystore/src/entities/platform/wasm/mls/hpke_private_key.rs   | 4 ++--
 keystore/src/entities/platform/wasm/mls/pending_group.rs      | 4 ++--
 keystore/src/entities/platform/wasm/mls/pending_message.rs    | 4 ++--
 keystore/src/entities/platform/wasm/mls/psk_bundle.rs         | 4 ++--
 keystore/src/entities/platform/wasm/mls/refresh_token.rs      | 4 ++--
 keystore/src/entities/platform/wasm/proteus/identity.rs       | 4 ++--
 keystore/src/entities/platform/wasm/proteus/prekey.rs         | 4 ++--
 23 files changed, 44 insertions(+), 44 deletions(-)

diff --git a/crypto-macros/src/entity_derive/derive_impl.rs b/crypto-macros/src/entity_derive/derive_impl.rs
index 9027794802..111a2c232a 100644
--- a/crypto-macros/src/entity_derive/derive_impl.rs
+++ b/crypto-macros/src/entity_derive/derive_impl.rs
@@ -38,7 +38,7 @@ impl KeyStoreEntityFlattened {
                 }
 
                 fn to_transaction_entity(self) -> crate::transaction::dynamic_dispatch::Entity {
-                    crate::transaction::dynamic_dispatch::Entity::#struct_name(self)
+                    crate::transaction::dynamic_dispatch::Entity::#struct_name(self.into())
                 }
             }
         }
diff --git a/crypto-macros/src/entity_derive_new/derive_impl.rs b/crypto-macros/src/entity_derive_new/derive_impl.rs
index 3596f111d8..1e7f377ad2 100644
--- a/crypto-macros/src/entity_derive_new/derive_impl.rs
+++ b/crypto-macros/src/entity_derive_new/derive_impl.rs
@@ -37,7 +37,7 @@ impl Entity {
                 const COLLECTION_NAME: &'static str = #collection_name;
 
                 fn to_transaction_entity(self) -> crate::transaction::dynamic_dispatch::Entity {
-                    crate::transaction::dynamic_dispatch::Entity::#struct_name(self)
+                    crate::transaction::dynamic_dispatch::Entity::#struct_name(self.into())
                 }
             }
         }
diff --git a/keystore/src/entities/platform/generic/general.rs b/keystore/src/entities/platform/generic/general.rs
index 0e578ac28b..3f477a37f0 100644
--- a/keystore/src/entities/platform/generic/general.rs
+++ b/keystore/src/entities/platform/generic/general.rs
@@ -36,7 +36,7 @@ impl EntityBase for ConsumerData {
     }
 
     fn to_transaction_entity(self) -> crate::transaction::dynamic_dispatch::Entity {
-        crate::transaction::dynamic_dispatch::Entity::ConsumerData(self)
+        crate::transaction::dynamic_dispatch::Entity::ConsumerData(self.into())
     }
 }
 
@@ -46,7 +46,7 @@ impl NewEntityBase for ConsumerData {
     const COLLECTION_NAME: &'static str = "consumer_data";
 
     fn to_transaction_entity(self) -> crate::transaction::dynamic_dispatch::Entity {
-        crate::transaction::dynamic_dispatch::Entity::ConsumerData(self)
+        crate::transaction::dynamic_dispatch::Entity::ConsumerData(self.into())
     }
 }
 
diff --git a/keystore/src/entities/platform/generic/mls/credential.rs b/keystore/src/entities/platform/generic/mls/credential.rs
index df2fac4422..efe9d9a999 100644
--- a/keystore/src/entities/platform/generic/mls/credential.rs
+++ b/keystore/src/entities/platform/generic/mls/credential.rs
@@ -169,7 +169,7 @@ impl EntityBase for StoredCredential {
     }
 
     fn to_transaction_entity(self) -> crate::transaction::dynamic_dispatch::Entity {
-        crate::transaction::dynamic_dispatch::Entity::StoredCredential(self)
+        crate::transaction::dynamic_dispatch::Entity::StoredCredential(self.into())
     }
 }
 
@@ -292,7 +292,7 @@ impl NewEntityBase for StoredCredential {
     const COLLECTION_NAME: &'static str = "mls_credentials";
 
     fn to_transaction_entity(self) -> crate::transaction::dynamic_dispatch::Entity {
-        crate::transaction::dynamic_dispatch::Entity::StoredCredential(self)
+        crate::transaction::dynamic_dispatch::Entity::StoredCredential(self.into())
     }
 }
 
diff --git a/keystore/src/entities/platform/generic/mls/e2ei_acme_ca.rs b/keystore/src/entities/platform/generic/mls/e2ei_acme_ca.rs
index 7bbad6ba7a..f8be2a53cb 100644
--- a/keystore/src/entities/platform/generic/mls/e2ei_acme_ca.rs
+++ b/keystore/src/entities/platform/generic/mls/e2ei_acme_ca.rs
@@ -34,7 +34,7 @@ impl EntityBase for E2eiAcmeCA {
     }
 
     fn to_transaction_entity(self) -> crate::transaction::dynamic_dispatch::Entity {
-        crate::transaction::dynamic_dispatch::Entity::E2eiAcmeCA(self)
+        crate::transaction::dynamic_dispatch::Entity::E2eiAcmeCA(self.into())
     }
 }
 
@@ -44,7 +44,7 @@ impl NewEntityBase for E2eiAcmeCA {
     const COLLECTION_NAME: &'static str = "e2ei_acme_ca";
 
     fn to_transaction_entity(self) -> crate::transaction::dynamic_dispatch::Entity {
-        crate::transaction::dynamic_dispatch::Entity::E2eiAcmeCA(self)
+        crate::transaction::dynamic_dispatch::Entity::E2eiAcmeCA(self.into())
     }
 }
 
diff --git a/keystore/src/entities/platform/generic/mls/encryption_keypair.rs b/keystore/src/entities/platform/generic/mls/encryption_keypair.rs
index 5ae436c978..5555f95082 100644
--- a/keystore/src/entities/platform/generic/mls/encryption_keypair.rs
+++ b/keystore/src/entities/platform/generic/mls/encryption_keypair.rs
@@ -111,7 +111,7 @@ impl EntityBase for StoredEncryptionKeyPair {
     }
 
     fn to_transaction_entity(self) -> crate::transaction::dynamic_dispatch::Entity {
-        crate::transaction::dynamic_dispatch::Entity::EncryptionKeyPair(self)
+        crate::transaction::dynamic_dispatch::Entity::EncryptionKeyPair(self.into())
     }
 }
 
@@ -180,7 +180,7 @@ impl NewEntityBase for StoredEncryptionKeyPair {
     const COLLECTION_NAME: &'static str = "mls_encryption_keypairs";
 
     fn to_transaction_entity(self) -> crate::transaction::dynamic_dispatch::Entity {
-        crate::transaction::dynamic_dispatch::Entity::EncryptionKeyPair(self)
+        crate::transaction::dynamic_dispatch::Entity::EncryptionKeyPair(self.into())
     }
 }
 
diff --git a/keystore/src/entities/platform/generic/mls/hpke_private_key.rs b/keystore/src/entities/platform/generic/mls/hpke_private_key.rs
index 2c6fd3dd9d..0167036fd7 100644
--- a/keystore/src/entities/platform/generic/mls/hpke_private_key.rs
+++ b/keystore/src/entities/platform/generic/mls/hpke_private_key.rs
@@ -111,7 +111,7 @@ impl EntityBase for StoredHpkePrivateKey {
     }
 
     fn to_transaction_entity(self) -> crate::transaction::dynamic_dispatch::Entity {
-        crate::transaction::dynamic_dispatch::Entity::HpkePrivateKey(self)
+        crate::transaction::dynamic_dispatch::Entity::HpkePrivateKey(self.into())
     }
 }
 
@@ -178,7 +178,7 @@ impl NewEntityBase for StoredHpkePrivateKey {
     const COLLECTION_NAME: &'static str = "mls_hpke_private_keys";
 
     fn to_transaction_entity(self) -> crate::transaction::dynamic_dispatch::Entity {
-        crate::transaction::dynamic_dispatch::Entity::HpkePrivateKey(self)
+        crate::transaction::dynamic_dispatch::Entity::HpkePrivateKey(self.into())
     }
 }
 
diff --git a/keystore/src/entities/platform/generic/mls/pending_group.rs b/keystore/src/entities/platform/generic/mls/pending_group.rs
index a1ddce80c5..ccf4fa5231 100644
--- a/keystore/src/entities/platform/generic/mls/pending_group.rs
+++ b/keystore/src/entities/platform/generic/mls/pending_group.rs
@@ -206,7 +206,7 @@ impl EntityBase for PersistedMlsPendingGroup {
     }
 
     fn to_transaction_entity(self) -> crate::transaction::dynamic_dispatch::Entity {
-        crate::transaction::dynamic_dispatch::Entity::PersistedMlsPendingGroup(self)
+        crate::transaction::dynamic_dispatch::Entity::PersistedMlsPendingGroup(self.into())
     }
 }
 
@@ -315,7 +315,7 @@ impl NewEntityBase for PersistedMlsPendingGroup {
     const COLLECTION_NAME: &'static str = "mls_pending_groups";
 
     fn to_transaction_entity(self) -> crate::transaction::dynamic_dispatch::Entity {
-        crate::transaction::dynamic_dispatch::Entity::PersistedMlsPendingGroup(self)
+        crate::transaction::dynamic_dispatch::Entity::PersistedMlsPendingGroup(self.into())
     }
 }
 
diff --git a/keystore/src/entities/platform/generic/mls/pending_message.rs b/keystore/src/entities/platform/generic/mls/pending_message.rs
index 94e7f80dc5..7101ad63af 100644
--- a/keystore/src/entities/platform/generic/mls/pending_message.rs
+++ b/keystore/src/entities/platform/generic/mls/pending_message.rs
@@ -158,7 +158,7 @@ impl EntityBase for MlsPendingMessage {
     }
 
     fn to_transaction_entity(self) -> crate::transaction::dynamic_dispatch::Entity {
-        crate::transaction::dynamic_dispatch::Entity::MlsPendingMessage(self)
+        crate::transaction::dynamic_dispatch::Entity::MlsPendingMessage(self.into())
     }
 }
 
@@ -212,7 +212,7 @@ impl NewEntityBase for MlsPendingMessage {
     const COLLECTION_NAME: &'static str = "mls_pending_messages";
 
     fn to_transaction_entity(self) -> crate::transaction::dynamic_dispatch::Entity {
-        crate::transaction::dynamic_dispatch::Entity::MlsPendingMessage(self)
+        crate::transaction::dynamic_dispatch::Entity::MlsPendingMessage(self.into())
     }
 }
 
diff --git a/keystore/src/entities/platform/generic/mls/psk_bundle.rs b/keystore/src/entities/platform/generic/mls/psk_bundle.rs
index 035ef8698a..b41aa1f934 100644
--- a/keystore/src/entities/platform/generic/mls/psk_bundle.rs
+++ b/keystore/src/entities/platform/generic/mls/psk_bundle.rs
@@ -101,7 +101,7 @@ impl EntityBase for StoredPskBundle {
     }
 
     fn to_transaction_entity(self) -> crate::transaction::dynamic_dispatch::Entity {
-        crate::transaction::dynamic_dispatch::Entity::PskBundle(self)
+        crate::transaction::dynamic_dispatch::Entity::PskBundle(self.into())
     }
 }
 
@@ -167,7 +167,7 @@ impl NewEntityBase for StoredPskBundle {
     const COLLECTION_NAME: &'static str = "mls_psk_bundles";
 
     fn to_transaction_entity(self) -> crate::transaction::dynamic_dispatch::Entity {
-        crate::transaction::dynamic_dispatch::Entity::PskBundle(self)
+        crate::transaction::dynamic_dispatch::Entity::PskBundle(self.into())
     }
 }
 
diff --git a/keystore/src/entities/platform/generic/proteus/identity.rs b/keystore/src/entities/platform/generic/proteus/identity.rs
index fc25bd0737..6cc6ff946c 100644
--- a/keystore/src/entities/platform/generic/proteus/identity.rs
+++ b/keystore/src/entities/platform/generic/proteus/identity.rs
@@ -98,7 +98,7 @@ impl EntityBase for ProteusIdentity {
     }
 
     fn to_transaction_entity(self) -> crate::transaction::dynamic_dispatch::Entity {
-        crate::transaction::dynamic_dispatch::Entity::ProteusIdentity(self)
+        crate::transaction::dynamic_dispatch::Entity::ProteusIdentity(self.into())
     }
 }
 
@@ -158,7 +158,7 @@ impl NewEntityBase for ProteusIdentity {
     const COLLECTION_NAME: &'static str = "proteus_identities";
 
     fn to_transaction_entity(self) -> crate::transaction::dynamic_dispatch::Entity {
-        crate::transaction::dynamic_dispatch::Entity::ProteusIdentity(self)
+        crate::transaction::dynamic_dispatch::Entity::ProteusIdentity(self.into())
     }
 }
 
diff --git a/keystore/src/entities/platform/generic/proteus/prekey.rs b/keystore/src/entities/platform/generic/proteus/prekey.rs
index 7cfacc5848..f8c00cc659 100644
--- a/keystore/src/entities/platform/generic/proteus/prekey.rs
+++ b/keystore/src/entities/platform/generic/proteus/prekey.rs
@@ -91,7 +91,7 @@ impl EntityBase for ProteusPrekey {
     }
 
     fn to_transaction_entity(self) -> crate::transaction::dynamic_dispatch::Entity {
-        crate::transaction::dynamic_dispatch::Entity::ProteusPrekey(self)
+        crate::transaction::dynamic_dispatch::Entity::ProteusPrekey(self.into())
     }
 }
 
@@ -151,7 +151,7 @@ impl NewEntityBase for ProteusPrekey {
     const COLLECTION_NAME: &'static str = "proteus_prekeys";
 
     fn to_transaction_entity(self) -> crate::transaction::dynamic_dispatch::Entity {
-        crate::transaction::dynamic_dispatch::Entity::ProteusPrekey(self)
+        crate::transaction::dynamic_dispatch::Entity::ProteusPrekey(self.into())
     }
 }
 
diff --git a/keystore/src/entities/platform/wasm/general.rs b/keystore/src/entities/platform/wasm/general.rs
index 23a6331ed2..a89389c6e4 100644
--- a/keystore/src/entities/platform/wasm/general.rs
+++ b/keystore/src/entities/platform/wasm/general.rs
@@ -16,7 +16,7 @@ impl EntityBase for ConsumerData {
     }
 
     fn to_transaction_entity(self) -> crate::transaction::dynamic_dispatch::Entity {
-        crate::transaction::dynamic_dispatch::Entity::ConsumerData(self)
+        crate::transaction::dynamic_dispatch::Entity::ConsumerData(self.into())
     }
 }
 
@@ -37,7 +37,7 @@ impl NewEntityBase for ConsumerData {
     const COLLECTION_NAME: &'static str = "consumer_data";
 
     fn to_transaction_entity(self) -> crate::transaction::dynamic_dispatch::Entity {
-        crate::transaction::dynamic_dispatch::Entity::ConsumerData(self)
+        crate::transaction::dynamic_dispatch::Entity::ConsumerData(self.into())
     }
 }
 
diff --git a/keystore/src/entities/platform/wasm/mls/credential.rs b/keystore/src/entities/platform/wasm/mls/credential.rs
index 361d24f9ae..39de2f4836 100644
--- a/keystore/src/entities/platform/wasm/mls/credential.rs
+++ b/keystore/src/entities/platform/wasm/mls/credential.rs
@@ -22,7 +22,7 @@ impl EntityBase for StoredCredential {
     }
 
     fn to_transaction_entity(self) -> crate::transaction::dynamic_dispatch::Entity {
-        crate::transaction::dynamic_dispatch::Entity::StoredCredential(self)
+        crate::transaction::dynamic_dispatch::Entity::StoredCredential(self.into())
     }
 }
 
@@ -91,7 +91,7 @@ impl NewEntityBase for StoredCredential {
     const COLLECTION_NAME: &'static str = "mls_credentials";
 
     fn to_transaction_entity(self) -> crate::transaction::dynamic_dispatch::Entity {
-        crate::transaction::dynamic_dispatch::Entity::StoredCredential(self)
+        crate::transaction::dynamic_dispatch::Entity::StoredCredential(self.into())
     }
 }
 
diff --git a/keystore/src/entities/platform/wasm/mls/e2ei_acme_ca.rs b/keystore/src/entities/platform/wasm/mls/e2ei_acme_ca.rs
index de95080983..c779fb58b5 100644
--- a/keystore/src/entities/platform/wasm/mls/e2ei_acme_ca.rs
+++ b/keystore/src/entities/platform/wasm/mls/e2ei_acme_ca.rs
@@ -15,7 +15,7 @@ impl EntityBase for E2eiAcmeCA {
     }
 
     fn to_transaction_entity(self) -> crate::transaction::dynamic_dispatch::Entity {
-        crate::transaction::dynamic_dispatch::Entity::E2eiAcmeCA(self)
+        crate::transaction::dynamic_dispatch::Entity::E2eiAcmeCA(self.into())
     }
 }
 
@@ -35,7 +35,7 @@ impl NewEntityBase for E2eiAcmeCA {
     const COLLECTION_NAME: &'static str = "e2ei_acme_ca";
 
     fn to_transaction_entity(self) -> crate::transaction::dynamic_dispatch::Entity {
-        crate::transaction::dynamic_dispatch::Entity::E2eiAcmeCA(self)
+        crate::transaction::dynamic_dispatch::Entity::E2eiAcmeCA(self.into())
     }
 }
 
diff --git a/keystore/src/entities/platform/wasm/mls/encryption_keypair.rs b/keystore/src/entities/platform/wasm/mls/encryption_keypair.rs
index a2d59af5cb..b1b815b8a6 100644
--- a/keystore/src/entities/platform/wasm/mls/encryption_keypair.rs
+++ b/keystore/src/entities/platform/wasm/mls/encryption_keypair.rs
@@ -21,7 +21,7 @@ impl EntityBase for StoredEncryptionKeyPair {
     }
 
     fn to_transaction_entity(self) -> crate::transaction::dynamic_dispatch::Entity {
-        crate::transaction::dynamic_dispatch::Entity::EncryptionKeyPair(self)
+        crate::transaction::dynamic_dispatch::Entity::EncryptionKeyPair(self.into())
     }
 }
 
@@ -74,7 +74,7 @@ impl NewEntityBase for StoredEncryptionKeyPair {
     const COLLECTION_NAME: &'static str = "mls_encryption_keypairs";
 
     fn to_transaction_entity(self) -> crate::transaction::dynamic_dispatch::Entity {
-        crate::transaction::dynamic_dispatch::Entity::EncryptionKeyPair(self)
+        crate::transaction::dynamic_dispatch::Entity::EncryptionKeyPair(self.into())
     }
 }
 
diff --git a/keystore/src/entities/platform/wasm/mls/hpke_private_key.rs b/keystore/src/entities/platform/wasm/mls/hpke_private_key.rs
index f2d1bd737f..03003e3569 100644
--- a/keystore/src/entities/platform/wasm/mls/hpke_private_key.rs
+++ b/keystore/src/entities/platform/wasm/mls/hpke_private_key.rs
@@ -22,7 +22,7 @@ impl EntityBase for StoredHpkePrivateKey {
     }
 
     fn to_transaction_entity(self) -> crate::transaction::dynamic_dispatch::Entity {
-        crate::transaction::dynamic_dispatch::Entity::HpkePrivateKey(self)
+        crate::transaction::dynamic_dispatch::Entity::HpkePrivateKey(self.into())
     }
 }
 
@@ -75,7 +75,7 @@ impl NewEntityBase for StoredHpkePrivateKey {
     const COLLECTION_NAME: &'static str = "mls_hpke_private_keys";
 
     fn to_transaction_entity(self) -> crate::transaction::dynamic_dispatch::Entity {
-        crate::transaction::dynamic_dispatch::Entity::HpkePrivateKey(self)
+        crate::transaction::dynamic_dispatch::Entity::HpkePrivateKey(self.into())
     }
 }
 
diff --git a/keystore/src/entities/platform/wasm/mls/pending_group.rs b/keystore/src/entities/platform/wasm/mls/pending_group.rs
index 1379355b10..c9e2ea475d 100644
--- a/keystore/src/entities/platform/wasm/mls/pending_group.rs
+++ b/keystore/src/entities/platform/wasm/mls/pending_group.rs
@@ -23,7 +23,7 @@ impl EntityBase for PersistedMlsPendingGroup {
     }
 
     fn to_transaction_entity(self) -> crate::transaction::dynamic_dispatch::Entity {
-        crate::transaction::dynamic_dispatch::Entity::PersistedMlsPendingGroup(self)
+        crate::transaction::dynamic_dispatch::Entity::PersistedMlsPendingGroup(self.into())
     }
 }
 
@@ -90,7 +90,7 @@ impl NewEntityBase for PersistedMlsPendingGroup {
     const COLLECTION_NAME: &'static str = "mls_pending_groups";
 
     fn to_transaction_entity(self) -> crate::transaction::dynamic_dispatch::Entity {
-        crate::transaction::dynamic_dispatch::Entity::PersistedMlsPendingGroup(self)
+        crate::transaction::dynamic_dispatch::Entity::PersistedMlsPendingGroup(self.into())
     }
 }
 
diff --git a/keystore/src/entities/platform/wasm/mls/pending_message.rs b/keystore/src/entities/platform/wasm/mls/pending_message.rs
index 4af774e43d..051fc449a6 100644
--- a/keystore/src/entities/platform/wasm/mls/pending_message.rs
+++ b/keystore/src/entities/platform/wasm/mls/pending_message.rs
@@ -24,7 +24,7 @@ impl EntityBase for MlsPendingMessage {
     }
 
     fn to_transaction_entity(self) -> crate::transaction::dynamic_dispatch::Entity {
-        crate::transaction::dynamic_dispatch::Entity::MlsPendingMessage(self)
+        crate::transaction::dynamic_dispatch::Entity::MlsPendingMessage(self.into())
     }
 }
 
@@ -112,7 +112,7 @@ impl NewEntityBase for MlsPendingMessage {
     const COLLECTION_NAME: &'static str = "mls_pending_messages";
 
     fn to_transaction_entity(self) -> crate::transaction::dynamic_dispatch::Entity {
-        crate::transaction::dynamic_dispatch::Entity::MlsPendingMessage(self)
+        crate::transaction::dynamic_dispatch::Entity::MlsPendingMessage(self.into())
     }
 }
 
diff --git a/keystore/src/entities/platform/wasm/mls/psk_bundle.rs b/keystore/src/entities/platform/wasm/mls/psk_bundle.rs
index b2bb96bfa9..5c94c224ec 100644
--- a/keystore/src/entities/platform/wasm/mls/psk_bundle.rs
+++ b/keystore/src/entities/platform/wasm/mls/psk_bundle.rs
@@ -22,7 +22,7 @@ impl EntityBase for StoredPskBundle {
     }
 
     fn to_transaction_entity(self) -> crate::transaction::dynamic_dispatch::Entity {
-        crate::transaction::dynamic_dispatch::Entity::PskBundle(self)
+        crate::transaction::dynamic_dispatch::Entity::PskBundle(self.into())
     }
 }
 
@@ -75,7 +75,7 @@ impl NewEntityBase for StoredPskBundle {
     const COLLECTION_NAME: &'static str = "mls_psk_bundles";
 
     fn to_transaction_entity(self) -> crate::transaction::dynamic_dispatch::Entity {
-        crate::transaction::dynamic_dispatch::Entity::PskBundle(self)
+        crate::transaction::dynamic_dispatch::Entity::PskBundle(self.into())
     }
 }
 
diff --git a/keystore/src/entities/platform/wasm/mls/refresh_token.rs b/keystore/src/entities/platform/wasm/mls/refresh_token.rs
index 5f7ac5a4ea..249297b8dd 100644
--- a/keystore/src/entities/platform/wasm/mls/refresh_token.rs
+++ b/keystore/src/entities/platform/wasm/mls/refresh_token.rs
@@ -16,7 +16,7 @@ impl EntityBase for E2eiRefreshToken {
     }
 
     fn to_transaction_entity(self) -> crate::transaction::dynamic_dispatch::Entity {
-        crate::transaction::dynamic_dispatch::Entity::E2eiRefreshToken(self)
+        crate::transaction::dynamic_dispatch::Entity::E2eiRefreshToken(self.into())
     }
 }
 
@@ -37,7 +37,7 @@ impl NewEntityBase for E2eiRefreshToken {
     const COLLECTION_NAME: &'static str = "e2ei_refresh_token";
 
     fn to_transaction_entity(self) -> crate::transaction::dynamic_dispatch::Entity {
-        crate::transaction::dynamic_dispatch::Entity::E2eiRefreshToken(self)
+        crate::transaction::dynamic_dispatch::Entity::E2eiRefreshToken(self.into())
     }
 }
 
diff --git a/keystore/src/entities/platform/wasm/proteus/identity.rs b/keystore/src/entities/platform/wasm/proteus/identity.rs
index cb0e627840..4ab31a352a 100644
--- a/keystore/src/entities/platform/wasm/proteus/identity.rs
+++ b/keystore/src/entities/platform/wasm/proteus/identity.rs
@@ -22,7 +22,7 @@ impl EntityBase for ProteusIdentity {
     }
 
     fn to_transaction_entity(self) -> crate::transaction::dynamic_dispatch::Entity {
-        crate::transaction::dynamic_dispatch::Entity::ProteusIdentity(self)
+        crate::transaction::dynamic_dispatch::Entity::ProteusIdentity(self.into())
     }
 }
 
@@ -81,7 +81,7 @@ impl NewEntityBase for ProteusIdentity {
     const COLLECTION_NAME: &'static str = "proteus_identities";
 
     fn to_transaction_entity(self) -> crate::transaction::dynamic_dispatch::Entity {
-        crate::transaction::dynamic_dispatch::Entity::ProteusIdentity(self)
+        crate::transaction::dynamic_dispatch::Entity::ProteusIdentity(self.into())
     }
 }
 
diff --git a/keystore/src/entities/platform/wasm/proteus/prekey.rs b/keystore/src/entities/platform/wasm/proteus/prekey.rs
index 6f4df017eb..25dae6edbd 100644
--- a/keystore/src/entities/platform/wasm/proteus/prekey.rs
+++ b/keystore/src/entities/platform/wasm/proteus/prekey.rs
@@ -22,7 +22,7 @@ impl EntityBase for ProteusPrekey {
     }
 
     fn to_transaction_entity(self) -> crate::transaction::dynamic_dispatch::Entity {
-        crate::transaction::dynamic_dispatch::Entity::ProteusPrekey(self)
+        crate::transaction::dynamic_dispatch::Entity::ProteusPrekey(self.into())
     }
 }
 
@@ -75,7 +75,7 @@ impl NewEntityBase for ProteusPrekey {
     const COLLECTION_NAME: &'static str = "proteus_prekeys";
 
     fn to_transaction_entity(self) -> crate::transaction::dynamic_dispatch::Entity {
-        crate::transaction::dynamic_dispatch::Entity::ProteusPrekey(self)
+        crate::transaction::dynamic_dispatch::Entity::ProteusPrekey(self.into())
     }
 }
 

From d4585973b7ecca48567639cc25b03e80d875690f Mon Sep 17 00:00:00 2001
From: Peter Goodspeed-Niklaus <peter.goodspeed-niklaus@wire.com>
Date: Fri, 19 Dec 2025 11:05:21 +0100
Subject: [PATCH 12/48] chore(keystore): update entity impls to track new trait
 requirements

Fascinatingly, we're allowed to simply omit the trait bounds where
they can be statically proven, i.e. on an associated type. We're
learning about Rust today!
---
 .../src/entity_derive_new/derive_impl.rs      | 25 ++++++++++--------
 keystore/src/entities/dummy_entity.rs         | 26 +++++++------------
 .../platform/generic/mls/pending_group.rs     | 15 ++++-------
 .../platform/generic/mls/pending_message.rs   |  9 +++----
 .../platform/wasm/mls/pending_group.rs        | 15 ++++-------
 .../platform/wasm/mls/pending_message.rs      |  9 +++----
 6 files changed, 42 insertions(+), 57 deletions(-)

diff --git a/crypto-macros/src/entity_derive_new/derive_impl.rs b/crypto-macros/src/entity_derive_new/derive_impl.rs
index 1e7f377ad2..34490d450c 100644
--- a/crypto-macros/src/entity_derive_new/derive_impl.rs
+++ b/crypto-macros/src/entity_derive_new/derive_impl.rs
@@ -150,19 +150,21 @@ impl Entity {
                     &self.#pk_field_name
                 }
 
-                async fn get_borrowed<Q>(conn: &mut Self::ConnectionType, key: &Q) -> crate::CryptoKeystoreResult<Option<Self>>
+                async fn get_borrowed(conn: &mut Self::ConnectionType, key: &Self::BorrowedPrimaryKey)
+                    -> crate::CryptoKeystoreResult<Option<Self>>
                 where
-                    Self::PrimaryKey: std::borrow::Borrow<Q>,
-                    Q: crate::traits::KeyType,
+                    for<'pk> &'pk Self::BorrowedPrimaryKey: crate::traits::KeyType,
                 {
+                    let key_bytes = <&Self::BorrowedPrimaryKey as crate::traits::KeyType>::bytes(&key);
+                    let key_bytes = key_bytes.as_ref();
                     #[cfg(target_family = "wasm")]
                     {
-                        conn.storage().new_get(key.bytes().as_ref()).await
+                        conn.storage().new_get(key_bytes).await
                     }
 
                     #[cfg(not(target_family = "wasm"))]
                     {
-                        crate::entities::platform::get_helper::<Self, _>(conn, #pk_column_name, key.bytes().as_ref(), |row| {
+                        crate::entities::platform::get_helper::<Self, _>(conn, #pk_column_name, key_bytes, |row| {
                             Ok(Self {
                                 #( #field_assignments, )*
                             })
@@ -246,22 +248,23 @@ impl Entity {
             #[cfg_attr(target_family = "wasm", async_trait::async_trait(?Send))]
             #[cfg_attr(not(target_family = "wasm"), async_trait::async_trait)]
             impl<'a> crate::traits::EntityDeleteBorrowed<'a> for #struct_name {
-                async fn delete_borrowed<Q>(
+                async fn delete_borrowed(
                     tx: &<Self as crate::traits::EntityDatabaseMutation<'a>>::Transaction,
-                    id: &Q,
+                    id: &<Self as crate::traits::BorrowPrimaryKey>::BorrowedPrimaryKey,
                 ) -> crate::CryptoKeystoreResult<bool>
                 where
-                    Self::PrimaryKey: std::borrow::Borrow<Q>,
-                    Q: crate::traits::KeyType
+                    for<'pk> &'pk <Self as crate::traits::BorrowPrimaryKey>::BorrowedPrimaryKey: crate::traits::KeyType,
                 {
+                    let id_bytes = <&<Self as crate::traits::BorrowPrimaryKey>::BorrowedPrimaryKey as crate::traits::KeyType>::bytes(&id);
+                    let id_bytes = id_bytes.as_ref();
                     #[cfg(target_family = "wasm")]
                     {
-                        tx.new_delete::<Self>(id.bytes().as_ref()).await
+                        tx.new_delete::<Self>(id_bytes).await
                     }
 
                     #[cfg(not(target_family = "wasm"))]
                     {
-                        crate::entities::platform::delete_helper::<Self>(tx, #id_column_name, id.bytes().as_ref()).await
+                        crate::entities::platform::delete_helper::<Self>(tx, #id_column_name, id_bytes).await
                     }
                 }
             }
diff --git a/keystore/src/entities/dummy_entity.rs b/keystore/src/entities/dummy_entity.rs
index 725f62b292..af75beb8f7 100644
--- a/keystore/src/entities/dummy_entity.rs
+++ b/keystore/src/entities/dummy_entity.rs
@@ -1,4 +1,4 @@
-use std::{borrow::Borrow, collections::HashSet, sync::LazyLock};
+use std::{collections::HashSet, sync::LazyLock};
 
 use async_lock::RwLock;
 use sha2::{Digest as _, Sha256};
@@ -123,11 +123,10 @@ impl BorrowPrimaryKey for DummyStoreValue {
         &[]
     }
 
-    async fn get_borrowed<Q>(_conn: &mut Self::ConnectionType, _key: &Q) -> CryptoKeystoreResult<Option<Self>>
-    where
-        Self::PrimaryKey: Borrow<Q>,
-        Q: KeyType,
-    {
+    async fn get_borrowed(
+        _conn: &mut Self::ConnectionType,
+        _key: &Self::BorrowedPrimaryKey,
+    ) -> CryptoKeystoreResult<Option<Self>> {
         Ok(None)
     }
 }
@@ -219,11 +218,10 @@ impl BorrowPrimaryKey for NewDummyStoreValue {
         &self.id
     }
 
-    async fn get_borrowed<Q>(_conn: &mut Self::ConnectionType, key: &Q) -> CryptoKeystoreResult<Option<Self>>
-    where
-        Self::PrimaryKey: Borrow<Q>,
-        Q: KeyType,
-    {
+    async fn get_borrowed(
+        _conn: &mut Self::ConnectionType,
+        key: &Self::BorrowedPrimaryKey,
+    ) -> CryptoKeystoreResult<Option<Self>> {
         let guard = NEW_DUMMY_STORE_IDS.read().await;
         let key = key.bytes();
         let key = key.as_ref();
@@ -271,11 +269,7 @@ impl<'a> EntityDatabaseMutation<'a> for NewDummyStoreValue {
 #[async_trait::async_trait]
 impl EntityDeleteBorrowed<'_> for NewDummyStoreValue {
     /// Delete an entity by a borrowed form of its primary key.
-    async fn delete_borrowed<Q>(_tx: &Self::Transaction, id: &Q) -> CryptoKeystoreResult<bool>
-    where
-        Self::PrimaryKey: Borrow<Q>,
-        Q: KeyType,
-    {
+    async fn delete_borrowed(_tx: &Self::Transaction, id: &Self::BorrowedPrimaryKey) -> CryptoKeystoreResult<bool> {
         let mut guard = NEW_DUMMY_STORE_IDS.write().await;
         let removed = guard.remove::<[u8]>(id.bytes().as_ref());
         Ok(removed)
diff --git a/keystore/src/entities/platform/generic/mls/pending_group.rs b/keystore/src/entities/platform/generic/mls/pending_group.rs
index ccf4fa5231..1803c653b6 100644
--- a/keystore/src/entities/platform/generic/mls/pending_group.rs
+++ b/keystore/src/entities/platform/generic/mls/pending_group.rs
@@ -348,11 +348,10 @@ impl BorrowPrimaryKey for PersistedMlsPendingGroup {
         &self.id
     }
 
-    async fn get_borrowed<Q>(conn: &mut Self::ConnectionType, key: &Q) -> CryptoKeystoreResult<Option<Self>>
-    where
-        Self::PrimaryKey: Borrow<Q>,
-        Q: KeyType,
-    {
+    async fn get_borrowed(
+        conn: &mut Self::ConnectionType,
+        key: &Self::BorrowedPrimaryKey,
+    ) -> CryptoKeystoreResult<Option<Self>> {
         get_helper::<Self, _>(conn, "id", key.bytes().as_ref(), Self::from_row).await
     }
 }
@@ -380,11 +379,7 @@ impl<'a> EntityDatabaseMutation<'a> for PersistedMlsPendingGroup {
 
 #[async_trait]
 impl<'a> EntityDeleteBorrowed<'a> for PersistedMlsPendingGroup {
-    async fn delete_borrowed<Q>(tx: &Self::Transaction, id: &Q) -> CryptoKeystoreResult<bool>
-    where
-        Self::PrimaryKey: Borrow<Q>,
-        Q: KeyType,
-    {
+    async fn delete_borrowed(tx: &Self::Transaction, id: &Self::BorrowedPrimaryKey) -> CryptoKeystoreResult<bool> {
         delete_helper::<Self>(tx, "id", id.bytes().as_ref()).await
     }
 }
diff --git a/keystore/src/entities/platform/generic/mls/pending_message.rs b/keystore/src/entities/platform/generic/mls/pending_message.rs
index 7101ad63af..daa568b576 100644
--- a/keystore/src/entities/platform/generic/mls/pending_message.rs
+++ b/keystore/src/entities/platform/generic/mls/pending_message.rs
@@ -250,11 +250,10 @@ impl BorrowPrimaryKey for MlsPendingMessage {
         &self.foreign_id
     }
 
-    async fn get_borrowed<Q>(conn: &mut Self::ConnectionType, key: &Q) -> CryptoKeystoreResult<Option<Self>>
-    where
-        Self::PrimaryKey: Borrow<Q>,
-        Q: KeyType,
-    {
+    async fn get_borrowed(
+        conn: &mut Self::ConnectionType,
+        key: &Self::BorrowedPrimaryKey,
+    ) -> CryptoKeystoreResult<Option<Self>> {
         panic!("cannot get `MlsPendingMessage` by primary key as it has no distinct primary key")
     }
 }
diff --git a/keystore/src/entities/platform/wasm/mls/pending_group.rs b/keystore/src/entities/platform/wasm/mls/pending_group.rs
index c9e2ea475d..db9563b7a6 100644
--- a/keystore/src/entities/platform/wasm/mls/pending_group.rs
+++ b/keystore/src/entities/platform/wasm/mls/pending_group.rs
@@ -123,11 +123,10 @@ impl BorrowPrimaryKey for PersistedMlsPendingGroup {
         &self.id
     }
 
-    async fn get_borrowed<Q>(conn: &mut Self::ConnectionType, key: &Q) -> CryptoKeystoreResult<Option<Self>>
-    where
-        Self::PrimaryKey: Borrow<Q>,
-        Q: KeyType,
-    {
+    async fn get_borrowed(
+        conn: &mut Self::ConnectionType,
+        key: &Self::BorrowedPrimaryKey,
+    ) -> CryptoKeystoreResult<Option<Self>> {
         conn.storage().new_get(key.bytes().as_ref()).await
     }
 }
@@ -151,11 +150,7 @@ impl<'a> EntityDatabaseMutation<'a> for PersistedMlsPendingGroup {
 
 #[async_trait(?Send)]
 impl<'a> EntityDeleteBorrowed<'a> for PersistedMlsPendingGroup {
-    async fn delete_borrowed<Q>(tx: &Self::Transaction, id: &Q) -> CryptoKeystoreResult<bool>
-    where
-        Self::PrimaryKey: Borrow<Q>,
-        Q: KeyType,
-    {
+    async fn delete_borrowed(tx: &Self::Transaction, id: &Self::BorrowedPrimaryKey) -> CryptoKeystoreResult<bool> {
         tx.new_delete::<Self>(id.bytes().as_ref()).await
     }
 }
diff --git a/keystore/src/entities/platform/wasm/mls/pending_message.rs b/keystore/src/entities/platform/wasm/mls/pending_message.rs
index 051fc449a6..c45d847bd5 100644
--- a/keystore/src/entities/platform/wasm/mls/pending_message.rs
+++ b/keystore/src/entities/platform/wasm/mls/pending_message.rs
@@ -153,11 +153,10 @@ impl BorrowPrimaryKey for MlsPendingMessage {
         &self.foreign_id
     }
 
-    async fn get_borrowed<Q>(conn: &mut Self::ConnectionType, key: &Q) -> CryptoKeystoreResult<Option<Self>>
-    where
-        Self::PrimaryKey: Borrow<Q>,
-        Q: KeyType,
-    {
+    async fn get_borrowed(
+        conn: &mut Self::ConnectionType,
+        key: &Self::BorrowedPrimaryKey,
+    ) -> CryptoKeystoreResult<Option<Self>> {
         panic!("cannot get `MlsPendingMessage` by primary key as it has no distinct primary key")
     }
 }

From 79ff1996c2ecc264f610530824911ae19a83afba Mon Sep 17 00:00:00 2001
From: Peter Goodspeed-Niklaus <peter.goodspeed-niklaus@wire.com>
Date: Fri, 19 Dec 2025 11:14:05 +0100
Subject: [PATCH 13/48] chore(keystore): rewrite `PersistedMlsGroupExt` in
 terms of new traits

---
 keystore/src/entities/mls.rs | 36 ++++++++++++++++++++++--------------
 1 file changed, 22 insertions(+), 14 deletions(-)

diff --git a/keystore/src/entities/mls.rs b/keystore/src/entities/mls.rs
index a909fe3bde..eb3db3fd1b 100644
--- a/keystore/src/entities/mls.rs
+++ b/keystore/src/entities/mls.rs
@@ -1,7 +1,10 @@
 use zeroize::Zeroize;
 
-use super::{Entity, EntityBase, EntityFindParams, EntityTransactionExt, StringEntityId};
-use crate::{CryptoKeystoreError, CryptoKeystoreResult, connection::TransactionWrapper};
+use crate::{
+    CryptoKeystoreError, CryptoKeystoreResult,
+    connection::TransactionWrapper,
+    traits::{BorrowPrimaryKey, Entity, EntityBase, KeyType, OwnedKeyType},
+};
 
 /// Entity representing a persisted `MlsGroup`
 #[derive(
@@ -28,27 +31,32 @@ pub struct PersistedMlsGroup {
 
 #[cfg_attr(target_family = "wasm", async_trait::async_trait(?Send))]
 #[cfg_attr(not(target_family = "wasm"), async_trait::async_trait)]
-pub trait PersistedMlsGroupExt: Entity {
+pub trait PersistedMlsGroupExt: Entity + BorrowPrimaryKey
+where
+    for<'a> &'a <Self as BorrowPrimaryKey>::BorrowedPrimaryKey: KeyType,
+{
     fn parent_id(&self) -> Option<&[u8]>;
 
-    async fn parent_group(
-        &self,
-        conn: &mut <Self as super::EntityBase>::ConnectionType,
-    ) -> CryptoKeystoreResult<Option<Self>> {
+    async fn parent_group(&self, conn: &mut Self::ConnectionType) -> CryptoKeystoreResult<Option<Self>> {
         let Some(parent_id) = self.parent_id() else {
             return Ok(None);
         };
 
-        <Self as super::Entity>::find_one(conn, &parent_id.into()).await
+        let parent_id = OwnedKeyType::from_bytes(parent_id)
+            .ok_or(CryptoKeystoreError::InvalidPrimaryKeyBytes(Self::COLLECTION_NAME))?;
+        Self::get(conn, &parent_id).await
     }
 
-    async fn child_groups(
-        &self,
-        conn: &mut <Self as super::EntityBase>::ConnectionType,
-    ) -> CryptoKeystoreResult<Vec<Self>> {
-        let entities = <Self as super::Entity>::find_all(conn, super::EntityFindParams::default()).await?;
+    async fn child_groups(&self, conn: &mut <Self as EntityBase>::ConnectionType) -> CryptoKeystoreResult<Vec<Self>> {
+        // A perfect opportunity for refactoring in WPB-20844
+        // when we do that, we no longer need varying implementations according to wasm or not,
+        // so both `parent_group` and this method should just be implemented directly on `PersistedMlsGroup`.
+        let entities = Self::load_all(conn).await?;
 
-        let id = self.id_raw();
+        // for whatever reason rustc needs each of these distinct bindings to prove to itself that the lifetimes work out
+        let id = self.borrow_primary_key();
+        let id = id.bytes();
+        let id = id.as_ref();
 
         Ok(entities
             .into_iter()

From df375cd0fc051a8eeb1510232e375ca5ab53aa30 Mon Sep 17 00:00:00 2001
From: Peter Goodspeed-Niklaus <peter.goodspeed-niklaus@wire.com>
Date: Fri, 19 Dec 2025 11:22:56 +0100
Subject: [PATCH 14/48] chore(keystore): update `proteus_traits::PreKeyStore`
 impl

---
 keystore/src/proteus.rs | 16 ++++++----------
 1 file changed, 6 insertions(+), 10 deletions(-)

diff --git a/keystore/src/proteus.rs b/keystore/src/proteus.rs
index 24b7bcaa90..70de143612 100644
--- a/keystore/src/proteus.rs
+++ b/keystore/src/proteus.rs
@@ -1,7 +1,6 @@
 use crate::{
-    CryptoKeystoreError, CryptoKeystoreResult,
-    connection::{Database, FetchFromDatabase},
-    entities::ProteusPrekey,
+    CryptoKeystoreError, CryptoKeystoreResult, connection::Database, entities::ProteusPrekey,
+    traits::FetchFromDatabase as _,
 };
 
 #[cfg_attr(target_family = "wasm", async_trait::async_trait(?Send))]
@@ -29,15 +28,12 @@ impl proteus_traits::PreKeyStore for Database {
         &mut self,
         id: proteus_traits::RawPreKeyId,
     ) -> Result<Option<proteus_traits::RawPreKey>, Self::Error> {
-        Ok(self
-            .find::<ProteusPrekey>(&id.to_le_bytes())
-            .await?
-            .map(|db_prekey| db_prekey.prekey.clone()))
+        self.get::<ProteusPrekey>(&id)
+            .await
+            .map(|db_prekey| db_prekey.map(|mut db_prekey| std::mem::take(&mut db_prekey.prekey)))
     }
 
     async fn remove(&mut self, id: proteus_traits::RawPreKeyId) -> Result<(), Self::Error> {
-        Database::remove::<ProteusPrekey, _>(self, id.to_le_bytes()).await?;
-
-        Ok(())
+        Database::remove::<ProteusPrekey>(self, &id).await
     }
 }

From c9edf61b92f2ca156f2fe2e292724fe50e678ac9 Mon Sep 17 00:00:00 2001
From: Peter Goodspeed-Niklaus <peter.goodspeed-niklaus@wire.com>
Date: Fri, 19 Dec 2025 11:27:23 +0100
Subject: [PATCH 15/48] chore(keystore): update mls-specific trait impls

---
 keystore/src/mls.rs | 94 ++++++++++++++++++++++++++-------------------
 1 file changed, 54 insertions(+), 40 deletions(-)

diff --git a/keystore/src/mls.rs b/keystore/src/mls.rs
index 3b56ffb8ce..2f898ca5bb 100644
--- a/keystore/src/mls.rs
+++ b/keystore/src/mls.rs
@@ -3,12 +3,12 @@ use openmls_basic_credential::SignatureKeyPair;
 use openmls_traits::key_store::{MlsEntity, MlsEntityId};
 
 use crate::{
-    CryptoKeystoreError, CryptoKeystoreResult, MissingKeyErrorKind,
-    connection::FetchFromDatabase,
+    CryptoKeystoreError, CryptoKeystoreResult, MissingKeyErrorKind, Sha256Hash,
     entities::{
-        EntityFindParams, PersistedMlsGroup, PersistedMlsPendingGroup, StoredCredential, StoredE2eiEnrollment,
-        StoredEncryptionKeyPair, StoredEpochEncryptionKeypair, StoredHpkePrivateKey, StoredKeypackage, StoredPskBundle,
+        PersistedMlsGroup, PersistedMlsPendingGroup, StoredCredential, StoredE2eiEnrollment, StoredEncryptionKeyPair,
+        StoredEpochEncryptionKeypair, StoredHpkePrivateKey, StoredKeypackage, StoredPskBundle,
     },
+    traits::FetchFromDatabase,
 };
 
 /// An interface for the specialized queries in the KeyStore
@@ -124,23 +124,19 @@ pub trait CryptoKeystoreMls: Sized {
 #[cfg_attr(not(target_family = "wasm"), async_trait::async_trait)]
 impl CryptoKeystoreMls for crate::Database {
     async fn mls_fetch_keypackages<V: MlsEntity>(&self, count: u32) -> CryptoKeystoreResult<Vec<V>> {
-        let reverse = !cfg!(target_family = "wasm");
-        let keypackages = self
-            .find_all::<StoredKeypackage>(EntityFindParams {
-                limit: Some(count),
-                offset: None,
-                reverse,
-            })
-            .await?;
-
+        let keypackages = self.load_all::<StoredKeypackage>().await?;
         Ok(keypackages
             .into_iter()
             .filter_map(|kpb| postcard::from_bytes(&kpb.keypackage).ok())
+            .take(count as _)
             .collect())
     }
 
     async fn mls_group_exists(&self, group_id: impl AsRef<[u8]> + Send) -> bool {
-        matches!(self.find::<PersistedMlsGroup>(group_id).await, Ok(Some(_)))
+        matches!(
+            self.get_borrowed::<PersistedMlsGroup>(group_id.as_ref()).await,
+            Ok(Some(_))
+        )
     }
 
     async fn mls_group_persist(
@@ -162,16 +158,20 @@ impl CryptoKeystoreMls for crate::Database {
     async fn mls_groups_restore(
         &self,
     ) -> CryptoKeystoreResult<std::collections::HashMap<Vec<u8>, (Option<Vec<u8>>, Vec<u8>)>> {
-        let groups = self.find_all::<PersistedMlsGroup>(EntityFindParams::default()).await?;
+        let groups = self.load_all::<PersistedMlsGroup>().await?;
         Ok(groups
             .into_iter()
-            .map(|group: PersistedMlsGroup| (group.id.clone(), (group.parent_id.clone(), group.state.clone())))
+            .map(|mut group: PersistedMlsGroup| {
+                let id = std::mem::take(&mut group.id);
+                let parent_id = std::mem::take(&mut group.parent_id);
+                let state = std::mem::take(&mut group.state);
+                (id, (parent_id, state))
+            })
             .collect())
     }
 
     async fn mls_group_delete(&self, group_id: impl AsRef<[u8]> + Send) -> CryptoKeystoreResult<()> {
-        self.remove::<PersistedMlsGroup, _>(group_id).await?;
-
+        self.remove_borrowed::<PersistedMlsGroup>(group_id.as_ref()).await?;
         Ok(())
     }
 
@@ -196,7 +196,7 @@ impl CryptoKeystoreMls for crate::Database {
         &self,
         group_id: impl AsRef<[u8]> + Send,
     ) -> CryptoKeystoreResult<(Vec<u8>, Vec<u8>)> {
-        self.find(group_id)
+        self.get_borrowed(group_id.as_ref())
             .await?
             .map(|r: PersistedMlsPendingGroup| (r.state.clone(), r.custom_configuration.clone()))
             .ok_or(CryptoKeystoreError::MissingKeyInStore(
@@ -205,7 +205,8 @@ impl CryptoKeystoreMls for crate::Database {
     }
 
     async fn mls_pending_groups_delete(&self, group_id: impl AsRef<[u8]> + Send) -> CryptoKeystoreResult<()> {
-        self.remove::<PersistedMlsPendingGroup, _>(group_id).await
+        self.remove_borrowed::<PersistedMlsPendingGroup>(group_id.as_ref())
+            .await
     }
 
     async fn save_e2ei_enrollment(&self, id: &[u8], content: &[u8]) -> CryptoKeystoreResult<()> {
@@ -219,13 +220,13 @@ impl CryptoKeystoreMls for crate::Database {
 
     async fn pop_e2ei_enrollment(&self, id: &[u8]) -> CryptoKeystoreResult<Vec<u8>> {
         // someone who has time could try to optimize this but honestly it's really on the cold path
-        let enrollment = self
-            .find::<StoredE2eiEnrollment>(id)
-            .await?
-            .ok_or(CryptoKeystoreError::MissingKeyInStore(
-                MissingKeyErrorKind::StoredE2eiEnrollment,
-            ))?;
-        self.remove::<StoredE2eiEnrollment, _>(id).await?;
+        let enrollment =
+            self.get_borrowed::<StoredE2eiEnrollment>(id)
+                .await?
+                .ok_or(CryptoKeystoreError::MissingKeyInStore(
+                    MissingKeyErrorKind::StoredE2eiEnrollment,
+                ))?;
+        self.remove_borrowed::<StoredE2eiEnrollment>(id).await?;
         Ok(enrollment.content.clone())
     }
 }
@@ -313,11 +314,12 @@ impl openmls_traits::key_store::OpenMlsKeyStore for crate::connection::Database
 
         match V::ID {
             MlsEntityId::GroupState => {
-                let group: PersistedMlsGroup = self.find(k).await.ok().flatten()?;
+                let group: PersistedMlsGroup = self.get_borrowed(k).await.ok().flatten()?;
                 deser(&group.state).ok()
             }
             MlsEntityId::SignatureKeyPair => {
-                let stored_credential = self.find::<StoredCredential>(k).await.ok().flatten()?;
+                let hash = Sha256Hash::from_existing_hash(k).ok()?;
+                let stored_credential = self.get::<StoredCredential>(&hash).await.ok().flatten()?;
                 let ciphersuite = Ciphersuite::try_from(stored_credential.ciphersuite).ok()?;
                 let signature_scheme = ciphersuite.signature_algorithm();
 
@@ -333,23 +335,26 @@ impl openmls_traits::key_store::OpenMlsKeyStore for crate::connection::Database
                 deser(&mls_keypair_serialized).ok()
             }
             MlsEntityId::KeyPackage => {
-                let kp: StoredKeypackage = self.find(k).await.ok().flatten()?;
+                let kp: StoredKeypackage = self.get_borrowed(k).await.ok().flatten()?;
                 deser(&kp.keypackage).ok()
             }
             MlsEntityId::HpkePrivateKey => {
-                let hpke_pk: StoredHpkePrivateKey = self.find(k).await.ok().flatten()?;
+                let hash = Sha256Hash::from_existing_hash(k).ok()?;
+                let hpke_pk: StoredHpkePrivateKey = self.get(&hash).await.ok().flatten()?;
                 deser(&hpke_pk.sk).ok()
             }
             MlsEntityId::PskBundle => {
-                let psk_bundle: StoredPskBundle = self.find(k).await.ok().flatten()?;
+                let hash = Sha256Hash::from_existing_hash(k).ok()?;
+                let psk_bundle: StoredPskBundle = self.get(&hash).await.ok().flatten()?;
                 deser(&psk_bundle.psk).ok()
             }
             MlsEntityId::EncryptionKeyPair => {
-                let kp: StoredEncryptionKeyPair = self.find(k).await.ok().flatten()?;
+                let hash = Sha256Hash::from_existing_hash(k).ok()?;
+                let kp: StoredEncryptionKeyPair = self.get(&hash).await.ok().flatten()?;
                 deser(&kp.sk).ok()
             }
             MlsEntityId::EpochEncryptionKeyPair => {
-                let kp: StoredEpochEncryptionKeypair = self.find(k).await.ok().flatten()?;
+                let kp: StoredEpochEncryptionKeypair = self.get_borrowed(k).await.ok().flatten()?;
                 deser(&kp.keypairs).ok()
             }
         }
@@ -357,16 +362,25 @@ impl openmls_traits::key_store::OpenMlsKeyStore for crate::connection::Database
 
     async fn delete<V: MlsEntity>(&self, k: &[u8]) -> Result<(), Self::Error> {
         match V::ID {
-            MlsEntityId::GroupState => self.remove::<PersistedMlsGroup, _>(k).await?,
+            MlsEntityId::GroupState => self.remove_borrowed::<PersistedMlsGroup>(k).await?,
             MlsEntityId::SignatureKeyPair => unimplemented!(
                 "Deleting a signature key pair should not be done through this API, any keypair should be deleted via
                 deleting a credential."
             ),
-            MlsEntityId::HpkePrivateKey => self.remove::<StoredHpkePrivateKey, _>(k).await?,
-            MlsEntityId::KeyPackage => self.remove::<StoredKeypackage, _>(k).await?,
-            MlsEntityId::PskBundle => self.remove::<StoredPskBundle, _>(k).await?,
-            MlsEntityId::EncryptionKeyPair => self.remove::<StoredEncryptionKeyPair, _>(k).await?,
-            MlsEntityId::EpochEncryptionKeyPair => self.remove::<StoredEpochEncryptionKeypair, _>(k).await?,
+            MlsEntityId::HpkePrivateKey => {
+                let hash = Sha256Hash::from_existing_hash(k)?;
+                self.remove::<StoredHpkePrivateKey>(&hash).await?
+            }
+            MlsEntityId::KeyPackage => self.remove_borrowed::<StoredKeypackage>(k).await?,
+            MlsEntityId::PskBundle => {
+                let hash = Sha256Hash::from_existing_hash(k)?;
+                self.remove::<StoredPskBundle>(&hash).await?
+            }
+            MlsEntityId::EncryptionKeyPair => {
+                let hash = Sha256Hash::from_existing_hash(k)?;
+                self.remove::<StoredEncryptionKeyPair>(&hash).await?
+            }
+            MlsEntityId::EpochEncryptionKeyPair => self.remove_borrowed::<StoredEpochEncryptionKeypair>(k).await?,
         }
 
         Ok(())

From a9c5cd9dbf5369df21c7c11b1289e3d50c4e6efb Mon Sep 17 00:00:00 2001
From: Peter Goodspeed-Niklaus <peter.goodspeed-niklaus@wire.com>
Date: Fri, 19 Dec 2025 11:31:39 +0100
Subject: [PATCH 16/48] chore(keystore): add `EntityId` constructor helpers

---
 .../transaction/dynamic_dispatch/entity_id.rs | 37 ++++++++++++++++---
 1 file changed, 31 insertions(+), 6 deletions(-)

diff --git a/keystore/src/transaction/dynamic_dispatch/entity_id.rs b/keystore/src/transaction/dynamic_dispatch/entity_id.rs
index 1d38313485..547c3bb09a 100644
--- a/keystore/src/transaction/dynamic_dispatch/entity_id.rs
+++ b/keystore/src/transaction/dynamic_dispatch/entity_id.rs
@@ -1,3 +1,5 @@
+use std::borrow::Cow;
+
 #[cfg(target_family = "wasm")]
 use crate::entities::E2eiRefreshToken;
 #[cfg(feature = "proteus-keystore")]
@@ -10,7 +12,7 @@ use crate::{
         StoredBufferedCommit, StoredCredential, StoredE2eiEnrollment, StoredEncryptionKeyPair,
         StoredEpochEncryptionKeypair, StoredHpkePrivateKey, StoredKeypackage, StoredPskBundle,
     },
-    traits::{Entity, EntityDatabaseMutation, KeyType as _, OwnedKeyType as _},
+    traits::{BorrowPrimaryKey, Entity, EntityDatabaseMutation, KeyType, OwnedKeyType as _},
     transaction::dynamic_dispatch::EntityType,
 };
 
@@ -29,17 +31,42 @@ impl EntityId {
             .ok_or(CryptoKeystoreError::InvalidPrimaryKeyBytes(self.typ.collection_name()))
     }
 
-    pub(crate) fn from_entity<E>(entity: &E) -> Self
+    fn from_key<E>(primary_key: Cow<'_, [u8]>) -> Self
     where
         E: Entity,
     {
         // assumption: nobody outside this crate will ever implement `Entity` on a foreign type
         let typ =
             EntityType::from_collection_name(E::COLLECTION_NAME).expect("all entities have a valid collection name");
-        let id = entity.primary_key().bytes().into_owned();
+        let id = primary_key.into_owned();
         Self { typ, id }
     }
 
+    pub(crate) fn from_entity<E>(entity: &E) -> Self
+    where
+        E: Entity,
+    {
+        Self::from_key::<E>(entity.primary_key().bytes())
+    }
+
+    pub(crate) fn from_primary_key<E>(primary_key: &E::PrimaryKey) -> Self
+    where
+        E: Entity,
+    {
+        Self::from_key::<E>(primary_key.bytes())
+    }
+
+    pub(crate) fn from_borrowed_primary_key<E>(primary_key: &E::BorrowedPrimaryKey) -> Self
+    where
+        E: Entity + BorrowPrimaryKey,
+    {
+        Self::from_key::<E>(primary_key.to_owned().bytes())
+    }
+
+    pub(crate) fn collection_name(&self) -> &'static str {
+        self.typ.collection_name()
+    }
+
     pub(crate) async fn execute_delete(&self, tx: &TransactionWrapper<'_>) -> CryptoKeystoreResult<bool> {
         match self.typ {
             EntityType::HpkePrivateKey => {
@@ -65,9 +92,7 @@ impl EntityId {
             EntityType::PersistedMlsPendingGroup => {
                 PersistedMlsPendingGroup::delete(tx, &self.primary_key::<PersistedMlsPendingGroup>()?).await
             }
-            EntityType::MlsPendingMessage => {
-                MlsPendingMessage::delete(tx, &self.primary_key::<MlsPendingMessage>()?).await
-            }
+            EntityType::MlsPendingMessage => MlsPendingMessage::delete_by_conversation_id(tx, &self.id).await,
             EntityType::StoredE2eiEnrollment => {
                 StoredE2eiEnrollment::delete(tx, &self.primary_key::<StoredE2eiEnrollment>()?).await
             }

From 70f337fb312dfdfa3533895170f24c0971c3efbc Mon Sep 17 00:00:00 2001
From: Peter Goodspeed-Niklaus <peter.goodspeed-niklaus@wire.com>
Date: Fri, 19 Dec 2025 11:33:13 +0100
Subject: [PATCH 17/48] chore(keystore): update transaction impls

Turns out that there's a real benefit to storing the dynamic-dispatch
entities the way we do, and that's a huge simplification of the
transaction stuff.
---
 keystore/src/transaction/mod.rs | 495 ++++++++++++++------------------
 1 file changed, 216 insertions(+), 279 deletions(-)

diff --git a/keystore/src/transaction/mod.rs b/keystore/src/transaction/mod.rs
index 0f955d9998..f9a9602a2a 100644
--- a/keystore/src/transaction/mod.rs
+++ b/keystore/src/transaction/mod.rs
@@ -1,38 +1,40 @@
 use std::{
-    collections::{HashMap, hash_map::Entry},
+    borrow::Cow,
+    collections::{HashMap, HashSet, hash_map::Entry},
     sync::Arc,
 };
 
 use async_lock::{RwLock, SemaphoreGuardArc};
 use itertools::Itertools;
-use zeroize::Zeroizing;
 
-#[cfg(feature = "proteus-keystore")]
-use crate::entities::proteus::*;
 use crate::{
     CryptoKeystoreError, CryptoKeystoreResult,
     connection::{Database, KeystoreDatabaseConnection},
-    entities::{ConsumerData, EntityBase, EntityFindParams, EntityTransactionExt, UniqueEntity, mls::*},
+    entities::{MlsPendingMessage, PersistedMlsGroupExt},
+    traits::{BorrowPrimaryKey, Entity, EntityBase as _, EntityDatabaseMutation, EntityDeleteBorrowed, KeyType},
     transaction::dynamic_dispatch::EntityId,
 };
 
-pub mod dynamic_dispatch;
+pub(crate) mod dynamic_dispatch;
 
-#[derive(Debug, Default, derive_more::Deref, derive_more::DerefMut)]
-struct InMemoryTable(HashMap<Vec<u8>, Zeroizing<Vec<u8>>>);
-
-type InMemoryCache = Arc<RwLock<HashMap<String, InMemoryTable>>>;
+/// table: primary key -> entity reference
+type InMemoryTable = HashMap<EntityId, dynamic_dispatch::Entity>;
+/// collection: collection name -> table
+type InMemoryCollection = Arc<RwLock<HashMap<&'static str, InMemoryTable>>>;
 
 /// This represents a transaction, where all operations will be done in memory and committed at the
 /// end
 #[derive(Debug, Clone)]
 pub(crate) struct KeystoreTransaction {
-    cache: InMemoryCache,
-    deleted: Arc<RwLock<Vec<EntityId>>>,
+    cache: InMemoryCollection,
+    deleted: Arc<RwLock<HashSet<EntityId>>>,
     _semaphore_guard: Arc<SemaphoreGuardArc>,
 }
 
 impl KeystoreTransaction {
+    /// Instantiate a new transaction.
+    ///
+    /// Requires a semaphore guard to ensure that only one exists at a time.
     pub(crate) async fn new(semaphore_guard: SemaphoreGuardArc) -> CryptoKeystoreResult<Self> {
         Ok(Self {
             cache: Default::default(),
@@ -41,357 +43,292 @@ impl KeystoreTransaction {
         })
     }
 
-    pub(crate) async fn save_mut<
-        E: crate::entities::Entity<ConnectionType = KeystoreDatabaseConnection> + EntityTransactionExt + Sync,
-    >(
-        &self,
-        mut entity: E,
-    ) -> CryptoKeystoreResult<E> {
-        entity.pre_save().await?;
+    /// Save an entity into this transaction.
+    ///
+    /// This is a multi-step process:
+    ///
+    /// - Adjust the entity by calling its [`pre_save()`][Entity::pre_save] method.
+    /// - Store the entity in an internal map.
+    /// - On [`Self::commit`], actually persist the entity into the supplied database.
+    pub(crate) async fn save<'a, E>(&self, mut entity: E) -> CryptoKeystoreResult<E::AutoGeneratedFields>
+    where
+        E: Entity<ConnectionType = KeystoreDatabaseConnection> + EntityDatabaseMutation<'a> + Send + Sync,
+    {
+        let auto_generated_fields = entity.pre_save().await?;
+
         let mut cache_guard = self.cache.write().await;
-        let table = cache_guard.entry(E::COLLECTION_NAME.to_string()).or_default();
-        let serialized = postcard::to_stdvec(&entity)?;
-        // Use merge_key() because `id_raw()` is not always unique for records.
-        // For `MlsPendingMessage` it's the id of the group it belongs to.
-        table.insert(entity.merge_key(), Zeroizing::new(serialized));
-        Ok(entity)
+        let table = cache_guard.entry(E::COLLECTION_NAME).or_default();
+        table.insert(EntityId::from_entity(&entity), entity.to_transaction_entity());
+
+        Ok(auto_generated_fields)
     }
 
-    pub(crate) async fn remove<
-        E: crate::entities::Entity<ConnectionType = KeystoreDatabaseConnection> + EntityTransactionExt,
-        S: AsRef<[u8]>,
-    >(
-        &self,
-        id: S,
-    ) -> CryptoKeystoreResult<()> {
+    async fn remove_by_entity_id<'a, E>(&self, entity_id: EntityId) -> CryptoKeystoreResult<()>
+    where
+        E: Entity + EntityDatabaseMutation<'a>,
+    {
+        // rm this entity from the set of added/modified items
+        // it might never touch the real db at all
         let mut cache_guard = self.cache.write().await;
-        if let Entry::Occupied(mut table) = cache_guard.entry(E::COLLECTION_NAME.to_string())
-            && let Entry::Occupied(cached_record) = table.get_mut().entry(id.as_ref().to_vec())
+        if let Entry::Occupied(mut table) = cache_guard.entry(E::COLLECTION_NAME)
+            && let Entry::Occupied(cached_record) = table.get_mut().entry(entity_id.clone())
         {
             cached_record.remove_entry();
         };
 
-        let mut deleted_list = self.deleted.write().await;
-        deleted_list.push(EntityId::from_collection_name(E::COLLECTION_NAME, id.as_ref())?);
+        // add this entity to the set of items which should be deleted from the persisted db
+        let mut deleted_set = self.deleted.write().await;
+        deleted_set.insert(entity_id);
         Ok(())
     }
 
-    pub(crate) async fn child_groups<E>(&self, entity: E, persisted_records: Vec<E>) -> CryptoKeystoreResult<Vec<E>>
+    /// Remove an entity by its primary key.
+    ///
+    /// Where the primary key has a distinct borrowed form, consider [`Self::remove_borrowed`].
+    ///
+    /// Note that this doesn't return whether or not anything was actually removed because
+    /// that won't happen until the transaction is committed.
+    pub(crate) async fn remove<'a, E>(&self, id: &E::PrimaryKey) -> CryptoKeystoreResult<()>
     where
-        E: crate::entities::Entity<ConnectionType = KeystoreDatabaseConnection> + PersistedMlsGroupExt + Sync,
+        E: Entity + EntityDatabaseMutation<'a>,
     {
-        // First get all raw groups from the cache, then deserialize them to enable filtering by there parent id
-        // matching `entity.id_raw()`.
-        let cached_records = self
-            .find_all_in_cache()
-            .await?
-            .into_iter()
-            .filter(|maybe_child: &E| {
+        let entity_id = EntityId::from_primary_key::<E>(id);
+        self.remove_by_entity_id::<E>(entity_id).await
+    }
+
+    /// Remove an entity by the borrowed form of its primary key.
+    ///
+    /// Note that this doesn't return whether or not anything was actually removed because
+    /// that won't happen until the transaction is committed.
+    pub(crate) async fn remove_borrowed<'a, E>(&self, id: &E::BorrowedPrimaryKey) -> CryptoKeystoreResult<()>
+    where
+        E: EntityDeleteBorrowed<'a> + BorrowPrimaryKey,
+    {
+        let entity_id = EntityId::from_borrowed_primary_key::<E>(id);
+        self.remove_by_entity_id::<E>(entity_id).await
+    }
+
+    pub(crate) async fn child_groups<E>(
+        &self,
+        entity: E,
+        persisted_records: impl IntoIterator<Item = E>,
+    ) -> CryptoKeystoreResult<Vec<E>>
+    where
+        E: Clone + Entity + BorrowPrimaryKey + PersistedMlsGroupExt + Send + Sync,
+        for<'pk> &'pk <E as BorrowPrimaryKey>::BorrowedPrimaryKey: KeyType,
+    {
+        // First get all raw groups from the cache, then filter by their parent id
+        let cached_records = self.find_all_in_cache::<E>().await;
+        let cached_records = cached_records
+            .iter()
+            .filter(|maybe_child| {
                 maybe_child
                     .parent_id()
-                    .map(|parent_id| parent_id == entity.id_raw())
+                    .map(|parent_id| parent_id == entity.borrow_primary_key().bytes().as_ref())
                     .unwrap_or_default()
             })
-            .collect();
+            .map(Arc::as_ref)
+            .map(Cow::Borrowed);
+
+        let persisted_records = persisted_records.into_iter().map(Cow::Owned);
 
-        Ok(self
-            .merge_records(cached_records, persisted_records, EntityFindParams::default())
-            .await)
+        Ok(self.merge_records(cached_records, persisted_records).await)
     }
 
-    pub(crate) async fn remove_pending_messages_by_conversation_id(
-        &self,
-        conversation_id: impl AsRef<[u8]> + Send,
-    ) -> CryptoKeystoreResult<()> {
-        // We cannot return an error from `retain()`, so we've got to do this dance with a mutable result.
-        let mut result = Ok(());
+    pub(crate) async fn remove_pending_messages_by_conversation_id(&self, conversation_id: impl AsRef<[u8]> + Send) {
+        let conversation_id = conversation_id.as_ref();
 
         let mut cache_guard = self.cache.write().await;
-        if let Entry::Occupied(mut table) = cache_guard.entry(MlsPendingMessage::COLLECTION_NAME.to_string()) {
-            table.get_mut().retain(|_key, record_bytes| {
-                postcard::from_bytes::<MlsPendingMessage>(record_bytes)
-                    .map(|pending_message| pending_message.foreign_id != conversation_id.as_ref())
-                    .inspect_err(|err| result = Err(err.clone()))
-                    .unwrap_or(false)
+        if let Entry::Occupied(mut table) = cache_guard.entry(MlsPendingMessage::COLLECTION_NAME) {
+            table.get_mut().retain(|_key, entity| {
+                let pending_message = entity
+                    .downcast::<MlsPendingMessage>()
+                    .expect("table for MlsPendingMessage contains only that type");
+                pending_message.foreign_id != conversation_id
             });
         }
 
-        let mut deleted_list = self.deleted.write().await;
-        deleted_list.push(EntityId::from_collection_name(
-            MlsPendingMessage::COLLECTION_NAME,
-            conversation_id.as_ref(),
-        )?);
-        result.map_err(Into::into)
+        let mut deleted_set = self.deleted.write().await;
+        deleted_set.insert(EntityId::from_borrowed_primary_key::<MlsPendingMessage>(
+            conversation_id,
+        ));
     }
 
     pub(crate) async fn find_pending_messages_by_conversation_id(
         &self,
         conversation_id: &[u8],
-        persisted_records: Vec<MlsPendingMessage>,
+        persisted_records: impl IntoIterator<Item = MlsPendingMessage>,
     ) -> CryptoKeystoreResult<Vec<MlsPendingMessage>> {
-        let cached_records = self
-            .find_all_in_cache::<MlsPendingMessage>()
-            .await?
-            .into_iter()
+        let persisted_records = persisted_records.into_iter().map(Cow::Owned);
+
+        let cached_records = self.find_all_in_cache::<MlsPendingMessage>().await;
+        let cached_records = cached_records
+            .iter()
             .filter(|pending_message| pending_message.foreign_id == conversation_id)
-            .collect();
-        let merged_records = self
-            .merge_records(cached_records, persisted_records, Default::default())
-            .await;
+            .map(Arc::as_ref)
+            .map(Cow::Borrowed);
+
+        let merged_records = self.merge_records(cached_records, persisted_records).await;
         Ok(merged_records)
     }
 
-    async fn find_in_cache<E>(&self, id: &[u8]) -> CryptoKeystoreResult<Option<E>>
+    async fn find_in_cache<E>(&self, entity_id: &EntityId) -> Option<Arc<E>>
     where
-        E: crate::entities::Entity<ConnectionType = KeystoreDatabaseConnection>,
+        E: Entity + Send + Sync,
     {
         let cache_guard = self.cache.read().await;
         cache_guard
             .get(E::COLLECTION_NAME)
-            .and_then(|table| {
-                table
-                    .get(id)
-                    .map(|record| -> CryptoKeystoreResult<_> { postcard::from_bytes::<E>(record).map_err(Into::into) })
-            })
-            .transpose()
+            .and_then(|table| table.get(entity_id).and_then(|entity| entity.downcast()))
     }
 
     /// The result of this function will have different contents for different scenarios:
     /// * `Some(Some(E))` - the transaction cache contains the record
     /// * `Some(None)` - the deletion of the record has been cached
     /// * `None` - there is no information about the record in the cache
-    pub(crate) async fn find<E>(&self, id: &[u8]) -> CryptoKeystoreResult<Option<Option<E>>>
+    async fn get_by_entity_id<E>(&self, entity_id: &EntityId) -> Option<Option<Arc<E>>>
     where
-        E: crate::entities::Entity<ConnectionType = KeystoreDatabaseConnection>,
+        E: Entity + Send + Sync,
     {
-        let maybe_cached_record = self.find_in_cache(id).await?;
-        if let Some(cached_record) = maybe_cached_record {
-            return Ok(Some(Some(cached_record)));
-        }
-
+        // when applying our transaction to the real database, we delete after inserting,
+        // so here we have to check for deletion before we check for existing values
         let deleted_list = self.deleted.read().await;
-        if deleted_list.contains(&EntityId::from_collection_name(E::COLLECTION_NAME, id)?) {
-            return Ok(Some(None));
+        if deleted_list.contains(entity_id) {
+            return Some(None);
         }
 
-        Ok(None)
+        self.find_in_cache::<E>(entity_id).await.map(Some)
     }
 
-    pub(crate) async fn find_unique<U: UniqueEntity<ConnectionType = KeystoreDatabaseConnection>>(
-        &self,
-    ) -> CryptoKeystoreResult<Option<U>> {
-        #[cfg(target_family = "wasm")]
-        let id = &U::ID;
-        #[cfg(not(target_family = "wasm"))]
-        let id = &[U::ID as u8];
-        let maybe_cached_record = self.find_in_cache::<U>(id).await?;
-        match maybe_cached_record {
-            Some(cached_record) => Ok(Some(cached_record)),
-            _ => {
-                // The deleted list doesn't have to be checked because unique entities don't implement
-                // deletion, just replace. So we can directly return None.
-                Ok(None)
-            }
-        }
+    /// The result of this function will have different contents for different scenarios:
+    /// * `Some(Some(E))` - the transaction cache contains the record
+    /// * `Some(None)` - the deletion of the record has been cached
+    /// * `None` - there is no information about the record in the cache
+    pub(crate) async fn get<E>(&self, id: &E::PrimaryKey) -> Option<Option<Arc<E>>>
+    where
+        E: Entity + Send + Sync,
+    {
+        let entity_id = EntityId::from_primary_key::<E>(id);
+        self.get_by_entity_id(&entity_id).await
     }
 
-    async fn find_all_in_cache<E: crate::entities::Entity<ConnectionType = KeystoreDatabaseConnection>>(
-        &self,
-    ) -> CryptoKeystoreResult<Vec<E>> {
+    /// The result of this function will have different contents for different scenarios:
+    /// * `Some(Some(E))` - the transaction cache contains the record
+    /// * `Some(None)` - the deletion of the record has been cached
+    /// * `None` - there is no information about the record in the cache
+    pub(crate) async fn get_borrowed<E>(&self, id: &E::BorrowedPrimaryKey) -> Option<Option<Arc<E>>>
+    where
+        E: Entity + BorrowPrimaryKey + Send + Sync,
+    {
+        let entity_id = EntityId::from_borrowed_primary_key::<E>(id);
+        self.get_by_entity_id(&entity_id).await
+    }
+
+    async fn find_all_in_cache<E>(&self) -> Vec<Arc<E>>
+    where
+        E: Entity + Send + Sync,
+    {
         let cache_guard = self.cache.read().await;
-        let cached_records = cache_guard
+        cache_guard
             .get(E::COLLECTION_NAME)
             .map(|table| {
                 table
                     .values()
-                    .map(|record| postcard::from_bytes::<E>(record).map_err(Into::into))
-                    .collect::<CryptoKeystoreResult<Vec<_>>>()
+                    .map(|record: &dynamic_dispatch::Entity| {
+                        record
+                            .downcast::<E>()
+                            .expect("all entries in this table are of this type")
+                            .clone()
+                    })
+                    .collect::<Vec<_>>()
             })
-            .transpose()?
-            .unwrap_or_default();
-        Ok(cached_records)
+            .unwrap_or_default()
     }
 
-    pub(crate) async fn find_all<E: crate::entities::Entity<ConnectionType = KeystoreDatabaseConnection>>(
-        &self,
-        persisted_records: Vec<E>,
-        params: EntityFindParams,
-    ) -> CryptoKeystoreResult<Vec<E>> {
-        let cached_records = self.find_all_in_cache().await?;
-        let merged_records = self.merge_records(cached_records, persisted_records, params).await;
+    pub(crate) async fn find_all<E>(&self, persisted_records: Vec<E>) -> CryptoKeystoreResult<Vec<E>>
+    where
+        E: Clone + Entity + Send + Sync,
+    {
+        let cached_records = self.find_all_in_cache().await;
+        let merged_records = self
+            .merge_records(
+                cached_records.iter().map(Arc::as_ref).map(Cow::Borrowed),
+                persisted_records.into_iter().map(Cow::Owned),
+            )
+            .await;
         Ok(merged_records)
     }
 
-    pub(crate) async fn find_many<E: crate::entities::Entity<ConnectionType = KeystoreDatabaseConnection>>(
-        &self,
-        persisted_records: Vec<E>,
-        ids: &[Vec<u8>],
-    ) -> CryptoKeystoreResult<Vec<E>> {
-        let records = self
-            .find_all(persisted_records, EntityFindParams::default())
-            .await?
-            .into_iter()
-            .filter(|record| ids.contains(&record.id_raw().to_vec()))
-            .collect();
-        Ok(records)
-    }
-
     /// Build a single list of unique records from two potentially overlapping lists.
     /// In case of overlap, records in `records_a` are prioritized.
     /// Identity from the perspective of this function is determined by the output of
-    /// [crate::entities::Entity::merge_key].
+    /// [Entity::merge_key].
     ///
     /// Further, the output list of records is built with respect to the provided [EntityFindParams]
     /// and the deleted records cached in this [Self] instance.
-    async fn merge_records<E: crate::entities::Entity<ConnectionType = KeystoreDatabaseConnection>>(
+    async fn merge_records<'a, E>(
         &self,
-        records_a: Vec<E>,
-        records_b: Vec<E>,
-        params: EntityFindParams,
-    ) -> Vec<E> {
-        let mut merged = records_a.into_iter().chain(records_b).unique_by(|e| e.merge_key());
-
+        records_a: impl IntoIterator<Item = Cow<'a, E>>,
+        records_b: impl IntoIterator<Item = Cow<'a, E>>,
+    ) -> Vec<E>
+    where
+        E: Clone + Entity,
+    {
         let deleted_records = self.deleted.read().await;
 
-        let merged: &mut dyn Iterator<Item = E> = if params.reverse { &mut merged.rev() } else { &mut merged };
-
-        merged
-            .filter(|record| !Self::record_is_in_deleted_list(record, &deleted_records))
-            .skip(params.offset.unwrap_or(0) as usize)
-            .take(params.limit.unwrap_or(u32::MAX) as usize)
+        records_a
+            .into_iter()
+            .chain(records_b)
+            .unique_by(|e| e.primary_key().bytes().into_owned())
+            .filter_map(|record| {
+                let id = EntityId::from_entity(record.as_ref());
+                (!deleted_records.contains(&id)).then_some(record.into_owned())
+            })
             .collect()
     }
 
-    fn record_is_in_deleted_list<E: crate::entities::Entity<ConnectionType = KeystoreDatabaseConnection>>(
-        record: &E,
-        deleted_records: &[EntityId],
-    ) -> bool {
-        let id = EntityId::from_collection_name(E::COLLECTION_NAME, record.id_raw());
-        let Ok(id) = id else { return false };
-        deleted_records.contains(&id)
-    }
-}
+    /// Persists all the operations in the database. It will effectively open a transaction
+    /// internally, perform all the buffered operations and commit.
+    pub(crate) async fn commit(&self, db: &Database) -> Result<(), CryptoKeystoreError> {
+        let conn = db.conn().await?;
+        let mut conn = conn.conn().await;
+
+        let cache = self.cache.read().await;
+        let deleted_ids = self.deleted.read().await;
+
+        let table_names_with_deletion = deleted_ids.iter().map(|entity_id| entity_id.collection_name());
+        let table_names_with_save = cache
+            .values()
+            .flat_map(|table| table.keys())
+            .map(|entity_id| entity_id.collection_name());
+        let mut tables = table_names_with_deletion
+            .chain(table_names_with_save)
+            .collect::<Vec<_>>();
+
+        if tables.is_empty() {
+            log::debug!("Empty transaction was committed.");
+            return Ok(());
+        }
 
-/// Persist all records cached in `$keystore_transaction` (first argument),
-/// using a transaction on `$db` (second argument).
-/// Use the provided types to read from the cache and write to the `$db`.
-///
-/// # Examples
-/// ```rust,ignore
-/// let transaction = KeystoreTransaction::new();
-/// let db = Connection::new();
-///
-/// // Commit records of all provided types
-/// commit_transaction!(
-///     transaction, db,
-///     [
-///         (identifier_01, StoredCredential),
-///         (identifier_02, StoredSignatureKeypair),
-///     ],
-/// );
-///
-/// // Commit records of provided types in the first list. Commit records of types in the second
-/// // list only if the "proteus-keystore" cargo feature is enabled.
-/// commit_transaction!(
-///     transaction, db,
-///     [
-///         (identifier_01, StoredCredential),
-///         (identifier_02, StoredSignatureKeypair),
-///     ],
-///     proteus_types: [
-///         (identifier_03, ProteusPrekey),
-///         (identifier_04, ProteusIdentity),
-///         (identifier_05, ProteusSession)
-///     ]
-/// );
-/// ```
-macro_rules! commit_transaction {
-    ($keystore_transaction:expr_2021, $db:expr_2021, [ $( ($records:ident, $entity:ty) ),*], proteus_types: [ $( ($conditional_records:ident, $conditional_entity:ty) ),*]) => {
-        #[cfg(feature = "proteus-keystore")]
-        commit_transaction!($keystore_transaction, $db, [ $( ($records, $entity) ),*], [ $( ($conditional_records, $conditional_entity) ),*]);
-
-        #[cfg(not(feature = "proteus-keystore"))]
-        commit_transaction!($keystore_transaction, $db, [ $( ($records, $entity) ),*]);
-    };
-     ($keystore_transaction:expr_2021, $db:expr_2021, $([ $( ($records:ident, $entity:ty) ),*]),*) => {
-            let cached_collections = ( $( $(
-            $keystore_transaction.find_all_in_cache::<$entity>().await?,
-                )* )* );
-
-             let ( $( $( $records, )* )* ) = cached_collections;
-
-            let conn = $db.conn().await?;
-            let mut conn = conn.conn().await;
-            let deleted_ids = $keystore_transaction.deleted.read().await;
-
-            let mut tables = Vec::new();
-            $( $(
-                if !$records.is_empty() {
-                    tables.push(<$entity>::COLLECTION_NAME);
-                }
-            )* )*
-
-            for deleted_id in deleted_ids.iter() {
-                tables.push(deleted_id.collection_name());
-            }
-
-            if tables.is_empty() {
-                log::debug!("Empty transaction was committed.");
-                return Ok(());
-            }
-
-            #[cfg(target_family = "wasm")]
-            let tx = conn.new_transaction(&tables).await?;
-            #[cfg(not(target_family = "wasm"))]
-            let tx = conn.transaction()?.into();
-
-             $( $(
-                if !$records.is_empty() {
-                    for record in $records {
-                        dynamic_dispatch::execute_save(&tx, &record.to_transaction_entity()).await?;
-                    }
-                }
-             )* )*
+        tables.sort_unstable();
+        tables.dedup();
 
+        // open a database transaction
+        #[cfg(target_family = "wasm")]
+        let tx = conn.new_transaction(&tables).await?;
+        #[cfg(not(target_family = "wasm"))]
+        let tx = conn.transaction()?.into();
 
-        for deleted_id in deleted_ids.iter() {
-            dynamic_dispatch::execute_delete(&tx, deleted_id).await?
+        for entity in cache.values().flat_map(|table| table.values()) {
+            entity.execute_save(&tx).await?;
         }
 
-         tx.commit_tx().await?;
-     };
-}
+        for deleted_id in deleted_ids.iter() {
+            deleted_id.execute_delete(&tx).await?;
+        }
 
-impl KeystoreTransaction {
-    /// Persists all the operations in the database. It will effectively open a transaction
-    /// internally, perform all the buffered operations and commit.
-    pub(crate) async fn commit(&self, db: &Database) -> Result<(), CryptoKeystoreError> {
-        commit_transaction!(
-            self, db,
-            [
-                (identifier_01, StoredCredential),
-                // (identifier_02, StoredSignatureKeypair),
-                (identifier_03, StoredHpkePrivateKey),
-                (identifier_04, StoredEncryptionKeyPair),
-                (identifier_05, StoredEpochEncryptionKeypair),
-                (identifier_06, StoredPskBundle),
-                (identifier_07, StoredKeypackage),
-                (identifier_08, PersistedMlsGroup),
-                (identifier_09, PersistedMlsPendingGroup),
-                (identifier_10, MlsPendingMessage),
-                (identifier_11, StoredE2eiEnrollment),
-                // (identifier_12, E2eiRefreshToken),
-                (identifier_13, E2eiAcmeCA),
-                (identifier_14, E2eiIntermediateCert),
-                (identifier_15, E2eiCrl),
-                (identifier_16, ConsumerData)
-            ],
-            proteus_types: [
-                (identifier_17, ProteusPrekey),
-                (identifier_18, ProteusIdentity),
-                (identifier_19, ProteusSession)
-            ]
-        );
+        // and commit everything
+        tx.commit_tx().await?;
 
         Ok(())
     }

From b1582a532f44054e34093579a21db16c8b9df96b Mon Sep 17 00:00:00 2001
From: Peter Goodspeed-Niklaus <peter.goodspeed-niklaus@wire.com>
Date: Fri, 19 Dec 2025 11:44:56 +0100
Subject: [PATCH 18/48] chore(keystore): update `trait FetchFromDatabase` for
 coherency

The trait used to be a bit loosely specified, since nobody was implementing it.
Now we know precisely what it needs and how to implement it.
---
 keystore/src/traits/fetch_from_database.rs | 25 ++++++++++++++++------
 1 file changed, 18 insertions(+), 7 deletions(-)

diff --git a/keystore/src/traits/fetch_from_database.rs b/keystore/src/traits/fetch_from_database.rs
index b2401c9544..01cfb3455c 100644
--- a/keystore/src/traits/fetch_from_database.rs
+++ b/keystore/src/traits/fetch_from_database.rs
@@ -1,9 +1,11 @@
+use std::borrow::Borrow;
+
 use async_trait::async_trait;
 
 use crate::{
     CryptoKeystoreResult,
     connection::KeystoreDatabaseConnection,
-    traits::{BorrowPrimaryKey, Entity, UniqueEntity, UniqueEntityExt},
+    traits::{BorrowPrimaryKey, Entity, KeyType, UniqueEntity, UniqueEntityExt},
 };
 
 /// Interface to fetch from the database either from the connection directly or through a
@@ -17,17 +19,17 @@ pub trait FetchFromDatabase: Send + Sync {
     /// Get an instance of `E` from the database by its primary key.
     async fn get<E>(&self, id: &<E as Entity>::PrimaryKey) -> CryptoKeystoreResult<Option<E>>
     where
-        E: Entity<ConnectionType = KeystoreDatabaseConnection>;
+        E: Entity<ConnectionType = KeystoreDatabaseConnection> + Clone + Send + Sync;
 
     /// Count the number of `E`s in the database.
     async fn count<E>(&self) -> CryptoKeystoreResult<u32>
     where
-        E: Entity<ConnectionType = KeystoreDatabaseConnection>;
+        E: Entity<ConnectionType = KeystoreDatabaseConnection> + Clone + Send + Sync;
 
     /// Load all `E`s from the database.
     async fn load_all<E>(&self) -> CryptoKeystoreResult<Vec<E>>
     where
-        E: Entity<ConnectionType = KeystoreDatabaseConnection>;
+        E: Entity<ConnectionType = KeystoreDatabaseConnection> + Clone + Send + Sync;
 
     /// Get an instance of `E` from the database by the borrowed form of its primary key.
     async fn get_borrowed<E>(
@@ -35,15 +37,24 @@ pub trait FetchFromDatabase: Send + Sync {
         id: &<E as BorrowPrimaryKey>::BorrowedPrimaryKey,
     ) -> CryptoKeystoreResult<Option<E>>
     where
-        E: BorrowPrimaryKey;
+        E: Entity<ConnectionType = KeystoreDatabaseConnection> + BorrowPrimaryKey + Clone + Send + Sync,
+        E::PrimaryKey: Borrow<E::BorrowedPrimaryKey>,
+        for<'a> &'a E::BorrowedPrimaryKey: KeyType;
 
     /// Get the requested unique entity from the database.
     async fn get_unique<'a, U>(&self) -> CryptoKeystoreResult<Option<U>>
     where
-        U: UniqueEntityExt<'a>;
+        U: UniqueEntityExt<'a> + Entity<ConnectionType = KeystoreDatabaseConnection> + Clone + Send + Sync,
+    {
+        self.get::<U>(&U::KEY).await
+    }
 
     /// Determine whether a unique entity is present in the database.
     async fn exists<'a, U>(&self) -> CryptoKeystoreResult<bool>
     where
-        U: UniqueEntityExt<'a>;
+        U: UniqueEntityExt<'a> + Entity<ConnectionType = KeystoreDatabaseConnection> + Clone + Send + Sync,
+    {
+        let count = self.count::<U>().await?;
+        Ok(count > 0)
+    }
 }

From fcbcc18d07f913f8a545b2f8e6b11287ce76e99e Mon Sep 17 00:00:00 2001
From: Peter Goodspeed-Niklaus <peter.goodspeed-niklaus@wire.com>
Date: Fri, 19 Dec 2025 11:46:26 +0100
Subject: [PATCH 19/48] chore(keystore): update proteus entity impls to use new
 traits

---
 keystore/src/entities/proteus.rs | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/keystore/src/entities/proteus.rs b/keystore/src/entities/proteus.rs
index 7b54278c3a..d8496b5c5a 100644
--- a/keystore/src/entities/proteus.rs
+++ b/keystore/src/entities/proteus.rs
@@ -1,6 +1,6 @@
 use zeroize::Zeroize;
 
-use crate::connection::FetchFromDatabase;
+use crate::traits::FetchFromDatabase as _;
 
 #[derive(core_crypto_macros::Debug, Clone, Zeroize, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
 #[zeroize(drop)]
@@ -75,7 +75,7 @@ impl ProteusPrekey {
             if id == limit {
                 return Err(crate::CryptoKeystoreError::NoFreePrekeyId);
             }
-            if conn.find::<Self>(&id.to_le_bytes()).await?.is_none() {
+            if conn.get::<Self>(&id).await?.is_none() {
                 break;
             }
             id += 1;

From df7fa6b45e159257120ee9dafab2aca3b15e186b Mon Sep 17 00:00:00 2001
From: Peter Goodspeed-Niklaus <peter.goodspeed-niklaus@wire.com>
Date: Fri, 19 Dec 2025 11:51:41 +0100
Subject: [PATCH 20/48] chore(keystore): update `Database` impls to use new
 traits

---
 keystore/src/connection/mod.rs                | 154 +++++++++---------
 .../src/connection/platform/generic/mod.rs    |   4 +-
 2 files changed, 81 insertions(+), 77 deletions(-)

diff --git a/keystore/src/connection/mod.rs b/keystore/src/connection/mod.rs
index 24d51d2756..c7445dfe96 100644
--- a/keystore/src/connection/mod.rs
+++ b/keystore/src/connection/mod.rs
@@ -1,5 +1,6 @@
-use std::{fmt, ops::Deref};
+use std::{borrow::Borrow, fmt, ops::Deref};
 
+use async_trait::async_trait;
 use sha2::{Digest as _, Sha256};
 use zeroize::{Zeroize, ZeroizeOnDrop};
 
@@ -29,8 +30,11 @@ use async_lock::{Mutex, MutexGuard, Semaphore};
 pub use self::platform::*;
 use crate::{
     CryptoKeystoreError, CryptoKeystoreResult,
-    entities::{Entity, EntityFindParams, EntityTransactionExt, MlsPendingMessage, StringEntityId, UniqueEntity},
-    transaction::KeystoreTransaction,
+    entities::{MlsPendingMessage, PersistedMlsGroupExt},
+    traits::{
+        BorrowPrimaryKey, Entity, EntityDatabaseMutation, EntityDeleteBorrowed, FetchFromDatabase, KeyType,
+    },
+    transaction::{KeystoreTransaction},
 };
 
 /// Limit on the length of a blob to be stored in the database.
@@ -143,19 +147,19 @@ const ALLOWED_CONCURRENT_TRANSACTIONS_COUNT: usize = 1;
 /// transaaction
 #[cfg_attr(target_family = "wasm", async_trait::async_trait(?Send))]
 #[cfg_attr(not(target_family = "wasm"), async_trait::async_trait)]
-pub trait FetchFromDatabase: Send + Sync {
+pub trait OldFetchFromDatabase: Send + Sync {
     async fn find<E: Entity<ConnectionType = KeystoreDatabaseConnection>>(
         &self,
         id: impl AsRef<[u8]> + Send,
     ) -> CryptoKeystoreResult<Option<E>>;
 
-    async fn find_unique<U: UniqueEntity<ConnectionType = KeystoreDatabaseConnection>>(
+    async fn find_unique<U: crate::entities::UniqueEntity<ConnectionType = KeystoreDatabaseConnection>>(
         &self,
     ) -> CryptoKeystoreResult<U>;
 
     async fn find_all<E: Entity<ConnectionType = KeystoreDatabaseConnection>>(
         &self,
-        params: EntityFindParams,
+        params: crate::entities::EntityFindParams,
     ) -> CryptoKeystoreResult<Vec<E>>;
 
     async fn find_many<E: Entity<ConnectionType = KeystoreDatabaseConnection>>(
@@ -275,9 +279,7 @@ impl Database {
         }
         Ok(())
     }
-}
 
-impl Database {
     /// Close this database and delete its contents.
     pub async fn wipe(&self) -> CryptoKeystoreResult<()> {
         self.take().await?.wipe().await
@@ -322,12 +324,11 @@ impl Database {
         Ok(())
     }
 
-    pub async fn child_groups<
-        E: Entity<ConnectionType = KeystoreDatabaseConnection> + crate::entities::PersistedMlsGroupExt + Sync,
-    >(
-        &self,
-        entity: E,
-    ) -> CryptoKeystoreResult<Vec<E>> {
+    pub async fn child_groups<'a, E>(&self, entity: E) -> CryptoKeystoreResult<Vec<E>>
+    where
+        E: Clone + Entity + EntityDatabaseMutation<'a> + BorrowPrimaryKey + PersistedMlsGroupExt + Send + Sync,
+        for<'pk> &'pk <E as BorrowPrimaryKey>::BorrowedPrimaryKey: KeyType,
+    {
         let mut conn = self.conn().await?;
         let persisted_records = entity.child_groups(conn.deref_mut()).await?;
 
@@ -338,29 +339,37 @@ impl Database {
         transaction.child_groups(entity, persisted_records).await
     }
 
-    pub async fn save<E: Entity<ConnectionType = KeystoreDatabaseConnection> + Sync + EntityTransactionExt>(
-        &self,
-        entity: E,
-    ) -> CryptoKeystoreResult<E> {
+    pub async fn save<'a, E>(&self, entity: E) -> CryptoKeystoreResult<E::AutoGeneratedFields>
+    where
+        E: Entity + EntityDatabaseMutation<'a> + Send + Sync,
+    {
         let transaction_guard = self.transaction.lock().await;
         let Some(transaction) = transaction_guard.as_ref() else {
             return Err(CryptoKeystoreError::MutatingOperationWithoutTransaction);
         };
-        transaction.save_mut(entity).await
+        transaction.save(entity).await
     }
 
-    pub async fn remove<
-        E: Entity<ConnectionType = KeystoreDatabaseConnection> + EntityTransactionExt,
-        S: AsRef<[u8]>,
-    >(
-        &self,
-        id: S,
-    ) -> CryptoKeystoreResult<()> {
+    pub async fn remove<'a, E>(&self, id: &E::PrimaryKey) -> CryptoKeystoreResult<()>
+    where
+        E: Entity + EntityDatabaseMutation<'a>,
+    {
+        let transaction_guard = self.transaction.lock().await;
+        let Some(transaction) = transaction_guard.as_ref() else {
+            return Err(CryptoKeystoreError::MutatingOperationWithoutTransaction);
+        };
+        transaction.remove::<E>(id).await
+    }
+
+    pub async fn remove_borrowed<'a, E>(&self, id: &E::BorrowedPrimaryKey) -> CryptoKeystoreResult<()>
+    where
+        E: Entity + EntityDatabaseMutation<'a> + BorrowPrimaryKey + EntityDeleteBorrowed<'a>,
+    {
         let transaction_guard = self.transaction.lock().await;
         let Some(transaction) = transaction_guard.as_ref() else {
             return Err(CryptoKeystoreError::MutatingOperationWithoutTransaction);
         };
-        transaction.remove::<E, S>(id).await
+        transaction.remove_borrowed::<E>(id).await
     }
 
     pub async fn find_pending_messages_by_conversation_id(
@@ -390,81 +399,76 @@ impl Database {
         };
         transaction
             .remove_pending_messages_by_conversation_id(conversation_id)
-            .await
+            .await;
+        Ok(())
     }
+
+
 }
 
-#[cfg_attr(target_family = "wasm", async_trait::async_trait(?Send))]
-#[cfg_attr(not(target_family = "wasm"), async_trait::async_trait)]
+#[cfg_attr(target_family = "wasm", async_trait(?Send))]
+#[cfg_attr(not(target_family = "wasm"), async_trait)]
 impl FetchFromDatabase for Database {
-    async fn find<E: Entity<ConnectionType = KeystoreDatabaseConnection>>(
-        &self,
-        id: impl AsRef<[u8]> + Send,
-    ) -> CryptoKeystoreResult<Option<E>> {
+    async fn get<E>(&self, id: &<E as Entity>::PrimaryKey) -> CryptoKeystoreResult<Option<E>>
+    where
+        E: Entity<ConnectionType = KeystoreDatabaseConnection> + Clone + Send + Sync,
+    {
         // If a transaction is in progress...
         if let Some(transaction) = self.transaction.lock().await.as_ref()
             //... and it has information about this entity, ...
-            && let Some(cached_record) = transaction.find::<E>(id.as_ref()).await?
-        {
-            // ... return that result
-            return Ok(cached_record);
+            && let Some(cached_record) = transaction.get(id).await {
+                return Ok(cached_record.map(Arc::unwrap_or_clone));
         }
 
         // Otherwise get it from the database
         let mut conn = self.conn().await?;
-        E::find_one(&mut conn, &id.as_ref().into()).await
+        E::get(&mut conn, id).await
     }
 
-    async fn find_unique<U: UniqueEntity>(&self) -> CryptoKeystoreResult<U> {
-        // If a transaction is in progress...
+    async fn get_borrowed<E>(&self, id: &<E as BorrowPrimaryKey>::BorrowedPrimaryKey) -> CryptoKeystoreResult<Option<E>>
+    where
+        E: Entity<ConnectionType = KeystoreDatabaseConnection> + BorrowPrimaryKey + Clone + Send + Sync,
+        E::PrimaryKey: Borrow<E::BorrowedPrimaryKey>,
+        for<'a> &'a E::BorrowedPrimaryKey: KeyType,
+    {
+                // If a transaction is in progress...
         if let Some(transaction) = self.transaction.lock().await.as_ref()
             //... and it has information about this entity, ...
-            && let Some(cached_record) = transaction.find_unique::<U>().await?
-        {
-            // ... return that result
-            return Ok(cached_record);
+            && let Some(cached_record) = transaction.get_borrowed(id).await {
+                return Ok(cached_record.map(Arc::unwrap_or_clone));
         }
+
         // Otherwise get it from the database
         let mut conn = self.conn().await?;
-        U::find_unique(&mut conn).await
+        E::get_borrowed(&mut conn, id).await
     }
 
-    async fn find_all<E: Entity<ConnectionType = KeystoreDatabaseConnection>>(
-        &self,
-        params: EntityFindParams,
-    ) -> CryptoKeystoreResult<Vec<E>> {
-        let mut conn = self.conn().await?;
-        let persisted_records = E::find_all(&mut conn, params.clone()).await?;
-
-        let transaction_guard = self.transaction.lock().await;
-        let Some(transaction) = transaction_guard.as_ref() else {
-            return Ok(persisted_records);
-        };
-        transaction.find_all(persisted_records, params).await
+    async fn count<E>(&self) -> CryptoKeystoreResult<u32>
+    where
+        E: Entity<ConnectionType = KeystoreDatabaseConnection> + Clone + Send + Sync,
+    {
+        if self.transaction.lock().await.is_some() {
+            // Unfortunately, we have to do this because of possible record id overlap
+            // between cache and db.
+            let count = self.load_all::<E>().await?.len();
+            Ok(count as _)
+        } else {
+            let mut conn = self.conn().await?;
+            E::count(&mut conn).await
+        }
     }
 
-    async fn find_many<E: Entity<ConnectionType = KeystoreDatabaseConnection>>(
-        &self,
-        ids: &[Vec<u8>],
-    ) -> CryptoKeystoreResult<Vec<E>> {
-        let entity_ids: Vec<StringEntityId> = ids.iter().map(|id| id.as_slice().into()).collect();
+    async fn load_all<E>(&self) -> CryptoKeystoreResult<Vec<E>>
+    where
+        E: Entity<ConnectionType = KeystoreDatabaseConnection> + Clone + Send + Sync,
+    {
         let mut conn = self.conn().await?;
-        let persisted_records = E::find_many(&mut conn, &entity_ids).await?;
+        let persisted_records = E::load_all(&mut conn).await?;
 
         let transaction_guard = self.transaction.lock().await;
         let Some(transaction) = transaction_guard.as_ref() else {
             return Ok(persisted_records);
         };
-        transaction.find_many(persisted_records, ids).await
-    }
-
-    async fn count<E: Entity<ConnectionType = KeystoreDatabaseConnection>>(&self) -> CryptoKeystoreResult<usize> {
-        if self.transaction.lock().await.is_some() {
-            // Unfortunately, we have to do this because of possible record id overlap
-            // between cache and db.
-            return Ok(self.find_all::<E>(Default::default()).await?.len());
-        };
-        let mut conn = self.conn().await?;
-        E::count(&mut conn).await
+        transaction.find_all(persisted_records).await
     }
 }
diff --git a/keystore/src/connection/platform/generic/mod.rs b/keystore/src/connection/platform/generic/mod.rs
index dfd1ac1713..8389a82c92 100644
--- a/keystore/src/connection/platform/generic/mod.rs
+++ b/keystore/src/connection/platform/generic/mod.rs
@@ -324,7 +324,7 @@ mod migration_test {
     use crate::{
         ConnectionType, Database, DatabaseKey,
         connection::{FetchFromDatabase, MigrationTarget},
-        entities::{EntityFindParams, StoredCredential},
+        entities::StoredCredential,
     };
 
     const DB: &[u8] = include_bytes!("../../../../../crypto-ffi/bindings/jvm/src/test/resources/db-v10002003.sqlite");
@@ -438,7 +438,7 @@ mod migration_test {
                 .await
                 .unwrap();
             let deduplicated_credentials = db
-                .find_all::<StoredCredential>(EntityFindParams::default())
+                .load_all::<StoredCredential>()
                 .await
                 .expect("deduplicated credentials");
 

From 3830119645f0ab363b44129a9e2c5b6a94c85441 Mon Sep 17 00:00:00 2001
From: Peter Goodspeed-Niklaus <peter.goodspeed-niklaus@wire.com>
Date: Fri, 19 Dec 2025 11:56:19 +0100
Subject: [PATCH 21/48] chore(keystore): `cargo +nightly fmt`

---
 keystore/src/connection/mod.rs | 20 +++++++++-----------
 keystore/src/entities/mls.rs   |  3 ++-
 2 files changed, 11 insertions(+), 12 deletions(-)

diff --git a/keystore/src/connection/mod.rs b/keystore/src/connection/mod.rs
index c7445dfe96..14223662d5 100644
--- a/keystore/src/connection/mod.rs
+++ b/keystore/src/connection/mod.rs
@@ -31,10 +31,8 @@ pub use self::platform::*;
 use crate::{
     CryptoKeystoreError, CryptoKeystoreResult,
     entities::{MlsPendingMessage, PersistedMlsGroupExt},
-    traits::{
-        BorrowPrimaryKey, Entity, EntityDatabaseMutation, EntityDeleteBorrowed, FetchFromDatabase, KeyType,
-    },
-    transaction::{KeystoreTransaction},
+    traits::{BorrowPrimaryKey, Entity, EntityDatabaseMutation, EntityDeleteBorrowed, FetchFromDatabase, KeyType},
+    transaction::KeystoreTransaction,
 };
 
 /// Limit on the length of a blob to be stored in the database.
@@ -402,8 +400,6 @@ impl Database {
             .await;
         Ok(())
     }
-
-
 }
 
 #[cfg_attr(target_family = "wasm", async_trait(?Send))]
@@ -416,8 +412,9 @@ impl FetchFromDatabase for Database {
         // If a transaction is in progress...
         if let Some(transaction) = self.transaction.lock().await.as_ref()
             //... and it has information about this entity, ...
-            && let Some(cached_record) = transaction.get(id).await {
-                return Ok(cached_record.map(Arc::unwrap_or_clone));
+            && let Some(cached_record) = transaction.get(id).await
+        {
+            return Ok(cached_record.map(Arc::unwrap_or_clone));
         }
 
         // Otherwise get it from the database
@@ -431,11 +428,12 @@ impl FetchFromDatabase for Database {
         E::PrimaryKey: Borrow<E::BorrowedPrimaryKey>,
         for<'a> &'a E::BorrowedPrimaryKey: KeyType,
     {
-                // If a transaction is in progress...
+        // If a transaction is in progress...
         if let Some(transaction) = self.transaction.lock().await.as_ref()
             //... and it has information about this entity, ...
-            && let Some(cached_record) = transaction.get_borrowed(id).await {
-                return Ok(cached_record.map(Arc::unwrap_or_clone));
+            && let Some(cached_record) = transaction.get_borrowed(id).await
+        {
+            return Ok(cached_record.map(Arc::unwrap_or_clone));
         }
 
         // Otherwise get it from the database
diff --git a/keystore/src/entities/mls.rs b/keystore/src/entities/mls.rs
index eb3db3fd1b..0e7adc25eb 100644
--- a/keystore/src/entities/mls.rs
+++ b/keystore/src/entities/mls.rs
@@ -53,7 +53,8 @@ where
         // so both `parent_group` and this method should just be implemented directly on `PersistedMlsGroup`.
         let entities = Self::load_all(conn).await?;
 
-        // for whatever reason rustc needs each of these distinct bindings to prove to itself that the lifetimes work out
+        // for whatever reason rustc needs each of these distinct bindings to prove to itself that the lifetimes work
+        // out
         let id = self.borrow_primary_key();
         let id = id.bytes();
         let id = id.as_ref();

From f3db970965cdc54d8f7d56f662469879a31395fd Mon Sep 17 00:00:00 2001
From: Peter Goodspeed-Niklaus <peter.goodspeed-niklaus@wire.com>
Date: Fri, 19 Dec 2025 13:14:48 +0100
Subject: [PATCH 22/48] refactor(keystore): break out primary key from entity
 type

So, I thought there were no more trait refactors incoming. But then
it turned out that there was a bug in the auto-implementations
of the unique entity helper on wasm only: `Entity` had an auto-impl
where `T: UniqueEntity` (so it could reference `T::KEY`), but
`UniqueEntity` had an auto-impl where `T: Entity` (so it could reference
`T::PrimaryKey`), and rustc was running out of stack trying to figure
out if a particular instance of `T` implemented both or neither of those.

So we broke out `PrimaryKey` from `Entity` which broke that loop,
and as a nice side effect means that there's now better symmetry
between `PrimaryKey` and `BorrowPrimaryKey`, and `Entity` and `EntityGetBorrowed`
(as well as `EntityDatabaseMutation` and `EntityDeleteBorrowed`).

But now I have to fix all the impls.
---
 keystore/src/traits/entity.rs                 | 31 ++-------
 .../src/traits/entity_database_mutation.rs    |  9 +--
 keystore/src/traits/fetch_from_database.rs    |  4 +-
 .../traits/item_encryption/decrypt_data.rs    |  2 +-
 keystore/src/traits/mod.rs                    |  4 +-
 keystore/src/traits/primary_key.rs            | 30 ++++++++
 keystore/src/traits/unique_entity.rs          | 68 ++++++++++---------
 7 files changed, 82 insertions(+), 66 deletions(-)
 create mode 100644 keystore/src/traits/primary_key.rs

diff --git a/keystore/src/traits/entity.rs b/keystore/src/traits/entity.rs
index ee817f35c3..acaa782f12 100644
--- a/keystore/src/traits/entity.rs
+++ b/keystore/src/traits/entity.rs
@@ -4,7 +4,10 @@ use async_trait::async_trait;
 
 use crate::{
     CryptoKeystoreResult,
-    traits::{EntityBase, KeyType, OwnedKeyType},
+    traits::{
+        EntityBase, KeyType, OwnedKeyType,
+        primary_key::{BorrowPrimaryKey, PrimaryKey},
+    },
 };
 
 /// Something which can be stored in our database.
@@ -12,19 +15,7 @@ use crate::{
 /// It has a primary key, which uniquely identifies it.
 #[cfg_attr(target_family = "wasm", async_trait(?Send))]
 #[cfg_attr(not(target_family = "wasm"), async_trait)]
-pub trait Entity: EntityBase {
-    /// Each distinct `PrimaryKey` uniquely identifies either 0 or 1 instance.
-    ///
-    /// This constraint should be enforced at the DB level.
-    type PrimaryKey: OwnedKeyType;
-
-    /// Get this entity's primary key.
-    ///
-    /// This must return an owned type, because there are some entities for which only owned primary keys are possible.
-    /// However, entities which have primary keys owned within the entity itself should consider also implementing
-    /// [`BorrowPrimaryKey`] for greater efficiency.
-    fn primary_key(&self) -> Self::PrimaryKey;
-
+pub trait Entity: EntityBase + PrimaryKey {
     /// Get an entity by its primary key.
     ///
     /// For entites whose primary key has a distinct borrowed type, it is best to implement this as a direct
@@ -44,19 +35,9 @@ pub trait Entity: EntityBase {
     async fn load_all(conn: &mut Self::ConnectionType) -> CryptoKeystoreResult<Vec<Self>>;
 }
 
-/// An extension trait which should be implemented for all entities whose primary key has a distinct borrowed form.
-///
-/// i.e. `String`, `Vec<u8>`, etc.
 #[cfg_attr(target_family = "wasm", async_trait(?Send))]
 #[cfg_attr(not(target_family = "wasm"), async_trait)]
-pub trait BorrowPrimaryKey: Entity {
-    type BorrowedPrimaryKey: ?Sized + ToOwned<Owned = Self::PrimaryKey>;
-
-    /// Borrow this entity's primary key without copying any data.
-    ///
-    /// This borrowed key has a lifetime tied to that of this entity.
-    fn borrow_primary_key(&self) -> &Self::BorrowedPrimaryKey;
-
+pub trait EntityGetBorrowed: Entity + BorrowPrimaryKey {
     /// Get an entity by a borrowed form of its primary key.
     ///
     /// The type signature here is somewhat complicated, but it breaks down simply: if our primary key is something
diff --git a/keystore/src/traits/entity_database_mutation.rs b/keystore/src/traits/entity_database_mutation.rs
index d8c599135c..be31407504 100644
--- a/keystore/src/traits/entity_database_mutation.rs
+++ b/keystore/src/traits/entity_database_mutation.rs
@@ -4,7 +4,7 @@ use async_trait::async_trait;
 
 use crate::{
     CryptoKeystoreResult,
-    traits::{BorrowPrimaryKey, Entity, KeyType},
+    traits::{BorrowPrimaryKey, Entity, KeyType, PrimaryKey},
 };
 
 /// Extend an [`Entity`] with db-mutating operations which can be performed when provided with a transaction.
@@ -42,7 +42,7 @@ pub trait EntityDatabaseMutation<'a>: Entity<ConnectionType = crate::connection:
     ///     <Self as EntityTransactionDeleteBorrowed<'a>>::delete_borrowed(tx, id).await
     /// }
     /// ```
-    async fn delete(tx: &Self::Transaction, id: &<Self as Entity>::PrimaryKey) -> CryptoKeystoreResult<bool>;
+    async fn delete(tx: &Self::Transaction, id: &Self::PrimaryKey) -> CryptoKeystoreResult<bool>;
 }
 
 /// Extend an [`Entity`] with db-mutating operations which can be performed when provided with a transaction.
@@ -53,10 +53,7 @@ pub trait EntityDeleteBorrowed<'a>: EntityDatabaseMutation<'a> + BorrowPrimaryKe
     ///
     /// The type signature here is somewhat complicated, but it breaks down simply: if our primary key is something
     /// like `Vec<u8>`, we want to be able to use this method even if what we have on hand is `&[u8]`.
-    async fn delete_borrowed(
-        tx: &<Self as EntityDatabaseMutation<'a>>::Transaction,
-        id: &<Self as BorrowPrimaryKey>::BorrowedPrimaryKey,
-    ) -> CryptoKeystoreResult<bool>
+    async fn delete_borrowed(tx: &Self::Transaction, id: &Self::BorrowedPrimaryKey) -> CryptoKeystoreResult<bool>
     where
         for<'pk> &'pk Self::BorrowedPrimaryKey: KeyType;
 }
diff --git a/keystore/src/traits/fetch_from_database.rs b/keystore/src/traits/fetch_from_database.rs
index 01cfb3455c..cd5436e6a1 100644
--- a/keystore/src/traits/fetch_from_database.rs
+++ b/keystore/src/traits/fetch_from_database.rs
@@ -5,7 +5,7 @@ use async_trait::async_trait;
 use crate::{
     CryptoKeystoreResult,
     connection::KeystoreDatabaseConnection,
-    traits::{BorrowPrimaryKey, Entity, KeyType, UniqueEntity, UniqueEntityExt},
+    traits::{BorrowPrimaryKey, Entity, KeyType, PrimaryKey, UniqueEntity, UniqueEntityExt},
 };
 
 /// Interface to fetch from the database either from the connection directly or through a
@@ -17,7 +17,7 @@ use crate::{
 #[cfg_attr(not(target_family = "wasm"), async_trait)]
 pub trait FetchFromDatabase: Send + Sync {
     /// Get an instance of `E` from the database by its primary key.
-    async fn get<E>(&self, id: &<E as Entity>::PrimaryKey) -> CryptoKeystoreResult<Option<E>>
+    async fn get<E>(&self, id: &E::PrimaryKey) -> CryptoKeystoreResult<Option<E>>
     where
         E: Entity<ConnectionType = KeystoreDatabaseConnection> + Clone + Send + Sync;
 
diff --git a/keystore/src/traits/item_encryption/decrypt_data.rs b/keystore/src/traits/item_encryption/decrypt_data.rs
index cdd01d9ad6..e83cd8b1a9 100644
--- a/keystore/src/traits/item_encryption/decrypt_data.rs
+++ b/keystore/src/traits/item_encryption/decrypt_data.rs
@@ -29,7 +29,7 @@ pub trait DecryptData: Entity {
     /// Decrypt some data, symmetrically to the process [`encrypt_data`][super::EncryptData::encrypt_data] uses.
     fn decrypt_data(
         cipher: &aes_gcm::Aes256Gcm,
-        primary_key: &<Self as Entity>::PrimaryKey,
+        primary_key: &Self::PrimaryKey,
         data: &[u8],
     ) -> CryptoKeystoreResult<Vec<u8>>;
 }
diff --git a/keystore/src/traits/mod.rs b/keystore/src/traits/mod.rs
index ae9d54a922..8c26414d41 100644
--- a/keystore/src/traits/mod.rs
+++ b/keystore/src/traits/mod.rs
@@ -9,12 +9,14 @@ mod entity_database_mutation;
 mod fetch_from_database;
 mod item_encryption;
 mod key_type;
+mod primary_key;
 mod unique_entity;
 
-pub use entity::{BorrowPrimaryKey, Entity};
+pub use entity::{Entity, EntityGetBorrowed};
 pub use entity_base::EntityBase;
 pub use entity_database_mutation::{EntityDatabaseMutation, EntityDeleteBorrowed};
 pub use fetch_from_database::FetchFromDatabase;
 pub use item_encryption::{DecryptData, Decryptable, Decrypting, EncryptData, Encrypting};
 pub use key_type::{KeyType, OwnedKeyType};
+pub use primary_key::{BorrowPrimaryKey, PrimaryKey};
 pub use unique_entity::{UniqueEntity, UniqueEntityExt, UniqueEntityImplementationHelper};
diff --git a/keystore/src/traits/primary_key.rs b/keystore/src/traits/primary_key.rs
new file mode 100644
index 0000000000..8aa218812b
--- /dev/null
+++ b/keystore/src/traits/primary_key.rs
@@ -0,0 +1,30 @@
+use async_trait::async_trait;
+
+use crate::traits::OwnedKeyType;
+
+/// Something which has a distinct primary key which can uniquely identify it.
+pub trait PrimaryKey {
+    /// Each distinct `PrimaryKey` uniquely identifies either 0 or 1 instance.
+    ///
+    /// This constraint should be enforced at the DB level.
+    type PrimaryKey: OwnedKeyType;
+
+    /// Get this entity's primary key.
+    ///
+    /// This must return an owned type, because there are some entities for which only owned primary keys are possible.
+    /// However, entities which have primary keys owned within the entity itself should consider also implementing
+    /// [`BorrowPrimaryKey`] for greater efficiency.
+    fn primary_key(&self) -> Self::PrimaryKey;
+}
+
+/// Something whose primary key can be borrowed as a distinct type.
+///
+/// i.e. `String`, `Vec<u8>`, etc.
+pub trait BorrowPrimaryKey: PrimaryKey {
+    type BorrowedPrimaryKey: ?Sized + ToOwned<Owned = Self::PrimaryKey>;
+
+    /// Borrow this entity's primary key without copying any data.
+    ///
+    /// This borrowed key has a lifetime tied to that of this entity.
+    fn borrow_primary_key(&self) -> &Self::BorrowedPrimaryKey;
+}
diff --git a/keystore/src/traits/unique_entity.rs b/keystore/src/traits/unique_entity.rs
index 3560081713..9d2bc3d7c7 100644
--- a/keystore/src/traits/unique_entity.rs
+++ b/keystore/src/traits/unique_entity.rs
@@ -1,6 +1,6 @@
 use async_trait::async_trait;
 #[cfg(not(target_family = "wasm"))]
-use rusqlite::{OptionalExtension as _, params};
+use rusqlite::{OptionalExtension as _, ToSql, params};
 #[cfg(target_family = "wasm")]
 use serde::de::DeserializeOwned;
 
@@ -10,14 +10,16 @@ use crate::entities::{count_helper, count_helper_tx, delete_helper, load_all_hel
 use crate::traits::{Decryptable, Decrypting, Encrypting, KeyType as _};
 use crate::{
     CryptoKeystoreResult,
-    connection::TransactionWrapper,
-    traits::{Entity, EntityBase, entity_database_mutation::EntityDatabaseMutation},
+    connection::{KeystoreDatabaseConnection, TransactionWrapper},
+    traits::{Entity, EntityBase, PrimaryKey, entity_database_mutation::EntityDatabaseMutation},
 };
 
 /// A unique entity can appear either 0 or 1 times in the database.
-pub trait UniqueEntity: EntityBase<ConnectionType = crate::connection::KeystoreDatabaseConnection> + Entity {
+pub trait UniqueEntity:
+    EntityBase<ConnectionType = crate::connection::KeystoreDatabaseConnection> + PrimaryKey
+{
     /// The id used as they key when storing this entity in a KV store.
-    const KEY: <Self as Entity>::PrimaryKey;
+    const KEY: Self::PrimaryKey;
 }
 
 /// Unique entities get some convenience methods implemented automatically.
@@ -96,6 +98,7 @@ where
 ///
 /// If you implement this trait, you get the following traits auto-implemented:
 ///
+/// - `PrimaryKey`
 /// - `UniqueEntity`
 /// - `Entity`
 /// - `EntityDatabaseMutation`
@@ -110,13 +113,29 @@ pub trait UniqueEntityImplementationHelper {
     fn content(&self) -> &[u8];
 }
 
+impl<T> PrimaryKey for T
+where
+    T: EntityBase<ConnectionType = KeystoreDatabaseConnection> + UniqueEntityImplementationHelper,
+{
+    // The old keystore trait used usize as the primary key type, but that would vary
+    // in width across various implementations and so is intentionally not a `KeyType`.
+    // So we distinguish betwen `u32` and `u64` according to whether or not we're on wasm.
+    #[cfg(target_family = "wasm")]
+    type PrimaryKey = u32;
+    #[cfg(not(target_family = "wasm"))]
+    type PrimaryKey = u64;
+
+    fn primary_key(&self) -> Self::PrimaryKey {
+        Self::KEY
+    }
+}
+
 #[cfg(target_family = "wasm")]
-#[async_trait(?Send)]
 impl<T> UniqueEntity for T
 where
     T: EntityBase<ConnectionType = crate::connection::KeystoreDatabaseConnection>
         + UniqueEntityImplementationHelper
-        + Entity<PrimaryKey = u32>,
+        + PrimaryKey<PrimaryKey = u32>,
 {
     const KEY: u32 = 0;
 }
@@ -131,15 +150,6 @@ where
         + Decryptable<'static>,
     <T as Decryptable<'static>>::DecryptableFrom: DeserializeOwned,
 {
-    // The old trait used usize as the primary key type, but that would vary
-    // in width across various implementations and so is intentionally not a `KeyType`.
-    // Instead we use `u32` which should be the same width on wasm.
-    type PrimaryKey = u32;
-
-    fn primary_key(&self) -> Self::PrimaryKey {
-        Self::KEY
-    }
-
     async fn get(conn: &mut Self::ConnectionType, key: &Self::PrimaryKey) -> CryptoKeystoreResult<Option<Self>> {
         conn.storage().new_get(key.bytes().as_ref()).await
     }
@@ -175,16 +185,17 @@ where
         tx.new_count::<Self>().await
     }
 
-    async fn delete(tx: &Self::Transaction, id: &<Self as Entity>::PrimaryKey) -> CryptoKeystoreResult<bool> {
+    async fn delete(tx: &Self::Transaction, id: &Self::PrimaryKey) -> CryptoKeystoreResult<bool> {
         tx.new_delete::<Self>(id.bytes().as_ref()).await
     }
 }
 
 #[cfg(not(target_family = "wasm"))]
-#[async_trait]
 impl<T> UniqueEntity for T
 where
-    T: EntityBase<ConnectionType = crate::connection::KeystoreDatabaseConnection> + UniqueEntityImplementationHelper,
+    T: EntityBase<ConnectionType = crate::connection::KeystoreDatabaseConnection>
+        + UniqueEntityImplementationHelper
+        + PrimaryKey<PrimaryKey = u64>,
 {
     const KEY: u64 = 0;
 }
@@ -193,18 +204,11 @@ where
 #[async_trait]
 impl<T> Entity for T
 where
-    T: EntityBase<ConnectionType = crate::connection::KeystoreDatabaseConnection> + UniqueEntityImplementationHelper,
+    T: EntityBase<ConnectionType = crate::connection::KeystoreDatabaseConnection>
+        + PrimaryKey
+        + UniqueEntityImplementationHelper,
+    <T as PrimaryKey>::PrimaryKey: ToSql,
 {
-    // The old trait used usize as the primary key type, not u64, but that would vary
-    // in width across various implementations and so is intentionally not a `KeyType`.
-    // Instead we use `u64` which should be the same width on the expected runtimes
-    // for non-wasm.
-    type PrimaryKey = u64;
-
-    fn primary_key(&self) -> Self::PrimaryKey {
-        Self::KEY
-    }
-
     async fn get(conn: &mut Self::ConnectionType, key: &Self::PrimaryKey) -> CryptoKeystoreResult<Option<Self>> {
         let conn = conn.conn().await;
         let mut statement = conn.prepare_cached(&format!(
@@ -233,8 +237,10 @@ where
 impl<'a, T> EntityDatabaseMutation<'a> for T
 where
     T: EntityBase<ConnectionType = crate::connection::KeystoreDatabaseConnection, AutoGeneratedFields = ()>
+        + UniqueEntity
         + UniqueEntityImplementationHelper
         + Sync,
+    <T as PrimaryKey>::PrimaryKey: ToSql,
 {
     type Transaction = TransactionWrapper<'a>;
 
@@ -251,7 +257,7 @@ where
         count_helper_tx::<Self>(tx).await
     }
 
-    async fn delete(tx: &Self::Transaction, id: &<Self as Entity>::PrimaryKey) -> CryptoKeystoreResult<bool> {
+    async fn delete(tx: &Self::Transaction, id: &Self::PrimaryKey) -> CryptoKeystoreResult<bool> {
         delete_helper::<Self>(tx, "id", id).await
     }
 }

From 823296e852efd4fd92a87db0000a9293b8047d34 Mon Sep 17 00:00:00 2001
From: Peter Goodspeed-Niklaus <peter.goodspeed-niklaus@wire.com>
Date: Fri, 19 Dec 2025 14:53:51 +0100
Subject: [PATCH 23/48] chore(keystore): update all entities to new traits

Just a bunch of busywork necessary to support the switch to `PrimaryKey`,
`BorrowPrimaryKey`, `EntityGetBorrowed`.
---
 .../src/entity_derive_new/derive_impl.rs      | 74 +++++++++----------
 keystore/src/connection/mod.rs                |  9 ++-
 keystore/src/entities/dummy_entity.rs         | 56 +++++++-------
 .../platform/generic/mls/credential.rs        | 10 ++-
 .../generic/mls/encryption_keypair.rs         |  8 +-
 .../platform/generic/mls/hpke_private_key.rs  |  8 +-
 .../platform/generic/mls/pending_group.rs     | 24 +++---
 .../platform/generic/mls/pending_message.rs   | 24 +++---
 .../platform/generic/mls/psk_bundle.rs        |  8 +-
 .../platform/generic/proteus/identity.rs      |  8 +-
 .../platform/generic/proteus/prekey.rs        |  8 +-
 .../entities/platform/wasm/mls/credential.rs  | 10 ++-
 .../platform/wasm/mls/encryption_keypair.rs   | 10 ++-
 .../platform/wasm/mls/hpke_private_key.rs     | 10 ++-
 .../platform/wasm/mls/pending_group.rs        | 27 ++++---
 .../platform/wasm/mls/pending_message.rs      | 38 ++++------
 .../entities/platform/wasm/mls/psk_bundle.rs  |  8 +-
 .../platform/wasm/mls/refresh_token.rs        | 36 ++++++++-
 .../platform/wasm/proteus/identity.rs         | 10 ++-
 .../entities/platform/wasm/proteus/prekey.rs  | 11 +--
 keystore/src/traits/fetch_from_database.rs    |  6 +-
 21 files changed, 236 insertions(+), 167 deletions(-)

diff --git a/crypto-macros/src/entity_derive_new/derive_impl.rs b/crypto-macros/src/entity_derive_new/derive_impl.rs
index 34490d450c..b255f9c27e 100644
--- a/crypto-macros/src/entity_derive_new/derive_impl.rs
+++ b/crypto-macros/src/entity_derive_new/derive_impl.rs
@@ -8,9 +8,10 @@ use crate::entity_derive_new::{Entity, column_type::ColumnType};
 impl quote::ToTokens for Entity {
     fn to_tokens(&self, tokens: &mut TokenStream) {
         tokens.extend(self.impl_entity_base());
+        tokens.extend(self.impl_primary_key());
         tokens.extend(self.impl_entity_generic());
         tokens.extend(self.impl_entity_wasm());
-        tokens.extend(self.impl_borrow_primary_key());
+        tokens.extend(self.impl_entity_get_borrowed());
         tokens.extend(self.impl_entity_database_mutation());
         tokens.extend(self.impl_entity_delete_borrowed());
         tokens.extend(self.impl_decrypting());
@@ -43,6 +44,35 @@ impl Entity {
         }
     }
 
+    /// `impl PrimaryKey for MyEntity` and `impl BorrowPrimaryKey for MyEntity`
+    fn impl_primary_key(&self) -> TokenStream {
+        let Self {
+            struct_name, id_column, ..
+        } = self;
+
+        let primary_key = id_column.column_type.owned();
+        let borrowed_primary_key = id_column.column_type.borrowed();
+        let pk_field_name = &id_column.field_name;
+
+        quote! {
+            impl crate::traits::PrimaryKey for #struct_name {
+                type PrimaryKey = #primary_key;
+
+                fn primary_key(&self) -> Self::PrimaryKey {
+                    self.#pk_field_name.clone()
+                }
+            }
+
+            impl crate::traits::BorrowPrimaryKey for #struct_name {
+                type BorrowedPrimaryKey = #borrowed_primary_key;
+
+                fn borrow_primary_key(&self) -> &Self::BorrowedPrimaryKey {
+                    &self.#pk_field_name
+                }
+            }
+        }
+    }
+
     /// `#[cfg(not(target_family = "wasm"))] impl Entity for MyEntity`
     fn impl_entity_generic(&self) -> TokenStream {
         let Self {
@@ -52,9 +82,6 @@ impl Entity {
             ..
         } = self;
 
-        let primary_key = id_column.column_type.owned();
-        let id_field_name = &id_column.field_name;
-
         let field_assignments = std::iter::once(id_column.field_assignment())
             .chain(other_columns.iter().map(|column| column.field_assignment()));
 
@@ -62,14 +89,8 @@ impl Entity {
             #[cfg(not(target_family = "wasm"))]
             #[::async_trait::async_trait]
             impl crate::traits::Entity for #struct_name {
-                type PrimaryKey = #primary_key;
-
-                fn primary_key(&self) -> #primary_key {
-                    self.#id_field_name.clone()
-                }
-
                 async fn get(conn: &mut Self::ConnectionType, key: &Self::PrimaryKey) -> crate::CryptoKeystoreResult<Option<Self>> {
-                    <Self as crate::traits::BorrowPrimaryKey>::get_borrowed(conn, key).await
+                    <Self as crate::traits::EntityGetBorrowed>::get_borrowed(conn, key).await
                 }
 
                 async fn count(conn: &mut Self::ConnectionType) -> crate::CryptoKeystoreResult<u32> {
@@ -89,25 +110,14 @@ impl Entity {
 
     /// `#[cfg(target_family = "wasm")] impl Entity for MyEntity`
     fn impl_entity_wasm(&self) -> TokenStream {
-        let Self {
-            struct_name, id_column, ..
-        } = self;
-
-        let primary_key = id_column.column_type.owned();
-        let id_field_name = &id_column.field_name;
+        let Self { struct_name, .. } = self;
 
         quote! {
             #[cfg(target_family = "wasm")]
             #[::async_trait::async_trait(?Send)]
             impl crate::traits::Entity for #struct_name {
-                type PrimaryKey = #primary_key;
-
-                fn primary_key(&self) -> #primary_key {
-                    self.#id_field_name.clone()
-                }
-
                 async fn get(conn: &mut Self::ConnectionType, key: &Self::PrimaryKey) -> crate::CryptoKeystoreResult<Option<Self>> {
-                    <Self as crate::traits::BorrowPrimaryKey>::get_borrowed(conn, key).await
+                    <Self as crate::traits::EntityGetBorrowed>::get_borrowed(conn, key).await
                 }
 
                 async fn count(conn: &mut Self::ConnectionType) -> crate::CryptoKeystoreResult<u32> {
@@ -121,8 +131,8 @@ impl Entity {
         }
     }
 
-    /// `impl BorrowPrimaryKey for MyEntity`
-    fn impl_borrow_primary_key(&self) -> TokenStream {
+    /// `impl EntityGetBorrowed for MyEntity`
+    fn impl_entity_get_borrowed(&self) -> TokenStream {
         let Self {
             struct_name,
             id_column,
@@ -130,8 +140,6 @@ impl Entity {
             ..
         } = self;
 
-        let borrowed_primary_key = id_column.column_type.borrowed();
-        let pk_field_name = &id_column.field_name;
         let pk_column_name = id_column
             .column_name
             .clone()
@@ -143,17 +151,9 @@ impl Entity {
         quote! {
             #[cfg_attr(target_family = "wasm", ::async_trait::async_trait(?Send))]
             #[cfg_attr(not(target_family = "wasm"), ::async_trait::async_trait)]
-            impl crate::traits::BorrowPrimaryKey for #struct_name {
-                type BorrowedPrimaryKey = #borrowed_primary_key;
-
-                fn borrow_primary_key(&self) -> &Self::BorrowedPrimaryKey {
-                    &self.#pk_field_name
-                }
-
+            impl crate::traits::EntityGetBorrowed for #struct_name {
                 async fn get_borrowed(conn: &mut Self::ConnectionType, key: &Self::BorrowedPrimaryKey)
                     -> crate::CryptoKeystoreResult<Option<Self>>
-                where
-                    for<'pk> &'pk Self::BorrowedPrimaryKey: crate::traits::KeyType,
                 {
                     let key_bytes = <&Self::BorrowedPrimaryKey as crate::traits::KeyType>::bytes(&key);
                     let key_bytes = key_bytes.as_ref();
diff --git a/keystore/src/connection/mod.rs b/keystore/src/connection/mod.rs
index 14223662d5..a31fdd6a0f 100644
--- a/keystore/src/connection/mod.rs
+++ b/keystore/src/connection/mod.rs
@@ -31,7 +31,10 @@ pub use self::platform::*;
 use crate::{
     CryptoKeystoreError, CryptoKeystoreResult,
     entities::{MlsPendingMessage, PersistedMlsGroupExt},
-    traits::{BorrowPrimaryKey, Entity, EntityDatabaseMutation, EntityDeleteBorrowed, FetchFromDatabase, KeyType},
+    traits::{
+        BorrowPrimaryKey, Entity, EntityDatabaseMutation, EntityDeleteBorrowed, EntityGetBorrowed, FetchFromDatabase,
+        KeyType,
+    },
     transaction::KeystoreTransaction,
 };
 
@@ -405,7 +408,7 @@ impl Database {
 #[cfg_attr(target_family = "wasm", async_trait(?Send))]
 #[cfg_attr(not(target_family = "wasm"), async_trait)]
 impl FetchFromDatabase for Database {
-    async fn get<E>(&self, id: &<E as Entity>::PrimaryKey) -> CryptoKeystoreResult<Option<E>>
+    async fn get<E>(&self, id: &E::PrimaryKey) -> CryptoKeystoreResult<Option<E>>
     where
         E: Entity<ConnectionType = KeystoreDatabaseConnection> + Clone + Send + Sync,
     {
@@ -424,7 +427,7 @@ impl FetchFromDatabase for Database {
 
     async fn get_borrowed<E>(&self, id: &<E as BorrowPrimaryKey>::BorrowedPrimaryKey) -> CryptoKeystoreResult<Option<E>>
     where
-        E: Entity<ConnectionType = KeystoreDatabaseConnection> + BorrowPrimaryKey + Clone + Send + Sync,
+        E: EntityGetBorrowed<ConnectionType = KeystoreDatabaseConnection> + Clone + Send + Sync,
         E::PrimaryKey: Borrow<E::BorrowedPrimaryKey>,
         for<'a> &'a E::BorrowedPrimaryKey: KeyType,
     {
diff --git a/keystore/src/entities/dummy_entity.rs b/keystore/src/entities/dummy_entity.rs
index af75beb8f7..fa6733f530 100644
--- a/keystore/src/entities/dummy_entity.rs
+++ b/keystore/src/entities/dummy_entity.rs
@@ -8,7 +8,7 @@ use crate::{
     entities::{Entity, EntityBase, EntityFindParams, StringEntityId},
     traits::{
         BorrowPrimaryKey, DecryptData as _, Decryptable, Decrypting, EncryptData as _, Encrypting, Entity as NewEntity,
-        EntityBase as NewEntityBase, KeyType, UniqueEntity,
+        EntityBase as NewEntityBase, EntityGetBorrowed, KeyType, PrimaryKey, UniqueEntity,
     },
 };
 #[cfg(not(target_family = "wasm"))]
@@ -90,15 +90,23 @@ impl Entity for DummyStoreValue {
     }
 }
 
-#[cfg_attr(target_family = "wasm", async_trait::async_trait(?Send))]
-#[cfg_attr(not(target_family = "wasm"), async_trait::async_trait)]
-impl NewEntity for DummyStoreValue {
+impl PrimaryKey for DummyStoreValue {
     type PrimaryKey = Vec<u8>;
-
-    fn primary_key(&self) -> Vec<u8> {
+    fn primary_key(&self) -> Self::PrimaryKey {
         Vec::new()
     }
+}
+
+impl BorrowPrimaryKey for DummyStoreValue {
+    type BorrowedPrimaryKey = [u8];
+    fn borrow_primary_key(&self) -> &Self::BorrowedPrimaryKey {
+        &[]
+    }
+}
 
+#[cfg_attr(target_family = "wasm", async_trait::async_trait(?Send))]
+#[cfg_attr(not(target_family = "wasm"), async_trait::async_trait)]
+impl NewEntity for DummyStoreValue {
     async fn get(_conn: &mut Self::ConnectionType, _key: &Self::PrimaryKey) -> CryptoKeystoreResult<Option<Self>> {
         Ok(None)
     }
@@ -116,13 +124,7 @@ impl NewEntity for DummyStoreValue {
 
 #[cfg_attr(target_family = "wasm", async_trait::async_trait(?Send))]
 #[cfg_attr(not(target_family = "wasm"), async_trait::async_trait)]
-impl BorrowPrimaryKey for DummyStoreValue {
-    type BorrowedPrimaryKey = [u8];
-
-    fn borrow_primary_key(&self) -> &Self::BorrowedPrimaryKey {
-        &[]
-    }
-
+impl EntityGetBorrowed for DummyStoreValue {
     async fn get_borrowed(
         _conn: &mut Self::ConnectionType,
         _key: &Self::BorrowedPrimaryKey,
@@ -132,7 +134,7 @@ impl BorrowPrimaryKey for DummyStoreValue {
 }
 
 impl UniqueEntity for DummyStoreValue {
-    const KEY: <Self as NewEntity>::PrimaryKey = Vec::new();
+    const KEY: Self::PrimaryKey = Vec::new();
 }
 
 #[derive(Debug, Clone, PartialEq, Eq)]
@@ -182,15 +184,23 @@ impl NewEntityBase for NewDummyStoreValue {
     }
 }
 
-#[cfg_attr(target_family = "wasm", async_trait::async_trait(?Send))]
-#[cfg_attr(not(target_family = "wasm"), async_trait::async_trait)]
-impl NewEntity for NewDummyStoreValue {
+impl PrimaryKey for NewDummyStoreValue {
     type PrimaryKey = Vec<u8>;
-
-    fn primary_key(&self) -> Vec<u8> {
+    fn primary_key(&self) -> Self::PrimaryKey {
         self.id.clone()
     }
+}
+
+impl BorrowPrimaryKey for NewDummyStoreValue {
+    type BorrowedPrimaryKey = [u8];
+    fn borrow_primary_key(&self) -> &Self::BorrowedPrimaryKey {
+        &self.id
+    }
+}
 
+#[cfg_attr(target_family = "wasm", async_trait::async_trait(?Send))]
+#[cfg_attr(not(target_family = "wasm"), async_trait::async_trait)]
+impl NewEntity for NewDummyStoreValue {
     async fn get(conn: &mut Self::ConnectionType, key: &Self::PrimaryKey) -> CryptoKeystoreResult<Option<Self>> {
         Self::get_borrowed(conn, key).await
     }
@@ -211,13 +221,7 @@ impl NewEntity for NewDummyStoreValue {
 
 #[cfg_attr(target_family = "wasm", async_trait::async_trait(?Send))]
 #[cfg_attr(not(target_family = "wasm"), async_trait::async_trait)]
-impl BorrowPrimaryKey for NewDummyStoreValue {
-    type BorrowedPrimaryKey = [u8];
-
-    fn borrow_primary_key(&self) -> &Self::BorrowedPrimaryKey {
-        &self.id
-    }
-
+impl EntityGetBorrowed for NewDummyStoreValue {
     async fn get_borrowed(
         _conn: &mut Self::ConnectionType,
         key: &Self::BorrowedPrimaryKey,
diff --git a/keystore/src/entities/platform/generic/mls/credential.rs b/keystore/src/entities/platform/generic/mls/credential.rs
index efe9d9a999..84a2a79bf5 100644
--- a/keystore/src/entities/platform/generic/mls/credential.rs
+++ b/keystore/src/entities/platform/generic/mls/credential.rs
@@ -16,7 +16,7 @@ use crate::{
     },
     traits::{
         BorrowPrimaryKey, Entity as NewEntity, EntityBase as NewEntityBase, EntityDatabaseMutation,
-        EntityDeleteBorrowed, KeyType,
+        EntityDeleteBorrowed, KeyType, PrimaryKey,
     },
 };
 
@@ -296,14 +296,16 @@ impl NewEntityBase for StoredCredential {
     }
 }
 
-#[async_trait]
-impl NewEntity for StoredCredential {
+impl PrimaryKey for StoredCredential {
     type PrimaryKey = Sha256Hash;
 
     fn primary_key(&self) -> Self::PrimaryKey {
         Sha256Hash::hash_from(&self.public_key)
     }
+}
 
+#[async_trait]
+impl NewEntity for StoredCredential {
     async fn get(conn: &mut Self::ConnectionType, key: &Self::PrimaryKey) -> CryptoKeystoreResult<Option<Self>> {
         let conn = conn.conn().await;
         let mut stmt = conn.prepare_cached(
@@ -387,7 +389,7 @@ impl<'a> EntityDatabaseMutation<'a> for StoredCredential {
         count_helper_tx::<Self>(tx).await
     }
 
-    async fn delete(tx: &Self::Transaction, id: &<Self as NewEntity>::PrimaryKey) -> CryptoKeystoreResult<bool> {
+    async fn delete(tx: &Self::Transaction, id: &Self::PrimaryKey) -> CryptoKeystoreResult<bool> {
         delete_helper::<Self>(tx, "public_key_sha256", id).await
     }
 }
diff --git a/keystore/src/entities/platform/generic/mls/encryption_keypair.rs b/keystore/src/entities/platform/generic/mls/encryption_keypair.rs
index 5555f95082..0d0a09a0a0 100644
--- a/keystore/src/entities/platform/generic/mls/encryption_keypair.rs
+++ b/keystore/src/entities/platform/generic/mls/encryption_keypair.rs
@@ -15,7 +15,7 @@ use crate::{
     },
     traits::{
         BorrowPrimaryKey, Entity as NewEntity, EntityBase as NewEntityBase, EntityDatabaseMutation,
-        EntityDeleteBorrowed, KeyType,
+        EntityDeleteBorrowed, KeyType, PrimaryKey,
     },
 };
 
@@ -184,14 +184,16 @@ impl NewEntityBase for StoredEncryptionKeyPair {
     }
 }
 
-#[async_trait]
-impl NewEntity for StoredEncryptionKeyPair {
+impl PrimaryKey for StoredEncryptionKeyPair {
     type PrimaryKey = Sha256Hash;
 
     fn primary_key(&self) -> Sha256Hash {
         Sha256Hash::hash_from(&self.pk)
     }
+}
 
+#[async_trait]
+impl NewEntity for StoredEncryptionKeyPair {
     async fn get(conn: &mut Self::ConnectionType, id: &Sha256Hash) -> CryptoKeystoreResult<Option<Self>> {
         get_helper::<Self, _>(conn, "pk_sha256", id, Self::from_row).await
     }
diff --git a/keystore/src/entities/platform/generic/mls/hpke_private_key.rs b/keystore/src/entities/platform/generic/mls/hpke_private_key.rs
index 0167036fd7..fdaf64719f 100644
--- a/keystore/src/entities/platform/generic/mls/hpke_private_key.rs
+++ b/keystore/src/entities/platform/generic/mls/hpke_private_key.rs
@@ -15,7 +15,7 @@ use crate::{
     },
     traits::{
         BorrowPrimaryKey, Entity as NewEntity, EntityBase as NewEntityBase, EntityDatabaseMutation,
-        EntityDeleteBorrowed, KeyType,
+        EntityDeleteBorrowed, KeyType, PrimaryKey,
     },
 };
 
@@ -182,14 +182,16 @@ impl NewEntityBase for StoredHpkePrivateKey {
     }
 }
 
-#[async_trait]
-impl NewEntity for StoredHpkePrivateKey {
+impl PrimaryKey for StoredHpkePrivateKey {
     type PrimaryKey = Sha256Hash;
 
     fn primary_key(&self) -> Sha256Hash {
         Sha256Hash::hash_from(&self.pk)
     }
+}
 
+#[async_trait]
+impl NewEntity for StoredHpkePrivateKey {
     async fn get(conn: &mut Self::ConnectionType, id: &Sha256Hash) -> CryptoKeystoreResult<Option<Self>> {
         get_helper::<Self, _>(conn, "pk_sha256", id, Self::from_row).await
     }
diff --git a/keystore/src/entities/platform/generic/mls/pending_group.rs b/keystore/src/entities/platform/generic/mls/pending_group.rs
index 1803c653b6..5b13c87a66 100644
--- a/keystore/src/entities/platform/generic/mls/pending_group.rs
+++ b/keystore/src/entities/platform/generic/mls/pending_group.rs
@@ -12,7 +12,7 @@ use crate::{
     },
     traits::{
         BorrowPrimaryKey, Entity as NewEntity, EntityBase as NewEntityBase, EntityDatabaseMutation,
-        EntityDeleteBorrowed, KeyType,
+        EntityDeleteBorrowed, EntityGetBorrowed, KeyType, PrimaryKey,
     },
 };
 
@@ -319,14 +319,24 @@ impl NewEntityBase for PersistedMlsPendingGroup {
     }
 }
 
-#[async_trait]
-impl NewEntity for PersistedMlsPendingGroup {
+impl PrimaryKey for PersistedMlsPendingGroup {
     type PrimaryKey = Vec<u8>;
 
     fn primary_key(&self) -> Self::PrimaryKey {
         self.id.clone()
     }
+}
+
+impl BorrowPrimaryKey for PersistedMlsPendingGroup {
+    type BorrowedPrimaryKey = [u8];
+
+    fn borrow_primary_key(&self) -> &[u8] {
+        &self.id
+    }
+}
 
+#[async_trait]
+impl NewEntity for PersistedMlsPendingGroup {
     async fn get(conn: &mut Self::ConnectionType, key: &Self::PrimaryKey) -> CryptoKeystoreResult<Option<Self>> {
         Self::get_borrowed(conn, key).await
     }
@@ -341,13 +351,7 @@ impl NewEntity for PersistedMlsPendingGroup {
 }
 
 #[async_trait]
-impl BorrowPrimaryKey for PersistedMlsPendingGroup {
-    type BorrowedPrimaryKey = [u8];
-
-    fn borrow_primary_key(&self) -> &[u8] {
-        &self.id
-    }
-
+impl EntityGetBorrowed for PersistedMlsPendingGroup {
     async fn get_borrowed(
         conn: &mut Self::ConnectionType,
         key: &Self::BorrowedPrimaryKey,
diff --git a/keystore/src/entities/platform/generic/mls/pending_message.rs b/keystore/src/entities/platform/generic/mls/pending_message.rs
index daa568b576..3c7d48e08b 100644
--- a/keystore/src/entities/platform/generic/mls/pending_message.rs
+++ b/keystore/src/entities/platform/generic/mls/pending_message.rs
@@ -12,7 +12,7 @@ use crate::{
     },
     traits::{
         BorrowPrimaryKey, Entity as NewEntity, EntityBase as NewEntityBase, EntityDatabaseMutation,
-        EntityDeleteBorrowed, KeyType,
+        EntityDeleteBorrowed, EntityGetBorrowed, KeyType, PrimaryKey,
     },
 };
 
@@ -222,13 +222,23 @@ impl NewEntityBase for MlsPendingMessage {
 ///
 /// However, we have to fake it here to support `KeystoreTransaction::remove_pending_messages_by_conversation_id`.
 /// This is temporary! Post WPB-20844, we should remove that whole API and also reset the primary key type here to `()`.
-#[async_trait]
-impl NewEntity for MlsPendingMessage {
+impl PrimaryKey for MlsPendingMessage {
     type PrimaryKey = Vec<u8>;
     fn primary_key(&self) -> Self::PrimaryKey {
         self.foreign_id.clone()
     }
+}
+
+impl BorrowPrimaryKey for MlsPendingMessage {
+    type BorrowedPrimaryKey = [u8];
+
+    fn borrow_primary_key(&self) -> &Self::BorrowedPrimaryKey {
+        &self.foreign_id
+    }
+}
 
+#[async_trait]
+impl NewEntity for MlsPendingMessage {
     async fn get(conn: &mut Self::ConnectionType, key: &Self::PrimaryKey) -> CryptoKeystoreResult<Option<Self>> {
         panic!("cannot get `MlsPendingMessage` by primary key as it has no distinct primary key")
     }
@@ -243,13 +253,7 @@ impl NewEntity for MlsPendingMessage {
 }
 
 #[async_trait]
-impl BorrowPrimaryKey for MlsPendingMessage {
-    type BorrowedPrimaryKey = [u8];
-
-    fn borrow_primary_key(&self) -> &Self::BorrowedPrimaryKey {
-        &self.foreign_id
-    }
-
+impl EntityGetBorrowed for MlsPendingMessage {
     async fn get_borrowed(
         conn: &mut Self::ConnectionType,
         key: &Self::BorrowedPrimaryKey,
diff --git a/keystore/src/entities/platform/generic/mls/psk_bundle.rs b/keystore/src/entities/platform/generic/mls/psk_bundle.rs
index b41aa1f934..ef75cbd7c8 100644
--- a/keystore/src/entities/platform/generic/mls/psk_bundle.rs
+++ b/keystore/src/entities/platform/generic/mls/psk_bundle.rs
@@ -10,7 +10,7 @@ use crate::{
         Entity, EntityBase, EntityFindParams, EntityIdStringExt, EntityTransactionExt, StoredPskBundle, StringEntityId,
         count_helper, count_helper_tx, delete_helper, get_helper, load_all_helper,
     },
-    traits::{Entity as NewEntity, EntityBase as NewEntityBase, EntityDatabaseMutation},
+    traits::{Entity as NewEntity, EntityBase as NewEntityBase, EntityDatabaseMutation, PrimaryKey},
 };
 
 #[async_trait::async_trait]
@@ -171,14 +171,16 @@ impl NewEntityBase for StoredPskBundle {
     }
 }
 
-#[async_trait]
-impl NewEntity for StoredPskBundle {
+impl PrimaryKey for StoredPskBundle {
     type PrimaryKey = Sha256Hash;
 
     fn primary_key(&self) -> Self::PrimaryKey {
         Sha256Hash::hash_from(&self.psk_id)
     }
+}
 
+#[async_trait]
+impl NewEntity for StoredPskBundle {
     async fn get(conn: &mut Self::ConnectionType, key: &Self::PrimaryKey) -> CryptoKeystoreResult<Option<Self>> {
         get_helper::<Self, _>(conn, "id_sha256", key, Self::from_row).await
     }
diff --git a/keystore/src/entities/platform/generic/proteus/identity.rs b/keystore/src/entities/platform/generic/proteus/identity.rs
index 6cc6ff946c..adffd5c519 100644
--- a/keystore/src/entities/platform/generic/proteus/identity.rs
+++ b/keystore/src/entities/platform/generic/proteus/identity.rs
@@ -8,7 +8,7 @@ use crate::{
         Entity, EntityBase, EntityFindParams, EntityTransactionExt, ProteusIdentity, StringEntityId, count_helper,
         count_helper_tx, load_all_helper,
     },
-    traits::{Entity as NewEntity, EntityBase as NewEntityBase, EntityDatabaseMutation, UniqueEntity},
+    traits::{Entity as NewEntity, EntityBase as NewEntityBase, EntityDatabaseMutation, PrimaryKey, UniqueEntity},
 };
 
 #[async_trait::async_trait]
@@ -166,11 +166,13 @@ impl UniqueEntity for ProteusIdentity {
     const KEY: () = ();
 }
 
-#[async_trait]
-impl NewEntity for ProteusIdentity {
+impl PrimaryKey for ProteusIdentity {
     type PrimaryKey = ();
     fn primary_key(&self) -> Self::PrimaryKey {}
+}
 
+#[async_trait]
+impl NewEntity for ProteusIdentity {
     async fn get(conn: &mut Self::ConnectionType, _key: &()) -> CryptoKeystoreResult<Option<Self>> {
         let conn = conn.conn().await;
         let mut stmt = conn.prepare_cached("SELECT rowid FROM proteus_identities ORDER BY rowid ASC LIMIT 1")?;
diff --git a/keystore/src/entities/platform/generic/proteus/prekey.rs b/keystore/src/entities/platform/generic/proteus/prekey.rs
index f8c00cc659..7d2f0ccc5e 100644
--- a/keystore/src/entities/platform/generic/proteus/prekey.rs
+++ b/keystore/src/entities/platform/generic/proteus/prekey.rs
@@ -8,7 +8,7 @@ use crate::{
         Entity, EntityBase, EntityFindParams, EntityTransactionExt, ProteusPrekey, StringEntityId, count_helper,
         count_helper_tx, delete_helper, get_helper, load_all_helper,
     },
-    traits::{Entity as NewEntity, EntityBase as NewEntityBase, EntityDatabaseMutation},
+    traits::{Entity as NewEntity, EntityBase as NewEntityBase, EntityDatabaseMutation, PrimaryKey},
 };
 
 #[async_trait::async_trait]
@@ -155,14 +155,16 @@ impl NewEntityBase for ProteusPrekey {
     }
 }
 
-#[async_trait]
-impl NewEntity for ProteusPrekey {
+impl PrimaryKey for ProteusPrekey {
     type PrimaryKey = u16;
 
     fn primary_key(&self) -> u16 {
         self.id
     }
+}
 
+#[async_trait]
+impl NewEntity for ProteusPrekey {
     async fn get(conn: &mut Self::ConnectionType, key: &u16) -> CryptoKeystoreResult<Option<Self>> {
         get_helper::<Self, _>(conn, "id", *key, Self::from_row).await
     }
diff --git a/keystore/src/entities/platform/wasm/mls/credential.rs b/keystore/src/entities/platform/wasm/mls/credential.rs
index 39de2f4836..92f16afdf8 100644
--- a/keystore/src/entities/platform/wasm/mls/credential.rs
+++ b/keystore/src/entities/platform/wasm/mls/credential.rs
@@ -7,7 +7,7 @@ use crate::{
     entities::{Entity, EntityBase, EntityFindParams, EntityTransactionExt, StoredCredential, StringEntityId},
     traits::{
         DecryptData, Decryptable, Decrypting, EncryptData, Encrypting, Entity as NewEntity,
-        EntityBase as NewEntityBase, EntityDatabaseMutation, KeyType as _,
+        EntityBase as NewEntityBase, EntityDatabaseMutation, KeyType as _, PrimaryKey,
     },
 };
 
@@ -95,14 +95,16 @@ impl NewEntityBase for StoredCredential {
     }
 }
 
-#[async_trait(?Send)]
-impl NewEntity for StoredCredential {
+impl PrimaryKey for StoredCredential {
     type PrimaryKey = Sha256Hash;
 
     fn primary_key(&self) -> Self::PrimaryKey {
         Sha256Hash::hash_from(&self.public_key)
     }
+}
 
+#[async_trait(?Send)]
+impl NewEntity for StoredCredential {
     async fn get(conn: &mut Self::ConnectionType, key: &Self::PrimaryKey) -> CryptoKeystoreResult<Option<Self>> {
         conn.storage().new_get(key.bytes().as_ref()).await
     }
@@ -138,7 +140,7 @@ impl<'a> EntityDatabaseMutation<'a> for StoredCredential {
         tx.new_count::<Self>().await
     }
 
-    async fn delete(tx: &Self::Transaction, id: &<Self as NewEntity>::PrimaryKey) -> CryptoKeystoreResult<bool> {
+    async fn delete(tx: &Self::Transaction, id: &Self::PrimaryKey) -> CryptoKeystoreResult<bool> {
         tx.new_delete::<Self>(id.bytes().as_ref()).await
     }
 }
diff --git a/keystore/src/entities/platform/wasm/mls/encryption_keypair.rs b/keystore/src/entities/platform/wasm/mls/encryption_keypair.rs
index b1b815b8a6..376aca7b4e 100644
--- a/keystore/src/entities/platform/wasm/mls/encryption_keypair.rs
+++ b/keystore/src/entities/platform/wasm/mls/encryption_keypair.rs
@@ -6,7 +6,7 @@ use crate::{
     entities::{Entity, EntityBase, EntityFindParams, EntityTransactionExt, StoredEncryptionKeyPair, StringEntityId},
     traits::{
         DecryptData, Decryptable, Decrypting, EncryptData, Encrypting, Entity as NewEntity,
-        EntityBase as NewEntityBase, EntityDatabaseMutation, KeyType,
+        EntityBase as NewEntityBase, EntityDatabaseMutation, KeyType, PrimaryKey,
     },
 };
 
@@ -78,14 +78,16 @@ impl NewEntityBase for StoredEncryptionKeyPair {
     }
 }
 
-#[async_trait(?Send)]
-impl NewEntity for StoredEncryptionKeyPair {
+impl PrimaryKey for StoredEncryptionKeyPair {
     type PrimaryKey = Sha256Hash;
 
-    fn primary_key(&self) -> Sha256Hash {
+    fn primary_key(&self) -> Self::PrimaryKey {
         Sha256Hash::hash_from(&self.pk)
     }
+}
 
+#[async_trait(?Send)]
+impl NewEntity for StoredEncryptionKeyPair {
     async fn get(conn: &mut Self::ConnectionType, id: &Sha256Hash) -> CryptoKeystoreResult<Option<Self>> {
         conn.storage().new_get(id.bytes().as_ref()).await
     }
diff --git a/keystore/src/entities/platform/wasm/mls/hpke_private_key.rs b/keystore/src/entities/platform/wasm/mls/hpke_private_key.rs
index 03003e3569..22ec5565fd 100644
--- a/keystore/src/entities/platform/wasm/mls/hpke_private_key.rs
+++ b/keystore/src/entities/platform/wasm/mls/hpke_private_key.rs
@@ -7,7 +7,7 @@ use crate::{
     entities::{Entity, EntityBase, EntityFindParams, EntityTransactionExt, StoredHpkePrivateKey, StringEntityId},
     traits::{
         DecryptData, Decryptable, Decrypting, EncryptData, Encrypting, Entity as NewEntity,
-        EntityBase as NewEntityBase, EntityDatabaseMutation, KeyType,
+        EntityBase as NewEntityBase, EntityDatabaseMutation, KeyType, PrimaryKey,
     },
 };
 
@@ -79,14 +79,16 @@ impl NewEntityBase for StoredHpkePrivateKey {
     }
 }
 
-#[async_trait(?Send)]
-impl NewEntity for StoredHpkePrivateKey {
+impl PrimaryKey for StoredHpkePrivateKey {
     type PrimaryKey = Sha256Hash;
 
-    fn primary_key(&self) -> Sha256Hash {
+    fn primary_key(&self) -> Self::PrimaryKey {
         Sha256Hash::hash_from(&self.pk)
     }
+}
 
+#[async_trait(?Send)]
+impl NewEntity for StoredHpkePrivateKey {
     async fn get(conn: &mut Self::ConnectionType, id: &Sha256Hash) -> CryptoKeystoreResult<Option<Self>> {
         conn.storage().new_get(id.bytes().as_ref()).await
     }
diff --git a/keystore/src/entities/platform/wasm/mls/pending_group.rs b/keystore/src/entities/platform/wasm/mls/pending_group.rs
index db9563b7a6..3ccdf1d0dc 100644
--- a/keystore/src/entities/platform/wasm/mls/pending_group.rs
+++ b/keystore/src/entities/platform/wasm/mls/pending_group.rs
@@ -1,5 +1,3 @@
-use std::borrow::Borrow;
-
 use async_trait::async_trait;
 use serde::{Deserialize, Serialize};
 
@@ -9,7 +7,8 @@ use crate::{
     entities::{Entity, EntityBase, EntityFindParams, EntityTransactionExt, PersistedMlsPendingGroup, StringEntityId},
     traits::{
         BorrowPrimaryKey, DecryptData, Decryptable, Decrypting, EncryptData, Encrypting, Entity as NewEntity,
-        EntityBase as NewEntityBase, EntityDatabaseMutation, EntityDeleteBorrowed, KeyType,
+        EntityBase as NewEntityBase, EntityDatabaseMutation, EntityDeleteBorrowed, EntityGetBorrowed, KeyType,
+        PrimaryKey,
     },
 };
 
@@ -94,14 +93,24 @@ impl NewEntityBase for PersistedMlsPendingGroup {
     }
 }
 
-#[async_trait(?Send)]
-impl NewEntity for PersistedMlsPendingGroup {
+impl PrimaryKey for PersistedMlsPendingGroup {
     type PrimaryKey = Vec<u8>;
 
     fn primary_key(&self) -> Self::PrimaryKey {
         self.id.clone()
     }
+}
+
+impl BorrowPrimaryKey for PersistedMlsPendingGroup {
+    type BorrowedPrimaryKey = [u8];
 
+    fn borrow_primary_key(&self) -> &Self::BorrowedPrimaryKey {
+        &self.id
+    }
+}
+
+#[async_trait(?Send)]
+impl NewEntity for PersistedMlsPendingGroup {
     async fn get(conn: &mut Self::ConnectionType, key: &Self::PrimaryKey) -> CryptoKeystoreResult<Option<Self>> {
         Self::get_borrowed(conn, key).await
     }
@@ -116,13 +125,7 @@ impl NewEntity for PersistedMlsPendingGroup {
 }
 
 #[async_trait(?Send)]
-impl BorrowPrimaryKey for PersistedMlsPendingGroup {
-    type BorrowedPrimaryKey = [u8];
-
-    fn borrow_primary_key(&self) -> &[u8] {
-        &self.id
-    }
-
+impl EntityGetBorrowed for PersistedMlsPendingGroup {
     async fn get_borrowed(
         conn: &mut Self::ConnectionType,
         key: &Self::BorrowedPrimaryKey,
diff --git a/keystore/src/entities/platform/wasm/mls/pending_message.rs b/keystore/src/entities/platform/wasm/mls/pending_message.rs
index c45d847bd5..e740b233fa 100644
--- a/keystore/src/entities/platform/wasm/mls/pending_message.rs
+++ b/keystore/src/entities/platform/wasm/mls/pending_message.rs
@@ -8,8 +8,8 @@ use crate::{
     connection::{KeystoreDatabaseConnection, TransactionWrapper},
     entities::{Entity, EntityBase, EntityFindParams, EntityTransactionExt, MlsPendingMessage, StringEntityId},
     traits::{
-        DecryptData, Decryptable, Decrypting, EncryptData, Encrypting, Entity as NewEntity,
-        EntityBase as NewEntityBase, EntityDatabaseMutation,
+        BorrowPrimaryKey, DecryptData, Decryptable, Decrypting, EncryptData, Encrypting, Entity as NewEntity,
+        EntityBase as NewEntityBase, EntityDatabaseMutation, PrimaryKey,
     },
 };
 
@@ -122,43 +122,37 @@ impl NewEntityBase for MlsPendingMessage {
 ///
 /// However, we have to fake it as a byte vector in this impl in order for encryption and decryption
 /// to work.
-#[async_trait(?Send)]
-impl NewEntity for MlsPendingMessage {
+impl PrimaryKey for MlsPendingMessage {
     type PrimaryKey = Vec<u8>;
 
-    fn primary_key(&self) -> Vec<u8> {
+    fn primary_key(&self) -> Self::PrimaryKey {
         self.foreign_id.clone()
     }
-
-    async fn get(_conn: &mut Self::ConnectionType, _key: &Self::PrimaryKey) -> CryptoKeystoreResult<Option<Self>> {
-        panic!("cannot get `MlsPendingMessage` by primary key as it has no distinct primary key")
-    }
-
-    async fn count(conn: &mut Self::ConnectionType) -> CryptoKeystoreResult<u32> {
-        conn.storage().new_count::<Self>().await
-    }
-
-    async fn load_all(conn: &mut Self::ConnectionType) -> CryptoKeystoreResult<Vec<Self>> {
-        conn.storage().new_get_all::<Self>().await
-    }
 }
 
 /// This implementation is purely to support `KeystoreTransaction::remove_pending_messages_by_conversation_id`;
 /// after WPB-20844, we should remove that whole API and also this implementation.
-#[async_trait(?Send)]
 impl BorrowPrimaryKey for MlsPendingMessage {
     type BorrowedPrimaryKey = [u8];
 
     fn borrow_primary_key(&self) -> &Self::BorrowedPrimaryKey {
         &self.foreign_id
     }
+}
 
-    async fn get_borrowed(
-        conn: &mut Self::ConnectionType,
-        key: &Self::BorrowedPrimaryKey,
-    ) -> CryptoKeystoreResult<Option<Self>> {
+#[async_trait(?Send)]
+impl NewEntity for MlsPendingMessage {
+    async fn get(_conn: &mut Self::ConnectionType, _key: &Self::PrimaryKey) -> CryptoKeystoreResult<Option<Self>> {
         panic!("cannot get `MlsPendingMessage` by primary key as it has no distinct primary key")
     }
+
+    async fn count(conn: &mut Self::ConnectionType) -> CryptoKeystoreResult<u32> {
+        conn.storage().new_count::<Self>().await
+    }
+
+    async fn load_all(conn: &mut Self::ConnectionType) -> CryptoKeystoreResult<Vec<Self>> {
+        conn.storage().new_get_all::<Self>().await
+    }
 }
 
 #[async_trait(?Send)]
diff --git a/keystore/src/entities/platform/wasm/mls/psk_bundle.rs b/keystore/src/entities/platform/wasm/mls/psk_bundle.rs
index 5c94c224ec..1b2a4b721f 100644
--- a/keystore/src/entities/platform/wasm/mls/psk_bundle.rs
+++ b/keystore/src/entities/platform/wasm/mls/psk_bundle.rs
@@ -7,7 +7,7 @@ use crate::{
     entities::{Entity, EntityBase, EntityFindParams, EntityTransactionExt, StoredPskBundle, StringEntityId},
     traits::{
         DecryptData, Decryptable, Decrypting, EncryptData, Encrypting, Entity as NewEntity,
-        EntityBase as NewEntityBase, EntityDatabaseMutation, KeyType as _,
+        EntityBase as NewEntityBase, EntityDatabaseMutation, KeyType as _, PrimaryKey,
     },
 };
 
@@ -79,14 +79,16 @@ impl NewEntityBase for StoredPskBundle {
     }
 }
 
-#[async_trait(?Send)]
-impl NewEntity for StoredPskBundle {
+impl PrimaryKey for StoredPskBundle {
     type PrimaryKey = Sha256Hash;
 
     fn primary_key(&self) -> Self::PrimaryKey {
         Sha256Hash::hash_from(&self.psk_id)
     }
+}
 
+#[async_trait(?Send)]
+impl NewEntity for StoredPskBundle {
     async fn get(conn: &mut Self::ConnectionType, key: &Self::PrimaryKey) -> CryptoKeystoreResult<Option<Self>> {
         conn.storage().new_get(key.bytes().as_ref()).await
     }
diff --git a/keystore/src/entities/platform/wasm/mls/refresh_token.rs b/keystore/src/entities/platform/wasm/mls/refresh_token.rs
index 249297b8dd..f674bac0de 100644
--- a/keystore/src/entities/platform/wasm/mls/refresh_token.rs
+++ b/keystore/src/entities/platform/wasm/mls/refresh_token.rs
@@ -1,8 +1,13 @@
+use serde::{Deserialize, Serialize};
+
 use crate::{
-    MissingKeyErrorKind,
+    CryptoKeystoreResult, MissingKeyErrorKind,
     connection::KeystoreDatabaseConnection,
     entities::{E2eiRefreshToken, EntityBase, UniqueEntity},
-    traits::{EntityBase as NewEntityBase, UniqueEntityImplementationHelper},
+    traits::{
+        DecryptData, Decryptable, Decrypting, EncryptData, Encrypting, EntityBase as NewEntityBase, UniqueEntity as _,
+        UniqueEntityImplementationHelper,
+    },
 };
 
 #[async_trait::async_trait(?Send)]
@@ -50,3 +55,30 @@ impl UniqueEntityImplementationHelper for E2eiRefreshToken {
         &self.content
     }
 }
+
+#[derive(Serialize, Deserialize)]
+pub struct E2eiRefreshTokenEncrypted {
+    content: Vec<u8>,
+}
+
+impl<'a> Encrypting<'a> for E2eiRefreshToken {
+    type EncryptedForm = E2eiRefreshTokenEncrypted;
+
+    fn encrypt(&'a self, cipher: &aes_gcm::Aes256Gcm) -> CryptoKeystoreResult<Self::EncryptedForm> {
+        let content = <Self as EncryptData>::encrypt_data(self, cipher, &self.content)?;
+        Ok(E2eiRefreshTokenEncrypted { content })
+    }
+}
+
+impl Decrypting<'static> for E2eiRefreshTokenEncrypted {
+    type DecryptedForm = E2eiRefreshToken;
+
+    fn decrypt(self, cipher: &aes_gcm::Aes256Gcm) -> CryptoKeystoreResult<Self::DecryptedForm> {
+        let content = <E2eiRefreshToken as DecryptData>::decrypt_data(cipher, &E2eiRefreshToken::KEY, &self.content)?;
+        Ok(E2eiRefreshToken { content })
+    }
+}
+
+impl Decryptable<'static> for E2eiRefreshToken {
+    type DecryptableFrom = E2eiRefreshTokenEncrypted;
+}
diff --git a/keystore/src/entities/platform/wasm/proteus/identity.rs b/keystore/src/entities/platform/wasm/proteus/identity.rs
index 4ab31a352a..730b14fdb7 100644
--- a/keystore/src/entities/platform/wasm/proteus/identity.rs
+++ b/keystore/src/entities/platform/wasm/proteus/identity.rs
@@ -7,7 +7,7 @@ use crate::{
     entities::{Entity, EntityBase, EntityFindParams, EntityTransactionExt, ProteusIdentity, StringEntityId},
     traits::{
         DecryptData, Decryptable, Decrypting, EncryptData, Encrypting, Entity as NewEntity,
-        EntityBase as NewEntityBase, EntityDatabaseMutation, KeyType as _, UniqueEntity,
+        EntityBase as NewEntityBase, EntityDatabaseMutation, KeyType as _, PrimaryKey, UniqueEntity,
     },
 };
 
@@ -89,13 +89,15 @@ impl UniqueEntity for ProteusIdentity {
     const KEY: [u8; 1] = [1];
 }
 
-#[async_trait(?Send)]
-impl NewEntity for ProteusIdentity {
+impl PrimaryKey for ProteusIdentity {
     type PrimaryKey = [u8; 1];
-    fn primary_key(&self) -> [u8; 1] {
+    fn primary_key(&self) -> Self::PrimaryKey {
         Self::KEY
     }
+}
 
+#[async_trait(?Send)]
+impl NewEntity for ProteusIdentity {
     async fn get(conn: &mut Self::ConnectionType, _key: &Self::PrimaryKey) -> CryptoKeystoreResult<Option<Self>> {
         let identity = Self::load_all(conn).await?.pop();
         Ok(identity)
diff --git a/keystore/src/entities/platform/wasm/proteus/prekey.rs b/keystore/src/entities/platform/wasm/proteus/prekey.rs
index 25dae6edbd..f74591f213 100644
--- a/keystore/src/entities/platform/wasm/proteus/prekey.rs
+++ b/keystore/src/entities/platform/wasm/proteus/prekey.rs
@@ -7,7 +7,7 @@ use crate::{
     entities::{Entity, EntityBase, EntityFindParams, EntityTransactionExt, ProteusPrekey, StringEntityId},
     traits::{
         DecryptData, Decryptable, Decrypting, EncryptData, Encrypting, Entity as NewEntity,
-        EntityBase as NewEntityBase, EntityDatabaseMutation, KeyType,
+        EntityBase as NewEntityBase, EntityDatabaseMutation, KeyType, PrimaryKey,
     },
 };
 
@@ -79,14 +79,15 @@ impl NewEntityBase for ProteusPrekey {
     }
 }
 
-#[async_trait(?Send)]
-impl NewEntity for ProteusPrekey {
+impl PrimaryKey for ProteusPrekey {
     type PrimaryKey = u16;
-
-    fn primary_key(&self) -> u16 {
+    fn primary_key(&self) -> Self::PrimaryKey {
         self.id
     }
+}
 
+#[async_trait(?Send)]
+impl NewEntity for ProteusPrekey {
     async fn get(conn: &mut Self::ConnectionType, key: &u16) -> CryptoKeystoreResult<Option<Self>> {
         conn.storage().new_get(key.bytes().as_ref()).await
     }
diff --git a/keystore/src/traits/fetch_from_database.rs b/keystore/src/traits/fetch_from_database.rs
index cd5436e6a1..7bd472b920 100644
--- a/keystore/src/traits/fetch_from_database.rs
+++ b/keystore/src/traits/fetch_from_database.rs
@@ -5,7 +5,9 @@ use async_trait::async_trait;
 use crate::{
     CryptoKeystoreResult,
     connection::KeystoreDatabaseConnection,
-    traits::{BorrowPrimaryKey, Entity, KeyType, PrimaryKey, UniqueEntity, UniqueEntityExt},
+    traits::{
+        BorrowPrimaryKey, Entity, EntityBase, EntityGetBorrowed, KeyType, PrimaryKey, UniqueEntity, UniqueEntityExt,
+    },
 };
 
 /// Interface to fetch from the database either from the connection directly or through a
@@ -37,7 +39,7 @@ pub trait FetchFromDatabase: Send + Sync {
         id: &<E as BorrowPrimaryKey>::BorrowedPrimaryKey,
     ) -> CryptoKeystoreResult<Option<E>>
     where
-        E: Entity<ConnectionType = KeystoreDatabaseConnection> + BorrowPrimaryKey + Clone + Send + Sync,
+        E: EntityGetBorrowed<ConnectionType = KeystoreDatabaseConnection> + Clone + Send + Sync,
         E::PrimaryKey: Borrow<E::BorrowedPrimaryKey>,
         for<'a> &'a E::BorrowedPrimaryKey: KeyType;
 

From 21f3dbb05a894c4b8c564af8ec347b660af01969 Mon Sep 17 00:00:00 2001
From: Peter Goodspeed-Niklaus <peter.goodspeed-niklaus@wire.com>
Date: Fri, 19 Dec 2025 14:55:23 +0100
Subject: [PATCH 24/48] chore(keystore): ensure wasm stuff is also updated

---
 keystore/src/entities/mls.rs                  | 11 +++--
 keystore/src/entities/mod.rs                  |  7 +--
 .../src/entities/platform/wasm/general.rs     | 34 +++++++++++++-
 .../platform/wasm/mls/e2ei_acme_ca.rs         | 34 +++++++++++++-
 keystore/src/traits/unique_entity.rs          | 45 ++++++++++++++++++-
 .../transaction/dynamic_dispatch/entity.rs    | 19 ++++----
 .../transaction/dynamic_dispatch/entity_id.rs |  6 +--
 7 files changed, 130 insertions(+), 26 deletions(-)

diff --git a/keystore/src/entities/mls.rs b/keystore/src/entities/mls.rs
index 0e7adc25eb..0a78b4b898 100644
--- a/keystore/src/entities/mls.rs
+++ b/keystore/src/entities/mls.rs
@@ -248,7 +248,7 @@ pub struct StoredE2eiEnrollment {
 #[cfg(target_family = "wasm")]
 #[async_trait::async_trait(?Send)]
 pub trait UniqueEntity:
-    EntityBase<ConnectionType = crate::connection::KeystoreDatabaseConnection>
+    crate::entities::EntityBase<ConnectionType = crate::connection::KeystoreDatabaseConnection>
     + serde::Serialize
     + serde::de::DeserializeOwned
 where
@@ -268,7 +268,7 @@ where
             .ok_or(CryptoKeystoreError::NotFound(Self::COLLECTION_NAME, "".to_string()))?)
     }
 
-    async fn find_all(conn: &mut Self::ConnectionType, _params: EntityFindParams) -> CryptoKeystoreResult<Vec<Self>> {
+    async fn find_all(conn: &mut Self::ConnectionType) -> CryptoKeystoreResult<Vec<Self>> {
         match Self::find_unique(conn).await {
             Ok(record) => Ok(vec![record]),
             Err(CryptoKeystoreError::NotFound(..)) => Ok(vec![]),
@@ -377,7 +377,10 @@ pub trait UniqueEntity: EntityBase<ConnectionType = crate::connection::KeystoreD
 #[cfg_attr(not(target_family = "wasm"), async_trait::async_trait)]
 impl<T> crate::entities::EntityTransactionExt for T
 where
-    T: crate::entities::Entity + UniqueEntity + Send + Sync,
+    T: crate::entities::Entity<ConnectionType = crate::connection::KeystoreDatabaseConnection>
+        + UniqueEntity
+        + Send
+        + Sync,
 {
     #[cfg(not(target_family = "wasm"))]
     async fn save(&self, tx: &TransactionWrapper<'_>) -> CryptoKeystoreResult<()> {
@@ -400,7 +403,7 @@ where
     #[cfg(target_family = "wasm")]
     async fn delete_fail_on_missing_id<'a>(
         _: &TransactionWrapper<'a>,
-        _id: StringEntityId<'a>,
+        _id: crate::entities::StringEntityId<'a>,
     ) -> CryptoKeystoreResult<()> {
         Err(CryptoKeystoreError::NotImplemented)
     }
diff --git a/keystore/src/entities/mod.rs b/keystore/src/entities/mod.rs
index b934ff11cb..39d0795576 100644
--- a/keystore/src/entities/mod.rs
+++ b/keystore/src/entities/mod.rs
@@ -245,13 +245,14 @@ cfg_if::cfg_if! {
         }
 
         #[async_trait::async_trait(?Send)]
-        impl<T: UniqueEntity + serde::Serialize + serde::de::DeserializeOwned> Entity for T {
+        impl<T> Entity for T
+        where T : UniqueEntity + crate::entities::EntityBase<ConnectionType = crate::connection::KeystoreDatabaseConnection> {
             fn id_raw(&self) -> &[u8] {
                 &Self::ID
             }
 
-            async fn find_all(conn: &mut Self::ConnectionType, params: EntityFindParams) -> CryptoKeystoreResult<Vec<Self>> {
-                    <Self as UniqueEntity>::find_all(conn, params).await
+            async fn find_all(conn: &mut Self::ConnectionType, _params: EntityFindParams) -> CryptoKeystoreResult<Vec<Self>> {
+                    <Self as UniqueEntity>::find_all(conn).await
                 }
 
             async fn find_one(conn: &mut Self::ConnectionType, _id: &StringEntityId) -> CryptoKeystoreResult<Option<Self>> {
diff --git a/keystore/src/entities/platform/wasm/general.rs b/keystore/src/entities/platform/wasm/general.rs
index a89389c6e4..250c290348 100644
--- a/keystore/src/entities/platform/wasm/general.rs
+++ b/keystore/src/entities/platform/wasm/general.rs
@@ -1,8 +1,11 @@
 use crate::{
-    MissingKeyErrorKind,
+    CryptoKeystoreResult, MissingKeyErrorKind,
     connection::KeystoreDatabaseConnection,
     entities::{ConsumerData, EntityBase, UniqueEntity},
-    traits::{EntityBase as NewEntityBase, UniqueEntityImplementationHelper},
+    traits::{
+        DecryptData, Decryptable, Decrypting, EncryptData, Encrypting, EntityBase as NewEntityBase, UniqueEntity as _,
+        UniqueEntityImplementationHelper,
+    },
 };
 
 #[async_trait::async_trait(?Send)]
@@ -50,3 +53,30 @@ impl UniqueEntityImplementationHelper for ConsumerData {
         &self.content
     }
 }
+
+#[derive(serde::Serialize, serde::Deserialize)]
+pub struct ConsumerDataEncrypted {
+    content: Vec<u8>,
+}
+
+impl<'a> Encrypting<'a> for ConsumerData {
+    type EncryptedForm = ConsumerDataEncrypted;
+
+    fn encrypt(&'a self, cipher: &aes_gcm::Aes256Gcm) -> CryptoKeystoreResult<Self::EncryptedForm> {
+        let content = <Self as EncryptData>::encrypt_data(self, cipher, &self.content)?;
+        Ok(ConsumerDataEncrypted { content })
+    }
+}
+
+impl Decrypting<'static> for ConsumerDataEncrypted {
+    type DecryptedForm = ConsumerData;
+
+    fn decrypt(self, cipher: &aes_gcm::Aes256Gcm) -> CryptoKeystoreResult<Self::DecryptedForm> {
+        let content = <ConsumerData as DecryptData>::decrypt_data(cipher, &ConsumerData::KEY, &self.content)?;
+        Ok(ConsumerData { content })
+    }
+}
+
+impl Decryptable<'static> for ConsumerData {
+    type DecryptableFrom = ConsumerDataEncrypted;
+}
diff --git a/keystore/src/entities/platform/wasm/mls/e2ei_acme_ca.rs b/keystore/src/entities/platform/wasm/mls/e2ei_acme_ca.rs
index c779fb58b5..fd34e2d4b3 100644
--- a/keystore/src/entities/platform/wasm/mls/e2ei_acme_ca.rs
+++ b/keystore/src/entities/platform/wasm/mls/e2ei_acme_ca.rs
@@ -1,8 +1,11 @@
 use crate::{
-    MissingKeyErrorKind,
+    CryptoKeystoreResult, MissingKeyErrorKind,
     connection::KeystoreDatabaseConnection,
     entities::{E2eiAcmeCA, EntityBase, UniqueEntity},
-    traits::{EntityBase as NewEntityBase, UniqueEntityImplementationHelper},
+    traits::{
+        DecryptData, Decryptable, Decrypting, EncryptData, Encrypting, EntityBase as NewEntityBase, UniqueEntity as _,
+        UniqueEntityImplementationHelper,
+    },
 };
 
 impl EntityBase for E2eiAcmeCA {
@@ -47,3 +50,30 @@ impl UniqueEntityImplementationHelper for E2eiAcmeCA {
         &self.content
     }
 }
+
+#[derive(serde::Serialize, serde::Deserialize)]
+pub struct E2eiAcmeCAEncrypted {
+    content: Vec<u8>,
+}
+
+impl<'a> Encrypting<'a> for E2eiAcmeCA {
+    type EncryptedForm = E2eiAcmeCAEncrypted;
+
+    fn encrypt(&'a self, cipher: &aes_gcm::Aes256Gcm) -> CryptoKeystoreResult<Self::EncryptedForm> {
+        let content = <Self as EncryptData>::encrypt_data(self, cipher, &self.content)?;
+        Ok(E2eiAcmeCAEncrypted { content })
+    }
+}
+
+impl Decrypting<'static> for E2eiAcmeCAEncrypted {
+    type DecryptedForm = E2eiAcmeCA;
+
+    fn decrypt(self, cipher: &aes_gcm::Aes256Gcm) -> CryptoKeystoreResult<Self::DecryptedForm> {
+        let content = <E2eiAcmeCA as DecryptData>::decrypt_data(cipher, &E2eiAcmeCA::KEY, &self.content)?;
+        Ok(E2eiAcmeCA { content })
+    }
+}
+
+impl Decryptable<'static> for E2eiAcmeCA {
+    type DecryptableFrom = E2eiAcmeCAEncrypted;
+}
diff --git a/keystore/src/traits/unique_entity.rs b/keystore/src/traits/unique_entity.rs
index 9d2bc3d7c7..eda1a2f587 100644
--- a/keystore/src/traits/unique_entity.rs
+++ b/keystore/src/traits/unique_entity.rs
@@ -44,8 +44,49 @@ pub trait UniqueEntityExt<'a>: UniqueEntity + EntityDatabaseMutation<'a> {
     async fn exists(conn: &mut Self::ConnectionType) -> CryptoKeystoreResult<bool>;
 }
 
-#[cfg_attr(target_family = "wasm", async_trait(?Send))]
-#[cfg_attr(not(target_family = "wasm"), async_trait)]
+// unfortunately we have to implement this trait twice, with nearly-identical but distinct bounds
+
+#[cfg(target_family = "wasm")]
+#[async_trait(?Send)]
+impl<'a, E> UniqueEntityExt<'a> for E
+where
+    E: UniqueEntity + EntityDatabaseMutation<'a> + Sync,
+{
+    /// Get this unique entity from the database.
+    async fn get_unique(conn: &mut Self::ConnectionType) -> CryptoKeystoreResult<Option<Self>> {
+        Self::get(conn, &Self::KEY).await
+    }
+
+    /// Set this unique entity into the database, replacing it if it already exists.
+    ///
+    /// Returns `true` if the entity previously existed and was replaced, or
+    /// `false` if it was not removed and this was a pure insertion.
+    async fn set_and_replace(&'a self, tx: &Self::Transaction) -> CryptoKeystoreResult<bool> {
+        let deleted = Self::delete(tx, &Self::KEY).await?;
+        self.save(tx).await?;
+        Ok(deleted)
+    }
+
+    /// Set this unique entity into the database if it does not already exist.
+    ///
+    /// Returns `true` if the entity was saved, or `false` if it aborted due to an already-existing entity.
+    async fn set_if_absent(&'a self, tx: &Self::Transaction) -> CryptoKeystoreResult<bool> {
+        let count = <Self as EntityDatabaseMutation>::count(tx).await?;
+        if count > 0 {
+            return Ok(false);
+        }
+        self.save(tx).await?;
+        Ok(true)
+    }
+
+    /// Returns whether or not the database contains an instance of this unique entity.
+    async fn exists(conn: &mut Self::ConnectionType) -> CryptoKeystoreResult<bool> {
+        <Self as Entity>::count(conn).await.map(|count| count > 0)
+    }
+}
+
+#[cfg(not(target_family = "wasm"))]
+#[async_trait]
 impl<'a, E> UniqueEntityExt<'a> for E
 where
     E: UniqueEntity + EntityDatabaseMutation<'a> + Sync,
diff --git a/keystore/src/transaction/dynamic_dispatch/entity.rs b/keystore/src/transaction/dynamic_dispatch/entity.rs
index d048ba6d90..2009484930 100644
--- a/keystore/src/transaction/dynamic_dispatch/entity.rs
+++ b/keystore/src/transaction/dynamic_dispatch/entity.rs
@@ -70,18 +70,20 @@ impl Entity {
             Entity::E2eiAcmeCA(e2ei_acme_ca) => e2ei_acme_ca.clone().downcast_arc(),
             Entity::E2eiIntermediateCert(e2ei_intermediate_cert) => e2ei_intermediate_cert.clone().downcast_arc(),
             Entity::E2eiCrl(e2ei_crl) => e2ei_crl.clone().downcast_arc(),
+            #[cfg(target_family = "wasm")]
+            Entity::E2eiRefreshToken(e2ei_refresh_token) => e2ei_refresh_token.clone().downcast_arc(),
+            #[cfg(feature = "proteus-keystore")]
             Entity::ProteusIdentity(proteus_identity) => proteus_identity.clone().downcast_arc(),
+            #[cfg(feature = "proteus-keystore")]
             Entity::ProteusPrekey(proteus_prekey) => proteus_prekey.clone().downcast_arc(),
+            #[cfg(feature = "proteus-keystore")]
             Entity::ProteusSession(proteus_session) => proteus_session.clone().downcast_arc(),
         }
     }
 
     pub(crate) async fn execute_save(&self, tx: &TransactionWrapper<'_>) -> CryptoKeystoreResult<()> {
         match self {
-            Entity::ConsumerData(consumer_data) => {
-                consumer_data.set_and_replace(tx).await?;
-                Ok(())
-            }
+            Entity::ConsumerData(consumer_data) => consumer_data.set_and_replace(tx).await.map(|_| ()),
             Entity::HpkePrivateKey(mls_hpke_private_key) => mls_hpke_private_key.save(tx).await,
             Entity::StoredKeypackage(mls_key_package) => mls_key_package.save(tx).await,
             Entity::PskBundle(mls_psk_bundle) => mls_psk_bundle.save(tx).await,
@@ -95,13 +97,10 @@ impl Entity {
             Entity::PersistedMlsPendingGroup(persisted_mls_pending_group) => persisted_mls_pending_group.save(tx).await,
             Entity::MlsPendingMessage(mls_pending_message) => mls_pending_message.save(tx).await,
             Entity::StoredE2eiEnrollment(e2ei_enrollment) => e2ei_enrollment.save(tx).await,
-            #[cfg(target_family = "wasm")]
-            Entity::E2eiRefreshToken(e2ei_refresh_token) => e2ei_refresh_token.replace(tx).await,
-            Entity::E2eiAcmeCA(e2ei_acme_ca) => {
-                e2ei_acme_ca.set_and_replace(tx).await?;
-                Ok(())
-            }
+            Entity::E2eiAcmeCA(e2ei_acme_ca) => e2ei_acme_ca.set_and_replace(tx).await.map(|_| ()),
             Entity::E2eiIntermediateCert(e2ei_intermediate_cert) => e2ei_intermediate_cert.save(tx).await,
+            #[cfg(target_family = "wasm")]
+            Entity::E2eiRefreshToken(e2ei_refresh_token) => e2ei_refresh_token.set_and_replace(tx).await.map(|_| ()),
             Entity::E2eiCrl(e2ei_crl) => e2ei_crl.save(tx).await,
             #[cfg(feature = "proteus-keystore")]
             Entity::ProteusSession(record) => record.save(tx).await,
diff --git a/keystore/src/transaction/dynamic_dispatch/entity_id.rs b/keystore/src/transaction/dynamic_dispatch/entity_id.rs
index 547c3bb09a..b0b24ef3cb 100644
--- a/keystore/src/transaction/dynamic_dispatch/entity_id.rs
+++ b/keystore/src/transaction/dynamic_dispatch/entity_id.rs
@@ -8,7 +8,7 @@ use crate::{
     CryptoKeystoreError, CryptoKeystoreResult,
     connection::TransactionWrapper,
     entities::{
-        E2eiAcmeCA, E2eiCrl, E2eiIntermediateCert, MlsPendingMessage, PersistedMlsGroup, PersistedMlsPendingGroup,
+        E2eiCrl, E2eiIntermediateCert, MlsPendingMessage, PersistedMlsGroup, PersistedMlsPendingGroup,
         StoredBufferedCommit, StoredCredential, StoredE2eiEnrollment, StoredEncryptionKeyPair,
         StoredEpochEncryptionKeypair, StoredHpkePrivateKey, StoredKeypackage, StoredPskBundle,
     },
@@ -27,7 +27,7 @@ impl EntityId {
     where
         E: Entity,
     {
-        <E as Entity>::PrimaryKey::from_bytes(&self.id)
+        E::PrimaryKey::from_bytes(&self.id)
             .ok_or(CryptoKeystoreError::InvalidPrimaryKeyBytes(self.typ.collection_name()))
     }
 
@@ -100,7 +100,7 @@ impl EntityId {
             EntityType::E2eiRefreshToken => {
                 E2eiRefreshToken::delete(tx, &self.primary_key::<E2eiRefreshToken>()?).await
             }
-            EntityType::E2eiAcmeCA => E2eiAcmeCA::delete(tx, &self.primary_key::<E2eiAcmeCA>()?).await,
+            EntityType::E2eiAcmeCA => Err(CryptoKeystoreError::NotImplemented),
             EntityType::E2eiIntermediateCert => {
                 E2eiIntermediateCert::delete(tx, &self.primary_key::<E2eiIntermediateCert>()?).await
             }

From b62eb8df7ea3b2de6e4dff94ae2c0039d18ed474 Mon Sep 17 00:00:00 2001
From: Peter Goodspeed-Niklaus <peter.goodspeed-niklaus@wire.com>
Date: Fri, 19 Dec 2025 16:09:58 +0100
Subject: [PATCH 25/48] chore: adjust other crates to use new
 `FetchFromDatabase`

---
 crypto/src/e2e_identity/pki_env.rs            |   8 +-
 crypto/src/group_store.rs                     |   4 +-
 .../decrypt/buffer_commit.rs                  |   6 +-
 crypto/src/mls/conversation/merge.rs          |   7 +-
 crypto/src/mls/conversation/persistence.rs    |   8 +-
 crypto/src/mls/conversation/renew.rs          |   4 +-
 crypto/src/mls/conversation/welcome.rs        |   7 +-
 .../src/mls/credential/credential_ref/find.rs |   7 +-
 .../credential/credential_ref/persistence.rs  |   9 +-
 crypto/src/mls/credential/crl.rs              |   4 +-
 crypto/src/mls/credential/persistence.rs      |  16 +--
 crypto/src/mls/session/key_package.rs         |  19 +--
 crypto/src/mls/session/mod.rs                 |   2 +-
 crypto/src/proteus.rs                         |  22 ++--
 crypto/src/test_utils/context.rs              |  17 +--
 .../transaction_context/conversation/mod.rs   |   4 +-
 .../transaction_context/e2e_identity/error.rs |   2 +
 .../e2e_identity/init_certificates.rs         |  11 +-
 crypto/src/transaction_context/mod.rs         |   6 +-
 crypto/src/transaction_context/test_utils.rs  |  22 ++--
 keystore-dump/src/main.rs                     |  32 ++---
 keystore/tests/mls.rs                         |   6 +-
 keystore/tests/z_entities.rs                  | 120 +++++++++---------
 23 files changed, 162 insertions(+), 181 deletions(-)

diff --git a/crypto/src/e2e_identity/pki_env.rs b/crypto/src/e2e_identity/pki_env.rs
index 4e4ea3cfea..670246fc23 100644
--- a/crypto/src/e2e_identity/pki_env.rs
+++ b/crypto/src/e2e_identity/pki_env.rs
@@ -1,8 +1,8 @@
 use std::collections::HashSet;
 
 use core_crypto_keystore::{
-    connection::FetchFromDatabase,
     entities::{E2eiAcmeCA, E2eiCrl, E2eiIntermediateCert},
+    traits::FetchFromDatabase,
 };
 use wire_e2e_identity::prelude::x509::revocation::{PkiEnvironment, PkiEnvironmentParams};
 use x509_cert::der::Decode;
@@ -33,7 +33,7 @@ impl IntoIterator for NewCrlDistributionPoints {
 
 pub(crate) async fn restore_pki_env(data_provider: &impl FetchFromDatabase) -> Result<Option<PkiEnvironment>> {
     let mut trust_roots = vec![];
-    let Ok(ta_raw) = data_provider.find_unique::<E2eiAcmeCA>().await else {
+    let Ok(Some(ta_raw)) = data_provider.get_unique::<E2eiAcmeCA>().await else {
         return Ok(None);
     };
 
@@ -42,7 +42,7 @@ pub(crate) async fn restore_pki_env(data_provider: &impl FetchFromDatabase) -> R
     );
 
     let intermediates = data_provider
-        .find_all::<E2eiIntermediateCert>(Default::default())
+        .load_all::<E2eiIntermediateCert>()
         .await
         .map_err(KeystoreError::wrap("finding intermediate certificates"))?
         .into_iter()
@@ -50,7 +50,7 @@ pub(crate) async fn restore_pki_env(data_provider: &impl FetchFromDatabase) -> R
         .collect::<Result<Vec<_>, _>>()?;
 
     let crls = data_provider
-        .find_all::<E2eiCrl>(Default::default())
+        .load_all::<E2eiCrl>()
         .await
         .map_err(KeystoreError::wrap("finding crls"))?
         .into_iter()
diff --git a/crypto/src/group_store.rs b/crypto/src/group_store.rs
index 9869980faf..6dadf2669c 100644
--- a/crypto/src/group_store.rs
+++ b/crypto/src/group_store.rs
@@ -1,6 +1,6 @@
 use std::sync::Arc;
 
-use core_crypto_keystore::connection::FetchFromDatabase;
+use core_crypto_keystore::traits::FetchFromDatabase;
 
 use crate::{ConversationId, KeystoreError, MlsConversation, RecursiveError, Result};
 
@@ -31,7 +31,7 @@ impl GroupStoreEntity for MlsConversation {
         keystore: &impl FetchFromDatabase,
     ) -> crate::Result<Option<Self>> {
         let result = keystore
-            .find::<Self::RawStoreValue>(id)
+            .get_borrowed::<Self::RawStoreValue>(id.as_ref())
             .await
             .map_err(KeystoreError::wrap("finding mls conversation from keystore by id"))?;
         let Some(store_value) = result else {
diff --git a/crypto/src/mls/conversation/conversation_guard/decrypt/buffer_commit.rs b/crypto/src/mls/conversation/conversation_guard/decrypt/buffer_commit.rs
index c0b0e360b1..36dc02561a 100644
--- a/crypto/src/mls/conversation/conversation_guard/decrypt/buffer_commit.rs
+++ b/crypto/src/mls/conversation/conversation_guard/decrypt/buffer_commit.rs
@@ -1,4 +1,4 @@
-use core_crypto_keystore::{connection::FetchFromDatabase as _, entities::StoredBufferedCommit};
+use core_crypto_keystore::{entities::StoredBufferedCommit, traits::FetchFromDatabase as _};
 use log::info;
 use openmls::framing::MlsMessageIn;
 use openmls_traits::OpenMlsCryptoProvider as _;
@@ -34,7 +34,7 @@ impl ConversationGuard {
         self.crypto_provider()
             .await?
             .keystore()
-            .find::<StoredBufferedCommit>(conversation.id())
+            .get_borrowed::<StoredBufferedCommit>(conversation.id().as_ref())
             .await
             .map(|option| option.map(StoredBufferedCommit::into_commit_data))
             .map_err(KeystoreError::wrap("attempting to retrieve buffered commit"))
@@ -69,7 +69,7 @@ impl ConversationGuard {
         self.crypto_provider()
             .await?
             .keystore()
-            .remove::<StoredBufferedCommit, _>(conversation.id())
+            .remove_borrowed::<StoredBufferedCommit>(conversation.id().as_ref())
             .await
             .map_err(KeystoreError::wrap("attempting to clear buffered commit"))
             .map_err(Into::into)
diff --git a/crypto/src/mls/conversation/merge.rs b/crypto/src/mls/conversation/merge.rs
index e912c3b349..4202a0cff9 100644
--- a/crypto/src/mls/conversation/merge.rs
+++ b/crypto/src/mls/conversation/merge.rs
@@ -10,7 +10,7 @@
 //! | 0 pend. Proposal  | ❌              | ✅              |
 //! | 1+ pend. Proposal | ❌              | ✅              |
 
-use core_crypto_keystore::entities::StoredEncryptionKeyPair;
+use core_crypto_keystore::{Sha256Hash, entities::StoredEncryptionKeyPair};
 use mls_crypto_provider::MlsCryptoProvider;
 use openmls_traits::OpenMlsCryptoProvider;
 
@@ -33,7 +33,10 @@ impl MlsConversation {
         // ..so if there's any, we clear them after the commit is merged
         for oln in &previous_own_leaf_nodes {
             let ek = oln.encryption_key().as_slice();
-            let _ = backend.key_store().remove::<StoredEncryptionKeyPair, _>(ek).await;
+            let _ = backend
+                .key_store()
+                .remove::<StoredEncryptionKeyPair>(&Sha256Hash::hash_from(ek))
+                .await;
         }
 
         client
diff --git a/crypto/src/mls/conversation/persistence.rs b/crypto/src/mls/conversation/persistence.rs
index 9f7c53aac7..859e14db99 100644
--- a/crypto/src/mls/conversation/persistence.rs
+++ b/crypto/src/mls/conversation/persistence.rs
@@ -1,10 +1,6 @@
 use std::collections::HashMap;
 
-use core_crypto_keystore::{
-    CryptoKeystoreMls as _,
-    connection::FetchFromDatabase as _,
-    entities::{EntityFindParams, PersistedMlsGroup},
-};
+use core_crypto_keystore::{CryptoKeystoreMls as _, entities::PersistedMlsGroup, traits::FetchFromDatabase as _};
 use mls_crypto_provider::Database;
 use openmls::group::{InnerState, MlsGroup};
 
@@ -50,7 +46,7 @@ impl MlsConversation {
     /// Effectively [`Database::mls_groups_restore`] but with better types
     pub(crate) async fn load_all(keystore: &Database) -> Result<HashMap<ConversationId, Self>> {
         let groups = keystore
-            .find_all::<PersistedMlsGroup>(EntityFindParams::default())
+            .load_all::<PersistedMlsGroup>()
             .await
             .map_err(KeystoreError::wrap("finding all persisted mls groups"))?;
         groups
diff --git a/crypto/src/mls/conversation/renew.rs b/crypto/src/mls/conversation/renew.rs
index 7106a36c37..0b975d92e7 100644
--- a/crypto/src/mls/conversation/renew.rs
+++ b/crypto/src/mls/conversation/renew.rs
@@ -1,4 +1,4 @@
-use core_crypto_keystore::entities::StoredEncryptionKeyPair;
+use core_crypto_keystore::{Sha256Hash, entities::StoredEncryptionKeyPair};
 use mls_crypto_provider::MlsCryptoProvider;
 use openmls::prelude::{LeafNode, LeafNodeIndex, Proposal, QueuedProposal, Sender, StagedCommit};
 use openmls_traits::OpenMlsCryptoProvider;
@@ -131,7 +131,7 @@ impl MlsConversation {
             // encryption key from the keystore otherwise we would have a leak
             backend
                 .key_store()
-                .remove::<StoredEncryptionKeyPair, _>(leaf_node.encryption_key().as_slice())
+                .remove::<StoredEncryptionKeyPair>(&Sha256Hash::hash_from(leaf_node.encryption_key().as_slice()))
                 .await
                 .map_err(KeystoreError::wrap("removing mls encryption keypair"))?;
         }
diff --git a/crypto/src/mls/conversation/welcome.rs b/crypto/src/mls/conversation/welcome.rs
index 72742b90e4..788b26990a 100644
--- a/crypto/src/mls/conversation/welcome.rs
+++ b/crypto/src/mls/conversation/welcome.rs
@@ -1,4 +1,4 @@
-use core_crypto_keystore::{connection::FetchFromDatabase, entities::PersistedMlsPendingGroup};
+use core_crypto_keystore::{entities::PersistedMlsPendingGroup, traits::FetchFromDatabase};
 use mls_crypto_provider::MlsCryptoProvider;
 use openmls::prelude::{MlsGroup, Welcome};
 use openmls_traits::OpenMlsCryptoProvider;
@@ -50,7 +50,10 @@ impl MlsConversation {
         let existing_conversation = mls_groups.get_fetch(&id, &backend.keystore(), None).await;
         let conversation_exists = existing_conversation.ok().flatten().is_some();
 
-        let pending_group = backend.key_store().find::<PersistedMlsPendingGroup>(id.as_ref()).await;
+        let pending_group = backend
+            .key_store()
+            .get_borrowed::<PersistedMlsPendingGroup>(id.as_ref())
+            .await;
         let pending_group_exists = pending_group.ok().flatten().is_some();
 
         if conversation_exists || pending_group_exists {
diff --git a/crypto/src/mls/credential/credential_ref/find.rs b/crypto/src/mls/credential/credential_ref/find.rs
index 073e53fec3..9c82d32e49 100644
--- a/crypto/src/mls/credential/credential_ref/find.rs
+++ b/crypto/src/mls/credential/credential_ref/find.rs
@@ -1,7 +1,4 @@
-use core_crypto_keystore::{
-    connection::FetchFromDatabase as _,
-    entities::{EntityFindParams, StoredCredential},
-};
+use core_crypto_keystore::{entities::StoredCredential, traits::FetchFromDatabase as _};
 use mls_crypto_provider::Database;
 use openmls::prelude::Credential as MlsCredential;
 use tls_codec::Deserialize as _;
@@ -67,7 +64,7 @@ impl CredentialRef {
         } = find_filters;
 
         let partial_credentials = database
-            .find_all::<StoredCredential>(EntityFindParams::default())
+            .load_all::<StoredCredential>()
             .await
             .map_err(KeystoreError::wrap("finding all credentials"))?
             .into_iter()
diff --git a/crypto/src/mls/credential/credential_ref/persistence.rs b/crypto/src/mls/credential/credential_ref/persistence.rs
index ce06299a5d..bc4a7bfeec 100644
--- a/crypto/src/mls/credential/credential_ref/persistence.rs
+++ b/crypto/src/mls/credential/credential_ref/persistence.rs
@@ -4,10 +4,7 @@
 //! useful to end users. Clients building on the CC API can't do anything useful with a full [`Credential`],
 //! and it's wasteful to transfer one across the FFI boundary.
 
-use core_crypto_keystore::{
-    connection::FetchFromDatabase as _,
-    entities::{EntityFindParams, StoredCredential},
-};
+use core_crypto_keystore::{Sha256Hash, entities::StoredCredential, traits::FetchFromDatabase as _};
 use mls_crypto_provider::Database;
 
 use super::{Error, Result};
@@ -20,7 +17,7 @@ impl CredentialRef {
     /// For loading a single credential, prefer [`Self::load`].
     pub(crate) async fn load_stored_credentials(database: &Database) -> Result<Vec<StoredCredential>> {
         let credentials = database
-            .find_all::<StoredCredential>(EntityFindParams::default())
+            .load_all::<StoredCredential>()
             .await
             .map_err(KeystoreError::wrap("finding all mls credentials"))?;
         Ok(credentials)
@@ -31,7 +28,7 @@ impl CredentialRef {
     /// Note that this does not attach the credential to any Session; it just does the data manipulation.
     pub(crate) async fn load(&self, database: &Database) -> Result<Credential> {
         database
-            .find::<StoredCredential>(self.public_key())
+            .get::<StoredCredential>(&Sha256Hash::hash_from(self.public_key()))
             .await
             .map_err(KeystoreError::wrap("finding credential"))?
             .ok_or(Error::CredentialNotFound)
diff --git a/crypto/src/mls/credential/crl.rs b/crypto/src/mls/credential/crl.rs
index a1d07c73ac..93e13a0ecc 100644
--- a/crypto/src/mls/credential/crl.rs
+++ b/crypto/src/mls/credential/crl.rs
@@ -1,6 +1,6 @@
 use std::collections::HashSet;
 
-use core_crypto_keystore::{connection::FetchFromDatabase, entities::E2eiCrl};
+use core_crypto_keystore::{entities::E2eiCrl, traits::FetchFromDatabase};
 use mls_crypto_provider::MlsCryptoProvider;
 use openmls::{
     group::MlsGroup,
@@ -75,7 +75,7 @@ pub(crate) async fn get_new_crl_distribution_points(
 
     let stored_crls = backend
         .key_store()
-        .find_all::<E2eiCrl>(Default::default())
+        .load_all::<E2eiCrl>()
         .await
         .map_err(KeystoreError::wrap("finding all e2e crl"))?;
     let stored_crl_dps: HashSet<&str> = stored_crls.iter().map(|crl| crl.distribution_point.as_str()).collect();
diff --git a/crypto/src/mls/credential/persistence.rs b/crypto/src/mls/credential/persistence.rs
index 590f6a4cf0..0b7742b25b 100644
--- a/crypto/src/mls/credential/persistence.rs
+++ b/crypto/src/mls/credential/persistence.rs
@@ -1,4 +1,4 @@
-use core_crypto_keystore::entities::StoredCredential;
+use core_crypto_keystore::{Sha256Hash, entities::StoredCredential};
 use mls_crypto_provider::Database;
 use tls_codec::Serialize as _;
 
@@ -6,14 +6,6 @@ use super::{Error, Result};
 use crate::{Credential, CredentialRef, KeystoreError};
 
 impl Credential {
-    /// Update all the fields that were updated by the DB during the save.
-    ///
-    /// [`<StoredCredential as
-    /// EntityTransactionExt>::pre_save`][core_crypto_keystore::entities::EntityTransactionExt::pre_save].
-    fn update_from(&mut self, stored: StoredCredential) {
-        self.earliest_validity = stored.created_at;
-    }
-
     /// Persist this credential into the database.
     ///
     /// Returns a reference which is stable over time and across the FFI boundary.
@@ -26,7 +18,7 @@ impl Credential {
             .tls_serialize_detached()
             .map_err(Error::tls_serialize("credential"))?;
 
-        let stored_credential = database
+        self.earliest_validity = database
             .save(StoredCredential {
                 session_id: self.client_id().to_owned().into_inner(),
                 credential: credential_data,
@@ -38,15 +30,13 @@ impl Credential {
             .await
             .map_err(KeystoreError::wrap("saving credential"))?;
 
-        self.update_from(stored_credential);
-
         Ok(CredentialRef::from_credential(self))
     }
 
     /// Delete this credential from the database
     pub(crate) async fn delete(self, database: &Database) -> Result<()> {
         database
-            .remove::<StoredCredential, _>(self.signature_key_pair.public())
+            .remove::<StoredCredential>(&Sha256Hash::hash_from(self.signature_key_pair.public()))
             .await
             .map_err(KeystoreError::wrap("deleting credential"))?;
 
diff --git a/crypto/src/mls/session/key_package.rs b/crypto/src/mls/session/key_package.rs
index 1f7b0376a3..3ae138adde 100644
--- a/crypto/src/mls/session/key_package.rs
+++ b/crypto/src/mls/session/key_package.rs
@@ -1,8 +1,9 @@
 use std::{sync::Arc, time::Duration};
 
 use core_crypto_keystore::{
-    connection::FetchFromDatabase,
-    entities::{EntityFindParams, StoredEncryptionKeyPair, StoredHpkePrivateKey, StoredKeypackage},
+    Sha256Hash,
+    entities::{StoredEncryptionKeyPair, StoredHpkePrivateKey, StoredKeypackage},
+    traits::FetchFromDatabase,
 };
 use openmls::prelude::{CryptoConfig, Lifetime};
 
@@ -14,10 +15,10 @@ use crate::{
 
 /// Default number of Keypackages a client generates the first time it's created
 #[cfg(not(test))]
-pub const INITIAL_KEYING_MATERIAL_COUNT: usize = 100;
+pub const INITIAL_KEYING_MATERIAL_COUNT: u32 = 100;
 /// Default number of Keypackages a client generates the first time it's created
 #[cfg(test)]
-pub const INITIAL_KEYING_MATERIAL_COUNT: usize = 10;
+pub const INITIAL_KEYING_MATERIAL_COUNT: u32 = 10;
 
 /// Default lifetime of all generated Keypackages. Matches the limit defined in openmls
 pub const KEYPACKAGE_DEFAULT_LIFETIME: Duration = Duration::from_secs(60 * 60 * 24 * 28 * 3); // ~3 months
@@ -88,7 +89,7 @@ impl Session {
         let stored_keypackages: Vec<StoredKeypackage> = self
             .crypto_provider
             .keystore()
-            .find_all(EntityFindParams::default())
+            .load_all()
             .await
             .map_err(KeystoreError::wrap("finding all keypackages"))?;
 
@@ -116,7 +117,7 @@ impl Session {
     pub(crate) async fn load_keypackage(&self, kp_ref: &KeypackageRef) -> Result<Option<Keypackage>> {
         self.crypto_provider
             .keystore()
-            .find::<StoredKeypackage>(kp_ref.hash_ref())
+            .get_borrowed::<StoredKeypackage>(kp_ref.hash_ref())
             .await
             .map_err(KeystoreError::wrap("loading keypackage from database"))?
             .map(|stored_keypackage| from_stored(&stored_keypackage))
@@ -135,13 +136,13 @@ impl Session {
         };
 
         let db = self.crypto_provider.keystore();
-        db.remove::<StoredKeypackage, _>(kp_ref.hash_ref())
+        db.remove_borrowed::<StoredKeypackage>(kp_ref.hash_ref())
             .await
             .map_err(KeystoreError::wrap("removing key package from keystore"))?;
-        db.remove::<StoredHpkePrivateKey, _>(kp.hpke_init_key().as_slice())
+        db.remove::<StoredHpkePrivateKey>(&Sha256Hash::hash_from(kp.hpke_init_key().as_slice()))
             .await
             .map_err(KeystoreError::wrap("removing private key from keystore"))?;
-        db.remove::<StoredEncryptionKeyPair, _>(kp.leaf_node().encryption_key().as_slice())
+        db.remove::<StoredEncryptionKeyPair>(&Sha256Hash::hash_from(kp.leaf_node().encryption_key().as_slice()))
             .await
             .map_err(KeystoreError::wrap("removing encryption keypair from keystore"))?;
 
diff --git a/crypto/src/mls/session/mod.rs b/crypto/src/mls/session/mod.rs
index 84724655d8..5184202553 100644
--- a/crypto/src/mls/session/mod.rs
+++ b/crypto/src/mls/session/mod.rs
@@ -324,7 +324,7 @@ impl Session {
 
 #[cfg(test)]
 mod tests {
-    use core_crypto_keystore::{connection::FetchFromDatabase as _, entities::*};
+    use core_crypto_keystore::{entities::*, traits::FetchFromDatabase as _};
     use mls_crypto_provider::MlsCryptoProvider;
 
     use super::*;
diff --git a/crypto/src/proteus.rs b/crypto/src/proteus.rs
index 3851d5a1f1..f4e9164ed6 100644
--- a/crypto/src/proteus.rs
+++ b/crypto/src/proteus.rs
@@ -2,8 +2,8 @@ use std::{collections::HashMap, sync::Arc};
 
 use core_crypto_keystore::{
     Database as CryptoKeystore,
-    connection::FetchFromDatabase,
     entities::{ProteusIdentity, ProteusSession},
+    traits::FetchFromDatabase,
 };
 use proteus_wasm::{
     keys::{IdentityKeyPair, PreKeyBundle},
@@ -73,8 +73,11 @@ impl GroupStoreEntity for ProteusConversationSession {
         identity: Option<Self::IdentityType>,
         keystore: &impl FetchFromDatabase,
     ) -> crate::Result<Option<Self>> {
+        let id = str::from_utf8(id.as_ref()).map_err(KeystoreError::wrap(
+            "converting id to string to fetch ProteusConversationSession",
+        ))?;
         let result = keystore
-            .find::<Self::RawStoreValue>(id)
+            .get_borrowed::<Self::RawStoreValue>(id)
             .await
             .map_err(KeystoreError::wrap("finding raw group store entity by id"))?;
         let Some(store_value) = result else {
@@ -198,7 +201,7 @@ impl ProteusCentral {
     /// errors)
     async fn load_or_create_identity(keystore: &CryptoKeystore) -> Result<IdentityKeyPair> {
         let Some(identity) = keystore
-            .find::<ProteusIdentity>(ProteusIdentity::ID)
+            .get_unique::<ProteusIdentity>()
             .await
             .map_err(KeystoreError::wrap("finding proteus identity"))?
         else {
@@ -238,10 +241,9 @@ impl ProteusCentral {
     ) -> Result<GroupStore<ProteusConversationSession>> {
         let mut proteus_sessions = GroupStore::new_with_limit(crate::group_store::ITEM_LIMIT * 2);
         for session in keystore
-            .find_all::<ProteusSession>(Default::default())
+            .load_all::<ProteusSession>()
             .await
             .map_err(KeystoreError::wrap("finding all proteus sessions"))?
-            .into_iter()
         {
             let proteus_session = Session::deserialise(identity.clone(), &session.session)
                 .map_err(ProteusError::wrap("deserializing session"))?;
@@ -381,7 +383,7 @@ impl ProteusCentral {
 
     /// Deletes a session in the store
     pub(crate) async fn session_delete(&mut self, keystore: &CryptoKeystore, session_id: &str) -> Result<()> {
-        if keystore.remove::<ProteusSession, _>(session_id).await.is_ok() {
+        if keystore.remove_borrowed::<ProteusSession>(session_id).await.is_ok() {
             let _ = self.proteus_sessions.remove(session_id.as_bytes());
         }
         Ok(())
@@ -505,9 +507,7 @@ impl ProteusCentral {
     /// If it cannot be found, one will be created.
     pub(crate) async fn last_resort_prekey(&self, keystore: &CryptoKeystore) -> Result<Vec<u8>> {
         let last_resort = if let Some(last_resort) = keystore
-            .find::<core_crypto_keystore::entities::ProteusPrekey>(
-                Self::last_resort_prekey_id().to_le_bytes().as_slice(),
-            )
+            .get::<core_crypto_keystore::entities::ProteusPrekey>(&Self::last_resort_prekey_id())
             .await
             .map_err(KeystoreError::wrap("finding proteus prekey"))?
         {
@@ -807,7 +807,7 @@ mod tests {
             gap_ids.dedup();
         }
         for gap_id in gap_ids.iter() {
-            keystore.remove::<ProteusPrekey, _>(gap_id.to_le_bytes()).await.unwrap();
+            keystore.remove::<ProteusPrekey>(gap_id).await.unwrap();
         }
 
         gap_ids.sort();
@@ -828,7 +828,7 @@ mod tests {
             gap_ids.dedup();
         }
         for gap_id in gap_ids.iter() {
-            keystore.remove::<ProteusPrekey, _>(gap_id.to_le_bytes()).await.unwrap();
+            keystore.remove::<ProteusPrekey>(gap_id).await.unwrap();
         }
 
         let potential_range = *ID_TEST_RANGE.end()..=(*ID_TEST_RANGE.end() * 2);
diff --git a/crypto/src/test_utils/context.rs b/crypto/src/test_utils/context.rs
index cc7d29b003..2edd486a5f 100644
--- a/crypto/src/test_utils/context.rs
+++ b/crypto/src/test_utils/context.rs
@@ -1,8 +1,9 @@
 use std::sync::Arc;
 
 use core_crypto_keystore::{
-    connection::FetchFromDatabase,
-    entities::{EntityFindParams, StoredCredential, StoredEncryptionKeyPair, StoredHpkePrivateKey, StoredKeypackage},
+    Sha256Hash,
+    entities::{StoredCredential, StoredEncryptionKeyPair, StoredHpkePrivateKey, StoredKeypackage},
+    traits::FetchFromDatabase,
 };
 use openmls::prelude::{Credential as MlsCredential, ExternalSender, HpkePublicKey, KeyPackage, SignaturePublicKey};
 use openmls_traits::{OpenMlsCryptoProvider, crypto::OpenMlsCrypto, types::SignatureScheme};
@@ -59,7 +60,7 @@ impl SessionContext {
             .await
             .unwrap()
             .key_store()
-            .find_all::<StoredKeypackage>(EntityFindParams::default())
+            .load_all::<StoredKeypackage>()
             .await
             .unwrap()
             .into_iter()
@@ -147,7 +148,7 @@ impl SessionContext {
             .keystore()
             .await
             .unwrap()
-            .find::<StoredHpkePrivateKey>(&skp.tls_serialize_detached().unwrap())
+            .get::<StoredHpkePrivateKey>(&Sha256Hash::hash_from(skp.tls_serialize_detached().unwrap()))
             .await
             .unwrap()
     }
@@ -158,14 +159,14 @@ impl SessionContext {
             .keystore()
             .await
             .unwrap()
-            .find_all::<StoredCredential>(EntityFindParams::default())
+            .load_all::<StoredCredential>()
             .await
             .unwrap()
             .into_iter()
             .find(|c| c.credential[..] == credential)
     }
 
-    pub async fn count_hpke_private_key(&self) -> usize {
+    pub async fn count_hpke_private_key(&self) -> u32 {
         self.transaction
             .keystore()
             .await
@@ -175,7 +176,7 @@ impl SessionContext {
             .unwrap()
     }
 
-    pub async fn count_encryption_keypairs(&self) -> usize {
+    pub async fn count_encryption_keypairs(&self) -> u32 {
         self.transaction
             .keystore()
             .await
@@ -185,7 +186,7 @@ impl SessionContext {
             .unwrap()
     }
 
-    pub async fn count_credentials_in_keystore(&self) -> usize {
+    pub async fn count_credentials_in_keystore(&self) -> u32 {
         self.transaction
             .keystore()
             .await
diff --git a/crypto/src/transaction_context/conversation/mod.rs b/crypto/src/transaction_context/conversation/mod.rs
index b34355a8e5..862b9ae445 100644
--- a/crypto/src/transaction_context/conversation/mod.rs
+++ b/crypto/src/transaction_context/conversation/mod.rs
@@ -6,7 +6,7 @@ pub mod external_sender;
 pub(crate) mod proposal;
 pub mod welcome;
 
-use core_crypto_keystore::{connection::FetchFromDatabase as _, entities::PersistedMlsPendingGroup};
+use core_crypto_keystore::{entities::PersistedMlsPendingGroup, traits::FetchFromDatabase as _};
 
 use super::{Error, Result, TransactionContext};
 use crate::{
@@ -39,7 +39,7 @@ impl TransactionContext {
     pub(crate) async fn pending_conversation(&self, id: &ConversationIdRef) -> Result<PendingConversation> {
         let keystore = self.keystore().await?;
         let Some(pending_group) = keystore
-            .find::<PersistedMlsPendingGroup>(id)
+            .get_borrowed::<PersistedMlsPendingGroup>(id.as_ref())
             .await
             .map_err(KeystoreError::wrap("finding persisted mls pending group"))?
         else {
diff --git a/crypto/src/transaction_context/e2e_identity/error.rs b/crypto/src/transaction_context/e2e_identity/error.rs
index fa744ebeb8..894ba44699 100644
--- a/crypto/src/transaction_context/e2e_identity/error.rs
+++ b/crypto/src/transaction_context/e2e_identity/error.rs
@@ -21,6 +21,8 @@ pub enum Error {
     PkiEnvironmentUnset,
     #[error("The certificate chain is invalid or not complete")]
     InvalidCertificateChain,
+    #[error("{0} not found")]
+    NotFound(&'static str),
     #[error(transparent)]
     X509Error(#[from] wire_e2e_identity::prelude::x509::RustyX509CheckError),
     #[error(transparent)]
diff --git a/crypto/src/transaction_context/e2e_identity/init_certificates.rs b/crypto/src/transaction_context/e2e_identity/init_certificates.rs
index ea62a9a31a..cb72c53c40 100644
--- a/crypto/src/transaction_context/e2e_identity/init_certificates.rs
+++ b/crypto/src/transaction_context/e2e_identity/init_certificates.rs
@@ -1,6 +1,6 @@
 use core_crypto_keystore::{
-    connection::FetchFromDatabase,
     entities::{E2eiAcmeCA, E2eiCrl, E2eiIntermediateCert},
+    traits::FetchFromDatabase,
 };
 use openmls_traits::OpenMlsCryptoProvider;
 use wire_e2e_identity::prelude::x509::{
@@ -42,7 +42,7 @@ impl TransactionContext {
                 .await
                 .map_err(RecursiveError::transaction("getting mls provider"))?
                 .keystore()
-                .find_unique::<E2eiAcmeCA>()
+                .get_unique::<E2eiAcmeCA>()
                 .await
                 .is_ok()
             {
@@ -133,9 +133,10 @@ impl TransactionContext {
             .await
             .map_err(RecursiveError::transaction("getting keystore"))?;
         let trust_anchor = keystore
-            .find_unique::<E2eiAcmeCA>()
+            .get_unique::<E2eiAcmeCA>()
             .await
-            .map_err(KeystoreError::wrap("finding acme ca"))?;
+            .map_err(KeystoreError::wrap("finding acme ca"))?
+            .ok_or(Error::NotFound("E2eiAcmeCA"))?;
         let trust_anchor = x509_cert::Certificate::from_der(&trust_anchor.content)?;
 
         // the `/federation` endpoint from smallstep repeats the root CA
@@ -212,7 +213,7 @@ impl TransactionContext {
             .map_err(RecursiveError::transaction("getting keystore"))?;
 
         let dirty = ks
-            .find::<E2eiCrl>(crl_dp.as_bytes())
+            .get::<E2eiCrl>(&crl_dp)
             .await
             .ok()
             .flatten()
diff --git a/crypto/src/transaction_context/mod.rs b/crypto/src/transaction_context/mod.rs
index 4b03df5098..46ba099eeb 100644
--- a/crypto/src/transaction_context/mod.rs
+++ b/crypto/src/transaction_context/mod.rs
@@ -6,7 +6,7 @@ use std::sync::Arc;
 #[cfg(feature = "proteus")]
 use async_lock::Mutex;
 use async_lock::{RwLock, RwLockReadGuardArc, RwLockWriteGuardArc};
-use core_crypto_keystore::{CryptoKeystoreError, connection::FetchFromDatabase, entities::ConsumerData};
+use core_crypto_keystore::{CryptoKeystoreError, entities::ConsumerData, traits::FetchFromDatabase as _};
 pub use error::{Error, Result};
 use mls_crypto_provider::{Database, MlsCryptoProvider};
 use openmls_traits::OpenMlsCryptoProvider as _;
@@ -292,8 +292,8 @@ impl TransactionContext {
     /// Get the data that has previously been set by [TransactionContext::set_data].
     /// This is meant to be used as a check point at the end of a transaction.
     pub async fn get_data(&self) -> Result<Option<Vec<u8>>> {
-        match self.keystore().await?.find_unique::<ConsumerData>().await {
-            Ok(data) => Ok(Some(data.into())),
+        match self.keystore().await?.get_unique::<ConsumerData>().await {
+            Ok(maybe_data) => Ok(maybe_data.map(Into::into)),
             Err(CryptoKeystoreError::NotFound(..)) => Ok(None),
             Err(err) => Err(KeystoreError::wrap("finding unique consumer data")(err).into()),
         }
diff --git a/crypto/src/transaction_context/test_utils.rs b/crypto/src/transaction_context/test_utils.rs
index b22b35fd93..74a88a276c 100644
--- a/crypto/src/transaction_context/test_utils.rs
+++ b/crypto/src/transaction_context/test_utils.rs
@@ -1,25 +1,25 @@
 use core_crypto_keystore::{
-    connection::FetchFromDatabase as _,
     entities::{
         MlsPendingMessage, PersistedMlsGroup, PersistedMlsPendingGroup, StoredCredential, StoredE2eiEnrollment,
         StoredEncryptionKeyPair, StoredEpochEncryptionKeypair, StoredHpkePrivateKey, StoredKeypackage, StoredPskBundle,
     },
+    traits::FetchFromDatabase as _,
 };
 
 use super::TransactionContext;
 
 #[derive(Debug, Clone, Eq, PartialEq)]
 pub struct EntitiesCount {
-    pub credential: usize,
-    pub encryption_keypair: usize,
-    pub epoch_encryption_keypair: usize,
-    pub enrollment: usize,
-    pub group: usize,
-    pub hpke_private_key: usize,
-    pub key_package: usize,
-    pub pending_group: usize,
-    pub pending_messages: usize,
-    pub psk_bundle: usize,
+    pub credential: u32,
+    pub encryption_keypair: u32,
+    pub epoch_encryption_keypair: u32,
+    pub enrollment: u32,
+    pub group: u32,
+    pub hpke_private_key: u32,
+    pub key_package: u32,
+    pub pending_group: u32,
+    pub pending_messages: u32,
+    pub psk_bundle: u32,
 }
 
 impl TransactionContext {
diff --git a/keystore-dump/src/main.rs b/keystore-dump/src/main.rs
index bbb005a62f..38401100b6 100644
--- a/keystore-dump/src/main.rs
+++ b/keystore-dump/src/main.rs
@@ -23,7 +23,7 @@ async fn main() -> anyhow::Result<()> {
     use chrono::TimeZone;
     use clap::Parser as _;
     use core_crypto_keystore::{
-        ConnectionType, Database as Keystore, DatabaseKey, connection::FetchFromDatabase, entities::*,
+        ConnectionType, Database as Keystore, DatabaseKey, entities::*, traits::FetchFromDatabase,
     };
     use openmls::prelude::TlsDeserializeTrait;
     use serde::ser::{SerializeMap, Serializer};
@@ -43,11 +43,7 @@ async fn main() -> anyhow::Result<()> {
     let mut json_map = json_serializer.serialize_map(None)?;
 
     let mut credentials: Vec<serde_json::Value> = vec![];
-    for cred in keystore
-        .find_all::<StoredCredential>(Default::default())
-        .await?
-        .into_iter()
-    {
+    for cred in keystore.load_all::<StoredCredential>().await?.into_iter() {
         let mls_credential = openmls::prelude::Credential::tls_deserialize(&mut cred.credential.as_slice())?;
         let mls_keypair = openmls_basic_credential::SignatureKeyPair::from_raw(
             core_crypto::Ciphersuite::try_from(cred.ciphersuite)
@@ -73,7 +69,7 @@ async fn main() -> anyhow::Result<()> {
     json_map.serialize_entry("mls_credentials", &credentials)?;
 
     let hpke_sks: Vec<openmls_traits::types::HpkePrivateKey> = keystore
-        .find_all::<StoredHpkePrivateKey>(Default::default())
+        .load_all::<StoredHpkePrivateKey>()
         .await?
         .into_iter()
         .map(|hpke_sk| postcard::from_bytes::<openmls_traits::types::HpkePrivateKey>(&hpke_sk.sk))
@@ -81,7 +77,7 @@ async fn main() -> anyhow::Result<()> {
     json_map.serialize_entry("mls_hpke_private_keys", &hpke_sks)?;
 
     let hpke_keypairs: Vec<openmls_traits::types::HpkeKeyPair> = keystore
-        .find_all::<StoredEncryptionKeyPair>(Default::default())
+        .load_all::<StoredEncryptionKeyPair>()
         .await?
         .into_iter()
         .map(|hpke_kp| postcard::from_bytes::<openmls_traits::types::HpkeKeyPair>(&hpke_kp.sk))
@@ -89,11 +85,7 @@ async fn main() -> anyhow::Result<()> {
     json_map.serialize_entry("mls_hpke_keypairs", &hpke_keypairs)?;
 
     let mut external_psks: std::collections::HashMap<String, openmls::schedule::psk::PskBundle> = Default::default();
-    for psk in keystore
-        .find_all::<StoredPskBundle>(Default::default())
-        .await?
-        .into_iter()
-    {
+    for psk in keystore.load_all::<StoredPskBundle>().await?.into_iter() {
         let mls_psk = postcard::from_bytes::<openmls::schedule::psk::PskBundle>(&psk.psk)?;
         external_psks.insert(hex::encode(&psk.psk_id), mls_psk);
     }
@@ -101,7 +93,7 @@ async fn main() -> anyhow::Result<()> {
     json_map.serialize_entry("external_psks", &external_psks)?;
 
     let keypackages: Vec<openmls::prelude::KeyPackage> = keystore
-        .find_all::<StoredKeypackage>(Default::default())
+        .load_all::<StoredKeypackage>()
         .await?
         .into_iter()
         .map(|kp| postcard::from_bytes::<openmls::prelude::KeyPackage>(&kp.keypackage))
@@ -109,7 +101,7 @@ async fn main() -> anyhow::Result<()> {
     json_map.serialize_entry("mls_keypackages", &keypackages)?;
 
     let e2ei_enrollments: Vec<core_crypto::E2eiEnrollment> = keystore
-        .find_all::<StoredE2eiEnrollment>(Default::default())
+        .load_all::<StoredE2eiEnrollment>()
         .await?
         .into_iter()
         .map(|enrollment| serde_json::from_slice::<core_crypto::E2eiEnrollment>(&enrollment.content))
@@ -117,7 +109,7 @@ async fn main() -> anyhow::Result<()> {
     json_map.serialize_entry("e2ei_enrollments", &e2ei_enrollments)?;
 
     let pgroups: Vec<openmls::prelude::MlsGroup> = keystore
-        .find_all::<PersistedMlsGroup>(Default::default())
+        .load_all::<PersistedMlsGroup>()
         .await?
         .into_iter()
         .map(|pgroup| core_crypto_keystore::deser::<openmls::prelude::MlsGroup>(&pgroup.state))
@@ -125,14 +117,14 @@ async fn main() -> anyhow::Result<()> {
     json_map.serialize_entry("mls_groups", &pgroups)?;
 
     let pegroups: Vec<openmls::prelude::MlsGroup> = keystore
-        .find_all::<PersistedMlsPendingGroup>(Default::default())
+        .load_all::<PersistedMlsPendingGroup>()
         .await?
         .into_iter()
         .map(|pgroup| core_crypto_keystore::deser::<openmls::prelude::MlsGroup>(&pgroup.state))
         .collect::<core_crypto_keystore::CryptoKeystoreResult<_>>()?;
     json_map.serialize_entry("mls_pending_groups", &pegroups)?;
 
-    if let Some(proteus_identity) = keystore.find::<ProteusIdentity>(ProteusIdentity::ID).await? {
+    if let Some(proteus_identity) = keystore.get_unique::<ProteusIdentity>().await? {
         let identity = {
             let sk = proteus_identity.sk_raw();
             let pk = proteus_identity.pk_raw();
@@ -141,7 +133,7 @@ async fn main() -> anyhow::Result<()> {
         json_map.serialize_entry("proteus_identity", &identity)?;
 
         let prekeys: Vec<proteus_wasm::keys::PreKey> = keystore
-            .find_all::<ProteusPrekey>(Default::default())
+            .load_all::<ProteusPrekey>()
             .await?
             .into_iter()
             .map(|pk| proteus_wasm::keys::PreKey::deserialise(&pk.prekey))
@@ -149,7 +141,7 @@ async fn main() -> anyhow::Result<()> {
         json_map.serialize_entry("proteus_prekeys", &prekeys)?;
 
         let proteus_sessions: Vec<proteus_wasm::session::Session<proteus_wasm::keys::IdentityKeyPair>> = keystore
-            .find_all::<ProteusSession>(Default::default())
+            .load_all::<ProteusSession>()
             .await?
             .into_iter()
             .map(|session| proteus_wasm::session::Session::deserialise(identity.clone(), &session.session))
diff --git a/keystore/tests/mls.rs b/keystore/tests/mls.rs
index f8bc5d92aa..82f2107933 100644
--- a/keystore/tests/mls.rs
+++ b/keystore/tests/mls.rs
@@ -59,7 +59,7 @@ mod tests {
 
     #[apply(all_storage_types)]
     pub async fn can_add_read_delete_credential_openmls_traits(context: KeystoreTestContext) {
-        use core_crypto_keystore::connection::FetchFromDatabase;
+        use core_crypto_keystore::{Sha256Hash, traits::FetchFromDatabase};
         use itertools::Itertools as _;
         use openmls_basic_credential::SignatureKeyPair;
 
@@ -92,7 +92,7 @@ mod tests {
 
         let (credential_from_store,) = backend
             .key_store()
-            .find_all::<StoredCredential>(Default::default())
+            .load_all::<StoredCredential>()
             .await
             .unwrap()
             .into_iter()
@@ -102,7 +102,7 @@ mod tests {
 
         backend
             .key_store()
-            .remove::<StoredCredential, _>(credential_from_store.public_key.clone())
+            .remove::<StoredCredential>(&Sha256Hash::hash_from(&credential_from_store.public_key))
             .await
             .unwrap();
     }
diff --git a/keystore/tests/z_entities.rs b/keystore/tests/z_entities.rs
index 75559223f4..e1409e1902 100644
--- a/keystore/tests/z_entities.rs
+++ b/keystore/tests/z_entities.rs
@@ -35,8 +35,6 @@ macro_rules! test_for_entity {
             crate::tests_impl::can_remove_entity::<$entity>(&store, entity).await;
 
             let ignore_count = pat_to_bool!($($ignore_entity_count)?);
-            let ignore_find_many = pat_to_bool!($($ignore_find_many)?);
-            crate::tests_impl::can_list_entities_with_find_many::<$entity>(&store, ignore_count, ignore_find_many).await;
             crate::tests_impl::can_list_entities_with_find_all::<$entity>(&store, ignore_count).await;
         }
     };
@@ -45,29 +43,39 @@ macro_rules! test_for_entity {
 #[cfg(test)]
 mod tests_impl {
     use core_crypto_keystore::{
-        connection::{FetchFromDatabase, KeystoreDatabaseConnection},
-        entities::{Entity, EntityFindParams, EntityTransactionExt, MlsPendingMessage, StoredCredential},
+        connection::KeystoreDatabaseConnection,
+        entities::{MlsPendingMessage, StoredCredential},
+        traits::{Entity, EntityDatabaseMutation, FetchFromDatabase as _, PrimaryKey as _},
     };
 
     use super::common::*;
     use crate::{ENTITY_COUNT, utils::EntityRandomUpdateExt};
 
-    pub(crate) async fn can_save_entity<
-        R: EntityRandomUpdateExt + Entity<ConnectionType = KeystoreDatabaseConnection> + EntityTransactionExt + Sync,
-    >(
-        store: &CryptoKeystore,
-    ) -> R {
+    pub(crate) async fn can_save_entity<'a, R>(store: &CryptoKeystore) -> R
+    where
+        R: Clone
+            + EntityRandomUpdateExt
+            + Entity<ConnectionType = KeystoreDatabaseConnection>
+            + EntityDatabaseMutation<'a>
+            + Send
+            + Sync,
+    {
         let entity = R::random();
         store.save(entity.clone()).await.unwrap();
         entity
     }
 
-    pub(crate) async fn can_find_entity<
-        R: EntityRandomUpdateExt + Entity<ConnectionType = KeystoreDatabaseConnection> + 'static + Sync,
-    >(
-        store: &CryptoKeystore,
-        entity: &R,
-    ) {
+    pub(crate) async fn can_find_entity<'a, R>(store: &CryptoKeystore, entity: &R)
+    where
+        R: Clone
+            + std::fmt::Debug
+            + Eq
+            + EntityRandomUpdateExt
+            + Entity<ConnectionType = KeystoreDatabaseConnection>
+            + EntityDatabaseMutation<'a>
+            + Send
+            + Sync,
+    {
         if let Some(pending_message) = entity.downcast::<MlsPendingMessage>() {
             let pending_message_from_store = store
                 .find_pending_messages_by_conversation_id(&pending_message.foreign_id)
@@ -78,71 +86,61 @@ mod tests_impl {
             assert_eq!(*pending_message, pending_message_from_store);
         } else if let Some(credential) = entity.downcast::<StoredCredential>() {
             let mut credential_from_store = store
-                .find::<StoredCredential>(&entity.merge_key())
+                .get::<StoredCredential>(&credential.primary_key())
                 .await
                 .unwrap()
                 .unwrap();
             credential_from_store.equalize();
             assert_eq!(*credential, credential_from_store);
         } else {
-            let mut entity_from_store = store.find::<R>(entity.id_raw()).await.unwrap().unwrap();
+            let primary_key = entity.primary_key();
+            let mut entity_from_store = store.get::<R>(&primary_key).await.unwrap().unwrap();
             entity_from_store.equalize();
             assert_eq!(*entity, entity_from_store);
         };
     }
 
-    pub(crate) async fn can_update_entity<
-        R: EntityRandomUpdateExt + Entity<ConnectionType = KeystoreDatabaseConnection> + EntityTransactionExt + Sync,
-    >(
-        store: &CryptoKeystore,
-        entity: &mut R,
-    ) {
+    pub(crate) async fn can_update_entity<'a, R>(store: &CryptoKeystore, entity: &mut R)
+    where
+        R: Clone
+            + std::fmt::Debug
+            + Eq
+            + EntityRandomUpdateExt
+            + Entity<ConnectionType = KeystoreDatabaseConnection>
+            + EntityDatabaseMutation<'a>
+            + Send
+            + Sync,
+    {
         entity.random_update();
         store.save(entity.clone()).await.unwrap();
-        let entity2: R = store.find(entity.id_raw()).await.unwrap().unwrap();
+        let entity2: R = store.get(&entity.primary_key()).await.unwrap().unwrap();
         assert_eq!(*entity, entity2);
     }
 
-    pub(crate) async fn can_remove_entity<
-        R: EntityRandomUpdateExt + Entity<ConnectionType = KeystoreDatabaseConnection> + EntityTransactionExt + Sync,
-    >(
-        store: &CryptoKeystore,
-        entity: R,
-    ) {
-        store.remove::<R, _>(entity.id_raw()).await.unwrap();
-        let entity2: Option<R> = store.find(entity.id_raw()).await.unwrap();
+    pub(crate) async fn can_remove_entity<'a, R>(store: &CryptoKeystore, entity: R)
+    where
+        R: Clone
+            + EntityRandomUpdateExt
+            + Entity<ConnectionType = KeystoreDatabaseConnection>
+            + EntityDatabaseMutation<'a>
+            + Send
+            + Sync,
+    {
+        store.remove::<R>(&entity.primary_key()).await.unwrap();
+        let entity2: Option<R> = store.get(&entity.primary_key()).await.unwrap();
         assert!(entity2.is_none());
     }
 
-    pub(crate) async fn can_list_entities_with_find_many<
-        R: EntityRandomUpdateExt + Entity<ConnectionType = KeystoreDatabaseConnection> + EntityTransactionExt + Sync,
-    >(
-        store: &CryptoKeystore,
-        ignore_entity_count: bool,
-        ignore_find_many: bool,
-    ) {
-        let mut ids: Vec<Vec<u8>> = vec![];
-        for _ in 0..ENTITY_COUNT {
-            let entity = R::random();
-            ids.push(entity.id_raw().to_vec());
-            store.save(entity).await.unwrap();
-        }
-
-        if !ignore_find_many {
-            let entities = store.find_many::<R>(&ids).await.unwrap();
-            if !ignore_entity_count {
-                assert_eq!(entities.len(), ENTITY_COUNT);
-            }
-        }
-    }
-
-    pub(crate) async fn can_list_entities_with_find_all<
-        R: EntityRandomUpdateExt + Entity<ConnectionType = KeystoreDatabaseConnection> + Sync,
-    >(
-        store: &CryptoKeystore,
-        ignore_entity_count: bool,
-    ) {
-        let entities = store.find_all::<R>(EntityFindParams::default()).await.unwrap();
+    pub(crate) async fn can_list_entities_with_find_all<'a, R>(store: &CryptoKeystore, ignore_entity_count: bool)
+    where
+        R: Clone
+            + EntityRandomUpdateExt
+            + Entity<ConnectionType = KeystoreDatabaseConnection>
+            + EntityDatabaseMutation<'a>
+            + Send
+            + Sync,
+    {
+        let entities = store.load_all::<R>().await.unwrap();
         if !ignore_entity_count {
             assert_eq!(entities.len(), ENTITY_COUNT);
         }

From f687d5635435deee089cc618ba9623e7a5d61513 Mon Sep 17 00:00:00 2001
From: Peter Goodspeed-Niklaus <peter.goodspeed-niklaus@wire.com>
Date: Fri, 19 Dec 2025 16:29:33 +0100
Subject: [PATCH 26/48] chore(keystore/tests): rollback on drop, not commit

If the open transaction contains invalid state (i.e. violating DB constraints),
we'd rather not panic within the Drop impl.
---
 keystore/tests/common.rs | 9 ++++++---
 1 file changed, 6 insertions(+), 3 deletions(-)

diff --git a/keystore/tests/common.rs b/keystore/tests/common.rs
index c6a980e6c4..0744ccdec3 100644
--- a/keystore/tests/common.rs
+++ b/keystore/tests/common.rs
@@ -61,13 +61,16 @@ impl KeystoreTestContext {
 impl Drop for KeystoreTestContext {
     fn drop(&mut self) {
         if let Some(store) = self.store.take() {
-            let commit_and_wipe = async move {
-                store.commit_transaction().await.expect("Could not commit transaction");
+            let rollback_and_wipe = async move {
+                store
+                    .rollback_transaction()
+                    .await
+                    .expect("could not rollback transaction");
                 store.wipe().await.expect("Could not wipe store");
             };
 
             #[cfg(not(target_family = "wasm"))]
-            futures_lite::future::block_on(commit_and_wipe);
+            futures_lite::future::block_on(rollback_and_wipe);
             #[cfg(target_family = "wasm")]
             wasm_bindgen_futures::spawn_local(commit_and_wipe);
         }

From f237bfe921d3a709f0b8e22341a8834ebe98b769 Mon Sep 17 00:00:00 2001
From: Peter Goodspeed-Niklaus <peter.goodspeed-niklaus@wire.com>
Date: Fri, 19 Dec 2025 16:43:33 +0100
Subject: [PATCH 27/48] chore(crypto-macros): preserve `no-upsert` error
 behavior

---
 crypto-macros/src/entity_derive_new/derive_impl.rs | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/crypto-macros/src/entity_derive_new/derive_impl.rs b/crypto-macros/src/entity_derive_new/derive_impl.rs
index b255f9c27e..8831e7ff2f 100644
--- a/crypto-macros/src/entity_derive_new/derive_impl.rs
+++ b/crypto-macros/src/entity_derive_new/derive_impl.rs
@@ -198,6 +198,10 @@ impl Entity {
             .map(|tokens| quote!(#tokens,))
             .collect::<TokenStream>();
 
+        let sql_map_err = (!upsert).then_some(quote! {
+            .map_err(|_| CryptoKeystoreError::AlreadyExists(Self::COLLECTION_NAME))
+        });
+
         quote! {
             #[cfg_attr(target_family = "wasm", async_trait::async_trait(?Send))]
             #[cfg_attr(not(target_family = "wasm"), async_trait::async_trait)]
@@ -213,7 +217,7 @@ impl Entity {
                         #[cfg(not(target_family = "wasm"))]
                         {
                             let mut stmt = tx.prepare_cached(#sql_statement)?;
-                            stmt.execute(rusqlite::params![#fields])?;
+                            stmt.execute(rusqlite::params![#fields])#sql_map_err?;
                             Ok(())
                         }
                 }

From 4a94c94f3de7cd891514f0d5c0c94dbcf022ea9b Mon Sep 17 00:00:00 2001
From: Peter Goodspeed-Niklaus <peter.goodspeed-niklaus@wire.com>
Date: Fri, 19 Dec 2025 16:55:33 +0100
Subject: [PATCH 28/48] chore(keystore/tests): restore stateful behavior within
 z-entities tests

---
 keystore/tests/z_entities.rs | 16 ++++++++++++++++
 1 file changed, 16 insertions(+)

diff --git a/keystore/tests/z_entities.rs b/keystore/tests/z_entities.rs
index e1409e1902..5188ea137e 100644
--- a/keystore/tests/z_entities.rs
+++ b/keystore/tests/z_entities.rs
@@ -35,6 +35,7 @@ macro_rules! test_for_entity {
             crate::tests_impl::can_remove_entity::<$entity>(&store, entity).await;
 
             let ignore_count = pat_to_bool!($($ignore_entity_count)?);
+            crate::tests_impl::insert_count_entities::<$entity>(&store).await;
             crate::tests_impl::can_list_entities_with_find_all::<$entity>(&store, ignore_count).await;
         }
     };
@@ -131,6 +132,21 @@ mod tests_impl {
         assert!(entity2.is_none());
     }
 
+    pub(super) async fn insert_count_entities<'a, R>(store: &CryptoKeystore)
+    where
+        R: Clone
+            + EntityRandomUpdateExt
+            + Entity<ConnectionType = KeystoreDatabaseConnection>
+            + EntityDatabaseMutation<'a>
+            + Send
+            + Sync,
+    {
+        for _ in 0..ENTITY_COUNT {
+            let entity = R::random();
+            store.save(entity).await.unwrap();
+        }
+    }
+
     pub(crate) async fn can_list_entities_with_find_all<'a, R>(store: &CryptoKeystore, ignore_entity_count: bool)
     where
         R: Clone

From 1f1474b1d82e3aafa9af7660200d20a9c10f1d07 Mon Sep 17 00:00:00 2001
From: Peter Goodspeed-Niklaus <peter.goodspeed-niklaus@wire.com>
Date: Fri, 19 Dec 2025 17:00:23 +0100
Subject: [PATCH 29/48] fixup! chore(keystore/tests): rollback on drop, not
 commit

---
 keystore/tests/common.rs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/keystore/tests/common.rs b/keystore/tests/common.rs
index 0744ccdec3..3b26a6e1ce 100644
--- a/keystore/tests/common.rs
+++ b/keystore/tests/common.rs
@@ -72,7 +72,7 @@ impl Drop for KeystoreTestContext {
             #[cfg(not(target_family = "wasm"))]
             futures_lite::future::block_on(rollback_and_wipe);
             #[cfg(target_family = "wasm")]
-            wasm_bindgen_futures::spawn_local(commit_and_wipe);
+            wasm_bindgen_futures::spawn_local(rollback_and_wipe);
         }
     }
 }

From 61dc48abddcb8ec5b9dcb2658019afb850d8571d Mon Sep 17 00:00:00 2001
From: Peter Goodspeed-Niklaus <peter.goodspeed-niklaus@wire.com>
Date: Fri, 19 Dec 2025 17:57:04 +0100
Subject: [PATCH 30/48] chore(keystore): use correct primary key for mls types

For some reason we're computing and storing a sha256 hash for these
types, but we're not using it as the primary key. OpenMLS uses the
raw id field as the primary key in a way that we can't modify,
so that's what has to be the pk. As for the hash: wasted space.
---
 .../generic/mls/encryption_keypair.rs         | 38 +++++++++++---
 .../platform/generic/mls/hpke_private_key.rs  | 38 +++++++++++---
 .../platform/generic/mls/psk_bundle.rs        | 38 ++++++++++++--
 .../platform/wasm/mls/encryption_keypair.rs   | 49 ++++++++++++++----
 .../platform/wasm/mls/hpke_private_key.rs     | 51 ++++++++++++++-----
 .../entities/platform/wasm/mls/psk_bundle.rs  | 51 ++++++++++++++-----
 keystore/src/mls.rs                           | 24 +++------
 7 files changed, 215 insertions(+), 74 deletions(-)

diff --git a/keystore/src/entities/platform/generic/mls/encryption_keypair.rs b/keystore/src/entities/platform/generic/mls/encryption_keypair.rs
index 0d0a09a0a0..bf5fb7f004 100644
--- a/keystore/src/entities/platform/generic/mls/encryption_keypair.rs
+++ b/keystore/src/entities/platform/generic/mls/encryption_keypair.rs
@@ -15,7 +15,7 @@ use crate::{
     },
     traits::{
         BorrowPrimaryKey, Entity as NewEntity, EntityBase as NewEntityBase, EntityDatabaseMutation,
-        EntityDeleteBorrowed, KeyType, PrimaryKey,
+        EntityDeleteBorrowed, EntityGetBorrowed, KeyType, PrimaryKey,
     },
 };
 
@@ -185,17 +185,25 @@ impl NewEntityBase for StoredEncryptionKeyPair {
 }
 
 impl PrimaryKey for StoredEncryptionKeyPair {
-    type PrimaryKey = Sha256Hash;
+    type PrimaryKey = Vec<u8>;
 
-    fn primary_key(&self) -> Sha256Hash {
-        Sha256Hash::hash_from(&self.pk)
+    fn primary_key(&self) -> Vec<u8> {
+        self.pk.clone()
+    }
+}
+
+impl BorrowPrimaryKey for StoredEncryptionKeyPair {
+    type BorrowedPrimaryKey = [u8];
+
+    fn borrow_primary_key(&self) -> &Self::BorrowedPrimaryKey {
+        &self.pk
     }
 }
 
 #[async_trait]
 impl NewEntity for StoredEncryptionKeyPair {
-    async fn get(conn: &mut Self::ConnectionType, id: &Sha256Hash) -> CryptoKeystoreResult<Option<Self>> {
-        get_helper::<Self, _>(conn, "pk_sha256", id, Self::from_row).await
+    async fn get(conn: &mut Self::ConnectionType, id: &Vec<u8>) -> CryptoKeystoreResult<Option<Self>> {
+        Self::get_borrowed(conn, id.as_slice()).await
     }
 
     async fn load_all(conn: &mut Self::ConnectionType) -> CryptoKeystoreResult<Vec<Self>> {
@@ -207,6 +215,13 @@ impl NewEntity for StoredEncryptionKeyPair {
     }
 }
 
+#[async_trait]
+impl EntityGetBorrowed for StoredEncryptionKeyPair {
+    async fn get_borrowed(conn: &mut Self::ConnectionType, id: &[u8]) -> CryptoKeystoreResult<Option<Self>> {
+        get_helper::<Self, _>(conn, "pk", id, Self::from_row).await
+    }
+}
+
 #[async_trait]
 impl<'a> EntityDatabaseMutation<'a> for StoredEncryptionKeyPair {
     type Transaction = TransactionWrapper<'a>;
@@ -222,7 +237,14 @@ impl<'a> EntityDatabaseMutation<'a> for StoredEncryptionKeyPair {
         count_helper_tx::<Self>(tx).await
     }
 
-    async fn delete(tx: &Self::Transaction, id: &Sha256Hash) -> CryptoKeystoreResult<bool> {
-        delete_helper::<Self>(tx, "pk_sha256", id).await
+    async fn delete(tx: &Self::Transaction, id: &Vec<u8>) -> CryptoKeystoreResult<bool> {
+        Self::delete_borrowed(tx, id.as_slice()).await
+    }
+}
+
+#[async_trait]
+impl<'a> EntityDeleteBorrowed<'a> for StoredEncryptionKeyPair {
+    async fn delete_borrowed(tx: &Self::Transaction, id: &[u8]) -> CryptoKeystoreResult<bool> {
+        delete_helper::<Self>(tx, "pk", id).await
     }
 }
diff --git a/keystore/src/entities/platform/generic/mls/hpke_private_key.rs b/keystore/src/entities/platform/generic/mls/hpke_private_key.rs
index fdaf64719f..7c6f06e7f8 100644
--- a/keystore/src/entities/platform/generic/mls/hpke_private_key.rs
+++ b/keystore/src/entities/platform/generic/mls/hpke_private_key.rs
@@ -15,7 +15,7 @@ use crate::{
     },
     traits::{
         BorrowPrimaryKey, Entity as NewEntity, EntityBase as NewEntityBase, EntityDatabaseMutation,
-        EntityDeleteBorrowed, KeyType, PrimaryKey,
+        EntityDeleteBorrowed, EntityGetBorrowed, KeyType, PrimaryKey,
     },
 };
 
@@ -183,17 +183,25 @@ impl NewEntityBase for StoredHpkePrivateKey {
 }
 
 impl PrimaryKey for StoredHpkePrivateKey {
-    type PrimaryKey = Sha256Hash;
+    type PrimaryKey = Vec<u8>;
 
-    fn primary_key(&self) -> Sha256Hash {
-        Sha256Hash::hash_from(&self.pk)
+    fn primary_key(&self) -> Vec<u8> {
+        self.pk.clone()
+    }
+}
+
+impl BorrowPrimaryKey for StoredHpkePrivateKey {
+    type BorrowedPrimaryKey = [u8];
+
+    fn borrow_primary_key(&self) -> &Self::BorrowedPrimaryKey {
+        &self.pk
     }
 }
 
 #[async_trait]
 impl NewEntity for StoredHpkePrivateKey {
-    async fn get(conn: &mut Self::ConnectionType, id: &Sha256Hash) -> CryptoKeystoreResult<Option<Self>> {
-        get_helper::<Self, _>(conn, "pk_sha256", id, Self::from_row).await
+    async fn get(conn: &mut Self::ConnectionType, id: &Vec<u8>) -> CryptoKeystoreResult<Option<Self>> {
+        Self::get_borrowed(conn, id.as_slice()).await
     }
 
     async fn load_all(conn: &mut Self::ConnectionType) -> CryptoKeystoreResult<Vec<Self>> {
@@ -205,6 +213,13 @@ impl NewEntity for StoredHpkePrivateKey {
     }
 }
 
+#[async_trait]
+impl EntityGetBorrowed for StoredHpkePrivateKey {
+    async fn get_borrowed(conn: &mut Self::ConnectionType, id: &[u8]) -> CryptoKeystoreResult<Option<Self>> {
+        get_helper::<Self, _>(conn, "pk", id, Self::from_row).await
+    }
+}
+
 #[async_trait]
 impl<'a> EntityDatabaseMutation<'a> for StoredHpkePrivateKey {
     type Transaction = TransactionWrapper<'a>;
@@ -220,7 +235,14 @@ impl<'a> EntityDatabaseMutation<'a> for StoredHpkePrivateKey {
         count_helper_tx::<Self>(tx).await
     }
 
-    async fn delete(tx: &Self::Transaction, id: &Sha256Hash) -> CryptoKeystoreResult<bool> {
-        delete_helper::<Self>(tx, "pk_sha256", id).await
+    async fn delete(tx: &Self::Transaction, id: &Vec<u8>) -> CryptoKeystoreResult<bool> {
+        Self::delete_borrowed(tx, id.as_slice()).await
+    }
+}
+
+#[async_trait]
+impl<'a> EntityDeleteBorrowed<'a> for StoredHpkePrivateKey {
+    async fn delete_borrowed(tx: &Self::Transaction, id: &[u8]) -> CryptoKeystoreResult<bool> {
+        delete_helper::<Self>(tx, "pk", id).await
     }
 }
diff --git a/keystore/src/entities/platform/generic/mls/psk_bundle.rs b/keystore/src/entities/platform/generic/mls/psk_bundle.rs
index ef75cbd7c8..23615394bc 100644
--- a/keystore/src/entities/platform/generic/mls/psk_bundle.rs
+++ b/keystore/src/entities/platform/generic/mls/psk_bundle.rs
@@ -10,7 +10,10 @@ use crate::{
         Entity, EntityBase, EntityFindParams, EntityIdStringExt, EntityTransactionExt, StoredPskBundle, StringEntityId,
         count_helper, count_helper_tx, delete_helper, get_helper, load_all_helper,
     },
-    traits::{Entity as NewEntity, EntityBase as NewEntityBase, EntityDatabaseMutation, PrimaryKey},
+    traits::{
+        BorrowPrimaryKey, Entity as NewEntity, EntityBase as NewEntityBase, EntityDatabaseMutation,
+        EntityDeleteBorrowed, EntityGetBorrowed, PrimaryKey,
+    },
 };
 
 #[async_trait::async_trait]
@@ -172,17 +175,25 @@ impl NewEntityBase for StoredPskBundle {
 }
 
 impl PrimaryKey for StoredPskBundle {
-    type PrimaryKey = Sha256Hash;
+    type PrimaryKey = Vec<u8>;
 
     fn primary_key(&self) -> Self::PrimaryKey {
-        Sha256Hash::hash_from(&self.psk_id)
+        self.psk_id.clone()
+    }
+}
+
+impl BorrowPrimaryKey for StoredPskBundle {
+    type BorrowedPrimaryKey = [u8];
+
+    fn borrow_primary_key(&self) -> &Self::BorrowedPrimaryKey {
+        &self.psk_id
     }
 }
 
 #[async_trait]
 impl NewEntity for StoredPskBundle {
     async fn get(conn: &mut Self::ConnectionType, key: &Self::PrimaryKey) -> CryptoKeystoreResult<Option<Self>> {
-        get_helper::<Self, _>(conn, "id_sha256", key, Self::from_row).await
+        Self::get_borrowed(conn, key.as_slice()).await
     }
 
     async fn count(conn: &mut Self::ConnectionType) -> CryptoKeystoreResult<u32> {
@@ -194,6 +205,16 @@ impl NewEntity for StoredPskBundle {
     }
 }
 
+#[async_trait]
+impl EntityGetBorrowed for StoredPskBundle {
+    async fn get_borrowed(
+        conn: &mut Self::ConnectionType,
+        key: &Self::BorrowedPrimaryKey,
+    ) -> CryptoKeystoreResult<Option<Self>> {
+        get_helper::<Self, _>(conn, "id", key, Self::from_row).await
+    }
+}
+
 #[async_trait]
 impl<'a> EntityDatabaseMutation<'a> for StoredPskBundle {
     type Transaction = TransactionWrapper<'a>;
@@ -210,6 +231,13 @@ impl<'a> EntityDatabaseMutation<'a> for StoredPskBundle {
     }
 
     async fn delete(tx: &Self::Transaction, id: &Self::PrimaryKey) -> CryptoKeystoreResult<bool> {
-        delete_helper::<Self>(tx, "id_sha256", id).await
+        Self::delete_borrowed(tx, id.as_slice()).await
+    }
+}
+
+#[async_trait]
+impl<'a> EntityDeleteBorrowed<'a> for StoredPskBundle {
+    async fn delete_borrowed(tx: &Self::Transaction, id: &Self::BorrowedPrimaryKey) -> CryptoKeystoreResult<bool> {
+        delete_helper::<Self>(tx, "id", id).await
     }
 }
diff --git a/keystore/src/entities/platform/wasm/mls/encryption_keypair.rs b/keystore/src/entities/platform/wasm/mls/encryption_keypair.rs
index 376aca7b4e..b38aa0682a 100644
--- a/keystore/src/entities/platform/wasm/mls/encryption_keypair.rs
+++ b/keystore/src/entities/platform/wasm/mls/encryption_keypair.rs
@@ -1,12 +1,12 @@
 use async_trait::async_trait;
 
 use crate::{
-    CryptoKeystoreResult, MissingKeyErrorKind, Sha256Hash,
+    CryptoKeystoreResult, MissingKeyErrorKind,
     connection::{DatabaseConnection, KeystoreDatabaseConnection, TransactionWrapper},
     entities::{Entity, EntityBase, EntityFindParams, EntityTransactionExt, StoredEncryptionKeyPair, StringEntityId},
     traits::{
-        DecryptData, Decryptable, Decrypting, EncryptData, Encrypting, Entity as NewEntity,
-        EntityBase as NewEntityBase, EntityDatabaseMutation, KeyType, PrimaryKey,
+        BorrowPrimaryKey, DecryptData, Decryptable, Decrypting, EncryptData, Encrypting, Entity as NewEntity,
+        EntityBase as NewEntityBase, EntityDatabaseMutation, EntityDeleteBorrowed, EntityGetBorrowed, PrimaryKey,
     },
 };
 
@@ -79,17 +79,25 @@ impl NewEntityBase for StoredEncryptionKeyPair {
 }
 
 impl PrimaryKey for StoredEncryptionKeyPair {
-    type PrimaryKey = Sha256Hash;
+    type PrimaryKey = Vec<u8>;
 
     fn primary_key(&self) -> Self::PrimaryKey {
-        Sha256Hash::hash_from(&self.pk)
+        self.pk.clone()
+    }
+}
+
+impl BorrowPrimaryKey for StoredEncryptionKeyPair {
+    type BorrowedPrimaryKey = [u8];
+
+    fn borrow_primary_key(&self) -> &Self::BorrowedPrimaryKey {
+        &self.pk
     }
 }
 
 #[async_trait(?Send)]
 impl NewEntity for StoredEncryptionKeyPair {
-    async fn get(conn: &mut Self::ConnectionType, id: &Sha256Hash) -> CryptoKeystoreResult<Option<Self>> {
-        conn.storage().new_get(id.bytes().as_ref()).await
+    async fn get(conn: &mut Self::ConnectionType, id: &Vec<u8>) -> CryptoKeystoreResult<Option<Self>> {
+        Self::get_borrowed(conn, id.as_slice()).await
     }
 
     async fn load_all(conn: &mut Self::ConnectionType) -> CryptoKeystoreResult<Vec<Self>> {
@@ -101,6 +109,16 @@ impl NewEntity for StoredEncryptionKeyPair {
     }
 }
 
+#[async_trait(?Send)]
+impl EntityGetBorrowed for StoredEncryptionKeyPair {
+    async fn get_borrowed(
+        conn: &mut <Self as NewEntityBase>::ConnectionType,
+        id: &[u8],
+    ) -> CryptoKeystoreResult<Option<Self>> {
+        conn.storage().new_get(id).await
+    }
+}
+
 #[async_trait(?Send)]
 impl<'a> EntityDatabaseMutation<'a> for StoredEncryptionKeyPair {
     type Transaction = TransactionWrapper<'a>;
@@ -113,8 +131,18 @@ impl<'a> EntityDatabaseMutation<'a> for StoredEncryptionKeyPair {
         tx.new_count::<Self>().await
     }
 
-    async fn delete(tx: &Self::Transaction, id: &Sha256Hash) -> CryptoKeystoreResult<bool> {
-        tx.new_delete::<Self>(id.bytes().as_ref()).await
+    async fn delete(tx: &Self::Transaction, id: &Vec<u8>) -> CryptoKeystoreResult<bool> {
+        Self::delete_borrowed(tx, id.as_slice()).await
+    }
+}
+
+#[async_trait(?Send)]
+impl<'a> EntityDeleteBorrowed<'a> for StoredEncryptionKeyPair {
+    async fn delete_borrowed(
+        tx: &<Self as EntityDatabaseMutation<'a>>::Transaction,
+        id: &[u8],
+    ) -> CryptoKeystoreResult<bool> {
+        tx.new_delete::<Self>(id).await
     }
 }
 
@@ -149,8 +177,7 @@ impl Decrypting<'static> for StoredEncryptionKeyPairDecrypt {
     fn decrypt(self, cipher: &aes_gcm::Aes256Gcm) -> CryptoKeystoreResult<Self::DecryptedForm> {
         let Self { pk, sk } = self;
 
-        let primary_key = Sha256Hash::hash_from(&pk);
-        let sk = <StoredEncryptionKeyPair as DecryptData>::decrypt_data(cipher, &primary_key, &sk)?;
+        let sk = <StoredEncryptionKeyPair as DecryptData>::decrypt_data(cipher, &pk, &sk)?;
 
         Ok(StoredEncryptionKeyPair { pk, sk })
     }
diff --git a/keystore/src/entities/platform/wasm/mls/hpke_private_key.rs b/keystore/src/entities/platform/wasm/mls/hpke_private_key.rs
index 22ec5565fd..fedcc6245f 100644
--- a/keystore/src/entities/platform/wasm/mls/hpke_private_key.rs
+++ b/keystore/src/entities/platform/wasm/mls/hpke_private_key.rs
@@ -2,12 +2,12 @@ use async_trait::async_trait;
 use serde::{Deserialize, Serialize};
 
 use crate::{
-    CryptoKeystoreResult, MissingKeyErrorKind, Sha256Hash,
+    CryptoKeystoreResult, MissingKeyErrorKind,
     connection::{DatabaseConnection, KeystoreDatabaseConnection, TransactionWrapper},
     entities::{Entity, EntityBase, EntityFindParams, EntityTransactionExt, StoredHpkePrivateKey, StringEntityId},
     traits::{
-        DecryptData, Decryptable, Decrypting, EncryptData, Encrypting, Entity as NewEntity,
-        EntityBase as NewEntityBase, EntityDatabaseMutation, KeyType, PrimaryKey,
+        BorrowPrimaryKey, DecryptData, Decryptable, Decrypting, EncryptData, Encrypting, Entity as NewEntity,
+        EntityBase as NewEntityBase, EntityDatabaseMutation, EntityDeleteBorrowed, EntityGetBorrowed, PrimaryKey,
     },
 };
 
@@ -80,17 +80,25 @@ impl NewEntityBase for StoredHpkePrivateKey {
 }
 
 impl PrimaryKey for StoredHpkePrivateKey {
-    type PrimaryKey = Sha256Hash;
+    type PrimaryKey = Vec<u8>;
 
-    fn primary_key(&self) -> Self::PrimaryKey {
-        Sha256Hash::hash_from(&self.pk)
+    fn primary_key(&self) -> Vec<u8> {
+        self.pk.clone()
+    }
+}
+
+impl BorrowPrimaryKey for StoredHpkePrivateKey {
+    type BorrowedPrimaryKey = [u8];
+
+    fn borrow_primary_key(&self) -> &Self::BorrowedPrimaryKey {
+        &self.pk
     }
 }
 
 #[async_trait(?Send)]
 impl NewEntity for StoredHpkePrivateKey {
-    async fn get(conn: &mut Self::ConnectionType, id: &Sha256Hash) -> CryptoKeystoreResult<Option<Self>> {
-        conn.storage().new_get(id.bytes().as_ref()).await
+    async fn get(conn: &mut Self::ConnectionType, id: &Vec<u8>) -> CryptoKeystoreResult<Option<Self>> {
+        Self::get_borrowed(conn, id.as_slice()).await
     }
 
     async fn load_all(conn: &mut Self::ConnectionType) -> CryptoKeystoreResult<Vec<Self>> {
@@ -102,6 +110,16 @@ impl NewEntity for StoredHpkePrivateKey {
     }
 }
 
+#[async_trait(?Send)]
+impl EntityGetBorrowed for StoredHpkePrivateKey {
+    async fn get_borrowed(
+        conn: &mut <Self as NewEntityBase>::ConnectionType,
+        id: &[u8],
+    ) -> CryptoKeystoreResult<Option<Self>> {
+        conn.storage().new_get(id).await
+    }
+}
+
 #[async_trait(?Send)]
 impl<'a> EntityDatabaseMutation<'a> for StoredHpkePrivateKey {
     type Transaction = TransactionWrapper<'a>;
@@ -114,8 +132,18 @@ impl<'a> EntityDatabaseMutation<'a> for StoredHpkePrivateKey {
         tx.new_count::<Self>().await
     }
 
-    async fn delete(tx: &Self::Transaction, id: &Sha256Hash) -> CryptoKeystoreResult<bool> {
-        tx.new_delete::<Self>(id.bytes().as_ref()).await
+    async fn delete(tx: &Self::Transaction, id: &Vec<u8>) -> CryptoKeystoreResult<bool> {
+        Self::delete_borrowed(tx, id.as_slice()).await
+    }
+}
+
+#[async_trait(?Send)]
+impl<'a> EntityDeleteBorrowed<'a> for StoredHpkePrivateKey {
+    async fn delete_borrowed(
+        tx: &<Self as EntityDatabaseMutation<'a>>::Transaction,
+        id: &[u8],
+    ) -> CryptoKeystoreResult<bool> {
+        tx.new_delete::<Self>(id).await
     }
 }
 
@@ -144,8 +172,7 @@ impl Decrypting<'static> for StoredHpkePrivateKeyDecrypt {
     type DecryptedForm = StoredHpkePrivateKey;
 
     fn decrypt(self, cipher: &aes_gcm::Aes256Gcm) -> CryptoKeystoreResult<Self::DecryptedForm> {
-        let primary_key = Sha256Hash::hash_from(&self.pk);
-        let sk = <StoredHpkePrivateKey as DecryptData>::decrypt_data(cipher, &primary_key, &self.sk)?;
+        let sk = <StoredHpkePrivateKey as DecryptData>::decrypt_data(cipher, &self.pk, &self.sk)?;
         Ok(StoredHpkePrivateKey { sk, pk: self.pk })
     }
 }
diff --git a/keystore/src/entities/platform/wasm/mls/psk_bundle.rs b/keystore/src/entities/platform/wasm/mls/psk_bundle.rs
index 1b2a4b721f..593b114063 100644
--- a/keystore/src/entities/platform/wasm/mls/psk_bundle.rs
+++ b/keystore/src/entities/platform/wasm/mls/psk_bundle.rs
@@ -2,12 +2,12 @@ use async_trait::async_trait;
 use serde::{Deserialize, Serialize};
 
 use crate::{
-    CryptoKeystoreResult, MissingKeyErrorKind, Sha256Hash,
+    CryptoKeystoreResult, MissingKeyErrorKind,
     connection::{DatabaseConnection, KeystoreDatabaseConnection, TransactionWrapper},
     entities::{Entity, EntityBase, EntityFindParams, EntityTransactionExt, StoredPskBundle, StringEntityId},
     traits::{
-        DecryptData, Decryptable, Decrypting, EncryptData, Encrypting, Entity as NewEntity,
-        EntityBase as NewEntityBase, EntityDatabaseMutation, KeyType as _, PrimaryKey,
+        BorrowPrimaryKey, DecryptData, Decryptable, Decrypting, EncryptData, Encrypting, Entity as NewEntity,
+        EntityBase as NewEntityBase, EntityDatabaseMutation, EntityDeleteBorrowed, EntityGetBorrowed, PrimaryKey,
     },
 };
 
@@ -80,17 +80,25 @@ impl NewEntityBase for StoredPskBundle {
 }
 
 impl PrimaryKey for StoredPskBundle {
-    type PrimaryKey = Sha256Hash;
+    type PrimaryKey = Vec<u8>;
 
-    fn primary_key(&self) -> Self::PrimaryKey {
-        Sha256Hash::hash_from(&self.psk_id)
+    fn primary_key(&self) -> Vec<u8> {
+        self.psk_id.clone()
+    }
+}
+
+impl BorrowPrimaryKey for StoredPskBundle {
+    type BorrowedPrimaryKey = [u8];
+
+    fn borrow_primary_key(&self) -> &[u8] {
+        &self.psk_id
     }
 }
 
 #[async_trait(?Send)]
 impl NewEntity for StoredPskBundle {
-    async fn get(conn: &mut Self::ConnectionType, key: &Self::PrimaryKey) -> CryptoKeystoreResult<Option<Self>> {
-        conn.storage().new_get(key.bytes().as_ref()).await
+    async fn get(conn: &mut Self::ConnectionType, key: &Vec<u8>) -> CryptoKeystoreResult<Option<Self>> {
+        Self::get_borrowed(conn, key.as_slice()).await
     }
 
     async fn count(conn: &mut Self::ConnectionType) -> CryptoKeystoreResult<u32> {
@@ -102,6 +110,16 @@ impl NewEntity for StoredPskBundle {
     }
 }
 
+#[async_trait(?Send)]
+impl EntityGetBorrowed for StoredPskBundle {
+    async fn get_borrowed(
+        conn: &mut <Self as NewEntityBase>::ConnectionType,
+        key: &[u8],
+    ) -> CryptoKeystoreResult<Option<Self>> {
+        conn.storage().new_get(key).await
+    }
+}
+
 #[async_trait(?Send)]
 impl<'a> EntityDatabaseMutation<'a> for StoredPskBundle {
     type Transaction = TransactionWrapper<'a>;
@@ -114,8 +132,18 @@ impl<'a> EntityDatabaseMutation<'a> for StoredPskBundle {
         tx.new_count::<Self>().await
     }
 
-    async fn delete(tx: &Self::Transaction, id: &Self::PrimaryKey) -> CryptoKeystoreResult<bool> {
-        tx.new_delete::<Self>(id.bytes().as_ref()).await
+    async fn delete(tx: &Self::Transaction, id: &Vec<u8>) -> CryptoKeystoreResult<bool> {
+        Self::delete_borrowed(tx, id.as_slice()).await
+    }
+}
+
+#[async_trait(?Send)]
+impl<'a> EntityDeleteBorrowed<'a> for StoredPskBundle {
+    async fn delete_borrowed(
+        tx: &<Self as EntityDatabaseMutation<'a>>::Transaction,
+        id: &[u8],
+    ) -> CryptoKeystoreResult<bool> {
+        tx.new_delete::<Self>(id).await
     }
 }
 
@@ -147,8 +175,7 @@ impl Decrypting<'static> for StoredPskBundleDecrypt {
     type DecryptedForm = StoredPskBundle;
 
     fn decrypt(self, cipher: &aes_gcm::Aes256Gcm) -> CryptoKeystoreResult<Self::DecryptedForm> {
-        let primary_key = Sha256Hash::hash_from(&self.psk_id);
-        let psk = <StoredPskBundle as DecryptData>::decrypt_data(cipher, &primary_key, &self.psk)?;
+        let psk = <StoredPskBundle as DecryptData>::decrypt_data(cipher, &self.psk_id, &self.psk)?;
         Ok(StoredPskBundle {
             psk_id: self.psk_id,
             psk,
diff --git a/keystore/src/mls.rs b/keystore/src/mls.rs
index 2f898ca5bb..c2dd12f25f 100644
--- a/keystore/src/mls.rs
+++ b/keystore/src/mls.rs
@@ -339,18 +339,15 @@ impl openmls_traits::key_store::OpenMlsKeyStore for crate::connection::Database
                 deser(&kp.keypackage).ok()
             }
             MlsEntityId::HpkePrivateKey => {
-                let hash = Sha256Hash::from_existing_hash(k).ok()?;
-                let hpke_pk: StoredHpkePrivateKey = self.get(&hash).await.ok().flatten()?;
+                let hpke_pk: StoredHpkePrivateKey = self.get_borrowed(k).await.ok().flatten()?;
                 deser(&hpke_pk.sk).ok()
             }
             MlsEntityId::PskBundle => {
-                let hash = Sha256Hash::from_existing_hash(k).ok()?;
-                let psk_bundle: StoredPskBundle = self.get(&hash).await.ok().flatten()?;
+                let psk_bundle: StoredPskBundle = self.get_borrowed(k).await.ok().flatten()?;
                 deser(&psk_bundle.psk).ok()
             }
             MlsEntityId::EncryptionKeyPair => {
-                let hash = Sha256Hash::from_existing_hash(k).ok()?;
-                let kp: StoredEncryptionKeyPair = self.get(&hash).await.ok().flatten()?;
+                let kp: StoredEncryptionKeyPair = self.get_borrowed(k).await.ok().flatten()?;
                 deser(&kp.sk).ok()
             }
             MlsEntityId::EpochEncryptionKeyPair => {
@@ -367,19 +364,10 @@ impl openmls_traits::key_store::OpenMlsKeyStore for crate::connection::Database
                 "Deleting a signature key pair should not be done through this API, any keypair should be deleted via
                 deleting a credential."
             ),
-            MlsEntityId::HpkePrivateKey => {
-                let hash = Sha256Hash::from_existing_hash(k)?;
-                self.remove::<StoredHpkePrivateKey>(&hash).await?
-            }
+            MlsEntityId::HpkePrivateKey => self.remove_borrowed::<StoredHpkePrivateKey>(k).await?,
             MlsEntityId::KeyPackage => self.remove_borrowed::<StoredKeypackage>(k).await?,
-            MlsEntityId::PskBundle => {
-                let hash = Sha256Hash::from_existing_hash(k)?;
-                self.remove::<StoredPskBundle>(&hash).await?
-            }
-            MlsEntityId::EncryptionKeyPair => {
-                let hash = Sha256Hash::from_existing_hash(k)?;
-                self.remove::<StoredEncryptionKeyPair>(&hash).await?
-            }
+            MlsEntityId::PskBundle => self.remove_borrowed::<StoredPskBundle>(k).await?,
+            MlsEntityId::EncryptionKeyPair => self.remove_borrowed::<StoredEncryptionKeyPair>(k).await?,
             MlsEntityId::EpochEncryptionKeyPair => self.remove_borrowed::<StoredEpochEncryptionKeypair>(k).await?,
         }
 

From 7e607b9b9854d39d6102d7be47df680a45846df9 Mon Sep 17 00:00:00 2001
From: Peter Goodspeed-Niklaus <peter.goodspeed-niklaus@wire.com>
Date: Mon, 22 Dec 2025 09:39:56 +0100
Subject: [PATCH 31/48] chore: adjust other crates to use new
 `FetchFromDatabase`

Update the previous version with this which fixes the correct primary key
---
 crypto/src/mls/conversation/merge.rs  | 7 ++-----
 crypto/src/mls/conversation/renew.rs  | 4 ++--
 crypto/src/mls/session/key_package.rs | 5 ++---
 crypto/src/test_utils/context.rs      | 3 +--
 4 files changed, 7 insertions(+), 12 deletions(-)

diff --git a/crypto/src/mls/conversation/merge.rs b/crypto/src/mls/conversation/merge.rs
index 4202a0cff9..22eee153ab 100644
--- a/crypto/src/mls/conversation/merge.rs
+++ b/crypto/src/mls/conversation/merge.rs
@@ -10,7 +10,7 @@
 //! | 0 pend. Proposal  | ❌              | ✅              |
 //! | 1+ pend. Proposal | ❌              | ✅              |
 
-use core_crypto_keystore::{Sha256Hash, entities::StoredEncryptionKeyPair};
+use core_crypto_keystore::entities::StoredEncryptionKeyPair;
 use mls_crypto_provider::MlsCryptoProvider;
 use openmls_traits::OpenMlsCryptoProvider;
 
@@ -33,10 +33,7 @@ impl MlsConversation {
         // ..so if there's any, we clear them after the commit is merged
         for oln in &previous_own_leaf_nodes {
             let ek = oln.encryption_key().as_slice();
-            let _ = backend
-                .key_store()
-                .remove::<StoredEncryptionKeyPair>(&Sha256Hash::hash_from(ek))
-                .await;
+            let _ = backend.key_store().remove_borrowed::<StoredEncryptionKeyPair>(ek).await;
         }
 
         client
diff --git a/crypto/src/mls/conversation/renew.rs b/crypto/src/mls/conversation/renew.rs
index 0b975d92e7..483a647f27 100644
--- a/crypto/src/mls/conversation/renew.rs
+++ b/crypto/src/mls/conversation/renew.rs
@@ -1,4 +1,4 @@
-use core_crypto_keystore::{Sha256Hash, entities::StoredEncryptionKeyPair};
+use core_crypto_keystore::entities::StoredEncryptionKeyPair;
 use mls_crypto_provider::MlsCryptoProvider;
 use openmls::prelude::{LeafNode, LeafNodeIndex, Proposal, QueuedProposal, Sender, StagedCommit};
 use openmls_traits::OpenMlsCryptoProvider;
@@ -131,7 +131,7 @@ impl MlsConversation {
             // encryption key from the keystore otherwise we would have a leak
             backend
                 .key_store()
-                .remove::<StoredEncryptionKeyPair>(&Sha256Hash::hash_from(leaf_node.encryption_key().as_slice()))
+                .remove_borrowed::<StoredEncryptionKeyPair>(leaf_node.encryption_key().as_slice())
                 .await
                 .map_err(KeystoreError::wrap("removing mls encryption keypair"))?;
         }
diff --git a/crypto/src/mls/session/key_package.rs b/crypto/src/mls/session/key_package.rs
index 3ae138adde..3f0a8fda55 100644
--- a/crypto/src/mls/session/key_package.rs
+++ b/crypto/src/mls/session/key_package.rs
@@ -1,7 +1,6 @@
 use std::{sync::Arc, time::Duration};
 
 use core_crypto_keystore::{
-    Sha256Hash,
     entities::{StoredEncryptionKeyPair, StoredHpkePrivateKey, StoredKeypackage},
     traits::FetchFromDatabase,
 };
@@ -139,10 +138,10 @@ impl Session {
         db.remove_borrowed::<StoredKeypackage>(kp_ref.hash_ref())
             .await
             .map_err(KeystoreError::wrap("removing key package from keystore"))?;
-        db.remove::<StoredHpkePrivateKey>(&Sha256Hash::hash_from(kp.hpke_init_key().as_slice()))
+        db.remove_borrowed::<StoredHpkePrivateKey>(kp.hpke_init_key().as_slice())
             .await
             .map_err(KeystoreError::wrap("removing private key from keystore"))?;
-        db.remove::<StoredEncryptionKeyPair>(&Sha256Hash::hash_from(kp.leaf_node().encryption_key().as_slice()))
+        db.remove_borrowed::<StoredEncryptionKeyPair>(kp.leaf_node().encryption_key().as_slice())
             .await
             .map_err(KeystoreError::wrap("removing encryption keypair from keystore"))?;
 
diff --git a/crypto/src/test_utils/context.rs b/crypto/src/test_utils/context.rs
index 2edd486a5f..40a25cd6d4 100644
--- a/crypto/src/test_utils/context.rs
+++ b/crypto/src/test_utils/context.rs
@@ -1,7 +1,6 @@
 use std::sync::Arc;
 
 use core_crypto_keystore::{
-    Sha256Hash,
     entities::{StoredCredential, StoredEncryptionKeyPair, StoredHpkePrivateKey, StoredKeypackage},
     traits::FetchFromDatabase,
 };
@@ -148,7 +147,7 @@ impl SessionContext {
             .keystore()
             .await
             .unwrap()
-            .get::<StoredHpkePrivateKey>(&Sha256Hash::hash_from(skp.tls_serialize_detached().unwrap()))
+            .get::<StoredHpkePrivateKey>(&skp.tls_serialize_detached().unwrap())
             .await
             .unwrap()
     }

From a08053edaf7c595276df9f68a38ed4310f55b00d Mon Sep 17 00:00:00 2001
From: Peter Goodspeed-Niklaus <peter.goodspeed-niklaus@wire.com>
Date: Mon, 22 Dec 2025 10:20:27 +0100
Subject: [PATCH 32/48] chore(crypto-macros): sql data type for hex fields is
 string

---
 crypto-macros/src/entity_derive_new/column.rs | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/crypto-macros/src/entity_derive_new/column.rs b/crypto-macros/src/entity_derive_new/column.rs
index ecc28c5944..ae26f2b77e 100644
--- a/crypto-macros/src/entity_derive_new/column.rs
+++ b/crypto-macros/src/entity_derive_new/column.rs
@@ -57,7 +57,11 @@ where
     pub(super) fn load_expression(&self) -> TokenStream {
         let column_name = self.sql_name();
 
-        let expr = quote!(row.get::<_, Vec<u8>>(#column_name)?);
+        let sql_data_type = match self.transformation {
+            None => quote!(Vec<u8>),
+            Some(FieldTransformation::Hex) => quote!(String),
+        };
+        let expr = quote!(row.get::<_, #sql_data_type>(#column_name)?);
 
         let expr = match self.transformation {
             None => expr,

From 1274cadd0cc176a8aa7d4ff278874cd2c308ce6b Mon Sep 17 00:00:00 2001
From: Peter Goodspeed-Niklaus <peter.goodspeed-niklaus@wire.com>
Date: Mon, 22 Dec 2025 12:57:56 +0100
Subject: [PATCH 33/48] chore: add test demonstrating storing and loading
 keypackages

We'd had an issue where keypackages couldn't be loaded due to
improper handling of hex-encoded keys, but this test now proves
that we've resolved that issue.
---
 crypto/src/mls/session/key_package.rs | 15 +++++++++++++++
 1 file changed, 15 insertions(+)

diff --git a/crypto/src/mls/session/key_package.rs b/crypto/src/mls/session/key_package.rs
index 3f0a8fda55..38ccf0f901 100644
--- a/crypto/src/mls/session/key_package.rs
+++ b/crypto/src/mls/session/key_package.rs
@@ -349,4 +349,19 @@ mod tests {
         })
         .await
     }
+
+    #[apply(all_cred_cipher)]
+    async fn can_store_and_load_key_packages(case: TestContext) {
+        let [cc] = case.sessions().await;
+
+        // generate a keypackage; automatically saves it
+        let kp = cc.new_keypackage(&case).await;
+
+        let all_keypackages = cc.session.get_keypackages().await.unwrap();
+        assert_eq!(all_keypackages[0], kp);
+
+        let kp_ref = kp.make_ref().unwrap();
+        let by_ref = cc.session.load_keypackage(&kp_ref).await.unwrap().unwrap();
+        assert_eq!(kp, by_ref);
+    }
 }

From efa9773448d4ab0855b583b53586a34c62adbca4 Mon Sep 17 00:00:00 2001
From: Peter Goodspeed-Niklaus <peter.goodspeed-niklaus@wire.com>
Date: Mon, 22 Dec 2025 14:10:32 +0100
Subject: [PATCH 34/48] chore(crypto-macros): search by hex value of key not
 key when appropriate

Turns out that when we have a field whose ID is a hex-encoded value,
the field is a string and the hex is lowercase. And what that means
is that we have to perform that transform when searching for that key!
---
 .../src/entity_derive_new/derive_impl.rs      | 29 +++++++++++++------
 1 file changed, 20 insertions(+), 9 deletions(-)

diff --git a/crypto-macros/src/entity_derive_new/derive_impl.rs b/crypto-macros/src/entity_derive_new/derive_impl.rs
index 8831e7ff2f..8a7a9846e4 100644
--- a/crypto-macros/src/entity_derive_new/derive_impl.rs
+++ b/crypto-macros/src/entity_derive_new/derive_impl.rs
@@ -3,7 +3,7 @@ use proc_macro2::{Span, TokenStream};
 use quote::quote;
 use syn::{Ident, Lifetime};
 
-use crate::entity_derive_new::{Entity, column_type::ColumnType};
+use crate::entity_derive_new::{Entity, FieldTransformation, column_type::ColumnType};
 
 impl quote::ToTokens for Entity {
     fn to_tokens(&self, tokens: &mut TokenStream) {
@@ -148,6 +148,11 @@ impl Entity {
         let field_assignments = std::iter::once(id_column.field_assignment())
             .chain(other_columns.iter().map(|column| column.field_assignment()));
 
+        let key_transform = match id_column.transformation {
+            None => None,
+            Some(FieldTransformation::Hex) => Some(quote! {let key = hex::encode(key);}),
+        };
+
         quote! {
             #[cfg_attr(target_family = "wasm", ::async_trait::async_trait(?Send))]
             #[cfg_attr(not(target_family = "wasm"), ::async_trait::async_trait)]
@@ -155,16 +160,17 @@ impl Entity {
                 async fn get_borrowed(conn: &mut Self::ConnectionType, key: &Self::BorrowedPrimaryKey)
                     -> crate::CryptoKeystoreResult<Option<Self>>
                 {
-                    let key_bytes = <&Self::BorrowedPrimaryKey as crate::traits::KeyType>::bytes(&key);
-                    let key_bytes = key_bytes.as_ref();
+                    let key = <&Self::BorrowedPrimaryKey as crate::traits::KeyType>::bytes(&key);
+                    let key = key.as_ref();
+                    #key_transform
                     #[cfg(target_family = "wasm")]
                     {
-                        conn.storage().new_get(key_bytes).await
+                        conn.storage().new_get(key).await
                     }
 
                     #[cfg(not(target_family = "wasm"))]
                     {
-                        crate::entities::platform::get_helper::<Self, _>(conn, #pk_column_name, key_bytes, |row| {
+                        crate::entities::platform::get_helper::<Self, _>(conn, #pk_column_name, key, |row| {
                             Ok(Self {
                                 #( #field_assignments, )*
                             })
@@ -247,6 +253,10 @@ impl Entity {
         } = self;
 
         let id_column_name = id_column.sql_name();
+        let key_transform = match id_column.transformation {
+            None => None,
+            Some(FieldTransformation::Hex) => Some(quote! {let key = hex::encode(key);}),
+        };
 
         quote! {
             #[cfg_attr(target_family = "wasm", async_trait::async_trait(?Send))]
@@ -259,16 +269,17 @@ impl Entity {
                 where
                     for<'pk> &'pk <Self as crate::traits::BorrowPrimaryKey>::BorrowedPrimaryKey: crate::traits::KeyType,
                 {
-                    let id_bytes = <&<Self as crate::traits::BorrowPrimaryKey>::BorrowedPrimaryKey as crate::traits::KeyType>::bytes(&id);
-                    let id_bytes = id_bytes.as_ref();
+                    let key = <&<Self as crate::traits::BorrowPrimaryKey>::BorrowedPrimaryKey as crate::traits::KeyType>::bytes(&id);
+                    let key = key.as_ref();
+                    #key_transform
                     #[cfg(target_family = "wasm")]
                     {
-                        tx.new_delete::<Self>(id_bytes).await
+                        tx.new_delete::<Self>(key).await
                     }
 
                     #[cfg(not(target_family = "wasm"))]
                     {
-                        crate::entities::platform::delete_helper::<Self>(tx, #id_column_name, id_bytes).await
+                        crate::entities::platform::delete_helper::<Self>(tx, #id_column_name, key).await
                     }
                 }
             }

From 1842952d1e322dfe363778298f5e97c7b17b293e Mon Sep 17 00:00:00 2001
From: Peter Goodspeed-Niklaus <peter.goodspeed-niklaus@wire.com>
Date: Mon, 22 Dec 2025 14:16:58 +0100
Subject: [PATCH 35/48] chore(keystore): add z_entities test for `E2eiAcmeCa`

---
 keystore/tests/z_entities.rs | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/keystore/tests/z_entities.rs b/keystore/tests/z_entities.rs
index 5188ea137e..6814cca8fd 100644
--- a/keystore/tests/z_entities.rs
+++ b/keystore/tests/z_entities.rs
@@ -189,6 +189,7 @@ mod tests {
     test_for_entity!(test_e2ei_intermediate_cert, E2eiIntermediateCert);
     test_for_entity!(test_e2ei_crl, E2eiCrl);
     test_for_entity!(test_e2ei_enrollment, StoredE2eiEnrollment ignore_update:true);
+    test_for_entity!(test_e2ei_acme_ca, E2eiAcmeCA ignore_entity_count:true ignore_find_many:true);
 
     cfg_if::cfg_if! {
         if #[cfg(feature = "proteus-keystore")] {
@@ -231,7 +232,7 @@ mod tests {
 #[cfg(test)]
 pub mod utils {
     use core_crypto_keystore::entities::{
-        MlsPendingMessage, PersistedMlsGroup, PersistedMlsPendingGroup, ProteusSession, StoredCredential,
+        E2eiAcmeCA, MlsPendingMessage, PersistedMlsGroup, PersistedMlsPendingGroup, ProteusSession, StoredCredential,
         StoredE2eiEnrollment, StoredEncryptionKeyPair, StoredEpochEncryptionKeypair, StoredHpkePrivateKey,
         StoredKeypackage, StoredPskBundle,
     };
@@ -351,6 +352,7 @@ pub mod utils {
     impl_entity_random_update_ext!(MlsPendingMessage, id_field = foreign_id, blob_fields = [message,]);
     impl_entity_random_update_ext!(StoredE2eiEnrollment, id_field = id, blob_fields = [content,]);
     impl_entity_random_update_ext!(StoredEpochEncryptionKeypair, id_field = id, blob_fields = [keypairs,]);
+    impl_entity_random_update_ext!(E2eiAcmeCA, blob_fields = [content,]);
 
     impl EntityRandomExt for core_crypto_keystore::entities::E2eiIntermediateCert {
         fn random() -> Self {

From 051202d2f787fc156de88cc3f1a3175888e394c2 Mon Sep 17 00:00:00 2001
From: Peter Goodspeed-Niklaus <peter.goodspeed-niklaus@wire.com>
Date: Mon, 22 Dec 2025 14:19:23 +0100
Subject: [PATCH 36/48] chore(CC): improve `fn e2ei_register_acme_ca`

1. Only fetch the database once
2. Given that we no longer raise errors on missing values,
   update the condition to produce an "already registered" error appropriately.
---
 .../e2e_identity/init_certificates.rs         | 25 +++++++------------
 1 file changed, 9 insertions(+), 16 deletions(-)

diff --git a/crypto/src/transaction_context/e2e_identity/init_certificates.rs b/crypto/src/transaction_context/e2e_identity/init_certificates.rs
index cb72c53c40..74ae66a5aa 100644
--- a/crypto/src/transaction_context/e2e_identity/init_certificates.rs
+++ b/crypto/src/transaction_context/e2e_identity/init_certificates.rs
@@ -36,18 +36,14 @@ impl TransactionContext {
     /// # Parameters
     /// * `trust_anchor_pem` - PEM certificate to anchor as a Trust Root
     pub async fn e2ei_register_acme_ca(&self, trust_anchor_pem: String) -> Result<()> {
-        {
-            if self
-                .mls_provider()
-                .await
-                .map_err(RecursiveError::transaction("getting mls provider"))?
-                .keystore()
-                .get_unique::<E2eiAcmeCA>()
-                .await
-                .is_ok()
-            {
-                return Err(Error::TrustAnchorAlreadyRegistered);
-            }
+        let database = self
+            .mls_provider()
+            .await
+            .map_err(RecursiveError::transaction("getting mls provider"))?
+            .keystore();
+
+        if matches!(database.get_unique::<E2eiAcmeCA>().await, Ok(Some(_))) {
+            return Err(Error::TrustAnchorAlreadyRegistered);
         }
 
         let pki_env = PkiEnvironment::init(PkiEnvironmentParams {
@@ -66,10 +62,7 @@ impl TransactionContext {
         // Save DER repr in keystore
         let cert_der = PkiEnvironment::encode_cert_to_der(&root_cert)?;
         let acme_ca = E2eiAcmeCA { content: cert_der };
-        self.mls_provider()
-            .await
-            .map_err(RecursiveError::transaction("getting mls provider"))?
-            .keystore()
+        database
             .save(acme_ca)
             .await
             .map_err(KeystoreError::wrap("saving acme ca"))?;

From dec9c5b44a19a53e03ad19f1bd0ccd2ba7c0914c Mon Sep 17 00:00:00 2001
From: Peter Goodspeed-Niklaus <peter.goodspeed-niklaus@wire.com>
Date: Mon, 22 Dec 2025 14:27:13 +0100
Subject: [PATCH 37/48] chore(CC): improve group-from-welcome error mapping

---
 crypto/src/mls/conversation/welcome.rs | 16 +++++++++-------
 1 file changed, 9 insertions(+), 7 deletions(-)

diff --git a/crypto/src/mls/conversation/welcome.rs b/crypto/src/mls/conversation/welcome.rs
index 788b26990a..93df24d222 100644
--- a/crypto/src/mls/conversation/welcome.rs
+++ b/crypto/src/mls/conversation/welcome.rs
@@ -38,13 +38,15 @@ impl MlsConversation {
     ) -> Result<Self> {
         let mls_group_config = configuration.as_openmls_default_configuration()?;
 
-        let group = MlsGroup::new_from_welcome(backend, &mls_group_config, welcome, None).await;
-
-        let group = match group {
-            Err(openmls::prelude::WelcomeError::NoMatchingKeyPackage)
-            | Err(openmls::prelude::WelcomeError::NoMatchingEncryptionKey) => return Err(Error::OrphanWelcome),
-            _ => group.map_err(MlsError::wrap("group could not be created from welcome"))?,
-        };
+        let group = MlsGroup::new_from_welcome(backend, &mls_group_config, welcome, None)
+            .await
+            .map_err(|err| {
+                use openmls::prelude::WelcomeError;
+                match err {
+                    WelcomeError::NoMatchingKeyPackage | WelcomeError::NoMatchingEncryptionKey => Error::OrphanWelcome,
+                    _ => MlsError::wrap("group could not be created from welcome")(err).into(),
+                }
+            })?;
 
         let id = ConversationId::from(group.group_id().as_slice());
         let existing_conversation = mls_groups.get_fetch(&id, &backend.keystore(), None).await;

From 614be9279dbc1e4c041690cbeb5f1a50f4595b87 Mon Sep 17 00:00:00 2001
From: Peter Goodspeed-Niklaus <peter.goodspeed-niklaus@wire.com>
Date: Mon, 22 Dec 2025 15:11:33 +0100
Subject: [PATCH 38/48] chore(keystore): query proteus identities properly

---
 keystore/src/entities/platform/generic/proteus/identity.rs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/keystore/src/entities/platform/generic/proteus/identity.rs b/keystore/src/entities/platform/generic/proteus/identity.rs
index adffd5c519..82b236186d 100644
--- a/keystore/src/entities/platform/generic/proteus/identity.rs
+++ b/keystore/src/entities/platform/generic/proteus/identity.rs
@@ -175,7 +175,7 @@ impl PrimaryKey for ProteusIdentity {
 impl NewEntity for ProteusIdentity {
     async fn get(conn: &mut Self::ConnectionType, _key: &()) -> CryptoKeystoreResult<Option<Self>> {
         let conn = conn.conn().await;
-        let mut stmt = conn.prepare_cached("SELECT rowid FROM proteus_identities ORDER BY rowid ASC LIMIT 1")?;
+        let mut stmt = conn.prepare_cached("SELECT sk, pk FROM proteus_identities ORDER BY rowid ASC LIMIT 1")?;
         stmt.query_one([], Self::from_row).optional().map_err(Into::into)
     }
 

From 449db4606a1eb5abbcdd8acda02f841a6a090a89 Mon Sep 17 00:00:00 2001
From: Peter Goodspeed-Niklaus <peter.goodspeed-niklaus@wire.com>
Date: Mon, 22 Dec 2025 16:53:07 +0100
Subject: [PATCH 39/48] chore(keystore): when saving entities remove their id
 from the deleted set

Without this behavior, if we have these operations

- Rm the entity with id X
- add a new entity with id X

then it gets added and then immediately deleted again, which is not helpful.
---
 keystore/src/transaction/mod.rs | 17 ++++++++++++++---
 1 file changed, 14 insertions(+), 3 deletions(-)

diff --git a/keystore/src/transaction/mod.rs b/keystore/src/transaction/mod.rs
index f9a9602a2a..8cee5c1811 100644
--- a/keystore/src/transaction/mod.rs
+++ b/keystore/src/transaction/mod.rs
@@ -49,6 +49,7 @@ impl KeystoreTransaction {
     ///
     /// - Adjust the entity by calling its [`pre_save()`][Entity::pre_save] method.
     /// - Store the entity in an internal map.
+    ///   - Remove the entity from the set of deleted entities, if it was there.
     /// - On [`Self::commit`], actually persist the entity into the supplied database.
     pub(crate) async fn save<'a, E>(&self, mut entity: E) -> CryptoKeystoreResult<E::AutoGeneratedFields>
     where
@@ -56,9 +57,19 @@ impl KeystoreTransaction {
     {
         let auto_generated_fields = entity.pre_save().await?;
 
-        let mut cache_guard = self.cache.write().await;
-        let table = cache_guard.entry(E::COLLECTION_NAME).or_default();
-        table.insert(EntityId::from_entity(&entity), entity.to_transaction_entity());
+        let entity_id = EntityId::from_entity(&entity);
+        {
+            // start by adding the entity
+            let mut cache_guard = self.cache.write().await;
+            let table = cache_guard.entry(E::COLLECTION_NAME).or_default();
+            table.insert(entity_id.clone(), entity.to_transaction_entity());
+        }
+        {
+            // at this point remove the entity from the set of deleted entities to ensure that
+            // this new data gets propagated
+            let mut cache_guard = self.deleted.write().await;
+            cache_guard.remove(&entity_id);
+        }
 
         Ok(auto_generated_fields)
     }

From 94ae4efbbfd5a61fdafac91c9c90384318858b4e Mon Sep 17 00:00:00 2001
From: Peter Goodspeed-Niklaus <peter.goodspeed-niklaus@wire.com>
Date: Tue, 23 Dec 2025 14:06:52 +0100
Subject: [PATCH 40/48] chore(keystore): handle `MlsPendingMessage`s properly

This entity type is really annoying because it just doesn't
have a single primary key per entity. So when we faked it as just
the conversation id, as soon as there was more than 1 pending message
per conversation, they started overwriting each other within the transaction.

So we've added a bunch of changes and workarounds to make it go.
I'm hopeful that in the future we can get rid of all or nearly all
of this, but for now, we're stuck with it.
---
 keystore/src/entities/mls.rs                  | 81 ++++++++++++++++++-
 .../platform/generic/mls/pending_message.rs   | 33 +-------
 .../platform/wasm/mls/pending_message.rs      | 38 +++------
 keystore/src/traits/item_encryption/aad.rs    |  7 ++
 .../traits/item_encryption/decrypt_data.rs    | 31 ++++++-
 .../traits/item_encryption/encrypt_data.rs    | 38 +++++++++
 keystore/src/traits/item_encryption/mod.rs    |  4 +-
 keystore/src/traits/mod.rs                    |  5 +-
 .../transaction/dynamic_dispatch/entity_id.rs | 13 ++-
 keystore/src/transaction/mod.rs               |  7 +-
 10 files changed, 187 insertions(+), 70 deletions(-)

diff --git a/keystore/src/entities/mls.rs b/keystore/src/entities/mls.rs
index 0a78b4b898..a8e163b6f7 100644
--- a/keystore/src/entities/mls.rs
+++ b/keystore/src/entities/mls.rs
@@ -1,9 +1,9 @@
-use zeroize::Zeroize;
+use zeroize::{Zeroize, ZeroizeOnDrop};
 
 use crate::{
     CryptoKeystoreError, CryptoKeystoreResult,
     connection::TransactionWrapper,
-    traits::{BorrowPrimaryKey, Entity, EntityBase, KeyType, OwnedKeyType},
+    traits::{BorrowPrimaryKey, Entity, EntityBase, KeyType, OwnedKeyType, PrimaryKey},
 };
 
 /// Entity representing a persisted `MlsGroup`
@@ -79,6 +79,76 @@ 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.
+///
+/// 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>,
+}
+
+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 From<&MlsPendingMessage> for MlsPendingMessagePrimaryKey {
+    fn from(value: &MlsPendingMessage) -> Self {
+        Self {
+            foreign_id: value.foreign_id.clone(),
+            message: value.message.clone(),
+        }
+    }
+}
+
+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()
+    }
+}
+
+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(),
+        })
+    }
+}
+
 /// Entity representing a buffered message
 #[derive(core_crypto_macros::Debug, Clone, PartialEq, Eq, Zeroize, serde::Serialize, serde::Deserialize)]
 #[zeroize(drop)]
@@ -88,6 +158,13 @@ pub struct MlsPendingMessage {
     pub message: Vec<u8>,
 }
 
+impl PrimaryKey for MlsPendingMessage {
+    type PrimaryKey = MlsPendingMessagePrimaryKey;
+    fn primary_key(&self) -> Self::PrimaryKey {
+        self.into()
+    }
+}
+
 /// Entity representing a buffered commit.
 ///
 /// There should always exist either 0 or 1 of these in the store per conversation.
diff --git a/keystore/src/entities/platform/generic/mls/pending_message.rs b/keystore/src/entities/platform/generic/mls/pending_message.rs
index 3c7d48e08b..900b939e47 100644
--- a/keystore/src/entities/platform/generic/mls/pending_message.rs
+++ b/keystore/src/entities/platform/generic/mls/pending_message.rs
@@ -12,7 +12,7 @@ use crate::{
     },
     traits::{
         BorrowPrimaryKey, Entity as NewEntity, EntityBase as NewEntityBase, EntityDatabaseMutation,
-        EntityDeleteBorrowed, EntityGetBorrowed, KeyType, PrimaryKey,
+        EntityDeleteBorrowed, EntityGetBorrowed, KeyType, OwnedKeyType, PrimaryKey,
     },
 };
 
@@ -216,27 +216,6 @@ impl NewEntityBase for MlsPendingMessage {
     }
 }
 
-/// Pending messages 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 it here to support `KeystoreTransaction::remove_pending_messages_by_conversation_id`.
-/// This is temporary! Post WPB-20844, we should remove that whole API and also reset the primary key type here to `()`.
-impl PrimaryKey for MlsPendingMessage {
-    type PrimaryKey = Vec<u8>;
-    fn primary_key(&self) -> Self::PrimaryKey {
-        self.foreign_id.clone()
-    }
-}
-
-impl BorrowPrimaryKey for MlsPendingMessage {
-    type BorrowedPrimaryKey = [u8];
-
-    fn borrow_primary_key(&self) -> &Self::BorrowedPrimaryKey {
-        &self.foreign_id
-    }
-}
-
 #[async_trait]
 impl NewEntity for MlsPendingMessage {
     async fn get(conn: &mut Self::ConnectionType, key: &Self::PrimaryKey) -> CryptoKeystoreResult<Option<Self>> {
@@ -252,16 +231,6 @@ impl NewEntity for MlsPendingMessage {
     }
 }
 
-#[async_trait]
-impl EntityGetBorrowed for MlsPendingMessage {
-    async fn get_borrowed(
-        conn: &mut Self::ConnectionType,
-        key: &Self::BorrowedPrimaryKey,
-    ) -> CryptoKeystoreResult<Option<Self>> {
-        panic!("cannot get `MlsPendingMessage` by primary key as it has no distinct primary key")
-    }
-}
-
 #[async_trait]
 impl<'a> EntityDatabaseMutation<'a> for MlsPendingMessage {
     type Transaction = TransactionWrapper<'a>;
diff --git a/keystore/src/entities/platform/wasm/mls/pending_message.rs b/keystore/src/entities/platform/wasm/mls/pending_message.rs
index e740b233fa..dae5f18bd8 100644
--- a/keystore/src/entities/platform/wasm/mls/pending_message.rs
+++ b/keystore/src/entities/platform/wasm/mls/pending_message.rs
@@ -8,8 +8,8 @@ use crate::{
     connection::{KeystoreDatabaseConnection, TransactionWrapper},
     entities::{Entity, EntityBase, EntityFindParams, EntityTransactionExt, MlsPendingMessage, StringEntityId},
     traits::{
-        BorrowPrimaryKey, DecryptData, Decryptable, Decrypting, EncryptData, Encrypting, Entity as NewEntity,
-        EntityBase as NewEntityBase, EntityDatabaseMutation, PrimaryKey,
+        DecryptWithExplicitEncryptionKey as _, Decryptable, Decrypting, EncryptWithExplicitEncryptionKey as _,
+        Encrypting, EncryptionKey, Entity as NewEntity, EntityBase as NewEntityBase, EntityDatabaseMutation,
     },
 };
 
@@ -116,30 +116,6 @@ impl NewEntityBase for MlsPendingMessage {
     }
 }
 
-/// Pending messages 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 it as a byte vector in this impl in order for encryption and decryption
-/// to work.
-impl PrimaryKey for MlsPendingMessage {
-    type PrimaryKey = Vec<u8>;
-
-    fn primary_key(&self) -> Self::PrimaryKey {
-        self.foreign_id.clone()
-    }
-}
-
-/// This implementation is purely to support `KeystoreTransaction::remove_pending_messages_by_conversation_id`;
-/// after WPB-20844, we should remove that whole API and also this implementation.
-impl BorrowPrimaryKey for MlsPendingMessage {
-    type BorrowedPrimaryKey = [u8];
-
-    fn borrow_primary_key(&self) -> &Self::BorrowedPrimaryKey {
-        &self.foreign_id
-    }
-}
-
 #[async_trait(?Send)]
 impl NewEntity for MlsPendingMessage {
     async fn get(_conn: &mut Self::ConnectionType, _key: &Self::PrimaryKey) -> CryptoKeystoreResult<Option<Self>> {
@@ -172,6 +148,12 @@ impl<'a> EntityDatabaseMutation<'a> for MlsPendingMessage {
     }
 }
 
+impl EncryptionKey for MlsPendingMessage {
+    fn encryption_key(&self) -> &[u8] {
+        &self.foreign_id
+    }
+}
+
 #[derive(Serialize)]
 pub struct MlsPendingMessageEncrypt<'a> {
     foreign_id: &'a [u8],
@@ -182,7 +164,7 @@ impl<'a> Encrypting<'a> for MlsPendingMessage {
     type EncryptedForm = MlsPendingMessageEncrypt<'a>;
 
     fn encrypt(&'a self, cipher: &aes_gcm::Aes256Gcm) -> CryptoKeystoreResult<Self::EncryptedForm> {
-        let message = <Self as EncryptData>::encrypt_data(self, cipher, &self.message)?;
+        let message = self.encrypt_data_with_encryption_key(cipher, &self.message)?;
         Ok(MlsPendingMessageEncrypt {
             foreign_id: &self.foreign_id,
             message,
@@ -200,7 +182,7 @@ impl Decrypting<'static> for MlsPendingMessageDecrypt {
     type DecryptedForm = MlsPendingMessage;
 
     fn decrypt(self, cipher: &aes_gcm::Aes256Gcm) -> CryptoKeystoreResult<Self::DecryptedForm> {
-        let message = <MlsPendingMessage as DecryptData>::decrypt_data(cipher, &self.foreign_id, &self.message)?;
+        let message = MlsPendingMessage::decrypt_data_with_encryption_key(cipher, &self.foreign_id, &self.message)?;
         Ok(MlsPendingMessage {
             foreign_id: self.foreign_id,
             message,
diff --git a/keystore/src/traits/item_encryption/aad.rs b/keystore/src/traits/item_encryption/aad.rs
index 9537bcc2de..25d0ce22d8 100644
--- a/keystore/src/traits/item_encryption/aad.rs
+++ b/keystore/src/traits/item_encryption/aad.rs
@@ -32,4 +32,11 @@ impl Aad {
         let id = primary_key.bytes().into_owned();
         Self { type_name, id }
     }
+
+    /// Don't use this unless you really have to! Prefer [`Self::from_primary_key`].
+    pub(super) fn from_encryption_key_bytes<E: Entity>(key_bytes: impl AsRef<[u8]>) -> Self {
+        let type_name = E::COLLECTION_NAME.as_bytes().to_vec();
+        let id = key_bytes.as_ref().to_owned();
+        Self { type_name, id }
+    }
 }
diff --git a/keystore/src/traits/item_encryption/decrypt_data.rs b/keystore/src/traits/item_encryption/decrypt_data.rs
index e83cd8b1a9..be33d539cb 100644
--- a/keystore/src/traits/item_encryption/decrypt_data.rs
+++ b/keystore/src/traits/item_encryption/decrypt_data.rs
@@ -1,7 +1,7 @@
 use super::aad::{AES_GCM_256_NONCE_SIZE, Aad};
 use crate::{
     CryptoKeystoreError, CryptoKeystoreResult,
-    traits::{Entity, EntityDeleteBorrowed, KeyType as _},
+    traits::{EncryptionKey, Entity, EntityDeleteBorrowed, KeyType as _},
 };
 
 fn decrypt_with_nonce_and_aad(
@@ -47,3 +47,32 @@ impl<E: Entity> DecryptData for E {
         decrypt_with_nonce_and_aad(cipher, msg, nonce, &aad)
     }
 }
+
+/// This trait uses an explicitly-set decryption key to decrypt some data.
+///
+/// This should rarely be used.
+pub trait DecryptWithExplicitEncryptionKey {
+    /// Decrypt some data with an encryption key (see [`EncryptionKey`]) instead of the instance's primary key.
+    fn decrypt_data_with_encryption_key(
+        cipher: &aes_gcm::Aes256Gcm,
+        encryption_key: &[u8],
+        data: &[u8],
+    ) -> CryptoKeystoreResult<Vec<u8>>;
+}
+
+impl<E> DecryptWithExplicitEncryptionKey for E
+where
+    E: Entity + EncryptionKey,
+{
+    fn decrypt_data_with_encryption_key(
+        cipher: &aes_gcm::Aes256Gcm,
+        encryption_key: &[u8],
+        data: &[u8],
+    ) -> CryptoKeystoreResult<Vec<u8>> {
+        let aad = Aad::from_encryption_key_bytes::<E>(encryption_key).serialize()?;
+        let (nonce, msg) = data
+            .split_at_checked(AES_GCM_256_NONCE_SIZE)
+            .ok_or(CryptoKeystoreError::AesGcmError)?;
+        decrypt_with_nonce_and_aad(cipher, msg, nonce, &aad)
+    }
+}
diff --git a/keystore/src/traits/item_encryption/encrypt_data.rs b/keystore/src/traits/item_encryption/encrypt_data.rs
index 55b214066a..0f558d091d 100644
--- a/keystore/src/traits/item_encryption/encrypt_data.rs
+++ b/keystore/src/traits/item_encryption/encrypt_data.rs
@@ -46,3 +46,41 @@ impl<E: Entity> EncryptData for E {
         encrypt_with_nonce_and_aad(cipher, data, &nonce_bytes, &aad)
     }
 }
+
+/// This trait is an hack enabling us to encrypt types for which we don't use the primary key in the AAD.
+///
+/// The only reason we'd ever want this is if the primary key is not what we actually use, and the only
+/// reason that would be the case is if we're faking a primary key where no such key really exists.
+///
+/// In other words, MLS pending messages.
+pub trait EncryptionKey {
+    /// Get the key bytes which are to be used as the encryption key for this data.
+    fn encryption_key(&self) -> &[u8];
+}
+
+/// This trait uses the explicitly-set encryption key to encrypt some data.
+///
+/// This should rarely be used.
+pub trait EncryptWithExplicitEncryptionKey {
+    /// Encrypt some data with an encryption key (see [`EncryptionKey`]) instead of the instance's primary key.
+    fn encrypt_data_with_encryption_key(
+        &self,
+        cipher: &aes_gcm::Aes256Gcm,
+        data: &[u8],
+    ) -> CryptoKeystoreResult<Vec<u8>>;
+}
+
+impl<E> EncryptWithExplicitEncryptionKey for E
+where
+    E: Entity + EncryptionKey,
+{
+    fn encrypt_data_with_encryption_key(
+        &self,
+        cipher: &aes_gcm::Aes256Gcm,
+        data: &[u8],
+    ) -> CryptoKeystoreResult<Vec<u8>> {
+        let aad = Aad::from_encryption_key_bytes::<E>(self.encryption_key()).serialize()?;
+        let nonce_bytes: [u8; AES_GCM_256_NONCE_SIZE] = rand::random();
+        encrypt_with_nonce_and_aad(cipher, data, &nonce_bytes, &aad)
+    }
+}
diff --git a/keystore/src/traits/item_encryption/mod.rs b/keystore/src/traits/item_encryption/mod.rs
index 258a23bd66..b7d8085a14 100644
--- a/keystore/src/traits/item_encryption/mod.rs
+++ b/keystore/src/traits/item_encryption/mod.rs
@@ -39,7 +39,7 @@ mod decrypting;
 mod encrypt_data;
 mod encrypting;
 
-pub use decrypt_data::DecryptData;
+pub use decrypt_data::{DecryptData, DecryptWithExplicitEncryptionKey};
 pub use decrypting::{Decryptable, Decrypting};
-pub use encrypt_data::EncryptData;
+pub use encrypt_data::{EncryptData, EncryptWithExplicitEncryptionKey, EncryptionKey};
 pub use encrypting::Encrypting;
diff --git a/keystore/src/traits/mod.rs b/keystore/src/traits/mod.rs
index 8c26414d41..e3aff4b4b5 100644
--- a/keystore/src/traits/mod.rs
+++ b/keystore/src/traits/mod.rs
@@ -16,7 +16,10 @@ pub use entity::{Entity, EntityGetBorrowed};
 pub use entity_base::EntityBase;
 pub use entity_database_mutation::{EntityDatabaseMutation, EntityDeleteBorrowed};
 pub use fetch_from_database::FetchFromDatabase;
-pub use item_encryption::{DecryptData, Decryptable, Decrypting, EncryptData, Encrypting};
+pub use item_encryption::{
+    DecryptData, DecryptWithExplicitEncryptionKey, Decryptable, Decrypting, EncryptData,
+    EncryptWithExplicitEncryptionKey, Encrypting, EncryptionKey,
+};
 pub use key_type::{KeyType, OwnedKeyType};
 pub use primary_key::{BorrowPrimaryKey, PrimaryKey};
 pub use unique_entity::{UniqueEntity, UniqueEntityExt, UniqueEntityImplementationHelper};
diff --git a/keystore/src/transaction/dynamic_dispatch/entity_id.rs b/keystore/src/transaction/dynamic_dispatch/entity_id.rs
index b0b24ef3cb..fea343fa36 100644
--- a/keystore/src/transaction/dynamic_dispatch/entity_id.rs
+++ b/keystore/src/transaction/dynamic_dispatch/entity_id.rs
@@ -1,3 +1,4 @@
+use core::fmt;
 use std::borrow::Cow;
 
 #[cfg(target_family = "wasm")]
@@ -22,6 +23,13 @@ pub(crate) struct EntityId {
     id: Vec<u8>,
 }
 
+impl fmt::Display for EntityId {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        let Self { typ, id } = self;
+        write!(f, "{typ:?}: {}", hex::encode(&id))
+    }
+}
+
 impl EntityId {
     fn primary_key<E>(&self) -> CryptoKeystoreResult<E::PrimaryKey>
     where
@@ -92,7 +100,10 @@ impl EntityId {
             EntityType::PersistedMlsPendingGroup => {
                 PersistedMlsPendingGroup::delete(tx, &self.primary_key::<PersistedMlsPendingGroup>()?).await
             }
-            EntityType::MlsPendingMessage => MlsPendingMessage::delete_by_conversation_id(tx, &self.id).await,
+            EntityType::MlsPendingMessage => {
+                let primary_key = self.primary_key::<MlsPendingMessage>()?;
+                MlsPendingMessage::delete_by_conversation_id(tx, &primary_key.foreign_id).await
+            }
             EntityType::StoredE2eiEnrollment => {
                 StoredE2eiEnrollment::delete(tx, &self.primary_key::<StoredE2eiEnrollment>()?).await
             }
diff --git a/keystore/src/transaction/mod.rs b/keystore/src/transaction/mod.rs
index 8cee5c1811..d1f5b2cad2 100644
--- a/keystore/src/transaction/mod.rs
+++ b/keystore/src/transaction/mod.rs
@@ -10,7 +10,7 @@ use itertools::Itertools;
 use crate::{
     CryptoKeystoreError, CryptoKeystoreResult,
     connection::{Database, KeystoreDatabaseConnection},
-    entities::{MlsPendingMessage, PersistedMlsGroupExt},
+    entities::{MlsPendingMessage, MlsPendingMessagePrimaryKey, PersistedMlsGroupExt},
     traits::{BorrowPrimaryKey, Entity, EntityBase as _, EntityDatabaseMutation, EntityDeleteBorrowed, KeyType},
     transaction::dynamic_dispatch::EntityId,
 };
@@ -158,10 +158,11 @@ impl KeystoreTransaction {
                 pending_message.foreign_id != conversation_id
             });
         }
+        drop(cache_guard);
 
         let mut deleted_set = self.deleted.write().await;
-        deleted_set.insert(EntityId::from_borrowed_primary_key::<MlsPendingMessage>(
-            conversation_id,
+        deleted_set.insert(EntityId::from_primary_key::<MlsPendingMessage>(
+            &MlsPendingMessagePrimaryKey::from_conversation_id(conversation_id),
         ));
     }
 

From 5fb34da11498722fdceb982b00d0b7876a664cba Mon Sep 17 00:00:00 2001
From: Peter Goodspeed-Niklaus <peter.goodspeed-niklaus@wire.com>
Date: Tue, 23 Dec 2025 14:25:30 +0100
Subject: [PATCH 41/48] chore(keystore): handle keys which are not byte vectors

For whatever reason the entity new derive macro was calling
`new_get` and `new_delete` with `String` or `&str` instances instead
of the expected `Vec<u8>` / `&[u8]`, and of course that was failing
type checking. But actually we'd prefer to have those functions available.
So broadening the acceptable types solves the problem for us (probably).
---
 keystore/src/connection/platform/wasm/storage/storage.rs     | 3 ++-
 keystore/src/connection/platform/wasm/storage/transaction.rs | 3 ++-
 2 files changed, 4 insertions(+), 2 deletions(-)

diff --git a/keystore/src/connection/platform/wasm/storage/storage.rs b/keystore/src/connection/platform/wasm/storage/storage.rs
index b3b3b13bfa..24ad862599 100644
--- a/keystore/src/connection/platform/wasm/storage/storage.rs
+++ b/keystore/src/connection/platform/wasm/storage/storage.rs
@@ -280,11 +280,12 @@ impl WasmEncryptedStorage {
     // After putting some thought into it, I'd prefer not to redesign the `Decrypting` trait, though.
     // There's always the chance that `serde_wasm_bindgen` will relax that restriction, at which point
     // we can just relax this bound and all will be good.
-    pub async fn new_get<'a, E>(&self, key: &[u8]) -> CryptoKeystoreResult<Option<E>>
+    pub async fn new_get<'a, E>(&self, key: impl AsRef<[u8]>) -> CryptoKeystoreResult<Option<E>>
     where
         E: NewEntity<ConnectionType = WasmConnection> + Decryptable<'a>,
         <E as Decryptable<'a>>::DecryptableFrom: DeserializeOwned,
     {
+        let key = key.as_ref();
         let js_value = match &self.storage {
             WasmStorageWrapper::Persistent(idb) => {
                 let transaction = idb.transaction(&[E::COLLECTION_NAME], TransactionMode::ReadOnly)?;
diff --git a/keystore/src/connection/platform/wasm/storage/transaction.rs b/keystore/src/connection/platform/wasm/storage/transaction.rs
index ab748ee380..c57f7931f8 100644
--- a/keystore/src/connection/platform/wasm/storage/transaction.rs
+++ b/keystore/src/connection/platform/wasm/storage/transaction.rs
@@ -150,10 +150,11 @@ impl WasmStorageTransaction<'_> {
     /// `BorrowPrimaryKey` or not, and without specialization, we can't just do the right thing
     /// and accept the more general form. But we do know the primary key and its borrowed form
     /// both implement `KeyType`, so it's always safe to accept a byte reference.
-    pub(crate) async fn new_delete<E>(&self, key: &[u8]) -> CryptoKeystoreResult<bool>
+    pub(crate) async fn new_delete<E>(&self, key: impl AsRef<[u8]>) -> CryptoKeystoreResult<bool>
     where
         E: NewEntity<ConnectionType = WasmConnection>,
     {
+        let key = key.as_ref();
         match self {
             WasmStorageTransaction::Persistent { tx, .. } => {
                 let query = JsValue::from(Uint8Array::from(key));

From f483541d8d2fff735c24457f01185ad2ce2075e6 Mon Sep 17 00:00:00 2001
From: Peter Goodspeed-Niklaus <peter.goodspeed-niklaus@wire.com>
Date: Tue, 23 Dec 2025 14:30:54 +0100
Subject: [PATCH 42/48] fixup! chore(keystore): handle `MlsPendingMessage`s
 properly

---
 keystore/src/entities/mls.rs | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/keystore/src/entities/mls.rs b/keystore/src/entities/mls.rs
index a8e163b6f7..a57d7d553b 100644
--- a/keystore/src/entities/mls.rs
+++ b/keystore/src/entities/mls.rs
@@ -83,8 +83,9 @@ pub struct PersistedMlsPendingGroup {
 /// 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.
+/// 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.
 ///
 /// This implementation is fairly inefficient and hopefully temporary. But it at least implements the correct semantics.
 #[derive(ZeroizeOnDrop)]

From 4bd96351fa9e436081cfecef4642e1133f304830 Mon Sep 17 00:00:00 2001
From: Peter Goodspeed-Niklaus <peter.goodspeed-niklaus@wire.com>
Date: Tue, 23 Dec 2025 14:34:28 +0100
Subject: [PATCH 43/48] fixup! chore(crypto-macros): search by hex value of key
 not key when appropriate

---
 crypto-macros/src/entity_derive_new/derive_impl.rs | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/crypto-macros/src/entity_derive_new/derive_impl.rs b/crypto-macros/src/entity_derive_new/derive_impl.rs
index 8a7a9846e4..d731ad66d4 100644
--- a/crypto-macros/src/entity_derive_new/derive_impl.rs
+++ b/crypto-macros/src/entity_derive_new/derive_impl.rs
@@ -148,6 +148,8 @@ impl Entity {
         let field_assignments = std::iter::once(id_column.field_assignment())
             .chain(other_columns.iter().map(|column| column.field_assignment()));
 
+        // if we ever add a second field transformation, we'll want this match pattern
+        #[allow(clippy::manual_map)]
         let key_transform = match id_column.transformation {
             None => None,
             Some(FieldTransformation::Hex) => Some(quote! {let key = hex::encode(key);}),
@@ -253,6 +255,8 @@ impl Entity {
         } = self;
 
         let id_column_name = id_column.sql_name();
+        // if we ever add a second field transformation, we'll want this match pattern
+        #[allow(clippy::manual_map)]
         let key_transform = match id_column.transformation {
             None => None,
             Some(FieldTransformation::Hex) => Some(quote! {let key = hex::encode(key);}),

From 3472eac3b6924295a5b7ca3f6e1fd098e8c3f6a7 Mon Sep 17 00:00:00 2001
From: Peter Goodspeed-Niklaus <peter.goodspeed-niklaus@wire.com>
Date: Tue, 23 Dec 2025 14:35:11 +0100
Subject: [PATCH 44/48] fixup! chore(keystore): handle `MlsPendingMessage`s
 properly

---
 keystore/src/transaction/dynamic_dispatch/entity_id.rs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/keystore/src/transaction/dynamic_dispatch/entity_id.rs b/keystore/src/transaction/dynamic_dispatch/entity_id.rs
index fea343fa36..a70f5c5531 100644
--- a/keystore/src/transaction/dynamic_dispatch/entity_id.rs
+++ b/keystore/src/transaction/dynamic_dispatch/entity_id.rs
@@ -26,7 +26,7 @@ pub(crate) struct EntityId {
 impl fmt::Display for EntityId {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         let Self { typ, id } = self;
-        write!(f, "{typ:?}: {}", hex::encode(&id))
+        write!(f, "{typ:?}: {}", hex::encode(id))
     }
 }
 

From f032f1c38b469cd46fb0647fbddecbf8c55b8cd4 Mon Sep 17 00:00:00 2001
From: Peter Goodspeed-Niklaus <peter.goodspeed-niklaus@wire.com>
Date: Tue, 23 Dec 2025 16:02:15 +0100
Subject: [PATCH 45/48] chore(crypto-macros): properly read optional types from
 sql

Turns out that rusqlite won't read null values as an empty byte vector,
so we have to add another column method governing how to generate
those get expressions.
---
 crypto-macros/src/entity_derive_new/column.rs |  2 +-
 .../src/entity_derive_new/column_type.rs      | 20 ++++++++++++++-----
 2 files changed, 16 insertions(+), 6 deletions(-)

diff --git a/crypto-macros/src/entity_derive_new/column.rs b/crypto-macros/src/entity_derive_new/column.rs
index ae26f2b77e..9eddf18700 100644
--- a/crypto-macros/src/entity_derive_new/column.rs
+++ b/crypto-macros/src/entity_derive_new/column.rs
@@ -58,7 +58,7 @@ where
         let column_name = self.sql_name();
 
         let sql_data_type = match self.transformation {
-            None => quote!(Vec<u8>),
+            None => self.column_type.get_as_type(),
             Some(FieldTransformation::Hex) => quote!(String),
         };
         let expr = quote!(row.get::<_, #sql_data_type>(#column_name)?);
diff --git a/crypto-macros/src/entity_derive_new/column_type.rs b/crypto-macros/src/entity_derive_new/column_type.rs
index 5cb3f06fee..cde35bf274 100644
--- a/crypto-macros/src/entity_derive_new/column_type.rs
+++ b/crypto-macros/src/entity_derive_new/column_type.rs
@@ -133,6 +133,9 @@ impl TryFrom<Type> for ColumnType {
 pub(super) trait EmitGetExpression {
     /// Emit an expression which wraps the input expression, appropriately parsing according to this column type.
     fn emit_get_expression(&self, input: TokenStream) -> TokenStream;
+
+    /// Emit an expression with the rust type which should be used in the rusqlite `get` expression
+    fn get_as_type(&self) -> TokenStream;
 }
 
 impl EmitGetExpression for IdColumnType {
@@ -142,17 +145,24 @@ impl EmitGetExpression for IdColumnType {
             Self::String => quote!(String::from_utf8(#input).map_err(|err| err.utf8_error())?),
         }
     }
+
+    fn get_as_type(&self) -> TokenStream {
+        quote!(Vec<u8>)
+    }
 }
 
 impl EmitGetExpression for ColumnType {
     fn emit_get_expression(&self, input: TokenStream) -> TokenStream {
         match self {
-            ColumnType::Bytes => input,
+            ColumnType::Bytes | ColumnType::OptionalBytes => input,
             ColumnType::String => quote!(String::from_utf8(#input).map_err(|err| err.utf8_error())?),
-            ColumnType::OptionalBytes => quote! {{
-                let data = #input;
-                (!data.is_empty()).then_some(data)
-            }},
+        }
+    }
+
+    fn get_as_type(&self) -> TokenStream {
+        match self {
+            ColumnType::Bytes | ColumnType::String => quote!(Vec<u8>),
+            ColumnType::OptionalBytes => quote!(Option<Vec<u8>>),
         }
     }
 }

From df37aa173cf3f4ef16304496fb1420ce01e4a594 Mon Sep 17 00:00:00 2001
From: Peter Goodspeed-Niklaus <peter.goodspeed-niklaus@wire.com>
Date: Tue, 23 Dec 2025 16:26:12 +0100
Subject: [PATCH 46/48] chore(keystore): don't panic if an entity with unknown
 id is used

---
 keystore/src/error.rs                         |  2 ++
 .../transaction/dynamic_dispatch/entity_id.rs | 14 +++++------
 keystore/src/transaction/mod.rs               | 23 +++++++++++--------
 3 files changed, 22 insertions(+), 17 deletions(-)

diff --git a/keystore/src/error.rs b/keystore/src/error.rs
index 8ff9477c1e..9ebad10122 100644
--- a/keystore/src/error.rs
+++ b/keystore/src/error.rs
@@ -154,6 +154,8 @@ pub enum CryptoKeystoreError {
     MigrationFailed(String),
     #[error("the provided bytes could not be interpreted as the primary key of {0}")]
     InvalidPrimaryKeyBytes(&'static str),
+    #[error("the entity {0} had an unknown collection name and could not be found")]
+    UnknownEntity(&'static str),
 }
 
 #[cfg(target_family = "wasm")]
diff --git a/keystore/src/transaction/dynamic_dispatch/entity_id.rs b/keystore/src/transaction/dynamic_dispatch/entity_id.rs
index a70f5c5531..1c484530f2 100644
--- a/keystore/src/transaction/dynamic_dispatch/entity_id.rs
+++ b/keystore/src/transaction/dynamic_dispatch/entity_id.rs
@@ -39,32 +39,30 @@ impl EntityId {
             .ok_or(CryptoKeystoreError::InvalidPrimaryKeyBytes(self.typ.collection_name()))
     }
 
-    fn from_key<E>(primary_key: Cow<'_, [u8]>) -> Self
+    fn from_key<E>(primary_key: Cow<'_, [u8]>) -> Option<Self>
     where
         E: Entity,
     {
-        // assumption: nobody outside this crate will ever implement `Entity` on a foreign type
-        let typ =
-            EntityType::from_collection_name(E::COLLECTION_NAME).expect("all entities have a valid collection name");
+        let typ = EntityType::from_collection_name(E::COLLECTION_NAME)?;
         let id = primary_key.into_owned();
-        Self { typ, id }
+        Some(Self { typ, id })
     }
 
-    pub(crate) fn from_entity<E>(entity: &E) -> Self
+    pub(crate) fn from_entity<E>(entity: &E) -> Option<Self>
     where
         E: Entity,
     {
         Self::from_key::<E>(entity.primary_key().bytes())
     }
 
-    pub(crate) fn from_primary_key<E>(primary_key: &E::PrimaryKey) -> Self
+    pub(crate) fn from_primary_key<E>(primary_key: &E::PrimaryKey) -> Option<Self>
     where
         E: Entity,
     {
         Self::from_key::<E>(primary_key.bytes())
     }
 
-    pub(crate) fn from_borrowed_primary_key<E>(primary_key: &E::BorrowedPrimaryKey) -> Self
+    pub(crate) fn from_borrowed_primary_key<E>(primary_key: &E::BorrowedPrimaryKey) -> Option<Self>
     where
         E: Entity + BorrowPrimaryKey,
     {
diff --git a/keystore/src/transaction/mod.rs b/keystore/src/transaction/mod.rs
index d1f5b2cad2..479c64b998 100644
--- a/keystore/src/transaction/mod.rs
+++ b/keystore/src/transaction/mod.rs
@@ -57,7 +57,7 @@ impl KeystoreTransaction {
     {
         let auto_generated_fields = entity.pre_save().await?;
 
-        let entity_id = EntityId::from_entity(&entity);
+        let entity_id = EntityId::from_entity(&entity).ok_or(CryptoKeystoreError::UnknownEntity(E::COLLECTION_NAME))?;
         {
             // start by adding the entity
             let mut cache_guard = self.cache.write().await;
@@ -103,7 +103,8 @@ impl KeystoreTransaction {
     where
         E: Entity + EntityDatabaseMutation<'a>,
     {
-        let entity_id = EntityId::from_primary_key::<E>(id);
+        let entity_id =
+            EntityId::from_primary_key::<E>(id).ok_or(CryptoKeystoreError::UnknownEntity(E::COLLECTION_NAME))?;
         self.remove_by_entity_id::<E>(entity_id).await
     }
 
@@ -115,7 +116,8 @@ impl KeystoreTransaction {
     where
         E: EntityDeleteBorrowed<'a> + BorrowPrimaryKey,
     {
-        let entity_id = EntityId::from_borrowed_primary_key::<E>(id);
+        let entity_id = EntityId::from_borrowed_primary_key::<E>(id)
+            .ok_or(CryptoKeystoreError::UnknownEntity(E::COLLECTION_NAME))?;
         self.remove_by_entity_id::<E>(entity_id).await
     }
 
@@ -161,9 +163,12 @@ impl KeystoreTransaction {
         drop(cache_guard);
 
         let mut deleted_set = self.deleted.write().await;
-        deleted_set.insert(EntityId::from_primary_key::<MlsPendingMessage>(
-            &MlsPendingMessagePrimaryKey::from_conversation_id(conversation_id),
-        ));
+        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"),
+        );
     }
 
     pub(crate) async fn find_pending_messages_by_conversation_id(
@@ -220,7 +225,7 @@ impl KeystoreTransaction {
     where
         E: Entity + Send + Sync,
     {
-        let entity_id = EntityId::from_primary_key::<E>(id);
+        let entity_id = EntityId::from_primary_key::<E>(id)?;
         self.get_by_entity_id(&entity_id).await
     }
 
@@ -232,7 +237,7 @@ impl KeystoreTransaction {
     where
         E: Entity + BorrowPrimaryKey + Send + Sync,
     {
-        let entity_id = EntityId::from_borrowed_primary_key::<E>(id);
+        let entity_id = EntityId::from_borrowed_primary_key::<E>(id)?;
         self.get_by_entity_id(&entity_id).await
     }
 
@@ -293,7 +298,7 @@ impl KeystoreTransaction {
             .chain(records_b)
             .unique_by(|e| e.primary_key().bytes().into_owned())
             .filter_map(|record| {
-                let id = EntityId::from_entity(record.as_ref());
+                let id = EntityId::from_entity(record.as_ref())?;
                 (!deleted_records.contains(&id)).then_some(record.into_owned())
             })
             .collect()

From 024a45486ff2903bf3d9851848b35d86319f8207 Mon Sep 17 00:00:00 2001
From: Peter Goodspeed-Niklaus <peter.goodspeed-niklaus@wire.com>
Date: Tue, 23 Dec 2025 16:38:14 +0100
Subject: [PATCH 47/48] chore(keystore): `ConsumerData` is a real entity

---
 keystore/src/transaction/dynamic_dispatch/entity_id.rs   | 3 ++-
 keystore/src/transaction/dynamic_dispatch/entity_type.rs | 9 ++++++---
 2 files changed, 8 insertions(+), 4 deletions(-)

diff --git a/keystore/src/transaction/dynamic_dispatch/entity_id.rs b/keystore/src/transaction/dynamic_dispatch/entity_id.rs
index 1c484530f2..b8525b37d0 100644
--- a/keystore/src/transaction/dynamic_dispatch/entity_id.rs
+++ b/keystore/src/transaction/dynamic_dispatch/entity_id.rs
@@ -9,7 +9,7 @@ use crate::{
     CryptoKeystoreError, CryptoKeystoreResult,
     connection::TransactionWrapper,
     entities::{
-        E2eiCrl, E2eiIntermediateCert, MlsPendingMessage, PersistedMlsGroup, PersistedMlsPendingGroup,
+        ConsumerData, E2eiCrl, E2eiIntermediateCert, MlsPendingMessage, PersistedMlsGroup, PersistedMlsPendingGroup,
         StoredBufferedCommit, StoredCredential, StoredE2eiEnrollment, StoredEncryptionKeyPair,
         StoredEpochEncryptionKeypair, StoredHpkePrivateKey, StoredKeypackage, StoredPskBundle,
     },
@@ -120,6 +120,7 @@ impl EntityId {
             EntityType::ProteusIdentity => ProteusIdentity::delete(tx, &self.primary_key::<ProteusIdentity>()?).await,
             #[cfg(feature = "proteus-keystore")]
             EntityType::ProteusPrekey => ProteusPrekey::delete(tx, &self.primary_key::<ProteusPrekey>()?).await,
+            EntityType::ConsumerData => ConsumerData::delete(tx, &self.primary_key::<ConsumerData>()?).await,
         }
     }
 }
diff --git a/keystore/src/transaction/dynamic_dispatch/entity_type.rs b/keystore/src/transaction/dynamic_dispatch/entity_type.rs
index af224af14d..8531f5535e 100644
--- a/keystore/src/transaction/dynamic_dispatch/entity_type.rs
+++ b/keystore/src/transaction/dynamic_dispatch/entity_type.rs
@@ -4,9 +4,9 @@ use crate::entities::E2eiRefreshToken;
 use crate::entities::{ProteusIdentity, ProteusPrekey, ProteusSession};
 use crate::{
     entities::{
-        E2eiAcmeCA, E2eiCrl, E2eiIntermediateCert, MlsPendingMessage, PersistedMlsGroup, PersistedMlsPendingGroup,
-        StoredBufferedCommit, StoredCredential, StoredE2eiEnrollment, StoredEncryptionKeyPair,
-        StoredEpochEncryptionKeypair, StoredHpkePrivateKey, StoredKeypackage, StoredPskBundle,
+        ConsumerData, E2eiAcmeCA, E2eiCrl, E2eiIntermediateCert, MlsPendingMessage, PersistedMlsGroup,
+        PersistedMlsPendingGroup, StoredBufferedCommit, StoredCredential, StoredE2eiEnrollment,
+        StoredEncryptionKeyPair, StoredEpochEncryptionKeypair, StoredHpkePrivateKey, StoredKeypackage, StoredPskBundle,
     },
     traits::EntityBase as _,
 };
@@ -35,6 +35,7 @@ pub(crate) enum EntityType {
     ProteusPrekey,
     #[cfg(feature = "proteus-keystore")]
     ProteusSession,
+    ConsumerData,
 }
 
 impl EntityType {
@@ -62,6 +63,7 @@ impl EntityType {
             ProteusPrekey::COLLECTION_NAME => Some(Self::ProteusPrekey),
             #[cfg(feature = "proteus-keystore")]
             ProteusSession::COLLECTION_NAME => Some(Self::ProteusSession),
+            ConsumerData::COLLECTION_NAME => Some(Self::ConsumerData),
             _ => None,
         }
     }
@@ -90,6 +92,7 @@ impl EntityType {
             #[cfg(feature = "proteus-keystore")]
             Self::ProteusSession => ProteusSession::COLLECTION_NAME,
             Self::HpkePrivateKey => StoredHpkePrivateKey::COLLECTION_NAME,
+            Self::ConsumerData => ConsumerData::COLLECTION_NAME,
         }
     }
 }

From fff17b2979bca9f7199b05e1616571a875719592 Mon Sep 17 00:00:00 2001
From: Peter Goodspeed-Niklaus <peter.goodspeed-niklaus@wire.com>
Date: Tue, 23 Dec 2025 17:35:28 +0100
Subject: [PATCH 48/48] chore(keystore): fix loading of generic unique entities

Turns out that we were loading from the wrong keys this whole time.
---
 keystore/src/traits/unique_entity.rs |  6 ++---
 keystore/tests/z_entities.rs         | 39 ++++++++++++++++++++++++++++
 2 files changed, 42 insertions(+), 3 deletions(-)

diff --git a/keystore/src/traits/unique_entity.rs b/keystore/src/traits/unique_entity.rs
index eda1a2f587..a0dc231f00 100644
--- a/keystore/src/traits/unique_entity.rs
+++ b/keystore/src/traits/unique_entity.rs
@@ -253,11 +253,11 @@ where
     async fn get(conn: &mut Self::ConnectionType, key: &Self::PrimaryKey) -> CryptoKeystoreResult<Option<Self>> {
         let conn = conn.conn().await;
         let mut statement = conn.prepare_cached(&format!(
-            "SELECT * FROM {collection_name} WHERE id = ?",
+            "SELECT content FROM {collection_name} WHERE id = ?",
             collection_name = Self::COLLECTION_NAME
         ))?;
         statement
-            .query_row([key], |row| Ok(Self::new(row.get(0)?)))
+            .query_row([key], |row| Ok(Self::new(row.get("content")?)))
             .optional()
             .map_err(Into::into)
     }
@@ -269,7 +269,7 @@ where
 
     /// Retrieve all entities of this type from the database.
     async fn load_all(conn: &mut Self::ConnectionType) -> CryptoKeystoreResult<Vec<Self>> {
-        load_all_helper::<Self, _>(conn, |row| Ok(Self::new(row.get(0)?))).await
+        load_all_helper::<Self, _>(conn, |row| Ok(Self::new(row.get("content")?))).await
     }
 }
 
diff --git a/keystore/tests/z_entities.rs b/keystore/tests/z_entities.rs
index 6814cca8fd..9e17aaaa76 100644
--- a/keystore/tests/z_entities.rs
+++ b/keystore/tests/z_entities.rs
@@ -227,6 +227,45 @@ mod tests {
         store.rollback_transaction().await.unwrap();
         store.new_transaction().await.unwrap();
     }
+
+    #[apply(all_storage_types)]
+    async fn can_save_and_load_consumer_data(context: KeystoreTestContext) {
+        use core_crypto_keystore::traits::FetchFromDatabase as _;
+
+        eprintln!("creating store");
+        let store = context.store();
+
+        eprintln!("checking consumer data before it exists");
+        assert!(!store.exists::<ConsumerData>().await.unwrap());
+        let consumer_data = store.get_unique::<ConsumerData>().await.unwrap();
+        assert!(consumer_data.is_none());
+
+        eprintln!("saving some consumer data");
+        const DATA: &[u8] = b"here is some arbitrary data";
+        store
+            .save(ConsumerData {
+                content: DATA.to_owned(),
+            })
+            .await
+            .unwrap();
+
+        // from transaction
+        eprintln!("checking retrieving consumer data from active transaction");
+        assert!(store.exists::<ConsumerData>().await.unwrap());
+        let consumer_data = store.get_unique::<ConsumerData>().await.unwrap().unwrap();
+        assert_eq!(consumer_data.content, DATA);
+
+        eprintln!("committing transaction");
+        store.commit_transaction().await.unwrap();
+        // don't forget to open a new (blank) transaction
+        store.new_transaction().await.unwrap();
+
+        // from storage (fallthrough)
+        eprintln!("checking retrieving consumer data from storage");
+        assert!(store.exists::<ConsumerData>().await.unwrap());
+        let consumer_data = store.get_unique::<ConsumerData>().await.unwrap().unwrap();
+        assert_eq!(consumer_data.content, DATA);
+    }
 }
 
 #[cfg(test)]