feat: implement set_client_data() [WPB-10919]

keystore/src/connection/platform/generic/migrations/V13__consumer_data.sql

@@ -0,0 +1,4 @@
+CREATE TABLE consumer_data (
+    id INTEGER PRIMARY KEY CHECK ( id = 0 ),
+    content BLOB
+);

keystore/src/connection/platform/wasm/migrations.rs

@@ -1,10 +1,10 @@
 use crate::connection::storage::{WasmEncryptedStorage, WasmStorageWrapper};
 use crate::connection::KeystoreDatabaseConnection;
 use crate::entities::{
-    E2eiAcmeCA, E2eiCrl, E2eiEnrollment, E2eiIntermediateCert, E2eiRefreshToken, Entity, EntityBase, MlsCredential,
-    MlsEncryptionKeyPair, MlsEpochEncryptionKeyPair, MlsHpkePrivateKey, MlsKeyPackage, MlsPendingMessage, MlsPskBundle,
-    MlsSignatureKeyPair, PersistedMlsGroup, PersistedMlsPendingGroup, ProteusIdentity, ProteusPrekey, ProteusSession,
-    UniqueEntity,
+    ConsumerData, E2eiAcmeCA, E2eiCrl, E2eiEnrollment, E2eiIntermediateCert, E2eiRefreshToken, Entity, EntityBase,
+    MlsCredential, MlsEncryptionKeyPair, MlsEpochEncryptionKeyPair, MlsHpkePrivateKey, MlsKeyPackage,
+    MlsPendingMessage, MlsPskBundle, MlsSignatureKeyPair, PersistedMlsGroup, PersistedMlsPendingGroup, ProteusIdentity,
+    ProteusPrekey, ProteusSession, UniqueEntity,
 };
 use crate::{CryptoKeystoreError, CryptoKeystoreResult};
 use idb::builder::{DatabaseBuilder, IndexBuilder, ObjectStoreBuilder};
@@ -33,11 +33,13 @@ const fn db_version_number(counter: u32) -> u32 {
 }
 
 const DB_VERSION_0: u32 = db_version_number(0);
+const DB_VERSION_1: u32 = db_version_number(1);
+const DB_VERSION_2: u32 = db_version_number(2);
 
 /// Open an existing idb database with the given name and key, and migrate it if needed.
 pub(crate) async fn open_and_migrate(name: &str, key: &str) -> CryptoKeystoreResult<Database> {
     /// Increment when adding a new migration.
-    const TARGET_VERSION: u32 = db_version_number(1);
+    const TARGET_VERSION: u32 = DB_VERSION_2;
     let factory = Factory::new()?;
 
     let open_existing = factory.open(name, None)?;
@@ -74,10 +76,28 @@ async fn do_migration_step(from: u32, name: &str, key: &str) -> CryptoKeystoreRe
         // The version that results from the latest migration must match TARGET_VERSION
         //      to ensure convergence of the while loop this is called from.
         0..=DB_VERSION_0 => migrate_to_version_1(name, key).await,
+        DB_VERSION_1 => migrate_to_version_2(name).await,
         _ => Err(CryptoKeystoreError::MigrationNotSupported(from)),
     }
 }
 
+/// Open IDB once with the new builder and close it, this will add the new object store.
+async fn migrate_to_version_2(name: &str) -> CryptoKeystoreResult<u32> {
+    let migrated_idb = get_builder_v2(name).build().await?;
+    migrated_idb.close();
+    Ok(DB_VERSION_2)
+}
+
+/// Add a new object store for the ConsumerData struct.
+fn get_builder_v2(name: &str) -> DatabaseBuilder {
+    let previous_builder = get_builder_v0(name);
+    previous_builder.version(DB_VERSION_2).add_object_store(
+        ObjectStoreBuilder::new(ConsumerData::COLLECTION_NAME)
+            .auto_increment(false)
+            .add_index(IndexBuilder::new("id".into(), KeyPath::new_single("id")).unique(true)),
+    )
+}
+
 /// With the current feature set of stable rust macros, we're not aware how to construct an
 /// identifier for each entity inside the macro.
 ///
@@ -137,13 +157,12 @@ macro_rules! migrate_entities_to_version_1 {
             }
 
             // The migration is complete and the version counter can be incremented.
