feat: implement set_client_data() [WPB-10919]

crypto-ffi/bindings/js/CoreCryptoContext.ts

@@ -52,6 +52,23 @@ export default class CoreCryptoContext {
         return new CoreCryptoContext(ctx);
     }
 
+    /**
+     * Set arbitrary data to be retrieved by {@link getData}.
+     * This is meant to be used as a check point at the end of a transaction.
+     * The data should be limited to a reasonable size.
+     */
+    async setData(data: Uint8Array): Promise<void> {
+        return await CoreCryptoError.asyncMapErr(this.#ctx.set_data(data));
+    }
+
+    /**
+     * Get data if it has previously been set by {@link setData}, or `undefined` otherwise.
+     * This is meant to be used as a check point at the end of a transaction.
+     */
+    async getData(): Promise<Uint8Array | undefined> {
+        return await CoreCryptoError.asyncMapErr(this.#ctx.get_data());
+    }
+
     /**
      * Use this after {@link CoreCrypto.deferredInit} when you have a clientId. It initializes MLS.
      *

crypto-ffi/bindings/js/test/CoreCrypto.test.js

@@ -225,6 +225,54 @@ test("can use groupInfo enums", async () => {
     await ctx.close();
 });
 
+test("Setting data persists to DB", async () => {
+    const [ctx, page] = await initBrowser();
+
+    const [firstResult, expectedSecondResult, secondResult] =
+        await page.evaluate(async () => {
+            const { CoreCrypto, Ciphersuite } = await import("./corecrypto.js");
+            const ciphersuite =
+                Ciphersuite.MLS_128_DHKEMX25519_AES128GCM_SHA256_Ed25519;
+
+            const client2Config = {
+                databaseName: "test",
+                key: "test",
+                ciphersuites: [ciphersuite],
+                clientId: "test",
+            };
+
+            const cc = await CoreCrypto.init(client2Config);
+
+            const text = "my message processing checkpoint";
+            const encoder = new TextEncoder();
+            const expectedSecondResult = encoder.encode(text);
+
+            let firstResult;
+            await cc.transaction(async (ctx) => {
+                firstResult = await ctx.getData();
+                await ctx.setData(expectedSecondResult);
+            });
+
+            let secondResult;
+            await cc.transaction(async (ctx) => {
+                secondResult = await ctx.getData();
+            });
+
+            // To be sure we're not obscuring the case in which firstResult would be null, as when it gets
+            // passed out of this closure, undefined becomes null.
+            firstResult = firstResult === null ? "null" : firstResult;
+
+            return [firstResult, expectedSecondResult, secondResult];
+        });
+
+    // Undefined becomes null.
+    expect(firstResult).toBe(null);
+    expect(secondResult).toEqual(expectedSecondResult);
+
+    await page.close();
+    await ctx.close();
+});
+
 test("Using invalid context throws error", async () => {
     const [ctx, page] = await initBrowser();
 

crypto-ffi/bindings/jvm/src/main/kotlin/com/wire/crypto/client/CoreCryptoContext.kt

@@ -39,6 +39,23 @@ class CoreCryptoContext(private val cc: CoreCryptoContext) {
         )
     }
 
+    /**
+     * Set arbitrary data to be retrieved by [getData].
+     * This is meant to be used as a check point at the end of a transaction.
+     * The data should be limited to a reasonable size.
+     */
+    suspend fun setData(data: ByteArray) {
+        cc.setData(data)
+    }
+
+    /**
+    * Get the data that has previously been set by [setData], or null if no data has been set.
+    * This is meant to be used as a check point at the end of a transaction.
+    */
+    suspend fun getData(): ByteArray? {
+        return cc.getData()
+    }
+
     /**
      * This is your entrypoint to initialize [com.wire.crypto.client.MLSClient] with a Basic Credential
      */

crypto-ffi/bindings/jvm/src/test/kotlin/com/wire/crypto/client/MLSTest.kt

@@ -40,6 +40,23 @@ class MLSTest {
         internal val carolId = "carol"
     }
 
+    @Test
+    fun set_client_data_persists() = runTest {
+        val cc = initCc()
+
+        val data = "my message processing checkpoint".toByteArray()
+
+        cc.transaction { ctx ->
+            assertThat(ctx.getData()).isNull()
+            ctx.setData(data)
+        }
+
+        cc.transaction { ctx ->
+            assertThat(ctx.getData()).isEqualTo(data)
+        }
+
+    }
+
     @Test
     fun externally_generated_ClientId_should_init_the_MLS_client() = runTest {
         val (alice, handle) = initCc().externallyGeneratedMlsClient()

crypto-ffi/src/generic/context/mod.rs

@@ -87,6 +87,16 @@ impl CoreCrypto {
 
 #[uniffi::export]
 impl CoreCryptoContext {
+    /// See [core_crypto::context::CentralContext::set_data].
+    pub async fn set_data(&self, data: Vec<u8>) -> CoreCryptoResult<()> {
+        self.context.set_data(data).await.map_err(Into::into)
+    }
+
+    /// See [core_crypto::context::CentralContext::get_data].
+    pub async fn get_data(&self) -> CoreCryptoResult<Option<Vec<u8>>> {
+        self.context.get_data().await.map_err(Into::into)
+    }
+
     /// See [core_crypto::context::CentralContext::mls_init]
     pub async fn mls_init(
         &self,

crypto-ffi/src/wasm/context/mod.rs

@@ -77,6 +77,39 @@ impl CoreCrypto {
 
 #[wasm_bindgen]
 impl CoreCryptoContext {
+    /// Returns: [`WasmCryptoResult<()>`]
+    ///
+    /// see [core_crypto::context::CentralContext::set_data]
+    pub fn set_data(&self, data: Box<[u8]>) -> Promise {
+        let context = self.inner.clone();
+        future_to_promise(
+            async move {
+                context.set_data(data.into()).await.map_err(CoreCryptoError::from)?;
+                WasmCryptoResult::Ok(JsValue::UNDEFINED)
+            }
+            .err_into(),
+        )
+    }
+
+    /// Returns: [`WasmCryptoResult<Option<js_sys::Uint8Array>>`]
+    ///
+    /// see [core_crypto::context::CentralContext::get_data]
+    pub fn get_data(&self) -> Promise {
+        let context = self.inner.clone();
+        future_to_promise(
+            async move {
+                let data = context.get_data().await.map_err(CoreCryptoError::from)?;
+                let result = if let Some(data) = data {
+                    JsValue::from(js_sys::Uint8Array::from(data.as_slice()))
+                } else {
+                    JsValue::UNDEFINED
+                };
+                WasmCryptoResult::Ok(result)
+            }
+            .err_into(),
+        )
+    }
+
     /// see [core_crypto::mls::context::CentralContext::mls_init]
     pub fn mls_init(&self, client_id: FfiClientId, ciphersuites: Box<[u16]>, nb_key_package: Option<u32>) -> Promise {
         let context = self.inner.clone();

crypto/src/context.rs

@@ -1,10 +1,6 @@
 //! This module contains the primitives to enable transactional support on a higher level within the
 //! [MlsCentral]. All mutating operations need to be done through a [CentralContext].
 
-use async_lock::{Mutex, RwLock, RwLockReadGuardArc, RwLockWriteGuardArc};
-use mls_crypto_provider::{CryptoKeystore, MlsCryptoProvider};
-use std::{ops::Deref, sync::Arc};
-
 use crate::mls::MlsCentral;
 #[cfg(feature = "proteus")]
 use crate::proteus::ProteusCentral;
@@ -13,6 +9,12 @@ use crate::{
     prelude::{Client, MlsConversation},
     CoreCrypto, CoreCryptoCallbacks, CryptoError, CryptoResult,
 };
+use async_lock::{Mutex, RwLock, RwLockReadGuardArc, RwLockWriteGuardArc};
+use core_crypto_keystore::connection::FetchFromDatabase;
+use core_crypto_keystore::entities::ConsumerData;
+use core_crypto_keystore::CryptoKeystoreError;
+use mls_crypto_provider::{CryptoKeystore, MlsCryptoProvider};
+use std::{ops::Deref, sync::Arc};
 
 /// This struct provides transactional support for Core Crypto.
 ///
@@ -164,4 +166,22 @@ impl CentralContext {
         *guard = ContextState::Invalid;
         rollback_result.map_err(Into::into)
     }
+
+    /// Set arbitrary data to be retrieved by [CentralContext::get_data].
+    /// This is meant to be used as a check point at the end of a transaction.
+    /// The data should be limited to a reasonable size.
+    pub async fn set_data(&self, data: Vec<u8>) -> CryptoResult<()> {
+        self.keystore().await?.save(ConsumerData::from(data)).await?;
+        Ok(())
+    }
+
+    /// Get the data that has previously been set by [CentralContext::set_data].
+    /// This is meant to be used as a check point at the end of a transaction.
+    pub async fn get_data(&self) -> CryptoResult<Option<Vec<u8>>> {
+        match self.keystore().await?.find_unique::<ConsumerData>().await {
+            Ok(data) => Ok(Some(data.into())),
+            Err(CryptoKeystoreError::NotFound(..)) => Ok(None),
+            Err(err) => Err(err.into()),
+        }
+    }
 }

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! {