resolve review comments

Author: NilanjanDaw

keymanager/Cargo.toml

@@ -4,4 +4,22 @@ members = [
     "workload_service/key_custody_core",
     "key_protection_service/key_custody_core",
     "km_common",
-]
+]
+
+[workspace.package]
+version = "0.1.0"
+edition = "2024"
+
+[workspace.dependencies]
+km_common = { path = "km_common" }
+uuid = { version = "1.20.0", features = ["v4"] }
+zeroize = "1.8.2"
+prost = "0.13"
+libc = "0.2.180"
+memmap2 = "0.9.9"
+thiserror = "2.0"
+clear_on_drop = "0.2"
+bssl-crypto = { path = "third_party/bssl-crypto" }
+hex = "0.4"
+prost-build = "0.13"
+protoc-bin-vendored = "3.2.0"

keymanager/key_protection_service/key_custody_core/Cargo.toml

@@ -1,16 +1,16 @@
 [package]
 name = "kps_key_custody_core"
-version = "0.1.0"
-edition = "2024"
+version.workspace = true
+edition.workspace = true
 
 [lib]
 crate-type = ["staticlib", "rlib"]
 
 [dependencies]
-km_common = { path = "../../km_common" }
-uuid = { version = "1.20.0", features = ["v4"] }
-zeroize = "1.8"
-prost = "0.13"
+km_common.workspace = true
+uuid.workspace = true
+zeroize.workspace = true
+prost.workspace = true
 
 [dev-dependencies]
-km_common = { path = "../../km_common", features = ["test-utils"] }
+km_common = { workspace = true, features = ["test-utils"] }

keymanager/key_protection_service/key_custody_core/src/lib.rs

@@ -152,6 +152,41 @@ pub unsafe extern "C" fn key_manager_destroy_kem_key(uuid_bytes: *const u8) -> i
     .unwrap_or(-1)
 }
 
+/// Internal function to decapsulate and reseal a shared secret.
+fn decap_and_seal_internal(
+    uuid: Uuid,
+    encapsulated_key: &[u8],
+    aad: &[u8],
+) -> Result<(Vec<u8>, Vec<u8>), i32> {
+    // Get key record from registry
+    let Some(key_record) = KEY_REGISTRY.get_key(&uuid) else {
+        Err(-1)? // Key not found
+    };
+
+    let KeySpec::KemWithBindingPub {
+        algo: hpke_algo,
+        binding_public_key,
+        ..
+    } = &key_record.meta.spec
+    else {
+        Err(-1)? // Invalid key type
+    };
+
+    let priv_key = key_record.get_private_key();
+
+    // Decapsulate
+    let shared_secret = match km_common::crypto::decaps(&priv_key, encapsulated_key) {
+        Ok(s) => s,
+        Err(_) => return Err(-3),
+    };
+
+    // Seal
+    match km_common::crypto::hpke_seal(binding_public_key, &shared_secret, aad, hpke_algo) {
+        Ok((enc, ct)) => Ok((enc.to_vec(), ct.to_vec())),
+        Err(_) => Err(-4),
+    }
+}
+
 /// Decapsulates a shared secret using a stored KEM key and immediately reseals it using the associated binding public key.
 ///
 /// ## Arguments
@@ -161,10 +196,10 @@ pub unsafe extern "C" fn key_manager_destroy_kem_key(uuid_bytes: *const u8) -> i
 /// * `aad` - A pointer to the Additional Authenticated Data (AAD) for the sealing operation.
 /// * `aad_len` - The length of the AAD.
 /// * `out_encapsulated_key` - A pointer to a buffer where the new encapsulated key will be written.
-/// * `out_encapsulated_key_len` - A pointer to a `usize` containing the size of `out_encapsulated_key`.
+/// * `out_encapsulated_key_len` - The size of `out_encapsulated_key`.
 ///   On success, updated with the actual size.
 /// * `out_ciphertext` - A pointer to a buffer where the sealed ciphertext will be written.
-/// * `out_ciphertext_len` - A pointer to a `usize` containing the size of `out_ciphertext`.
+/// * `out_ciphertext_len` -The size of `out_ciphertext`.
 ///   On success, updated with the actual size.
 ///
 /// ## Safety
@@ -208,7 +243,7 @@ pub unsafe extern "C" fn key_manager_decap_and_seal(
         }
 
         // Convert to Safe Types
