[KeyManager] FFI for Enumerate keys (#665)

Implements FFI for the Enumerate Keys API in the Key Custody Core (KCC). It exposes the necessary Rust functions to allow the Go Workload Service to list active KEM keys.

## Dependency
- Created on top of PR #652 (Key Generation KOL changes)

## Key Changes
1.  **Key Custody Core (`key_custody_core`)**:
    -   Added `key_manager_enumerate_kem_keys` FFI function.
    -   Implemented `enumerate_kem_keys` in `KeyRegistry` to safely expose key metadata.
    -   Added `KmHpkeAlgorithm` struct for C compatibility.

2.  **C Header (`kps_key_custody_core.h`)**:
    -   Exported `key_manager_enumerate_kem_keys` definition.

## Verification
-   **Rust Tests**:
    -   Ran `cargo check` and `cargo test` in `keymanager/key_protection_service/key_custody_core`.
    -   Verified that the new FFI function compiles and integrates with the existing system.

Author: atulpatildbz

keymanager/key_protection_service/key_custody_core/cbindgen.toml

@@ -4,7 +4,7 @@ pragma_once = false
 autogen_warning = "/* Auto-generated by cbindgen. DO NOT EDIT. */"
 cpp_compat = true
 no_includes = true
-style = "tag"
+style = "type"
 documentation = false
 usize_is_size_t = true
 sys_includes = ["stddef.h", "stdint.h"]
@@ -15,7 +15,7 @@ parse_deps = false
 clean = true
 
 [export]
-item_types = ["functions"]
-include = ["key_manager_generate_kem_keypair"]
+item_types = ["functions", "structs", "constants", "enums"]
+include = ["key_manager_generate_kem_keypair", "key_manager_destroy_kem_key", "key_manager_enumerate_kem_keys", "KpsKeyInfo", "MAX_ALGORITHM_LEN", "MAX_PUBLIC_KEY_LEN"]
 
 

keymanager/key_protection_service/key_custody_core/include/kps_key_custody_core.h

@@ -6,6 +6,21 @@
 #include <stddef.h>
 #include <stdint.h>
 
+#define MAX_ALGORITHM_LEN 128
+
+#define MAX_PUBLIC_KEY_LEN 2048
+
+typedef struct {
+  uint8_t uuid[16];
+  uint8_t algorithm[MAX_ALGORITHM_LEN];
+  size_t algorithm_len;
+  uint8_t pub_key[MAX_PUBLIC_KEY_LEN];
+  size_t pub_key_len;
+  uint8_t binding_pub_key[MAX_PUBLIC_KEY_LEN];
+  size_t binding_pub_key_len;
+  uint64_t remaining_lifespan_secs;
+} KpsKeyInfo;
+
 #ifdef __cplusplus
 extern "C" {
 #endif // __cplusplus
@@ -21,6 +36,11 @@ int32_t key_manager_generate_kem_keypair(const uint8_t *algo_ptr,
 
 int32_t key_manager_destroy_kem_key(const uint8_t *uuid_bytes);
 
+int32_t key_manager_enumerate_kem_keys(KpsKeyInfo *out_entries,
+                                       size_t max_entries,
+                                       size_t offset,
+                                       bool *out_has_more);
+
 int32_t key_manager_decap_and_seal(const uint8_t *uuid_bytes,
                                    const uint8_t *encapsulated_key,
                                    size_t encapsulated_key_len,

keymanager/key_protection_service/key_custody_core/src/lib.rs

@@ -3,10 +3,10 @@ use km_common::crypto::PublicKey;
 use km_common::key_types::{KeyRecord, KeyRegistry, KeySpec};
 use prost::Message;
 use std::slice;
+use std::sync::atomic::AtomicBool;
 use std::sync::Arc;
 use std::sync::LazyLock;
-use std::sync::atomic::AtomicBool;
-use std::time::Duration;
+use std::time::{Duration, Instant};
 use uuid::Uuid;
 
 static KEY_REGISTRY: LazyLock<KeyRegistry> = LazyLock::new(|| {
@@ -196,6 +196,125 @@ fn decap_and_seal_internal(
     }
 }
 
+pub const MAX_ALGORITHM_LEN: usize = 128;
+pub const MAX_PUBLIC_KEY_LEN: usize = 2048;
+
+#[repr(C)]
+pub struct KpsKeyInfo {
+    pub uuid: [u8; 16],
+    pub algorithm: [u8; MAX_ALGORITHM_LEN],
+    pub algorithm_len: usize,
+    pub pub_key: [u8; MAX_PUBLIC_KEY_LEN],
+    pub pub_key_len: usize,
+    pub binding_pub_key: [u8; MAX_PUBLIC_KEY_LEN],
+    pub binding_pub_key_len: usize,
+    pub remaining_lifespan_secs: u64,
+}
+
+impl Default for KpsKeyInfo {
+    fn default() -> Self {
+        KpsKeyInfo {
+            uuid: [0; 16],
+            algorithm: [0; MAX_ALGORITHM_LEN],
+            algorithm_len: 0,
+            pub_key: [0; MAX_PUBLIC_KEY_LEN],
+            pub_key_len: 0,
+            binding_pub_key: [0; MAX_PUBLIC_KEY_LEN],
+            binding_pub_key_len: 0,
+            remaining_lifespan_secs: 0,
+        }
+    }
+}
+
+fn enumerate_kem_keys_internal(
+    entries: &mut [KpsKeyInfo],
+    offset: usize,
+) -> Result<(usize, bool), i32> {
+    let (metas, total_count) = KEY_REGISTRY.list_all_keys(offset, entries.len());
+    let count = metas.len();
+    let has_more = offset + count < total_count;
+
+    for (entry, meta) in entries.iter_mut().zip(metas.into_iter()) {
+        let KeySpec::KemWithBindingPub {
+            algo,
+            kem_public_key: pub_key,
+            binding_public_key: binding_pub_key,
+            ..
+        } = &meta.spec
+        else {
+            return Err(-1); // Implementation error, KPS should only contain KEM keys.
+        };
+
+        let algo_bytes = algo.encode_to_vec();
+
+        if pub_key.as_bytes().len() > MAX_PUBLIC_KEY_LEN || algo_bytes.len() > MAX_ALGORITHM_LEN {
+            debug_assert!(
+                false,
+                "Implementation error: Key size exceeds buffer limits! (algo={}, pub={})",
+                algo_bytes.len(),
+                pub_key.as_bytes().len()
+            );
+            return Err(-2); // Buffer Limit Exceeded
+        }
+        if binding_pub_key.as_bytes().len() > MAX_PUBLIC_KEY_LEN {
+            debug_assert!(
+                false,
+                "Implementation error: Binding Key size exceeds buffer limits! (bpk={})",
+                binding_pub_key.as_bytes().len()
+            );
+            return Err(-2);
+        }
+
+        let now = Instant::now();
+        let remaining = meta.delete_after.saturating_duration_since(now).as_secs();
+
+        *entry = KpsKeyInfo::default();
+
+        entry.uuid.copy_from_slice(meta.id.as_bytes());
+
+        entry.algorithm[..algo_bytes.len()].copy_from_slice(&algo_bytes);
+        entry.algorithm_len = algo_bytes.len();
+
+        entry.pub_key[..pub_key.as_bytes().len()].copy_from_slice(pub_key.as_bytes());
+        entry.pub_key_len = pub_key.as_bytes().len();
+
+        entry.binding_pub_key[..binding_pub_key.as_bytes().len()]
+            .copy_from_slice(binding_pub_key.as_bytes());
+        entry.binding_pub_key_len = binding_pub_key.as_bytes().len();
+
+        entry.remaining_lifespan_secs = remaining;
+    }
+
+    Ok((count, has_more))
+}
+
+#[unsafe(no_mangle)]
+pub unsafe extern "C" fn key_manager_enumerate_kem_keys(
+    out_entries: *mut KpsKeyInfo,
+    max_entries: usize,
+    offset: usize,
+    out_has_more: Option<&mut bool>,
+) -> i32 {
+    std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
+        if out_entries.is_null() {
+            return -1;
+        }
+
+        let entries = unsafe { slice::from_raw_parts_mut(out_entries, max_entries) };
+
+        match enumerate_kem_keys_internal(entries, offset) {
+            Ok((count, has_more)) => {
+                if let Some(has_more_ref) = out_has_more {
+                    *has_more_ref = has_more;
+                }
+                count as i32
+            }
+            Err(e) => e,
+        }
+    }))
+    .unwrap_or(-1)
+}
+
 /// Decapsulates a shared secret using a stored KEM key and immediately reseals it using the associated binding public key.
 ///
 /// ## Arguments
@@ -498,6 +617,150 @@ mod tests {
         assert_eq!(result, -1);
     }
 
+    #[test]
+    fn test_enumerate_kem_keys_null_pointers() {
+        let result = unsafe { key_manager_enumerate_kem_keys(std::ptr::null_mut(), 10, 0, None) };
+        assert_eq!(result, -1);
+    }
+
+    #[test]
+    fn test_enumerate_kem_keys_after_generate() {
+        let binding_pubkey = [7u8; 32];
+        let mut uuid_bytes = [0u8; 16];
+        let mut pubkey_bytes = [0u8; 32];
+        let pubkey_len: usize = 32;
+        let algo = HpkeAlgorithm {
+            kem: KemAlgorithm::DhkemX25519HkdfSha256 as i32,
+            kdf: KdfAlgorithm::HkdfSha256 as i32,
+            aead: AeadAlgorithm::Aes256Gcm as i32,
+        };
+        // MUST encode to bytes
+        let algo_bytes = algo.encode_to_vec();
+
+        // Generate a key first.
+        let rc = unsafe {
+            key_manager_generate_kem_keypair(
+                algo_bytes.as_ptr(),
+                algo_bytes.len(),
+                binding_pubkey.as_ptr(),
+                binding_pubkey.len(),
+                3600,
+                uuid_bytes.as_mut_ptr(),
+                pubkey_bytes.as_mut_ptr(),
+                pubkey_len,
+            )
+        };
+        assert_eq!(rc, 0);
+
+        // Enumerate.
+        let mut entries: Vec<KpsKeyInfo> = Vec::with_capacity(256);
+        // Initialize with default/zero values. Note: Arrays are larger now.
+        entries.resize_with(100, || KpsKeyInfo {
+            uuid: [0; 16],
+            algorithm: [0; MAX_ALGORITHM_LEN],
+            algorithm_len: 0,
+            pub_key: [0; MAX_PUBLIC_KEY_LEN],
+            pub_key_len: 0,
+            binding_pub_key: [0; MAX_PUBLIC_KEY_LEN],
+            binding_pub_key_len: 0,
+            remaining_lifespan_secs: 0,
+        });
+        let mut has_more = false;
+
+        let rc = unsafe {
+            // max_entries=100, offset=0
+            key_manager_enumerate_kem_keys(
+                entries.as_mut_ptr(),
+                entries.len(),
+                0,
+                Some(&mut has_more),
+            )
+        };
+        assert!(rc >= 1);
+        let count = rc as usize;
+
+        // Find our key in the results.
+        let mut found = false;
+        for i in 0..count {
+            if entries[i].uuid == uuid_bytes {
+                found = true;
+                let encoded_algo = &entries[i].algorithm[..entries[i].algorithm_len];
+                let decoded_algo = HpkeAlgorithm::decode(encoded_algo).unwrap();
+                assert_eq!(decoded_algo.kem, KemAlgorithm::DhkemX25519HkdfSha256 as i32);
+                assert_eq!(entries[i].pub_key_len, 32);
+                assert!(entries[i].remaining_lifespan_secs > 0);
+                break;
+            }
+        }
+        assert!(found, "generated key not found in enumerate results");
+    }
+
+    #[test]
+    fn test_enumerate_kem_keys_has_more() {
+        // Assume there is at least one key from initialization/other tests or we'll generate one
+        let mut uuid_bytes = [0u8; 16];
+        let mut pubkey_bytes = [0u8; 32];
+        let algo = HpkeAlgorithm {
+            kem: KemAlgorithm::DhkemX25519HkdfSha256 as i32,
+            kdf: KdfAlgorithm::HkdfSha256 as i32,
+            aead: AeadAlgorithm::Aes256Gcm as i32,
+        };
+        let algo_bytes = algo.encode_to_vec();
+
+        // Let's explicitly generate a key so we know there's at least one in the registry
+        unsafe {
+            key_manager_generate_kem_keypair(
+                algo_bytes.as_ptr(),
+                algo_bytes.len(),
+                [7u8; 32].as_ptr(), // fake binding key
+                32,
+                3600,
+                uuid_bytes.as_mut_ptr(),
+                pubkey_bytes.as_mut_ptr(),
+                32,
+            );
+        }
+
+        let mut entries: Vec<KpsKeyInfo> = Vec::with_capacity(256);
+        entries.resize_with(100, || KpsKeyInfo {
+            uuid: [0; 16],
+            algorithm: [0; MAX_ALGORITHM_LEN],
+            algorithm_len: 0,
+            pub_key: [0; MAX_PUBLIC_KEY_LEN],
+            pub_key_len: 0,
+            binding_pub_key: [0; MAX_PUBLIC_KEY_LEN],
+            binding_pub_key_len: 0,
+            remaining_lifespan_secs: 0,
+        });
+
+        // 1. Ask for 0 entries. We should get has_more = true.
+        let mut has_more = false;
+        let rc = unsafe {
+            key_manager_enumerate_kem_keys(entries.as_mut_ptr(), 0, 0, Some(&mut has_more))
+        };
+        assert_eq!(rc, 0);
+        assert!(
+            has_more,
+            "has_more should be true when max_entries is 0 and keys exist"
+        );
+
+        // 2. Ask for 100 entries (which should cover all generated keys). has_more = false.
+        has_more = true; // reset to true to ensure it gets set to false
+        let rc = unsafe {
+            key_manager_enumerate_kem_keys(
+                entries.as_mut_ptr(),
+                entries.len(),
+                0,
+                Some(&mut has_more),
+            )
+        };
+        assert!(rc >= 1);
+        assert!(
+            !has_more,
+            "has_more should be false when all keys are retrieved"
+        );
+    }
+
     #[test]
     fn test_decap_and_seal_success() {
         // 1. Setup binding key (receiver for seal)
@@ -531,9 +794,8 @@ mod tests {
         // 3. Generate a "client" ciphertext/encapsulation targeting KEM key.
         let aad = b"test_aad";
         // We use `encap` to act as the client to generate a valid encapsulation
-        let kem_pub_key_obj = PublicKey::try_from(kem_pubkey_bytes.to_vec()).unwrap();
-        let (client_shared_secret, client_enc) =
-            km_common::crypto::encap(&kem_pub_key_obj).unwrap();
+        let pub_key_obj = PublicKey::try_from(kem_pubkey_bytes.to_vec()).unwrap();
+        let (client_shared_secret, client_enc) = km_common::crypto::encap(&pub_key_obj).unwrap();
 
         // Step 3: Call `decap_and_seal`.
         let mut out_enc_key = [0u8; 32];
@@ -704,10 +966,9 @@ mod tests {
         assert_eq!(res, 0, "Setup failed: key generation returned error");
 
         // 3. Generate valid client encapsulation
-        let kem_pub_key_obj = PublicKey::try_from(kem_pubkey_bytes.to_vec()).unwrap();
+        let pub_key_obj = PublicKey::try_from(kem_pubkey_bytes.to_vec()).unwrap();
         let pt = km_common::crypto::secret_box::SecretBox::new(b"secret".to_vec());
-        let (client_enc, _) =
-            km_common::crypto::hpke_seal(&kem_pub_key_obj, &pt, b"", &algo).unwrap();
+        let (client_enc, _) = km_common::crypto::hpke_seal(&pub_key_obj, &pt, b"", &algo).unwrap();
 
         // 4. Call with small output buffers
         let mut out_enc_key = [0u8; 31]; // Small

keymanager/km_common/build.rs

@@ -7,7 +7,7 @@ fn main() -> Result<()> {
     }
 
     let mut config = prost_build::Config::new();
-    config.type_attribute("HpkeAlgorithm", "#[repr(C)]");
+
     config.compile_protos(&["proto/algorithms.proto"], &["proto/"])?;
 
     Ok(())