-            const MIGRATING_TO: u32 = db_version_number(1);
             let factory = Factory::new()?;
-            let open_request = factory.open($name, Some(MIGRATING_TO))?;
+            let open_request = factory.open($name, Some(DB_VERSION_1))?;
             let idb = open_request.await?;
             idb.close();
 
-            Ok(MIGRATING_TO)
+            Ok(DB_VERSION_1)
         }
     };
 }

keystore/src/entities/general.rs

@@ -0,0 +1,22 @@
+/// Consumers of this library can use this to specify data to be persisted at the end of
+/// a transaction.
+#[derive(Debug, Clone, PartialEq, Eq)]
+#[cfg_attr(
+    any(target_family = "wasm", feature = "serde"),
+    derive(serde::Serialize, serde::Deserialize)
+)]
+pub struct ConsumerData {
+    pub content: Vec<u8>,
+}
+
+impl From<Vec<u8>> for ConsumerData {
+    fn from(content: Vec<u8>) -> Self {
+        Self { content }
+    }
+}
+
+impl From<ConsumerData> for Vec<u8> {
+    fn from(consumer_data: ConsumerData) -> Self {
+        consumer_data.content
+    }
+}

keystore/src/entities/mls.rs

@@ -227,6 +227,18 @@ where
         }
     }
 
+    async fn find_one(conn: &mut Self::ConnectionType) -> CryptoKeystoreResult<Option<Self>> {
+        match Self::find_unique(conn).await {
+            Ok(record) => Ok(Some(record)),
+            Err(CryptoKeystoreError::NotFound(_, _)) => Ok(None),
+            Err(err) => Err(err),
+        }
+    }
+
+    async fn count(conn: &mut Self::ConnectionType) -> CryptoKeystoreResult<usize> {
+        conn.storage().count(Self::COLLECTION_NAME).await
+    }
+
     async fn replace<'a>(&'a self, transaction: &TransactionWrapper<'a>) -> CryptoKeystoreResult<()> {
         transaction.save(self.clone()).await?;
         Ok(())
@@ -267,6 +279,22 @@ pub trait UniqueEntity: Entity<ConnectionType = crate::connection::KeystoreDatab
         }
     }
 
