fix codebase to use secretbox, tighten security

Author: NilanjanDaw

keymanager/key_protection_service/key_custody_core/src/lib.rs

@@ -171,28 +171,40 @@ pub unsafe extern "C" fn key_manager_decap_and_seal(
     aad: *const u8,
     aad_len: usize,
     out_encapsulated_key: *mut u8,
-    out_encapsulated_key_len: *mut usize,
+    out_encapsulated_key_len: usize,
     out_ciphertext: *mut u8,
-    out_ciphertext_len: *mut usize,
+    out_ciphertext_len: usize,
 ) -> i32 {
     if uuid_bytes.is_null()
         || encapsulated_key.is_null()
         || encapsulated_key_len == 0
         || out_encapsulated_key.is_null()
-        || out_encapsulated_key_len.is_null()
+        || out_encapsulated_key_len == 0
         || out_ciphertext.is_null()
-        || out_ciphertext_len.is_null()
+        || out_ciphertext_len == 0
     {
         return -1;
     }
 
-    let uuid = unsafe {
-        let mut bytes = [0u8; 16];
-        std::ptr::copy_nonoverlapping(uuid_bytes, bytes.as_mut_ptr(), 16);
-        Uuid::from_bytes(bytes)
+    // Convert to Safe Types
+    let uuid = unsafe { std::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 {
+        unsafe { slice::from_raw_parts(aad, aad_len) }
+    } else {
+        &[]
+    };
+    let out_encapsulated_key =
+        unsafe { slice::from_raw_parts_mut(out_encapsulated_key, out_encapsulated_key_len) };
+    let out_ciphertext = unsafe { slice::from_raw_parts_mut(out_ciphertext, out_ciphertext_len) };
+
+    let uuid_val = match Uuid::from_slice(uuid) {
+        Ok(u) => u,
+        Err(_) => return -1,
     };
 
-    let key_record = match KEY_REGISTRY.get_key(&uuid) {
+    // Get key record from registry
+    let key_record = match KEY_REGISTRY.get_key(&uuid_val) {
         Some(record) => record,
         None => return -1,
     };
@@ -208,61 +220,39 @@ pub unsafe extern "C" fn key_manager_decap_and_seal(
 
     let _kem_algo = match km_common::algorithms::KemAlgorithm::try_from(hpke_algo.kem) {
         Ok(k) => k,
-        Err(_) => return -1,
+        Err(_) => return -1, // Invalid KEM algorithm
     };
 
-    let enc_key_slice = unsafe { slice::from_raw_parts(encapsulated_key, encapsulated_key_len) };
-
-    // Convert Vault bytes to SecretBox for active usage
-    let secret =
-        km_common::crypto::secret_box::SecretBox::new(key_record.private_key.as_bytes().to_vec());
-    let priv_key = km_common::crypto::PrivateKey::from_secret(secret);
+    let priv_key = key_record.get_private_key();
 
     // Decapsulate
-    let mut shared_secret = match km_common::crypto::decaps(&priv_key, enc_key_slice) {
+    let shared_secret = match km_common::crypto::decaps(&priv_key, enc_key_slice) {
         Ok(s) => s,
         Err(_) => return -3,
     };
 
-    let aad_slice = if !aad.is_null() && aad_len > 0 {
-        unsafe { slice::from_raw_parts(aad, aad_len) }
-    } else {
-        &[]
-    };
-
     // Seal
-    let (new_enc_key, sealed_ciphertext) = match km_common::crypto::hpke_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(_) => {
-            shared_secret.zeroize();
             return -4;
         }
     };
-    shared_secret.zeroize();
 
     // Copy outputs
-    unsafe {
-        let enc_len_req = new_enc_key.len();
-        let ct_len_req = sealed_ciphertext.len();
-        let enc_len_avail = *out_encapsulated_key_len;
-        let ct_len_avail = *out_ciphertext_len;
-
-        if enc_len_avail < enc_len_req || ct_len_avail < ct_len_req {
-            return -2;
-        }
-
-        std::ptr::copy_nonoverlapping(new_enc_key.as_ptr(), out_encapsulated_key, enc_len_req);
-        *out_encapsulated_key_len = enc_len_req;
+    let enc_len_req = enc_key.len();
+    let ct_len_req = sealed_ciphertext.len();
 
-        std::ptr::copy_nonoverlapping(sealed_ciphertext.as_ptr(), out_ciphertext, ct_len_req);
-        *out_ciphertext_len = ct_len_req;
+    if out_encapsulated_key_len != enc_len_req || out_ciphertext_len != ct_len_req {
+        return -2;
     }
-
+    out_encapsulated_key.copy_from_slice(enc_key.as_slice());
+    out_ciphertext.copy_from_slice(sealed_ciphertext.as_slice());
     0
 }
 
@@ -485,14 +475,14 @@ mod tests {
         // 2. Generate KEM key in registry
         let mut uuid_bytes = [0u8; 16];
         let mut kem_pubkey_bytes = [0u8; 32];
-        let mut kem_pubkey_len = 32;
+        let kem_pubkey_len = 32;
         let algo = HpkeAlgorithm {
             kem: KemAlgorithm::DhkemX25519HkdfSha256 as i32,
             kdf: KdfAlgorithm::HkdfSha256 as i32,
             aead: AeadAlgorithm::Aes256Gcm as i32,
         };
 
-        let result = unsafe {
+        unsafe {
             key_manager_generate_kem_keypair(
                 algo,
                 binding_pk.as_bytes().as_ptr(),
@@ -515,9 +505,9 @@ mod tests {
 
         // Step 3: Call `decap_and_seal`.
         let mut out_enc_key = [0u8; 32];
-        let mut out_enc_key_len = 32;
+        let out_enc_key_len = 32;
         let mut out_ct = [0u8; 48]; // 32 bytes secret + 16 tag
-        let mut out_ct_len = 48;
+        let out_ct_len = 48;
 
         let result = unsafe {
             key_manager_decap_and_seal(
@@ -527,9 +517,9 @@ mod tests {
                 aad.as_ptr(),
                 aad.len(),
                 out_enc_key.as_mut_ptr(),
-                &mut out_enc_key_len,
+                out_enc_key_len,
                 out_ct.as_mut_ptr(),
-                &mut out_ct_len,
+                out_ct_len,
             )
         };
 
@@ -545,10 +535,7 @@ mod tests {
         // 5. Verify the recovered secret matches what decaps would produce
         let key_record = KEY_REGISTRY.get_key(&Uuid::from_bytes(uuid_bytes)).unwrap();
 
-        let secret = km_common::crypto::secret_box::SecretBox::new(
-            key_record.private_key.as_bytes().to_vec(),
-        );
-        let priv_key = km_common::crypto::PrivateKey::from_secret(secret);
+        let priv_key = key_record.get_private_key();
 
         let expected_shared_secret =
             km_common::crypto::decaps(&priv_key, &client_enc).expect("decaps failed");
@@ -574,9 +561,9 @@ mod tests {
     #[test]
     fn test_decap_and_seal_invalid_uuid() {
         let mut out_enc_key = [0u8; 32];
-        let mut out_enc_key_len = 32;
+        let out_enc_key_len = 32;
         let mut out_ct = [0u8; 48];
-        let mut out_ct_len = 48;
+        let out_ct_len = 48;
 
         let result = unsafe {
             key_manager_decap_and_seal(
@@ -586,9 +573,9 @@ mod tests {
                 std::ptr::null(),
                 0,
                 out_enc_key.as_mut_ptr(),
-                &mut out_enc_key_len,
+                out_enc_key_len,
                 out_ct.as_mut_ptr(),
-                &mut out_ct_len,
+                out_ct_len,
             )
         };
 
@@ -598,7 +585,7 @@ mod tests {
     #[test]
     fn test_decap_and_seal_null_args() {
         let mut out_enc_key = [0u8; 32];
-        let mut out_enc_key_len = 32;
+        let out_enc_key_len = 32;
 
         let result = unsafe {
             key_manager_decap_and_seal(
@@ -608,9 +595,9 @@ mod tests {
                 std::ptr::null(),
                 0,
                 out_enc_key.as_mut_ptr(),
-                &mut out_enc_key_len,
-                std::ptr::null_mut(),
+                out_enc_key_len,
                 std::ptr::null_mut(),
+                0,
             )
         };
 
@@ -626,7 +613,7 @@ mod tests {
         // 2. Generate KEM key
         let mut uuid_bytes = [0u8; 16];
         let mut kem_pubkey_bytes = [0u8; 32];
-        let mut kem_pubkey_len = 32;
+        let kem_pubkey_len = 32;
         let algo = HpkeAlgorithm {
             kem: KemAlgorithm::DhkemX25519HkdfSha256 as i32,
             kdf: KdfAlgorithm::HkdfSha256 as i32,
@@ -640,16 +627,16 @@ mod tests {
                 3600,
                 uuid_bytes.as_mut_ptr(),
                 kem_pubkey_bytes.as_mut_ptr(),
-                &mut kem_pubkey_len,
+                kem_pubkey_len,
             )
         };
         assert_eq!(res, 0);
 
         // 3. Call with invalid encapsulated key (wrong length for X25519)
         let mut out_enc_key = [0u8; 32];
-        let mut out_enc_key_len = 32;
+        let out_enc_key_len = 32;
         let mut out_ct = [0u8; 48];
-        let mut out_ct_len = 48;
+        let out_ct_len = 48;
 
         let result = unsafe {
             key_manager_decap_and_seal(
@@ -659,9 +646,9 @@ mod tests {
                 std::ptr::null(),
                 0,
                 out_enc_key.as_mut_ptr(),
-                &mut out_enc_key_len,
+                out_enc_key_len,
                 out_ct.as_mut_ptr(),
-                &mut out_ct_len,
+                out_ct_len,
             )
         };
 
@@ -677,7 +664,7 @@ mod tests {
         // 2. Generate KEM key
         let mut uuid_bytes = [0u8; 16];
         let mut kem_pubkey_bytes = [0u8; 32];
-        let mut kem_pubkey_len = 32;
+        let kem_pubkey_len = 32;
         let algo = HpkeAlgorithm {
             kem: KemAlgorithm::DhkemX25519HkdfSha256 as i32,
             kdf: KdfAlgorithm::HkdfSha256 as i32,
@@ -691,7 +678,7 @@ mod tests {
                 3600,
                 uuid_bytes.as_mut_ptr(),
                 kem_pubkey_bytes.as_mut_ptr(),
-                &mut kem_pubkey_len,
+                kem_pubkey_len,
             );
         }
 
@@ -703,9 +690,9 @@ mod tests {
 
         // 4. Call with small output buffers
         let mut out_enc_key = [0u8; 31]; // Small
-        let mut out_enc_key_len = 31;
+        let out_enc_key_len = 31;
         let mut out_ct = [0u8; 47]; // Small
-        let mut out_ct_len = 47;
+        let out_ct_len = 47;
 
         let result = unsafe {
             key_manager_decap_and_seal(
@@ -715,9 +702,9 @@ mod tests {
                 std::ptr::null(),
                 0,
                 out_enc_key.as_mut_ptr(),
-                &mut out_enc_key_len,
+                out_enc_key_len,
                 out_ct.as_mut_ptr(),
-                &mut out_ct_len,
+                out_ct_len,
             )
         };
 

keymanager/km_common/src/crypto.rs

@@ -1,9 +1,8 @@
 use crate::algorithms::{AeadAlgorithm, HpkeAlgorithm, KdfAlgorithm, KemAlgorithm};
 pub mod secret_box;
 use crate::crypto::secret_box::SecretBox;
-use bssl_crypto::hkdf;
 #[cfg(any(test, feature = "test-utils"))]
-use bssl_crypto::{aead, aead::Aead};
+use bssl_crypto::{aead, aead::Aead, hkdf};
 use clear_on_drop::clear_stack_on_return;
 use thiserror::Error;
 
@@ -93,6 +92,12 @@ pub enum PrivateKey {
     X25519(X25519PrivateKey),
 }
 
+impl From<SecretBox> for PrivateKey {
+    fn from(secret: SecretBox) -> Self {
+        PrivateKey::X25519(X25519PrivateKey(secret))
+    }
+}
+
 impl From<PrivateKey> for SecretBox {
     fn from(key: PrivateKey) -> SecretBox {
         match key {

keymanager/km_common/src/key_types.rs

@@ -47,7 +47,7 @@ pub struct KeyMetadata {
 pub struct KeyRecord {
     pub meta: KeyMetadata,
     /// memfd_secrets backed secret key-material
-    pub private_key: Vault,
+    private_key: Vault,
 }
 
 pub type KeyHandle = Uuid;
@@ -102,6 +102,11 @@ impl KeyRegistry {
 }
 
 impl KeyRecord {
+    /// Returns the private key material.
+    pub fn get_private_key(&self) -> crypto::PrivateKey {
+        crypto::PrivateKey::from(self.private_key.get_secret())
+    }
+
     /// Creates a new long-term Binding key.
     pub fn create_binding_key(
         algo: HpkeAlgorithm,