Improve performance, fix bug with standard EDEK parsing (#114)

Cargo.toml

@@ -24,14 +24,15 @@ bytes = { version = "1.4.0", features = ["serde"] }
 convert_case = "0.6.0"
 futures = "0.3.29"
 hmac = { version = "0.12.1", features = ["std"] }
-ironcore-documents = "0.1"
+ironcore-documents = "0.2.1"
 itertools = "0.12"
 ndarray = "0.15.6"
 ndarray-rand = "0.14.0"
 protobuf = { version = "3.3", features = ["with-bytes"] }
 rand = "0.8.5"
 rand_chacha = "0.3.1"
 rand_distr = "0.4.3"
+rayon = "1.10.0"
 regex = "1.10.2"
 reqwest = { version = "0.12", default-features = false, features = [
     "json",
@@ -42,8 +43,8 @@ serde = { version = "1.0.163", features = ["derive"] }
 serde_json = { version = "1.0.96", features = ["float_roundtrip"] }
 thiserror = "1.0.50"
 uniffi = { git = "https://github.com/mozilla/uniffi-rs", features = [
-    "cli",
     "tokio",
+    "cli",
 ] }
 z85 = "3.0.5"
 
@@ -86,7 +87,7 @@ path = "uniffi-bindgen.rs"
 # 6.9M vs 1.5M in initial testing. Can further have `strip` (the Unix utility) run on it to save ~0.2 MB more.
 # WARNING: be careful changing this, since downstream integration tests (in "core/tests") depend on this profile.
 [profile.release]
-opt-level = 3       # Check 'z' sometimes in case much smaller. Initial testing showed 0.1M size difference but 25% performance hit, so 3 is better.
+opt-level = 3       # Our crypto basically requires opt-level 3 for reasonable speeds
 lto = true          # Enable Link Time Optimization
 codegen-units = 1   # Reduce number of codegen units to increase optimizations.
 panic = 'abort'     # Abort on panic

benches/ironcore_alloy_bench.rs

@@ -205,8 +205,11 @@ fn tsp_benches(c: &mut Criterion) {
     c.bench_function("TSP - batch encrypt 10 documents, 10 fields, 10B", |b| {
         b.to_async(Runtime::new().unwrap()).iter_batched(
             || {
-                PlaintextDocuments((0..10).fold(HashMap::new(), |mut acc, i| {
-                    let doc = generate_plaintext(10, 10, &mut rng);
+                let num_documents = 10;
+                let num_fields = 10;
+                let field_size = 10;
+                PlaintextDocuments((0..num_documents).fold(HashMap::new(), |mut acc, i| {
+                    let doc = generate_plaintext(field_size, num_fields, &mut rng);
                     acc.insert(DocumentId(format!("doc{}", i)), doc);
                     acc
                 }))

kotlin/benchmarks/src/SaasShieldBenchmark.kt

@@ -58,12 +58,17 @@ class SaasShieldBenchmark {
         }
     }
 
+    @TearDown
+    fun tearDown() {
+        saasShieldSdk.close()
+    }
+
     fun generatePlaintextDocument(bytesPerField: Int, numFields: Int): PlaintextDocument {
         val documentMap = HashMap<String, PlaintextBytes>()
         for (i in 1..numFields) {
             val byteArray = ByteArray(bytesPerField)
             kotlin.random.Random.nextBytes(byteArray)
-            documentMap.put("doc" + i, byteArray)
+            documentMap.put("field" + i, byteArray)
         }
         return documentMap
     }

kotlin/build.gradle.kts

@@ -10,8 +10,8 @@ group = "com.ironcorelabs"
 
 plugins {
     // Apply the org.jetbrains.kotlin.jvm Plugin to add support for Kotlin.
-    kotlin("jvm") version "1.8.10"
-    id("org.jetbrains.dokka") version "1.9.0"
+    kotlin("jvm") version "1.9.23"
+    id("org.jetbrains.dokka") version "1.9.20"
 
     // Apply the java-library plugin for API and implementation separation.
     `java-library`
@@ -23,8 +23,8 @@ plugins {
 
     // benchmark deps
     java
-    id("org.jetbrains.kotlinx.benchmark") version "0.4.9"
-    kotlin("plugin.allopen") version "1.9.20"
+    id("org.jetbrains.kotlinx.benchmark") version "0.4.10"
+    kotlin("plugin.allopen") version "1.9.23"
 }
 
 tasks.withType<org.jetbrains.kotlin.gradle.tasks.KotlinCompile>().configureEach {
@@ -55,8 +55,8 @@ dependencies {
     testImplementation("org.jetbrains.kotlin:kotlin-test-junit5")
     // Use the JUnit 5 integration.
     testImplementation("org.junit.jupiter:junit-jupiter-engine:5.9.1")
-    implementation("net.java.dev.jna:jna:5.13.0")
-    implementation("org.jetbrains.kotlinx:kotlinx-coroutines-core:1.7.3")
+    implementation("net.java.dev.jna:jna:5.14.0")
+    implementation("org.jetbrains.kotlinx:kotlinx-coroutines-core:1.8.0")
     implementation("org.jetbrains.kotlin:kotlin-scripting-jvm")
     benchmarksImplementation("org.jetbrains.kotlinx:kotlinx-benchmark-runtime:0.4.9")
     benchmarksImplementation(sourceSets.main.get().output + sourceSets.main.get().runtimeClasspath)
@@ -129,7 +129,7 @@ benchmark {
     targets {
         register(benchmarks) {
             this as JvmBenchmarkTarget
-            jmhVersion = "1.21"
+            jmhVersion = "1.37"
         }
     }
 }

rust-toolchain.toml

@@ -1,4 +1,4 @@
 [toolchain]
 profile = "default"
-channel = "1.75.0"
+channel = "1.77.2"
 components = ["rust-src", "rust-analyzer", "llvm-tools"]

src/saas_shield/config.rs

@@ -1,5 +1,6 @@
-use crate::errors::AlloyError;
-use crate::tenant_security_client::{ApiKey, TenantSecurityClient};
+use crate::tenant_security_client::TenantSecurityClient;
+use crate::{errors::AlloyError, tenant_security_client::ApiKey};
+use reqwest::header::{HeaderMap, HeaderValue};
 use std::sync::Arc;
 
 /// Configuration for the SaaS Shield SDKs. Sets the TSP domain/URI and API key to be used for SaaS Shield operations.
@@ -21,17 +22,25 @@ impl SaasShieldConfiguration {
         accept_invalid_certs: bool,
         approximation_factor: Option<f32>,
     ) -> Result<Arc<Self>, AlloyError> {
+        let parsed_api_key = ApiKey::try_from(api_key)?;
+        let default_headers = {
+            let mut headers: HeaderMap = HeaderMap::default();
+            headers.insert("Content-Type", HeaderValue::from_static("application/json"));
+            let mut auth_header: HeaderValue = format!("cmk {}", parsed_api_key.0)
+                .parse()
+                .expect("Invalid API_KEY");
+            auth_header.set_sensitive(true);
+            headers.insert("Authorization", auth_header);
+            headers
+        };
         let reqwest_client = reqwest::Client::builder()
+            .default_headers(default_headers)
             .danger_accept_invalid_certs(accept_invalid_certs)
             .build()
             .expect("Failed to create http client. This means there is a system misconfiguration.");
         Ok(Arc::new(Self {
             approximation_factor,
-            tenant_security_client: Arc::new(TenantSecurityClient::new(
-                tsp_uri,
-                ApiKey::try_from(api_key)?,
-                reqwest_client,
-            )),
+            tenant_security_client: Arc::new(TenantSecurityClient::new(tsp_uri, reqwest_client)),
         }))
     }
 }