Add deterministic batch functionality (#95)

kotlin/src/test/kotlin/com/ironcorelabs/ironcore_alloy/IroncoreAlloyTest.kt

@@ -179,7 +179,6 @@ class IroncoreAlloyTest {
         }
     }
 
-
     @Test
     fun sdkStandardAttachedRoundtrip() {
         val plaintextDocument = "My data".toByteArray()
@@ -243,6 +242,28 @@ class IroncoreAlloyTest {
         }
     }
 
+    @Test
+    fun sdkBatchRoundtripDeterministic() {
+        val plaintextInput = "My data".toByteArray()
+        val field = PlaintextField(plaintextInput, "", "")
+        val badField = PlaintextField("My data".toByteArray(), "bad_path", "bad_path")
+        val metadata = AlloyMetadata.newSimple("tenant")
+        val plaintextFields = mapOf("doc" to field, "badDoc" to badField)
+        runBlocking {
+            val encrypted = sdk.deterministic().encryptBatch(plaintextFields, metadata)
+            assertEquals(encrypted.successes.size, 1)
+            assertEquals(encrypted.failures.size, 1)
+            assertEquals(
+                    encrypted.failures.get("badDoc")?.message,
+                    "msg=Provided secret path `bad_path` does not exist in the deterministic configuration."
+            )
+            val decrypted = sdk.deterministic().decryptBatch(encrypted.successes, metadata)
+            assertEquals(decrypted.successes.size, 1)
+            assertEquals(decrypted.failures.size, 0)
+            assertContentEquals(decrypted.successes.get("doc")?.plaintextField, plaintextInput)
+        }
+    }
+
     @Test
     fun sdkStandardDecryptWrongType() {
         val err =

python/ironcore-alloy/tests/test_ironcore_alloy.py

@@ -198,6 +198,35 @@ async def test_encrypt_with_existing_edek(self):
         decrypted = await self.sdk.standard().decrypt(encrypted2, metadata)
         assert decrypted == document2
 
+    @pytest.mark.asyncio
+    async def test_deterministic_batch_roundtrip(self):
+        plaintext_input = b"foobar"
+        field = PlaintextField(
+            plaintext_field=plaintext_input,
+            secret_path="",
+            derivation_path="",
+        )
+        bad_field = PlaintextField(
+            plaintext_field=plaintext_input,
+            secret_path="bad_path",
+            derivation_path="bad_path",
+        )
+        fields = {"doc": field, "bad_doc": bad_field}
+        metadata = AlloyMetadata.new_simple("tenant")
+        encrypted = await self.sdk.deterministic().encrypt_batch(fields, metadata)
+        assert len(encrypted.successes) == 1
+        assert len(encrypted.failures) == 1
+        assert (
+            encrypted.failures["bad_doc"].msg  # type: ignore
+            == "Provided secret path `bad_path` does not exist in the deterministic configuration."
+        )
+        decrypted = await self.sdk.deterministic().decrypt_batch(
+            encrypted.successes, metadata
+        )
+        assert len(decrypted.successes) == 1
+        assert len(decrypted.failures) == 0
+        assert decrypted.successes["doc"].plaintext_field == plaintext_input
+
     @pytest.mark.asyncio
     async def test_decrypt_deterministic_metadata(self):
         field = EncryptedField(
@@ -299,7 +328,7 @@ async def test_rotate_deterministic_failures(self):
         assert len(rotated.failures) == 1
         assert (
             "Provided secret path `wrong_path` does not exist"
-            in rotated.failures["doc"].msg
+            in rotated.failures["doc"].msg  # type: ignore
         )
 
     @pytest.mark.skip(reason="need seeded client")

src/deterministic.rs

@@ -1,15 +1,12 @@
 use crate::{
-    errors::AlloyError,
-    util::{self, BatchResult},
-    AlloyMetadata, DerivationPath, EncryptedBytes, FieldId, PlaintextBytes, Secret, SecretPath,
-    TenantId,
+    create_batch_result_struct, errors::AlloyError, util, AlloyMetadata, DerivationPath,
+    EncryptedBytes, FieldId, PlaintextBytes, Secret, SecretPath, TenantId,
 };
 use aes_gcm::KeyInit;
 use aes_siv::siv::Aes256Siv;
 use bytes::Bytes;
 use ironcore_documents::v5::key_id_header::KeyIdHeader;
 use std::collections::HashMap;
-use uniffi::custom_newtype;
 
 #[derive(Debug, Clone, uniffi::Record)]
 pub struct EncryptedField {
@@ -27,26 +24,13 @@ pub struct PlaintextField {
 pub type PlaintextFields = HashMap<FieldId, PlaintextField>;
 pub type EncryptedFields = HashMap<FieldId, EncryptedField>;
 pub type GenerateQueryResult = HashMap<FieldId, Vec<EncryptedField>>;
-
-#[derive(Debug, Clone, uniffi::Record)]
-pub struct DeterministicRotateResult {
-    pub successes: HashMap<FieldId, EncryptedField>,
-    pub failures: HashMap<FieldId, AlloyError>,
-}
-
-impl From<BatchResult<EncryptedField>> for DeterministicRotateResult {
-    fn from(value: BatchResult<EncryptedField>) -> Self {
-        Self {
-            successes: value.successes,
-            failures: value.failures,
-        }
-    }
-}
+create_batch_result_struct!(DeterministicRotateResult, EncryptedField, FieldId);
+create_batch_result_struct!(DeterministicEncryptBatchResult, EncryptedField, FieldId);
+create_batch_result_struct!(DeterministicDecryptBatchResult, PlaintextField, FieldId);
 
 /// Key used for deterministic operations.
 #[derive(Debug, Clone)]
-pub struct DeterministicEncryptionKey(pub Vec<u8>);
-custom_newtype!(DeterministicEncryptionKey, Vec<u8>);
+pub(crate) struct DeterministicEncryptionKey(pub Vec<u8>);
 
 impl DeterministicEncryptionKey {
     /// A way to generate a key from the secret, tenant_id and derivation_path. This is done in the context of
@@ -74,12 +58,29 @@ pub trait DeterministicFieldOps {
         plaintext_field: PlaintextField,
         metadata: &AlloyMetadata,
     ) -> Result<EncryptedField, AlloyError>;
+    /// Deterministically encrypt the provided fields with the provided metadata.
+    /// Because the fields are encrypted deterministically with each call, the result will be the same for repeated calls.
+    /// This allows for exact matches and indexing of the encrypted field, but comes with some security considerations.
+    /// If you don't need to support these use cases, we recommend using `standard` encryption instead.
+    async fn encrypt_batch(
+        &self,
+        fields: PlaintextFields,
+        metadata: &AlloyMetadata,
+    ) -> Result<DeterministicEncryptBatchResult, AlloyError>;
     /// Decrypt a field that was deterministically encrypted with the provided metadata.
     async fn decrypt(
         &self,
         encrypted_field: EncryptedField,
         metadata: &AlloyMetadata,
     ) -> Result<PlaintextField, AlloyError>;
+    /// Decrypt each of the fields that were deterministically encrypted with the provided metadata.
+    /// Note that because the metadata is shared between the fields, they all must correspond to the
+    /// same tenant ID.
+    async fn decrypt_batch(
+        &self,
+        encrypted_fields: EncryptedFields,
+        metadata: &AlloyMetadata,
+    ) -> Result<DeterministicDecryptBatchResult, AlloyError>;
     /// Encrypt each plaintext field with any Current and InRotation keys for the provided secret path.
     /// The resulting encrypted fields should be used in tandem when querying the data store.
     async fn generate_query_field_values(

src/errors.rs

@@ -33,6 +33,7 @@ pub enum AlloyError {
         msg: String,
     },
 }
+
 impl std::fmt::Display for AlloyError {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         match self {

src/saas_shield/deterministic.rs

@@ -4,7 +4,8 @@ use super::{
 };
 
 use crate::deterministic::{
-    decrypt_internal, encrypt_internal, DeterministicEncryptionKey, DeterministicFieldOps,
+    decrypt_internal, encrypt_internal, DeterministicDecryptBatchResult,
+    DeterministicEncryptBatchResult, DeterministicEncryptionKey, DeterministicFieldOps,
     DeterministicRotateResult, EncryptedField, EncryptedFields, GenerateQueryResult,
     PlaintextField, PlaintextFields,
 };
@@ -76,14 +77,49 @@ impl DeterministicFieldOps for SaasShieldDeterministicClient {
         )
     }
 
+    /// Deterministically encrypt the provided fields with the provided metadata.
+    /// Because the fields are encrypted deterministically with each call, the result will be the same for repeated calls.
+    /// This allows for exact matches and indexing of the encrypted field, but comes with some security considerations.
+    /// If you don't need to support these use cases, we recommend using `standard` encryption instead.
+    async fn encrypt_batch(
+        &self,
+        plaintext_fields: PlaintextFields,
+        metadata: &AlloyMetadata,
+    ) -> Result<DeterministicEncryptBatchResult, AlloyError> {
+        let paths = plaintext_fields
+            .values()
+            .map(|field| (field.secret_path.clone(), field.derivation_path.clone()))
+            .collect_vec();
+        let all_keys = derive_keys_many_paths(
+            &self.tenant_security_client,
+            metadata,
+            paths,
+            SecretType::Deterministic,
+        )
+        .await?;
+        let encrypt_field = |plaintext_field: PlaintextField| {
+            let new_current_key = all_keys.get_key_for_path(
+                &plaintext_field.secret_path,
+                &plaintext_field.derivation_path,
+                DeriveKeyChoice::Current,
+            )?;
+            let key_id_header = Self::create_key_id_header(new_current_key.tenant_secret_id.0);
+            encrypt_internal(
+                DeterministicEncryptionKey(new_current_key.derived_key.0.clone()),
+                key_id_header,
+                plaintext_field,
+            )
+        };
+        Ok(collection_to_batch_result(plaintext_fields, encrypt_field).into())
+    }
+
     /// Decrypt a field that was deterministically encrypted with the provided metadata.
     async fn decrypt(
         &self,
         encrypted_field: EncryptedField,
         metadata: &AlloyMetadata,
     ) -> Result<PlaintextField, AlloyError> {
-        let (key_id, ciphertext) =
-            Self::decompose_key_id_header(encrypted_field.encrypted_field.clone())?;
+        let (key_id, ciphertext) = Self::decompose_key_id_header(encrypted_field.encrypted_field)?;
         let paths = [(
             encrypted_field.secret_path.clone(),
             [encrypted_field.derivation_path.clone()].into(),
@@ -118,6 +154,43 @@ impl DeterministicFieldOps for SaasShieldDeterministicClient {
         }
     }
 
+    /// Decrypt each of the fields that were deterministically encrypted with the provided metadata.
+    /// Note that because the metadata is shared between the fields, they all must correspond to the
+    /// same tenant ID.
+    async fn decrypt_batch(
+        &self,
+        encrypted_fields: EncryptedFields,
+        metadata: &AlloyMetadata,
+    ) -> Result<DeterministicDecryptBatchResult, AlloyError> {
+        let paths = encrypted_fields
+            .values()
+            .map(|field| (field.secret_path.clone(), field.derivation_path.clone()))
+            .collect_vec();
+        let all_keys = derive_keys_many_paths(
+            &self.tenant_security_client,
+            metadata,
+            paths,
+            SecretType::Deterministic,
+        )
+        .await?;
+        let decrypt_field = |encrypted_field: EncryptedField| {
+            let (original_key_id, ciphertext) =
+                Self::decompose_key_id_header(encrypted_field.encrypted_field)?;
+            let original_key = all_keys.get_key_for_path(
+                &encrypted_field.secret_path,
+                &encrypted_field.derivation_path,
+                DeriveKeyChoice::Specific(original_key_id),
+            )?;
+            decrypt_internal(
+                DeterministicEncryptionKey(original_key.derived_key.0.clone()),
+                ciphertext,
+                encrypted_field.secret_path.clone(),
+                encrypted_field.derivation_path.clone(),
+            )
+        };
+        Ok(collection_to_batch_result(encrypted_fields, decrypt_field).into())
+    }
+
     /// Encrypt each plaintext field with any Current and InRotation keys for the provided secret path.
     /// The resulting encrypted fields should be used in tandem when querying the data store.
     async fn generate_query_field_values(

src/standalone/deterministic.rs

@@ -1,6 +1,7 @@
 use super::config::RotatableSecret;
 use crate::deterministic::{
-    decrypt_internal, encrypt_internal, DeterministicEncryptionKey, DeterministicFieldOps,
+    decrypt_internal, encrypt_internal, DeterministicDecryptBatchResult,
+    DeterministicEncryptBatchResult, DeterministicEncryptionKey, DeterministicFieldOps,
     DeterministicRotateResult, EncryptedField, EncryptedFields, GenerateQueryResult,
     PlaintextField, PlaintextFields,
 };
@@ -63,8 +64,7 @@ impl StandaloneDeterministicClient {
         encrypted_field: EncryptedField,
         tenant_id: &TenantId,
     ) -> Result<PlaintextField, AlloyError> {
-        let (key_id, ciphertext) =
-            Self::decompose_key_id_header(encrypted_field.encrypted_field.clone())?;
+        let (key_id, ciphertext) = Self::decompose_key_id_header(encrypted_field.encrypted_field)?;
         let secret = self
             .config
             .get(&encrypted_field.secret_path)
@@ -121,6 +121,21 @@ impl DeterministicFieldOps for StandaloneDeterministicClient {
         self.encrypt_sync(plaintext_field, &metadata.tenant_id)
     }
 
+    /// Deterministically encrypt the provided fields with the provided metadata.
+    /// Because the fields are encrypted deterministically with each call, the result will be the same for repeated calls.
+    /// This allows for exact matches and indexing of the encrypted field, but comes with some security considerations.
+    /// If you don't need to support these use cases, we recommend using `standard` encryption instead.
+    async fn encrypt_batch(
+        &self,
+        plaintext_fields: PlaintextFields,
+        metadata: &AlloyMetadata,
+    ) -> Result<DeterministicEncryptBatchResult, AlloyError> {
+        let encrypt_field = |plaintext_field: PlaintextField| {
+            self.encrypt_sync(plaintext_field, &metadata.tenant_id)
+        };
+        Ok(collection_to_batch_result(plaintext_fields, encrypt_field).into())
+    }
+
     /// Decrypt a field that was deterministically encrypted with the provided metadata.
     async fn decrypt(
         &self,
@@ -130,6 +145,20 @@ impl DeterministicFieldOps for StandaloneDeterministicClient {
         self.decrypt_sync(encrypted_field, &metadata.tenant_id)
     }
 
+    /// Decrypt each of the fields that were deterministically encrypted with the provided metadata.
+    /// Note that because the metadata is shared between the fields, they all must correspond to the
+    /// same tenant ID.
+    async fn decrypt_batch(
+        &self,
+        encrypted_fields: EncryptedFields,
+        metadata: &AlloyMetadata,
+    ) -> Result<DeterministicDecryptBatchResult, AlloyError> {
+        let decrypt_field = |encrypted_field: EncryptedField| {
+            self.decrypt_sync(encrypted_field, &metadata.tenant_id)
+        };
+        Ok(collection_to_batch_result(encrypted_fields, decrypt_field).into())
+    }
+
     /// Encrypt each plaintext field with any Current and InRotation keys for the provided secret path.
     /// The resulting encrypted fields should be used in tandem when querying the data store.
     async fn generate_query_field_values(

src/standalone/standard_attached.rs

@@ -21,7 +21,7 @@ impl StandaloneAttachedStandardClient {
     }
 }
 
-#[uniffi::export(async_runtime = "tokio")]
+#[uniffi::export]
 impl StandardAttachedDocumentOps for StandaloneAttachedStandardClient {
     /// Encrypt a field with the provided metadata.
     /// A DEK (document encryption key) will be generated and encrypted using a derived key.

src/standard.rs

@@ -1,7 +1,5 @@
 use crate::{
-    alloy_client_trait::AlloyClient,
-    errors::AlloyError,
-    util::{get_rng, BatchResult},
+    alloy_client_trait::AlloyClient, create_batch_result_struct, errors::AlloyError, util::get_rng,
     AlloyMetadata, EncryptedBytes, FieldId, PlaintextBytes, TenantId,
 };
 use ironcore_documents::{
@@ -72,20 +70,7 @@ pub struct EncryptedDocument {
     pub document: HashMap<FieldId, EncryptedBytes>,
 }
 
-#[derive(Debug, Clone, uniffi::Record)]
-pub struct RekeyEdeksBatchResult {
-    pub successes: HashMap<String, EdekWithKeyIdHeader>,
-    pub failures: HashMap<String, AlloyError>,
-}
-
-impl From<BatchResult<EdekWithKeyIdHeader>> for RekeyEdeksBatchResult {
-    fn from(value: BatchResult<EdekWithKeyIdHeader>) -> Self {
-        Self {
-            successes: value.successes,
-            failures: value.failures,
-        }
-    }
-}
+create_batch_result_struct!(RekeyEdeksBatchResult, EdekWithKeyIdHeader, FieldId);
 
 /// API for encrypting and decrypting documents using our standard encryption. This class of encryption is the most
 /// broadly useful and secure. If you don't have a need to match on or preserve the distance properties of the