-        let uuid = unsafe { std::slice::from_raw_parts(uuid_bytes, 16) };
+        let uuid_slice = unsafe { slice::from_raw_parts(uuid_bytes, 16) };
         let enc_key_slice =
             unsafe { slice::from_raw_parts(encapsulated_key, encapsulated_key_len) };
         let aad_slice = if !aad.is_null() && aad_len > 0 {
@@ -218,65 +253,26 @@ pub unsafe extern "C" fn key_manager_decap_and_seal(
         };
         let out_encapsulated_key_slice =
             unsafe { slice::from_raw_parts_mut(out_encapsulated_key, out_encapsulated_key_len) };
-        let out_ciphertext =
+        let out_ciphertext_slice =
             unsafe { slice::from_raw_parts_mut(out_ciphertext, out_ciphertext_len) };
 
         let uuid = match Uuid::from_slice(uuid_slice) {
             Ok(u) => u,
             Err(_) => return -1,
         };
 
-        // Get key record from registry
-        let key_record = match KEY_REGISTRY.get_key(&uuid_val) {
-            Some(record) => record,
-            None => return -1,
-        };
-
-        let (hpke_algo, binding_public_key) = match &key_record.meta.spec {
-            KeySpec::KemWithBindingPub {
-                algo,
-                binding_public_key,
-                ..
-            } => (algo, binding_public_key),
-            _ => return -1, // Wrong key type
-        };
-
-        let _kem_algo = match km_common::algorithms::KemAlgorithm::try_from(hpke_algo.kem) {
-            Ok(k) => k,
-            Err(_) => return -1, // Invalid KEM algorithm
-        };
-
-        let priv_key = key_record.get_private_key();
-
-        // Decapsulate
-        let shared_secret = match km_common::crypto::decaps(&priv_key, enc_key_slice) {
-            Ok(s) => s,
-            Err(_) => return -3,
-        };
-
-        // Seal
-        let (enc_key, sealed_ciphertext) = match km_common::crypto::hpke_seal(
-            binding_public_key,
-            &shared_secret,
-            aad_slice,
-            hpke_algo,
-        ) {
-            Ok(res) => res,
-            Err(_) => {
-                return -4;
+        // Call Safe Internal Function
+        match decap_and_seal_internal(uuid, enc_key_slice, aad_slice) {
+            Ok((enc, ct)) => {
+                if out_encapsulated_key_len != enc.len() || out_ciphertext_len != ct.len() {
+                    return -2;
+                }
+                out_encapsulated_key_slice.copy_from_slice(&enc);
+                out_ciphertext_slice.copy_from_slice(&ct);
+                0 // Success
             }
-        };
-
-        // Copy outputs
-        let enc_len_req = enc_key.len();
-        let ct_len_req = sealed_ciphertext.len();
-
-        if out_encapsulated_key_len != enc_len_req || out_ciphertext_len != ct_len_req {
-            return -2;
+            Err(e) => e,
         }
-        out_encapsulated_key.copy_from_slice(enc_key.as_slice());
-        out_ciphertext.copy_from_slice(sealed_ciphertext.as_slice());
-        0
     }))
     .unwrap_or(-1)
 }
@@ -462,10 +458,11 @@ mod tests {
             kdf: KdfAlgorithm::HkdfSha256 as i32,
             aead: AeadAlgorithm::Aes256Gcm as i32,
         };