+    async fn find_one(conn: &mut Self::ConnectionType) -> CryptoKeystoreResult<Option<Self>> {
+        match Self::find_unique(conn).await {
+            Ok(record) => Ok(Some(record)),
+            Err(CryptoKeystoreError::NotFound(_, _)) => Ok(None),
+            Err(err) => Err(err),
+        }
+    }
+
+    async fn count(conn: &mut Self::ConnectionType) -> CryptoKeystoreResult<usize> {
+        Ok(
+            conn.query_row(&format!("SELECT COUNT(*) FROM {}", Self::COLLECTION_NAME), [], |r| {
+                r.get(0)
+            })?,
+        )
+    }
+
     fn content(&self) -> &[u8];
 
     async fn replace(&self, transaction: &TransactionWrapper<'_>) -> CryptoKeystoreResult<()> {

keystore/src/entities/mod.rs

@@ -14,7 +14,10 @@
 // You should have received a copy of the GNU General Public License
 // along with this program. If not, see http://www.gnu.org/licenses/.
 
+pub(crate) mod general;
 pub(crate) mod mls;
+
+pub use self::general::*;
 pub use self::mls::*;
 
 cfg_if::cfg_if! {

keystore/src/entities/platform/generic/general.rs

@@ -0,0 +1,49 @@
+use crate::{
+    connection::KeystoreDatabaseConnection,
+    entities::{ConsumerData, Entity, EntityBase, EntityFindParams, StringEntityId, UniqueEntity},
+    CryptoKeystoreResult, MissingKeyErrorKind,
+};
+
+impl Entity for ConsumerData {
+    fn id_raw(&self) -> &[u8] {
+        &[Self::ID as u8]
+    }
+}
+
+#[async_trait::async_trait]
+impl UniqueEntity for ConsumerData {
+    fn new(content: Vec<u8>) -> Self {
+        Self { content }
+    }
+
+    fn content(&self) -> &[u8] {
+        &self.content
+    }
+}
+
+#[async_trait::async_trait]
+impl EntityBase for ConsumerData {
+    type ConnectionType = KeystoreDatabaseConnection;
+    type AutoGeneratedFields = ();
+    const COLLECTION_NAME: &'static str = "consumer_data";
+
+    fn to_missing_key_err_kind() -> MissingKeyErrorKind {
+        MissingKeyErrorKind::ConsumerData
+    }
+
+    fn to_transaction_entity(self) -> crate::transaction::dynamic_dispatch::Entity {
+        crate::transaction::dynamic_dispatch::Entity::ConsumerData(self)
+    }
+
+    async fn find_all(conn: &mut Self::ConnectionType, params: EntityFindParams) -> CryptoKeystoreResult<Vec<Self>> {
+        <Self as UniqueEntity>::find_all(conn, params).await
+    }
+
+    async fn find_one(conn: &mut Self::ConnectionType, _id: &StringEntityId) -> CryptoKeystoreResult<Option<Self>> {
+        <Self as UniqueEntity>::find_one(conn).await
+    }
+
+    async fn count(conn: &mut Self::ConnectionType) -> CryptoKeystoreResult<usize> {
+        <Self as UniqueEntity>::count(conn).await
+    }
+}

keystore/src/entities/platform/generic/mod.rs

@@ -14,7 +14,9 @@
 // You should have received a copy of the GNU General Public License
 // along with this program. If not, see http://www.gnu.org/licenses/.
 
+mod general;
 mod mls;
+
 pub use self::mls::*;
 
 cfg_if::cfg_if! {

keystore/src/entities/platform/wasm/general.rs

@@ -0,0 +1,56 @@
+use crate::connection::DatabaseConnection;
+use crate::entities::Entity;
+use crate::{
+    connection::KeystoreDatabaseConnection,
+    entities::{ConsumerData, EntityBase, EntityFindParams, StringEntityId, UniqueEntity},
+    CryptoKeystoreResult, MissingKeyErrorKind,
+};
+
+#[cfg_attr(target_family = "wasm", async_trait::async_trait(?Send))]
+#[cfg_attr(not(target_family = "wasm"), async_trait::async_trait)]
+impl EntityBase for ConsumerData {
+    type ConnectionType = KeystoreDatabaseConnection;
+    type AutoGeneratedFields = ();
+    const COLLECTION_NAME: &'static str = "consumer_data";
+
+    fn to_missing_key_err_kind() -> MissingKeyErrorKind {
+        MissingKeyErrorKind::ConsumerData
+    }
+
+    fn to_transaction_entity(self) -> crate::transaction::dynamic_dispatch::Entity {
+        crate::transaction::dynamic_dispatch::Entity::ConsumerData(self)
+    }
+
+    async fn find_all(conn: &mut Self::ConnectionType, params: EntityFindParams) -> CryptoKeystoreResult<Vec<Self>> {
+        <Self as UniqueEntity>::find_all(conn, params).await
+    }
+
+    async fn find_one(conn: &mut Self::ConnectionType, _id: &StringEntityId) -> CryptoKeystoreResult<Option<Self>> {
+        <Self as UniqueEntity>::find_one(conn).await
+    }
+
+    async fn count(conn: &mut Self::ConnectionType) -> CryptoKeystoreResult<usize> {
+        <Self as UniqueEntity>::count(conn).await
+    }
+}
+
+impl Entity for ConsumerData {
+    fn id_raw(&self) -> &[u8] {
+        &Self::ID
+    }
+
+    fn encrypt(&mut self, cipher: &aes_gcm::Aes256Gcm) -> CryptoKeystoreResult<()> {
+        self.content = self.encrypt_data(cipher, self.content.as_slice())?;
+        Self::ConnectionType::check_buffer_size(self.content.len())?;
+        Ok(())
+    }
+
+    fn decrypt(&mut self, cipher: &aes_gcm::Aes256Gcm) -> CryptoKeystoreResult<()> {
+        self.content = self.decrypt_data(cipher, self.content.as_slice())?;
+        Ok(())
+    }
+}
+
+#[cfg_attr(target_family = "wasm", async_trait::async_trait(?Send))]
+#[cfg_attr(not(target_family = "wasm"), async_trait::async_trait)]
+impl UniqueEntity for ConsumerData {}

keystore/src/entities/platform/wasm/mod.rs

@@ -14,7 +14,9 @@
 // You should have received a copy of the GNU General Public License
 // along with this program. If not, see http://www.gnu.org/licenses/.
 
+mod general;
 mod mls;
+
 pub use self::mls::*;
 
 cfg_if::cfg_if! {

keystore/src/error.rs

@@ -17,6 +17,8 @@
 /// Error to represent when a key is not present in the KeyStore
 #[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
 pub enum MissingKeyErrorKind {
+    #[error("Consumer Data")]
+    ConsumerData,
     #[error("MLS KeyPackageBundle")]
     MlsKeyPackageBundle,
     #[error("MLS SignatureKeyPair")]

keystore/src/transaction/dynamic_dispatch.rs

@@ -3,17 +3,18 @@
 
 use crate::connection::TransactionWrapper;
 use crate::entities::{
-    E2eiAcmeCA, E2eiCrl, E2eiEnrollment, E2eiIntermediateCert, E2eiRefreshToken, EntityBase, EntityTransactionExt,
-    MlsCredential, MlsEncryptionKeyPair, MlsEpochEncryptionKeyPair, MlsHpkePrivateKey, MlsKeyPackage,
-    MlsPendingMessage, MlsPskBundle, MlsSignatureKeyPair, PersistedMlsGroup, PersistedMlsPendingGroup, StringEntityId,
-    UniqueEntity,
+    ConsumerData, E2eiAcmeCA, E2eiCrl, E2eiEnrollment, E2eiIntermediateCert, E2eiRefreshToken, EntityBase,
+    EntityTransactionExt, MlsCredential, MlsEncryptionKeyPair, MlsEpochEncryptionKeyPair, MlsHpkePrivateKey,
+    MlsKeyPackage, MlsPendingMessage, MlsPskBundle, MlsSignatureKeyPair, PersistedMlsGroup, PersistedMlsPendingGroup,
+    StringEntityId, UniqueEntity,
 };
 #[cfg(feature = "proteus-keystore")]
 use crate::entities::{ProteusIdentity, ProteusPrekey, ProteusSession};
 use crate::{CryptoKeystoreError, CryptoKeystoreResult};
 
 #[derive(Debug)]
 pub enum Entity {
+    ConsumerData(ConsumerData),
     SignatureKeyPair(MlsSignatureKeyPair),
     HpkePrivateKey(MlsHpkePrivateKey),
     KeyPackage(MlsKeyPackage),
@@ -146,6 +147,7 @@ impl EntityId {
 
 pub async fn execute_save(tx: &TransactionWrapper<'_>, entity: &Entity) -> CryptoKeystoreResult<()> {
     match entity {
+        Entity::ConsumerData(consumer_data) => consumer_data.replace(tx).await,
         Entity::SignatureKeyPair(mls_signature_key_pair) => mls_signature_key_pair.save(tx).await,
         Entity::HpkePrivateKey(mls_hpke_private_key) => mls_hpke_private_key.save(tx).await,
         Entity::KeyPackage(mls_key_package) => mls_key_package.save(tx).await,

keystore/src/transaction/mod.rs

@@ -3,7 +3,7 @@ pub mod dynamic_dispatch;
 use crate::entities::mls::*;
 #[cfg(feature = "proteus-keystore")]
 use crate::entities::proteus::*;
-use crate::entities::{EntityBase, EntityFindParams, EntityTransactionExt, UniqueEntity};
+use crate::entities::{ConsumerData, EntityBase, EntityFindParams, EntityTransactionExt, UniqueEntity};
 use crate::transaction::dynamic_dispatch::EntityId;
 use crate::{
     connection::{Connection, DatabaseConnection, FetchFromDatabase, KeystoreDatabaseConnection},
@@ -341,12 +341,13 @@ impl KeystoreTransaction {
                 (identifier_12, E2eiRefreshToken),
                 (identifier_13, E2eiAcmeCA),
                 (identifier_14, E2eiIntermediateCert),
-                (identifier_15, E2eiCrl)
+                (identifier_15, E2eiCrl),
+                (identifier_16, ConsumerData)
             ],
             proteus_types: [
-                (identifier_16, ProteusPrekey),
-                (identifier_17, ProteusIdentity),
-                (identifier_18, ProteusSession)
+                (identifier_17, ProteusPrekey),
+                (identifier_18, ProteusIdentity),
+                (identifier_19, ProteusSession)
             ]
         );