-
+        let algo_bytes = algo.encode_to_vec();
         unsafe {
             let res = key_manager_generate_kem_keypair(
-                algo,
+                algo_bytes.as_ptr(),
+                algo_bytes.len(),
                 binding_pubkey.as_ptr(),
                 binding_pubkey.len(),
                 3600,
@@ -513,10 +510,11 @@ mod tests {
             kdf: KdfAlgorithm::HkdfSha256 as i32,
             aead: AeadAlgorithm::Aes256Gcm as i32,
         };
-
+        let algo_bytes = algo.encode_to_vec();
         unsafe {
             key_manager_generate_kem_keypair(
-                algo,
+                algo_bytes.as_ptr(),
+                algo_bytes.len(),
                 binding_pk.as_bytes().as_ptr(),
                 binding_pk.as_bytes().len(),
                 3600,
@@ -632,9 +630,11 @@ mod tests {
             kdf: KdfAlgorithm::HkdfSha256 as i32,
             aead: AeadAlgorithm::Aes256Gcm as i32,
         };
+        let algo_bytes = algo.encode_to_vec();
         let res = unsafe {
             key_manager_generate_kem_keypair(
-                algo,
+                algo_bytes.as_ptr(),
+                algo_bytes.len(),
                 binding_pk.as_bytes().as_ptr(),
                 binding_pk.as_bytes().len(),
                 3600,
@@ -683,9 +683,12 @@ mod tests {
             kdf: KdfAlgorithm::HkdfSha256 as i32,
             aead: AeadAlgorithm::Aes256Gcm as i32,
         };
+
+        let algo_bytes = algo.encode_to_vec();
         let res = unsafe {
             key_manager_generate_kem_keypair(
-                algo,
+                algo_bytes.as_ptr(),
+                algo_bytes.len(),
                 binding_pk.as_bytes().as_ptr(),
                 binding_pk.as_bytes().len(),
                 3600,

keymanager/km_common/Cargo.toml

@@ -1,24 +1,24 @@
 [package]
 name = "km_common"
-version = "0.1.0"
-edition = "2024"
+version.workspace = true
+edition.workspace = true
 
 [dependencies]
-libc = "0.2.180"
-memmap2 = "0.9.9"
-prost = "0.13"
-uuid = { version = "1.20.0", features = ["v4", "serde"] }
-zeroize = { version = "1.8.2", features = ["derive"] }
-thiserror = "2.0"
-bssl-crypto = { path = "../third_party/bssl-crypto" }
-clear_on_drop = "0.2"
+libc.workspace = true
+memmap2.workspace = true
+prost.workspace = true
+uuid = { workspace = true, features = ["serde"] }
+zeroize = { workspace = true, features = ["derive"] }
+thiserror.workspace = true
+bssl-crypto.workspace = true
+clear_on_drop.workspace = true
 
 [features]
 test-utils = []
 
 [dev-dependencies]
-hex = "0.4"
+hex.workspace = true
 
 [build-dependencies]
-prost-build = "0.13"
-protoc-bin-vendored = "3.2.0"
+prost-build.workspace = true
+protoc-bin-vendored.workspace = true

keymanager/workload_service/key_custody_core/Cargo.toml

@@ -1,13 +1,13 @@
 [package]
 name = "ws_key_custody_core"
-version = "0.1.0"
-edition = "2024"
+version.workspace = true
+edition.workspace = true
 
 [lib]
 crate-type = ["staticlib", "rlib"]
 
 [dependencies]
-km_common = { path = "../../km_common" }
-uuid = { version = "1.20.0", features = ["v4"] }
-zeroize = "1.8"
-prost = "0.13"
+km_common.workspace = true
+uuid.workspace = true
+zeroize.workspace = true
+prost.workspace = true

keymanager/workload_service/key_custody_core/src/lib.rs

@@ -1,6 +1,6 @@
 use km_common::algorithms::HpkeAlgorithm;
-use km_common::crypto::secret_box::SecretBox;
 use km_common::crypto::PublicKey;
+use km_common::crypto::secret_box::SecretBox;
 use km_common::key_types::{KeyRecord, KeyRegistry, KeySpec};
 use prost::Message;
 use std::slice;
@@ -62,7 +62,6 @@ fn generate_binding_keypair_internal(
 /// * `0` on success.
 /// * `-1` if an error occurred during key generation.
 /// * `-2` if the `out_pubkey` buffer size does not match the key size.
-
 #[unsafe(no_mangle)]
 pub unsafe extern "C" fn key_manager_generate_binding_keypair(
     algo_ptr: *const u8,
@@ -134,15 +133,9 @@ pub unsafe extern "C" fn key_manager_destroy_binding_key(uuid_bytes: *const u8)
 }
 
 /// Internal function to decrypt a ciphertext using a stored binding key.
-fn open_internal(
-    uuid: Uuid,
-    enc: &[u8],
-    ciphertext: &[u8],
-    aad: &[u8],
-) -> Result<SecretBox, i32> {
-    let record = match KEY_REGISTRY.get_key(&uuid) {
-        Some(r) => r,
-        None => return Err(-1),
+fn open_internal(uuid: Uuid, enc: &[u8], ciphertext: &[u8], aad: &[u8]) -> Result<SecretBox, i32> {
+    let Some(record) = KEY_REGISTRY.get_key(&uuid) else {
+        Err(-1)? // Key not found
     };
 
     let algo = match &record.meta.spec {
@@ -376,10 +369,11 @@ mod tests {
             kdf: KdfAlgorithm::HkdfSha256 as i32,
             aead: AeadAlgorithm::Aes256Gcm as i32,
         };
-
+        let algo_bytes = algo.encode_to_vec();
         unsafe {
             let res = key_manager_generate_binding_keypair(
-                algo,
+                algo_bytes.as_ptr(),
+                algo_bytes.len(),
                 3600,
                 uuid_bytes.as_mut_ptr(),
                 pubkey_bytes.as_mut_ptr(),
@@ -421,9 +415,11 @@ mod tests {
         };
 
         // 1. Generate keypair
+        let algo_bytes = algo.encode_to_vec();
         unsafe {
             key_manager_generate_binding_keypair(
-                algo,
+                algo_bytes.as_ptr(),
+                algo_bytes.len(),
                 3600,
                 uuid_bytes.as_mut_ptr(),
                 pubkey_bytes.as_mut_ptr(),
@@ -490,9 +486,11 @@ mod tests {
             kdf: KdfAlgorithm::HkdfSha256 as i32,
             aead: AeadAlgorithm::Aes256Gcm as i32,
         };
+        let algo_bytes = algo.encode_to_vec();
         let res = unsafe {
             key_manager_generate_binding_keypair(
-                algo,
+                algo_bytes.as_ptr(),
+                algo_bytes.len(),
                 3600,
                 uuid_bytes.as_mut_ptr(),
                 pubkey_bytes.as_mut_ptr(),