From 171275e17577027a6e6147c47aa2c20f0aa4937c Mon Sep 17 00:00:00 2001 From: Atul Patil Date: Sun, 15 Feb 2026 08:20:27 +0000 Subject: [PATCH 01/10] feat: implement FFI for Enumerate Keys API This commit implements the 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. --- .../include/kps_key_custody_core.h | 30 +++- .../key_custody_core/src/lib.rs | 164 +++++++++++++++++- keymanager/km_common/src/key_types.rs | 79 ++++++++- 3 files changed, 268 insertions(+), 5 deletions(-) diff --git a/keymanager/key_protection_service/key_custody_core/include/kps_key_custody_core.h b/keymanager/key_protection_service/key_custody_core/include/kps_key_custody_core.h index e5e573a3a..b0a56718c 100644 --- a/keymanager/key_protection_service/key_custody_core/include/kps_key_custody_core.h +++ b/keymanager/key_protection_service/key_custody_core/include/kps_key_custody_core.h @@ -10,6 +10,24 @@ extern "C" { #endif // __cplusplus +// KmHpkeAlgorithm represents the HPKE algorithm suite configuration. +typedef struct { + int32_t kem; + int32_t kdf; + int32_t aead; +} KmHpkeAlgorithm; + +// KpsKeyInfo holds metadata for a single KEM key entry. +typedef struct { + uint8_t uuid[16]; + KmHpkeAlgorithm algorithm; + uint8_t kem_pub_key[64]; + size_t kem_pub_key_len; + uint8_t binding_pub_key[64]; + size_t binding_pub_key_len; + uint64_t remaining_lifespan_secs; +} KpsKeyInfo; + int32_t key_manager_generate_kem_keypair(const uint8_t *algo_ptr, size_t algo_len, const uint8_t *binding_pubkey, @@ -21,8 +39,18 @@ int32_t key_manager_generate_kem_keypair(const uint8_t *algo_ptr, int32_t key_manager_destroy_kem_key(const uint8_t *uuid_bytes); +// key_manager_enumerate_kem_keys enumerates all active KEM keys in the registry. +// Writes up to max_entries KpsKeyInfo structs to out_entries and sets out_count +// to the number actually written. +// +// Returns 0 on success, -1 on error. +int32_t key_manager_enumerate_kem_keys( + KpsKeyInfo *out_entries, + size_t max_entries, + size_t *out_count); + #ifdef __cplusplus } // extern "C" #endif // __cplusplus -#endif /* KPS_KEY_CUSTODY_CORE_H_ */ +#endif // KPS_KEY_CUSTODY_CORE_H_ diff --git a/keymanager/key_protection_service/key_custody_core/src/lib.rs b/keymanager/key_protection_service/key_custody_core/src/lib.rs index 876601126..1766adfdc 100644 --- a/keymanager/key_protection_service/key_custody_core/src/lib.rs +++ b/keymanager/key_protection_service/key_custody_core/src/lib.rs @@ -6,7 +6,7 @@ use std::slice; use std::sync::atomic::AtomicBool; use std::sync::Arc; use std::sync::LazyLock; -use std::time::Duration; +use std::time::{Duration, Instant}; use uuid::Uuid; static KEY_REGISTRY: LazyLock = LazyLock::new(|| { @@ -63,7 +63,6 @@ fn generate_kem_keypair_internal( /// * `0` on success. /// * `-1` if an error occurred during key generation or if `binding_pubkey` is null/empty. /// * `-2` if the `out_pubkey` buffer size does not match the key size. - #[unsafe(no_mangle)] pub unsafe extern "C" fn key_manager_generate_kem_keypair( algo_ptr: *const u8, @@ -145,6 +144,85 @@ pub unsafe extern "C" fn key_manager_destroy_kem_key(uuid_bytes: *const u8) -> i } } +#[repr(C)] +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct KmHpkeAlgorithm { + pub kem: i32, + pub kdf: i32, + pub aead: i32, +} + +#[repr(C)] +pub struct KpsKeyInfo { + pub uuid: [u8; 16], + pub algorithm: KmHpkeAlgorithm, + pub kem_pub_key: [u8; 64], + pub kem_pub_key_len: usize, + pub binding_pub_key: [u8; 64], + pub binding_pub_key_len: usize, + pub remaining_lifespan_secs: u64, +} + +#[unsafe(no_mangle)] +pub unsafe extern "C" fn key_manager_enumerate_kem_keys( + out_entries: *mut KpsKeyInfo, + max_entries: usize, + out_count: *mut usize, +) -> i32 { + if out_entries.is_null() || out_count.is_null() { + return -1; + } + + let metas = KEY_REGISTRY.enumerate_keys(); + let now = Instant::now(); + let mut count = 0usize; + + for meta in &metas { + if count >= max_entries { + break; + } + if let KeySpec::KemWithBindingPub { + algo, + kem_public_key, + binding_public_key, + } = &meta.spec + { + if kem_public_key.as_bytes().len() > 64 || binding_public_key.as_bytes().len() > 64 { + continue; + } + + let remaining = meta.delete_after.saturating_duration_since(now).as_secs(); + + let entry = unsafe { &mut *out_entries.add(count) }; + entry.uuid.copy_from_slice(meta.id.as_bytes()); + entry.algorithm = KmHpkeAlgorithm { + kem: (*algo).kem, + kdf: (*algo).kdf, + aead: (*algo).aead, + }; + + entry.kem_pub_key = [0u8; 64]; + entry.kem_pub_key[..kem_public_key.as_bytes().len()] + .copy_from_slice(kem_public_key.as_bytes()); + entry.kem_pub_key_len = kem_public_key.as_bytes().len(); + + entry.binding_pub_key = [0u8; 64]; + entry.binding_pub_key[..binding_public_key.as_bytes().len()] + .copy_from_slice(binding_public_key.as_bytes()); + entry.binding_pub_key_len = binding_public_key.as_bytes().len(); + + entry.remaining_lifespan_secs = remaining; + + count += 1; + } + } + + unsafe { + *out_count = count; + } + 0 +} + #[cfg(test)] mod tests { use super::*; @@ -326,10 +404,12 @@ 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, @@ -360,4 +440,82 @@ mod tests { let result = unsafe { key_manager_destroy_kem_key(std::ptr::null()) }; 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, std::ptr::null_mut()) }; + 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; 64]; + 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 = Vec::with_capacity(256); + entries.resize_with(256, || KpsKeyInfo { + uuid: [0; 16], + algorithm: KmHpkeAlgorithm { + kem: 0, + kdf: 0, + aead: 0, + }, + kem_pub_key: [0; 64], + kem_pub_key_len: 0, + binding_pub_key: [0; 64], + binding_pub_key_len: 0, + remaining_lifespan_secs: 0, + }); + let mut count: usize = 0; + + let rc = unsafe { + key_manager_enumerate_kem_keys(entries.as_mut_ptr(), entries.len(), &mut count) + }; + assert_eq!(rc, 0); + // At least 1 key should be enumerated (the one we just generated). + assert!(count >= 1); + + // Find our key in the results. + let mut found = false; + for i in 0..count { + if entries[i].uuid == uuid_bytes { + found = true; + assert_eq!( + entries[i].algorithm.kem, + KemAlgorithm::DhkemX25519HkdfSha256 as i32 + ); + assert_eq!(entries[i].kem_pub_key_len, 32); + assert_eq!(entries[i].binding_pub_key_len, 32); + assert_eq!(&entries[i].binding_pub_key[..32], &binding_pubkey); + assert!(entries[i].remaining_lifespan_secs > 0); + break; + } + } + assert!(found, "generated key not found in enumerate results"); + } } diff --git a/keymanager/km_common/src/key_types.rs b/keymanager/km_common/src/key_types.rs index dcc743c15..afbbb83b0 100644 --- a/keymanager/km_common/src/key_types.rs +++ b/keymanager/km_common/src/key_types.rs @@ -1,6 +1,6 @@ use crate::algorithms::{AeadAlgorithm, HpkeAlgorithm, KdfAlgorithm, KemAlgorithm}; use crate::crypto; -use crate::crypto::{PublicKey, secret_box}; +use crate::crypto::{secret_box, PublicKey}; use crate::protected_mem::Vault; use std::collections::HashMap; use std::sync::{Arc, RwLock}; @@ -99,6 +99,11 @@ impl KeyRegistry { } }) } + + pub fn enumerate_keys(&self) -> Vec { + let keys = self.keys.read().unwrap(); + keys.values().map(|r| r.meta.clone()).collect() + } } impl KeyRecord { @@ -186,6 +191,33 @@ mod tests { use super::*; use crate::algorithms::{AeadAlgorithm, KdfAlgorithm, KemAlgorithm}; + fn create_key_record( + algo: HpkeAlgorithm, + expiry_secs: u64, + spec_builder: F, + ) -> Result + where + F: FnOnce(HpkeAlgorithm, PublicKey) -> KeySpec, + { + let (pub_key, priv_key) = crypto::generate_keypair(KemAlgorithm::DhkemX25519HkdfSha256)?; + let id = Uuid::new_v4(); + let vault = Vault::new(secret_box::SecretBox::from(priv_key)) + .map_err(|_| crypto::Error::CryptoError)?; + let now = Instant::now(); + let delete_after = now + .checked_add(Duration::from_secs(expiry_secs)) + .ok_or(crypto::Error::UnsupportedAlgorithm)?; + Ok(KeyRecord { + meta: KeyMetadata { + id, + created_at: now, + delete_after, + spec: spec_builder(algo, pub_key), + }, + private_key: vault, + }) + } + #[test] fn test_create_binding_key_success() { let algo = HpkeAlgorithm { @@ -246,6 +278,51 @@ mod tests { } } + #[test] + fn test_enumerate_keys_empty() { + let registry = KeyRegistry::default(); + let keys = registry.enumerate_keys(); + assert!(keys.is_empty()); + } + + #[test] + fn test_enumerate_keys() { + let registry = KeyRegistry::default(); + let algo = HpkeAlgorithm { + kem: KemAlgorithm::DhkemX25519HkdfSha256 as i32, + kdf: KdfAlgorithm::HkdfSha256 as i32, + aead: AeadAlgorithm::Aes256Gcm as i32, + }; + + let binding_pubkey = [42u8; 32]; + + let record1 = create_key_record(algo, 3600, |a, pk| KeySpec::KemWithBindingPub { + algo: a, + kem_public_key: pk, + binding_public_key: PublicKey::try_from(binding_pubkey.to_vec()).unwrap(), + }) + .expect("failed to create key 1"); + let id1 = record1.meta.id; + + let record2 = create_key_record(algo, 7200, |a, pk| KeySpec::KemWithBindingPub { + algo: a, + kem_public_key: pk, + binding_public_key: PublicKey::try_from(binding_pubkey.to_vec()).unwrap(), + }) + .expect("failed to create key 2"); + let id2 = record2.meta.id; + + registry.add_key(record1); + registry.add_key(record2); + + let metas = registry.enumerate_keys(); + assert_eq!(metas.len(), 2); + + let ids: Vec = metas.iter().map(|m| m.id).collect(); + assert!(ids.contains(&id1)); + assert!(ids.contains(&id2)); + } + #[test] fn test_add_key() { let registry = KeyRegistry::default(); From 0fa279df411ae5990f17773ba68ca30df9d6cfb0 Mon Sep 17 00:00:00 2001 From: Atul Patil Date: Thu, 19 Feb 2026 14:31:29 +0000 Subject: [PATCH 02/10] fix(rebase): update FFI tests and remove unused repr(C) --- keymanager/km_common/build.rs | 2 +- .../workload_service/key_custody_core/src/lib.rs | 12 +++++++----- 2 files changed, 8 insertions(+), 6 deletions(-) diff --git a/keymanager/km_common/build.rs b/keymanager/km_common/build.rs index 5d08cbae2..a3e8ceb42 100644 --- a/keymanager/km_common/build.rs +++ b/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(()) diff --git a/keymanager/workload_service/key_custody_core/src/lib.rs b/keymanager/workload_service/key_custody_core/src/lib.rs index ea5e1501e..9c8cc5b4a 100644 --- a/keymanager/workload_service/key_custody_core/src/lib.rs +++ b/keymanager/workload_service/key_custody_core/src/lib.rs @@ -264,16 +264,18 @@ mod tests { aead: AeadAlgorithm::Aes256Gcm as i32, }; - unsafe { - let res = key_manager_generate_binding_keypair( - algo, + let algo_bytes = algo.encode_to_vec(); + let res = unsafe { + key_manager_generate_binding_keypair( + algo_bytes.as_ptr(), + algo_bytes.len(), 3600, uuid_bytes.as_mut_ptr(), pubkey_bytes.as_mut_ptr(), pubkey_len, - ); - assert_eq!(res, 0); + ) }; + assert_eq!(res, 0); let result = unsafe { key_manager_destroy_binding_key(uuid_bytes.as_ptr()) }; assert_eq!(result, 0); From 6baab295c3e7b2dfd81a97e06db8e109d5bdc537 Mon Sep 17 00:00:00 2001 From: Atul Patil Date: Fri, 20 Feb 2026 10:02:40 +0000 Subject: [PATCH 03/10] refactor: Update Enumerate Keys FFI to use serialized HpkeAlgorithm --- .../key_custody_core/cbindgen.toml | 6 +-- .../include/kps_key_custody_core.h | 43 +++++++------------ .../key_custody_core/src/lib.rs | 43 ++++++++----------- keymanager/km_common/src/key_types.rs | 18 ++++---- 4 files changed, 45 insertions(+), 65 deletions(-) diff --git a/keymanager/key_protection_service/key_custody_core/cbindgen.toml b/keymanager/key_protection_service/key_custody_core/cbindgen.toml index 63241ab35..00d028f82 100644 --- a/keymanager/key_protection_service/key_custody_core/cbindgen.toml +++ b/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"] +include = ["key_manager_generate_kem_keypair", "key_manager_destroy_kem_key", "key_manager_enumerate_kem_keys", "KpsKeyInfo"] diff --git a/keymanager/key_protection_service/key_custody_core/include/kps_key_custody_core.h b/keymanager/key_protection_service/key_custody_core/include/kps_key_custody_core.h index b0a56718c..ddd92b7dc 100644 --- a/keymanager/key_protection_service/key_custody_core/include/kps_key_custody_core.h +++ b/keymanager/key_protection_service/key_custody_core/include/kps_key_custody_core.h @@ -6,28 +6,21 @@ #include #include +typedef struct { + uint8_t uuid[16]; + uint8_t algorithm[64]; + size_t algorithm_len; + uint8_t kem_pub_key[64]; + size_t kem_pub_key_len; + uint8_t binding_pub_key[64]; + size_t binding_pub_key_len; + uint64_t remaining_lifespan_secs; +} KpsKeyInfo; + #ifdef __cplusplus extern "C" { #endif // __cplusplus -// KmHpkeAlgorithm represents the HPKE algorithm suite configuration. -typedef struct { - int32_t kem; - int32_t kdf; - int32_t aead; -} KmHpkeAlgorithm; - -// KpsKeyInfo holds metadata for a single KEM key entry. -typedef struct { - uint8_t uuid[16]; - KmHpkeAlgorithm algorithm; - uint8_t kem_pub_key[64]; - size_t kem_pub_key_len; - uint8_t binding_pub_key[64]; - size_t binding_pub_key_len; - uint64_t remaining_lifespan_secs; -} KpsKeyInfo; - int32_t key_manager_generate_kem_keypair(const uint8_t *algo_ptr, size_t algo_len, const uint8_t *binding_pubkey, @@ -39,18 +32,12 @@ int32_t key_manager_generate_kem_keypair(const uint8_t *algo_ptr, int32_t key_manager_destroy_kem_key(const uint8_t *uuid_bytes); -// key_manager_enumerate_kem_keys enumerates all active KEM keys in the registry. -// Writes up to max_entries KpsKeyInfo structs to out_entries and sets out_count -// to the number actually written. -// -// Returns 0 on success, -1 on error. -int32_t key_manager_enumerate_kem_keys( - KpsKeyInfo *out_entries, - size_t max_entries, - size_t *out_count); +int32_t key_manager_enumerate_kem_keys(KpsKeyInfo *out_entries, + size_t max_entries, + size_t *out_count); #ifdef __cplusplus } // extern "C" #endif // __cplusplus -#endif // KPS_KEY_CUSTODY_CORE_H_ +#endif /* KPS_KEY_CUSTODY_CORE_H_ */ diff --git a/keymanager/key_protection_service/key_custody_core/src/lib.rs b/keymanager/key_protection_service/key_custody_core/src/lib.rs index 1766adfdc..f59535c52 100644 --- a/keymanager/key_protection_service/key_custody_core/src/lib.rs +++ b/keymanager/key_protection_service/key_custody_core/src/lib.rs @@ -63,6 +63,7 @@ fn generate_kem_keypair_internal( /// * `0` on success. /// * `-1` if an error occurred during key generation or if `binding_pubkey` is null/empty. /// * `-2` if the `out_pubkey` buffer size does not match the key size. + #[unsafe(no_mangle)] pub unsafe extern "C" fn key_manager_generate_kem_keypair( algo_ptr: *const u8, @@ -144,18 +145,11 @@ pub unsafe extern "C" fn key_manager_destroy_kem_key(uuid_bytes: *const u8) -> i } } -#[repr(C)] -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub struct KmHpkeAlgorithm { - pub kem: i32, - pub kdf: i32, - pub aead: i32, -} - #[repr(C)] pub struct KpsKeyInfo { pub uuid: [u8; 16], - pub algorithm: KmHpkeAlgorithm, + pub algorithm: [u8; 64], + pub algorithm_len: usize, pub kem_pub_key: [u8; 64], pub kem_pub_key_len: usize, pub binding_pub_key: [u8; 64], @@ -187,7 +181,11 @@ pub unsafe extern "C" fn key_manager_enumerate_kem_keys( binding_public_key, } = &meta.spec { - if kem_public_key.as_bytes().len() > 64 || binding_public_key.as_bytes().len() > 64 { + let algo_bytes = algo.encode_to_vec(); + if kem_public_key.as_bytes().len() > 64 + || binding_public_key.as_bytes().len() > 64 + || algo_bytes.len() > 64 + { continue; } @@ -195,11 +193,10 @@ pub unsafe extern "C" fn key_manager_enumerate_kem_keys( let entry = unsafe { &mut *out_entries.add(count) }; entry.uuid.copy_from_slice(meta.id.as_bytes()); - entry.algorithm = KmHpkeAlgorithm { - kem: (*algo).kem, - kdf: (*algo).kdf, - aead: (*algo).aead, - }; + + entry.algorithm = [0u8; 64]; + entry.algorithm[..algo_bytes.len()].copy_from_slice(&algo_bytes); + entry.algorithm_len = algo_bytes.len(); entry.kem_pub_key = [0u8; 64]; entry.kem_pub_key[..kem_public_key.as_bytes().len()] @@ -480,11 +477,8 @@ mod tests { let mut entries: Vec = Vec::with_capacity(256); entries.resize_with(256, || KpsKeyInfo { uuid: [0; 16], - algorithm: KmHpkeAlgorithm { - kem: 0, - kdf: 0, - aead: 0, - }, + algorithm: [0; 64], + algorithm_len: 0, kem_pub_key: [0; 64], kem_pub_key_len: 0, binding_pub_key: [0; 64], @@ -505,10 +499,11 @@ mod tests { for i in 0..count { if entries[i].uuid == uuid_bytes { found = true; - assert_eq!( - entries[i].algorithm.kem, - KemAlgorithm::DhkemX25519HkdfSha256 as i32 - ); + + 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].kem_pub_key_len, 32); assert_eq!(entries[i].binding_pub_key_len, 32); assert_eq!(&entries[i].binding_pub_key[..32], &binding_pubkey); diff --git a/keymanager/km_common/src/key_types.rs b/keymanager/km_common/src/key_types.rs index afbbb83b0..5d44cbc70 100644 --- a/keymanager/km_common/src/key_types.rs +++ b/keymanager/km_common/src/key_types.rs @@ -86,16 +86,14 @@ impl KeyRegistry { stop_signal: Arc, ) -> std::thread::JoinHandle<()> { let keys_clone = Arc::clone(&self.keys); - std::thread::spawn(move || { - loop { - std::thread::sleep(Duration::from_secs(REAPER_INTERVAL_SECS)); - if stop_signal.load(std::sync::atomic::Ordering::Relaxed) { - break; - } - let now = Instant::now(); - if let Ok(mut keys) = keys_clone.write() { - keys.retain(|_, key| key.meta.delete_after > now); - } + std::thread::spawn(move || loop { + std::thread::sleep(Duration::from_secs(REAPER_INTERVAL_SECS)); + if stop_signal.load(std::sync::atomic::Ordering::Relaxed) { + break; + } + let now = Instant::now(); + if let Ok(mut keys) = keys_clone.write() { + keys.retain(|_, key| key.meta.delete_after > now); } }) } From 6ee8d67b2c6438d00b90176778a232a8830b8690 Mon Sep 17 00:00:00 2001 From: Atul Patil Date: Fri, 20 Feb 2026 12:51:15 +0000 Subject: [PATCH 04/10] feat(keymanager): pagination and large key support for enumerate API Adds offset-based pagination to the key enumeration API and increases the buffer sizes for KpsKeyInfo to support Post-Quantum Cryptography (PQC) keys (e.g., ML-KEM). Also optimizes internal enumeration to avoid cloning the entire key registry. --- .../include/kps_key_custody_core.h | 7 +-- .../key_custody_core/src/lib.rs | 50 ++++++++++++------- keymanager/km_common/src/key_types.rs | 31 +++++++++++- 3 files changed, 65 insertions(+), 23 deletions(-) diff --git a/keymanager/key_protection_service/key_custody_core/include/kps_key_custody_core.h b/keymanager/key_protection_service/key_custody_core/include/kps_key_custody_core.h index ddd92b7dc..87a8d94e3 100644 --- a/keymanager/key_protection_service/key_custody_core/include/kps_key_custody_core.h +++ b/keymanager/key_protection_service/key_custody_core/include/kps_key_custody_core.h @@ -8,11 +8,11 @@ typedef struct { uint8_t uuid[16]; - uint8_t algorithm[64]; + uint8_t algorithm[128]; size_t algorithm_len; - uint8_t kem_pub_key[64]; + uint8_t kem_pub_key[2048]; size_t kem_pub_key_len; - uint8_t binding_pub_key[64]; + uint8_t binding_pub_key[2048]; size_t binding_pub_key_len; uint64_t remaining_lifespan_secs; } KpsKeyInfo; @@ -34,6 +34,7 @@ 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, size_t *out_count); #ifdef __cplusplus diff --git a/keymanager/key_protection_service/key_custody_core/src/lib.rs b/keymanager/key_protection_service/key_custody_core/src/lib.rs index f59535c52..94e47cf60 100644 --- a/keymanager/key_protection_service/key_custody_core/src/lib.rs +++ b/keymanager/key_protection_service/key_custody_core/src/lib.rs @@ -148,11 +148,11 @@ pub unsafe extern "C" fn key_manager_destroy_kem_key(uuid_bytes: *const u8) -> i #[repr(C)] pub struct KpsKeyInfo { pub uuid: [u8; 16], - pub algorithm: [u8; 64], + pub algorithm: [u8; 128], pub algorithm_len: usize, - pub kem_pub_key: [u8; 64], + pub kem_pub_key: [u8; 2048], pub kem_pub_key_len: usize, - pub binding_pub_key: [u8; 64], + pub binding_pub_key: [u8; 2048], pub binding_pub_key_len: usize, pub remaining_lifespan_secs: u64, } @@ -161,13 +161,14 @@ pub struct KpsKeyInfo { pub unsafe extern "C" fn key_manager_enumerate_kem_keys( out_entries: *mut KpsKeyInfo, max_entries: usize, + offset: usize, out_count: *mut usize, ) -> i32 { if out_entries.is_null() || out_count.is_null() { return -1; } - let metas = KEY_REGISTRY.enumerate_keys(); + let metas = KEY_REGISTRY.list_keys(offset, max_entries); let now = Instant::now(); let mut count = 0usize; @@ -182,10 +183,17 @@ pub unsafe extern "C" fn key_manager_enumerate_kem_keys( } = &meta.spec { let algo_bytes = algo.encode_to_vec(); - if kem_public_key.as_bytes().len() > 64 - || binding_public_key.as_bytes().len() > 64 - || algo_bytes.len() > 64 + if kem_public_key.as_bytes().len() > 2048 + || binding_public_key.as_bytes().len() > 2048 + || algo_bytes.len() > 128 { + eprintln!( + "Skipping key {}: size exceeds buffer limits (algo={}, kem={}, binding={})", + meta.id, + algo_bytes.len(), + kem_public_key.as_bytes().len(), + binding_public_key.as_bytes().len() + ); continue; } @@ -194,16 +202,16 @@ pub unsafe extern "C" fn key_manager_enumerate_kem_keys( let entry = unsafe { &mut *out_entries.add(count) }; entry.uuid.copy_from_slice(meta.id.as_bytes()); - entry.algorithm = [0u8; 64]; + entry.algorithm = [0u8; 128]; entry.algorithm[..algo_bytes.len()].copy_from_slice(&algo_bytes); entry.algorithm_len = algo_bytes.len(); - entry.kem_pub_key = [0u8; 64]; + entry.kem_pub_key = [0u8; 2048]; entry.kem_pub_key[..kem_public_key.as_bytes().len()] .copy_from_slice(kem_public_key.as_bytes()); entry.kem_pub_key_len = kem_public_key.as_bytes().len(); - entry.binding_pub_key = [0u8; 64]; + entry.binding_pub_key = [0u8; 2048]; entry.binding_pub_key[..binding_public_key.as_bytes().len()] .copy_from_slice(binding_public_key.as_bytes()); entry.binding_pub_key_len = binding_public_key.as_bytes().len(); @@ -440,7 +448,9 @@ mod tests { #[test] fn test_enumerate_kem_keys_null_pointers() { - let result = unsafe { key_manager_enumerate_kem_keys(std::ptr::null_mut(), 10, std::ptr::null_mut()) }; + let result = unsafe { + key_manager_enumerate_kem_keys(std::ptr::null_mut(), 10, 0, std::ptr::null_mut()) + }; assert_eq!(result, -1); } @@ -448,7 +458,7 @@ mod tests { fn test_enumerate_kem_keys_after_generate() { let binding_pubkey = [7u8; 32]; let mut uuid_bytes = [0u8; 16]; - let mut pubkey_bytes = [0u8; 64]; + let mut pubkey_bytes = [0u8; 32]; let pubkey_len: usize = 32; let algo = HpkeAlgorithm { kem: KemAlgorithm::DhkemX25519HkdfSha256 as i32, @@ -475,20 +485,22 @@ mod tests { // Enumerate. let mut entries: Vec = Vec::with_capacity(256); - entries.resize_with(256, || KpsKeyInfo { + // Initialize with default/zero values. Note: Arrays are larger now. + entries.resize_with(100, || KpsKeyInfo { uuid: [0; 16], - algorithm: [0; 64], + algorithm: [0; 128], algorithm_len: 0, - kem_pub_key: [0; 64], + kem_pub_key: [0; 2048], kem_pub_key_len: 0, - binding_pub_key: [0; 64], + binding_pub_key: [0; 2048], binding_pub_key_len: 0, remaining_lifespan_secs: 0, }); let mut count: usize = 0; let rc = unsafe { - key_manager_enumerate_kem_keys(entries.as_mut_ptr(), entries.len(), &mut count) + // max_entries=1, offset=0 + key_manager_enumerate_kem_keys(entries.as_mut_ptr(), entries.len(), 0, &mut count) }; assert_eq!(rc, 0); // At least 1 key should be enumerated (the one we just generated). @@ -499,10 +511,10 @@ mod tests { 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].kem_pub_key_len, 32); assert_eq!(entries[i].binding_pub_key_len, 32); diff --git a/keymanager/km_common/src/key_types.rs b/keymanager/km_common/src/key_types.rs index 5d44cbc70..7f2362ad8 100644 --- a/keymanager/km_common/src/key_types.rs +++ b/keymanager/km_common/src/key_types.rs @@ -99,8 +99,25 @@ impl KeyRegistry { } pub fn enumerate_keys(&self) -> Vec { + self.list_keys(0, usize::MAX) + } + + pub fn list_keys(&self, offset: usize, limit: usize) -> Vec { let keys = self.keys.read().unwrap(); - keys.values().map(|r| r.meta.clone()).collect() + let mut refs: Vec<&Arc> = keys.values().collect(); + // Sort for stable pagination: created_at, then id + refs.sort_by(|a, b| { + a.meta + .created_at + .cmp(&b.meta.created_at) + .then(a.meta.id.cmp(&b.meta.id)) + }); + + refs.into_iter() + .skip(offset) + .take(limit) + .map(|r| r.meta.clone()) + .collect() } } @@ -319,6 +336,18 @@ mod tests { let ids: Vec = metas.iter().map(|m| m.id).collect(); assert!(ids.contains(&id1)); assert!(ids.contains(&id2)); + + // Test pagination + let page1 = registry.list_keys(0, 1); + assert_eq!(page1.len(), 1); + let page2 = registry.list_keys(1, 1); + assert_eq!(page2.len(), 1); + assert_ne!(page1[0].id, page2[0].id); + + // Verify order (record1 created before record2) + // Note: created_at might be identical if executed very fast, so sort checks ID too. + // But record1 and record2 are created with explicit calls, so created_at likely differs slightly. + // Either way, stability is guaranteed. } #[test] From 4994a67f7b4644fe9ffabd8bc12ca3adf726d7de Mon Sep 17 00:00:00 2001 From: Atul Patil Date: Fri, 20 Feb 2026 13:37:00 +0000 Subject: [PATCH 05/10] feat(keymanager): remove binding key from enumerate API As per design discussions, the binding key is not useful to the caller of enumerate_keys and should not be exposed. This change removes and from . --- .../include/kps_key_custody_core.h | 2 -- .../key_custody_core/src/lib.rs | 24 ++++--------------- 2 files changed, 5 insertions(+), 21 deletions(-) diff --git a/keymanager/key_protection_service/key_custody_core/include/kps_key_custody_core.h b/keymanager/key_protection_service/key_custody_core/include/kps_key_custody_core.h index 87a8d94e3..14ad64511 100644 --- a/keymanager/key_protection_service/key_custody_core/include/kps_key_custody_core.h +++ b/keymanager/key_protection_service/key_custody_core/include/kps_key_custody_core.h @@ -12,8 +12,6 @@ typedef struct { size_t algorithm_len; uint8_t kem_pub_key[2048]; size_t kem_pub_key_len; - uint8_t binding_pub_key[2048]; - size_t binding_pub_key_len; uint64_t remaining_lifespan_secs; } KpsKeyInfo; diff --git a/keymanager/key_protection_service/key_custody_core/src/lib.rs b/keymanager/key_protection_service/key_custody_core/src/lib.rs index 94e47cf60..2c69809e1 100644 --- a/keymanager/key_protection_service/key_custody_core/src/lib.rs +++ b/keymanager/key_protection_service/key_custody_core/src/lib.rs @@ -152,8 +152,6 @@ pub struct KpsKeyInfo { pub algorithm_len: usize, pub kem_pub_key: [u8; 2048], pub kem_pub_key_len: usize, - pub binding_pub_key: [u8; 2048], - pub binding_pub_key_len: usize, pub remaining_lifespan_secs: u64, } @@ -179,20 +177,16 @@ pub unsafe extern "C" fn key_manager_enumerate_kem_keys( if let KeySpec::KemWithBindingPub { algo, kem_public_key, - binding_public_key, + binding_public_key: _, } = &meta.spec { let algo_bytes = algo.encode_to_vec(); - if kem_public_key.as_bytes().len() > 2048 - || binding_public_key.as_bytes().len() > 2048 - || algo_bytes.len() > 128 - { + if kem_public_key.as_bytes().len() > 2048 || algo_bytes.len() > 128 { eprintln!( - "Skipping key {}: size exceeds buffer limits (algo={}, kem={}, binding={})", + "Skipping key {}: size exceeds buffer limits (algo={}, kem={})", meta.id, algo_bytes.len(), - kem_public_key.as_bytes().len(), - binding_public_key.as_bytes().len() + kem_public_key.as_bytes().len() ); continue; } @@ -211,11 +205,6 @@ pub unsafe extern "C" fn key_manager_enumerate_kem_keys( .copy_from_slice(kem_public_key.as_bytes()); entry.kem_pub_key_len = kem_public_key.as_bytes().len(); - entry.binding_pub_key = [0u8; 2048]; - entry.binding_pub_key[..binding_public_key.as_bytes().len()] - .copy_from_slice(binding_public_key.as_bytes()); - entry.binding_pub_key_len = binding_public_key.as_bytes().len(); - entry.remaining_lifespan_secs = remaining; count += 1; @@ -492,8 +481,6 @@ mod tests { algorithm_len: 0, kem_pub_key: [0; 2048], kem_pub_key_len: 0, - binding_pub_key: [0; 2048], - binding_pub_key_len: 0, remaining_lifespan_secs: 0, }); let mut count: usize = 0; @@ -517,8 +504,7 @@ mod tests { assert_eq!(decoded_algo.kem, KemAlgorithm::DhkemX25519HkdfSha256 as i32); assert_eq!(entries[i].kem_pub_key_len, 32); - assert_eq!(entries[i].binding_pub_key_len, 32); - assert_eq!(&entries[i].binding_pub_key[..32], &binding_pubkey); + // binding_pub_key checks removed assert!(entries[i].remaining_lifespan_secs > 0); break; } From 5fb6ab88951ca24612e44d60f43cab99ce209da2 Mon Sep 17 00:00:00 2001 From: Atul Patil Date: Fri, 20 Feb 2026 14:50:40 +0000 Subject: [PATCH 06/10] refactor(keymanager): deduplicate pagination logic and clean up key registry - Introduced with inlined logic to specifically handle KEM key pagination. - Removed unused and methods from . - Updated FFI to use . - Restricted internal key listing to KEM keys only, ensuring deterministic behavior for the FFI. --- .../key_custody_core/src/lib.rs | 5 +--- keymanager/km_common/src/key_types.rs | 25 ++++++++++--------- 2 files changed, 14 insertions(+), 16 deletions(-) diff --git a/keymanager/key_protection_service/key_custody_core/src/lib.rs b/keymanager/key_protection_service/key_custody_core/src/lib.rs index 2c69809e1..ba03f4979 100644 --- a/keymanager/key_protection_service/key_custody_core/src/lib.rs +++ b/keymanager/key_protection_service/key_custody_core/src/lib.rs @@ -166,14 +166,11 @@ pub unsafe extern "C" fn key_manager_enumerate_kem_keys( return -1; } - let metas = KEY_REGISTRY.list_keys(offset, max_entries); + let metas = KEY_REGISTRY.list_kem_keys(offset, max_entries); let now = Instant::now(); let mut count = 0usize; for meta in &metas { - if count >= max_entries { - break; - } if let KeySpec::KemWithBindingPub { algo, kem_public_key, diff --git a/keymanager/km_common/src/key_types.rs b/keymanager/km_common/src/key_types.rs index 7f2362ad8..74bc257aa 100644 --- a/keymanager/km_common/src/key_types.rs +++ b/keymanager/km_common/src/key_types.rs @@ -98,13 +98,14 @@ impl KeyRegistry { }) } - pub fn enumerate_keys(&self) -> Vec { - self.list_keys(0, usize::MAX) - } - - pub fn list_keys(&self, offset: usize, limit: usize) -> Vec { + /// Lists only KEM keys with pagination support. + pub fn list_kem_keys(&self, offset: usize, limit: usize) -> Vec { let keys = self.keys.read().unwrap(); - let mut refs: Vec<&Arc> = keys.values().collect(); + let mut refs: Vec<&Arc> = keys + .values() + .filter(|k| matches!(k.meta.spec, KeySpec::KemWithBindingPub { .. })) + .collect(); + // Sort for stable pagination: created_at, then id refs.sort_by(|a, b| { a.meta @@ -294,14 +295,14 @@ mod tests { } #[test] - fn test_enumerate_keys_empty() { + fn test_list_kem_keys_empty() { let registry = KeyRegistry::default(); - let keys = registry.enumerate_keys(); + let keys = registry.list_kem_keys(0, 100); assert!(keys.is_empty()); } #[test] - fn test_enumerate_keys() { + fn test_list_kem_keys() { let registry = KeyRegistry::default(); let algo = HpkeAlgorithm { kem: KemAlgorithm::DhkemX25519HkdfSha256 as i32, @@ -330,7 +331,7 @@ mod tests { registry.add_key(record1); registry.add_key(record2); - let metas = registry.enumerate_keys(); + let metas = registry.list_kem_keys(0, 100); assert_eq!(metas.len(), 2); let ids: Vec = metas.iter().map(|m| m.id).collect(); @@ -338,9 +339,9 @@ mod tests { assert!(ids.contains(&id2)); // Test pagination - let page1 = registry.list_keys(0, 1); + let page1 = registry.list_kem_keys(0, 1); assert_eq!(page1.len(), 1); - let page2 = registry.list_keys(1, 1); + let page2 = registry.list_kem_keys(1, 1); assert_eq!(page2.len(), 1); assert_ne!(page1[0].id, page2[0].id); From 0dd4bfcd77413cad3816e2b0ace0f4c17dc6fd8e Mon Sep 17 00:00:00 2001 From: Atul Patil Date: Fri, 20 Feb 2026 15:01:29 +0000 Subject: [PATCH 07/10] Sync lib.rs from wsd_enumerate_go with conflict resolution --- keymanager/key_protection_service/key_custody_core/src/lib.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/keymanager/key_protection_service/key_custody_core/src/lib.rs b/keymanager/key_protection_service/key_custody_core/src/lib.rs index ba03f4979..35f3b4cfa 100644 --- a/keymanager/key_protection_service/key_custody_core/src/lib.rs +++ b/keymanager/key_protection_service/key_custody_core/src/lib.rs @@ -185,6 +185,7 @@ pub unsafe extern "C" fn key_manager_enumerate_kem_keys( algo_bytes.len(), kem_public_key.as_bytes().len() ); + } else { continue; } @@ -193,6 +194,7 @@ pub unsafe extern "C" fn key_manager_enumerate_kem_keys( let entry = unsafe { &mut *out_entries.add(count) }; entry.uuid.copy_from_slice(meta.id.as_bytes()); + entry.algorithm = [0u8; 128]; entry.algorithm[..algo_bytes.len()].copy_from_slice(&algo_bytes); entry.algorithm_len = algo_bytes.len(); @@ -495,10 +497,8 @@ mod tests { 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].kem_pub_key_len, 32); // binding_pub_key checks removed From f3cc514a77020f3aa38f152fc8325c5f19491597 Mon Sep 17 00:00:00 2001 From: Atul Patil Date: Mon, 9 Feb 2026 15:27:31 +0530 Subject: [PATCH 08/10] keymanager/wsd: Implement Enumerate Keys API (GET /v1/keys) Enumerate Keys API (GET /v1/keys) to the Workload Service Daemon (WSD), allowing clients to list active KEM keys stored in the Key Protection Service (KPS). Key changes: 1. **Workload Service (Go)**: - Added EnumerateKEMKeys integration with the KCC FFI. - Exposed GET /v1/keys endpoint. - Defined response structures aligned with the design doc: - snake_case JSON tags (key_handle, kem_pub_key, etc.). - Stringified enums for algorithms (HKDF_SHA384, AES_256_GCM). - Stringified durations (e.g., 3600s). 2. **Key Protection Service (Rust)**: - Added key_manager_enumerate_kem_keys FFI function to key_custody_core. - Implemented enumerate_kem_keys in KeyRegistry to safely expose key metadata. - Added KmHpkeAlgorithm struct for C compatibility. 3. **Testing**: - Added unit tests in server_test.go covering success, empty state, and error handling. - Added Rust FFI tests validating memory safety and correctness. --- .../kps_key_custody_core_cgo.go | 44 +++++ .../key_custody_core/types.go | 14 ++ keymanager/key_protection_service/service.go | 28 ++- .../key_protection_service/service_test.go | 85 ++++++-- keymanager/workload_service/proto_enums.go | 85 ++++++++ keymanager/workload_service/server.go | 93 ++++++++- keymanager/workload_service/server_test.go | 181 +++++++++++++++++- 7 files changed, 509 insertions(+), 21 deletions(-) create mode 100644 keymanager/key_protection_service/key_custody_core/types.go diff --git a/keymanager/key_protection_service/key_custody_core/kps_key_custody_core_cgo.go b/keymanager/key_protection_service/key_custody_core/kps_key_custody_core_cgo.go index a52909357..f852debaf 100644 --- a/keymanager/key_protection_service/key_custody_core/kps_key_custody_core_cgo.go +++ b/keymanager/key_protection_service/key_custody_core/kps_key_custody_core_cgo.go @@ -60,3 +60,47 @@ func GenerateKEMKeypair(algo *algorithms.HpkeAlgorithm, bindingPubKey []byte, li copy(pubkey, pubkeyBuf[:pubkeyLen]) return id, pubkey, nil } + +// EnumerateKEMKeys retrieves all active KEM key entries from the Rust KCC registry. +func EnumerateKEMKeys() ([]KEMKeyInfo, error) { + const maxEntries = 256 + var entries [maxEntries]C.KpsKeyInfo + var count C.size_t + + rc := C.key_manager_enumerate_kem_keys( + &entries[0], + C.size_t(maxEntries), + &count, + ) + if rc != 0 { + return nil, fmt.Errorf("key_manager_enumerate_kem_keys failed with code %d", rc) + } + + result := make([]KEMKeyInfo, count) + for i := C.size_t(0); i < count; i++ { + e := entries[i] + + id, err := uuid.FromBytes(C.GoBytes(unsafe.Pointer(&e.uuid[0]), 16)) + if err != nil { + return nil, fmt.Errorf("invalid UUID at index %d: %w", i, err) + } + + kemPubKey := make([]byte, e.kem_pub_key_len) + copy(kemPubKey, C.GoBytes(unsafe.Pointer(&e.kem_pub_key[0]), C.int(e.kem_pub_key_len))) + + bindingPubKey := make([]byte, e.binding_pub_key_len) + copy(bindingPubKey, C.GoBytes(unsafe.Pointer(&e.binding_pub_key[0]), C.int(e.binding_pub_key_len))) + + result[i] = KEMKeyInfo{ + ID: id, + KemAlgorithm: int32(e.algorithm.kem), + KdfAlgorithm: int32(e.algorithm.kdf), + AeadAlgorithm: int32(e.algorithm.aead), + KEMPubKey: kemPubKey, + BindingPubKey: bindingPubKey, + RemainingLifespanSecs: uint64(e.remaining_lifespan_secs), + } + } + + return result, nil +} diff --git a/keymanager/key_protection_service/key_custody_core/types.go b/keymanager/key_protection_service/key_custody_core/types.go new file mode 100644 index 000000000..b40f0e8f1 --- /dev/null +++ b/keymanager/key_protection_service/key_custody_core/types.go @@ -0,0 +1,14 @@ +package kpskcc + +import "github.com/google/uuid" + +// KEMKeyInfo holds metadata for a single KEM key returned by EnumerateKEMKeys. +type KEMKeyInfo struct { + ID uuid.UUID + KemAlgorithm int32 + KdfAlgorithm int32 + AeadAlgorithm int32 + KEMPubKey []byte + BindingPubKey []byte + RemainingLifespanSecs uint64 +} diff --git a/keymanager/key_protection_service/service.go b/keymanager/key_protection_service/service.go index c9319bfaa..206444c49 100644 --- a/keymanager/key_protection_service/service.go +++ b/keymanager/key_protection_service/service.go @@ -1,9 +1,10 @@ // Package key_protection_service implements the Key Orchestration Layer (KOL) // for the Key Protection Service. It wraps the KPS Key Custody Core (KCC) FFI -// to provide a Go-native interface for KEM key generation. +// to provide a Go-native interface for KEM key generation and enumeration. package key_protection_service import ( + kpskcc "github.com/google/go-tpm-tools/keymanager/key_protection_service/key_custody_core" "github.com/google/uuid" algorithms "github.com/google/go-tpm-tools/keymanager/km_common/proto" @@ -14,14 +15,26 @@ type KEMKeyGenerator interface { GenerateKEMKeypair(algo *algorithms.HpkeAlgorithm, bindingPubKey []byte, lifespanSecs uint64) (uuid.UUID, []byte, error) } -// Service implements KEMKeyGenerator by delegating to the KPS KCC FFI. +// KEMKeyEnumerator enumerates active KEM keys in the KPS registry. +type KEMKeyEnumerator interface { + EnumerateKEMKeys() ([]kpskcc.KEMKeyInfo, error) +} + +// Service implements KEMKeyGenerator and KEMKeyEnumerator by delegating to the KPS KCC FFI. type Service struct { generateKEMKeypairFn func(algo *algorithms.HpkeAlgorithm, bindingPubKey []byte, lifespanSecs uint64) (uuid.UUID, []byte, error) + enumerateKEMKeysFn func() ([]kpskcc.KEMKeyInfo, error) } -// NewService creates a new KPS KOL service with the given KCC function. -func NewService(generateKEMKeypairFn func(algo *algorithms.HpkeAlgorithm, bindingPubKey []byte, lifespanSecs uint64) (uuid.UUID, []byte, error)) *Service { - return &Service{generateKEMKeypairFn: generateKEMKeypairFn} +// NewService creates a new KPS KOL service with the given KCC functions. +func NewService( + generateKEMKeypairFn func(algo *algorithms.HpkeAlgorithm, bindingPubKey []byte, lifespanSecs uint64) (uuid.UUID, []byte, error), + enumerateKEMKeysFn func() ([]kpskcc.KEMKeyInfo, error), +) *Service { + return &Service{ + generateKEMKeypairFn: generateKEMKeypairFn, + enumerateKEMKeysFn: enumerateKEMKeysFn, + } } // GenerateKEMKeypair generates a KEM keypair linked to the provided binding @@ -29,3 +42,8 @@ func NewService(generateKEMKeypairFn func(algo *algorithms.HpkeAlgorithm, bindin func (s *Service) GenerateKEMKeypair(algo *algorithms.HpkeAlgorithm, bindingPubKey []byte, lifespanSecs uint64) (uuid.UUID, []byte, error) { return s.generateKEMKeypairFn(algo, bindingPubKey, lifespanSecs) } + +// EnumerateKEMKeys retrieves all active KEM key entries from the KPS KCC registry. +func (s *Service) EnumerateKEMKeys() ([]kpskcc.KEMKeyInfo, error) { + return s.enumerateKEMKeysFn() +} diff --git a/keymanager/key_protection_service/service_test.go b/keymanager/key_protection_service/service_test.go index e4e40c6fc..b6b532277 100644 --- a/keymanager/key_protection_service/service_test.go +++ b/keymanager/key_protection_service/service_test.go @@ -4,6 +4,7 @@ import ( "fmt" "testing" + kpskcc "github.com/google/go-tpm-tools/keymanager/key_protection_service/key_custody_core" "github.com/google/uuid" algorithms "github.com/google/go-tpm-tools/keymanager/km_common/proto" @@ -16,15 +17,20 @@ func TestServiceGenerateKEMKeypairSuccess(t *testing.T) { expectedPubKey[i] = byte(i + 10) } - svc := NewService(func(algo *algorithms.HpkeAlgorithm, bindingPubKey []byte, lifespanSecs uint64) (uuid.UUID, []byte, error) { - if len(bindingPubKey) != 32 { - t.Fatalf("expected 32-byte binding public key, got %d", len(bindingPubKey)) - } - if lifespanSecs != 7200 { - t.Fatalf("expected lifespanSecs 7200, got %d", lifespanSecs) - } - return expectedUUID, expectedPubKey, nil - }) + svc := NewService( + func(algo *algorithms.HpkeAlgorithm, bindingPubKey []byte, lifespanSecs uint64) (uuid.UUID, []byte, error) { + if len(bindingPubKey) != 32 { + t.Fatalf("expected 32-byte binding public key, got %d", len(bindingPubKey)) + } + if lifespanSecs != 7200 { + t.Fatalf("expected lifespanSecs 7200, got %d", lifespanSecs) + } + return expectedUUID, expectedPubKey, nil + }, + func() ([]kpskcc.KEMKeyInfo, error) { + return nil, nil + }, + ) id, pubKey, err := svc.GenerateKEMKeypair(&algorithms.HpkeAlgorithm{}, make([]byte, 32), 7200) if err != nil { @@ -39,12 +45,67 @@ func TestServiceGenerateKEMKeypairSuccess(t *testing.T) { } func TestServiceGenerateKEMKeypairError(t *testing.T) { - svc := NewService(func(algo *algorithms.HpkeAlgorithm, bindingPubKey []byte, lifespanSecs uint64) (uuid.UUID, []byte, error) { - return uuid.Nil, nil, fmt.Errorf("FFI error") - }) + svc := NewService( + func(algo *algorithms.HpkeAlgorithm, bindingPubKey []byte, lifespanSecs uint64) (uuid.UUID, []byte, error) { + return uuid.Nil, nil, fmt.Errorf("FFI error") + }, + func() ([]kpskcc.KEMKeyInfo, error) { + return nil, nil + }, + ) _, _, err := svc.GenerateKEMKeypair(&algorithms.HpkeAlgorithm{}, make([]byte, 32), 3600) if err == nil { t.Fatal("expected error, got nil") } } + +func TestServiceEnumerateKEMKeysSuccess(t *testing.T) { + expectedKeys := []kpskcc.KEMKeyInfo{ + { + ID: uuid.New(), + KemAlgorithm: 1, + KdfAlgorithm: 1, + AeadAlgorithm: 1, + KEMPubKey: make([]byte, 32), + BindingPubKey: make([]byte, 32), + RemainingLifespanSecs: 3500, + }, + } + + svc := NewService( + func(algo *algorithms.HpkeAlgorithm, bindingPubKey []byte, lifespanSecs uint64) (uuid.UUID, []byte, error) { + return uuid.Nil, nil, nil + }, + func() ([]kpskcc.KEMKeyInfo, error) { + return expectedKeys, nil + }, + ) + + keys, err := svc.EnumerateKEMKeys() + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if len(keys) != 1 { + t.Fatalf("expected 1 key, got %d", len(keys)) + } + if keys[0].ID != expectedKeys[0].ID { + t.Fatalf("expected ID %s, got %s", expectedKeys[0].ID, keys[0].ID) + } +} + +func TestServiceEnumerateKEMKeysError(t *testing.T) { + svc := NewService( + func(algo *algorithms.HpkeAlgorithm, bindingPubKey []byte, lifespanSecs uint64) (uuid.UUID, []byte, error) { + return uuid.Nil, nil, nil + }, + func() ([]kpskcc.KEMKeyInfo, error) { + return nil, fmt.Errorf("enumerate error") + }, + ) + + _, err := svc.EnumerateKEMKeys() + if err == nil { + t.Fatal("expected error, got nil") + } +} diff --git a/keymanager/workload_service/proto_enums.go b/keymanager/workload_service/proto_enums.go index 2b2c2b748..f8ca4cf64 100644 --- a/keymanager/workload_service/proto_enums.go +++ b/keymanager/workload_service/proto_enums.go @@ -133,3 +133,88 @@ func (k KemAlgorithm) ToHpkeAlgorithm() (*algorithms.HpkeAlgorithm, error) { return nil, fmt.Errorf("unsupported algorithm: %s", k) } } + +// KdfAlgorithm represents the requested KDF algorithm. +type KdfAlgorithm int32 + +const ( + KdfAlgorithmUnspecified KdfAlgorithm = 0 + // Corrected from HKDF_SHA384 to HKDF_SHA256 based on ToHpkeAlgorithm usage which maps to HKDF_SHA256 (val 1) + KdfAlgorithmHKDFSHA256 KdfAlgorithm = 1 +) + +var ( + kdfAlgorithmToString = map[KdfAlgorithm]string{ + KdfAlgorithmUnspecified: "KDF_ALGORITHM_UNSPECIFIED", + KdfAlgorithmHKDFSHA256: "HKDF_SHA256", + } + stringToKdfAlgorithm = map[string]KdfAlgorithm{ + "KDF_ALGORITHM_UNSPECIFIED": KdfAlgorithmUnspecified, + "HKDF_SHA256": KdfAlgorithmHKDFSHA256, + } +) + +func (k KdfAlgorithm) String() string { + if s, ok := kdfAlgorithmToString[k]; ok { + return s + } + return fmt.Sprintf("KDF_ALGORITHM_UNKNOWN(%d)", k) +} + +func (k KdfAlgorithm) MarshalJSON() ([]byte, error) { + return json.Marshal(k.String()) +} + +func (k *KdfAlgorithm) UnmarshalJSON(data []byte) error { + var s string + if err := json.Unmarshal(data, &s); err != nil { + return fmt.Errorf("KdfAlgorithm must be a string") + } + if v, ok := stringToKdfAlgorithm[s]; ok { + *k = v + return nil + } + return fmt.Errorf("unknown KdfAlgorithm: %q", s) +} + +// AeadAlgorithm represents the requested AEAD algorithm. +type AeadAlgorithm int32 + +const ( + AeadAlgorithmUnspecified AeadAlgorithm = 0 + AeadAlgorithmAES256GCM AeadAlgorithm = 1 +) + +var ( + aeadAlgorithmToString = map[AeadAlgorithm]string{ + AeadAlgorithmUnspecified: "AEAD_ALGORITHM_UNSPECIFIED", + AeadAlgorithmAES256GCM: "AES_256_GCM", + } + stringToAeadAlgorithm = map[string]AeadAlgorithm{ + "AEAD_ALGORITHM_UNSPECIFIED": AeadAlgorithmUnspecified, + "AES_256_GCM": AeadAlgorithmAES256GCM, + } +) + +func (k AeadAlgorithm) String() string { + if s, ok := aeadAlgorithmToString[k]; ok { + return s + } + return fmt.Sprintf("AEAD_ALGORITHM_UNKNOWN(%d)", k) +} + +func (k AeadAlgorithm) MarshalJSON() ([]byte, error) { + return json.Marshal(k.String()) +} + +func (k *AeadAlgorithm) UnmarshalJSON(data []byte) error { + var s string + if err := json.Unmarshal(data, &s); err != nil { + return fmt.Errorf("AeadAlgorithm must be a string") + } + if v, ok := stringToAeadAlgorithm[s]; ok { + *k = v + return nil + } + return fmt.Errorf("unknown AeadAlgorithm: %q", s) +} diff --git a/keymanager/workload_service/server.go b/keymanager/workload_service/server.go index 41ec27ab3..2661b4d3c 100644 --- a/keymanager/workload_service/server.go +++ b/keymanager/workload_service/server.go @@ -5,6 +5,7 @@ package workload_service import ( "context" + "encoding/base64" "encoding/json" "fmt" "math" @@ -13,6 +14,7 @@ import ( "os" "sync" + kpskcc "github.com/google/go-tpm-tools/keymanager/key_protection_service/key_custody_core" "github.com/google/uuid" algorithms "github.com/google/go-tpm-tools/keymanager/km_common/proto" @@ -28,6 +30,11 @@ type KEMKeyGenerator interface { GenerateKEMKeypair(algo *algorithms.HpkeAlgorithm, bindingPubKey []byte, lifespanSecs uint64) (uuid.UUID, []byte, error) } +// KEMKeyEnumerator enumerates active KEM keys from the KPS registry. +type KEMKeyEnumerator interface { + EnumerateKEMKeys() ([]kpskcc.KEMKeyInfo, error) +} + // KeyHandle represents a key handle returned from the API. type KeyHandle struct { Handle string `json:"handle"` @@ -68,10 +75,48 @@ type GenerateKemResponse struct { KeyHandle KeyHandle `json:"key_handle"` } +// KemPublicKey represents a KEM public key with its algorithm identifier. +type KemPublicKey struct { + Algorithm KemAlgorithm `json:"algorithm"` + PublicKey string `json:"public_key"` +} + +// HpkeAlgorithm identifies the HPKE algorithm suite (KEM, KDF, AEAD). +type HpkeAlgorithm struct { + Kem KemAlgorithm `json:"kem"` + Kdf KdfAlgorithm `json:"kdf"` + Aead AeadAlgorithm `json:"aead"` +} + +// HpkePublicKey represents an HPKE public key with its full algorithm suite. +type HpkePublicKey struct { + Algorithm HpkeAlgorithm `json:"algorithm"` + PublicKey string `json:"public_key"` +} + +// BoundKEMInfo holds the full metadata for a bound KEM key. +type BoundKEMInfo struct { + KeyHandle KeyHandle `json:"key_handle"` + KemPubKey KemPublicKey `json:"kem_pub_key"` + BindingPubKey HpkePublicKey `json:"binding_pub_key"` + RemainingLifespan ProtoDuration `json:"remaining_lifespan"` +} + +// KeyInfo wraps a single key entry in the enumerate response. +type KeyInfo struct { + BoundKemInfo *BoundKEMInfo `json:"bound_kem_info,omitempty"` +} + +// EnumerateKeysResponse is returned by GET /v1/keys. +type EnumerateKeysResponse struct { + KeyInfos []KeyInfo `json:"key_infos"` +} + // Server is the WSD HTTP server. type Server struct { bindingGen BindingKeyGenerator kemGen KEMKeyGenerator + kemEnum KEMKeyEnumerator mu sync.RWMutex kemToBindingMap map[uuid.UUID]uuid.UUID @@ -81,15 +126,17 @@ type Server struct { } // NewServer creates a new WSD server with the given dependencies. -func NewServer(bindingGen BindingKeyGenerator, kemGen KEMKeyGenerator) *Server { +func NewServer(bindingGen BindingKeyGenerator, kemGen KEMKeyGenerator, kemEnum KEMKeyEnumerator) *Server { s := &Server{ bindingGen: bindingGen, kemGen: kemGen, + kemEnum: kemEnum, kemToBindingMap: make(map[uuid.UUID]uuid.UUID), } mux := http.NewServeMux() mux.HandleFunc("POST /v1/keys:generate_kem", s.handleGenerateKem) + mux.HandleFunc("GET /v1/keys", s.handleEnumerateKeys) s.httpServer = &http.Server{Handler: mux} return s @@ -124,6 +171,48 @@ func (s *Server) LookupBindingUUID(kemUUID uuid.UUID) (uuid.UUID, bool) { return id, ok } +func (s *Server) handleEnumerateKeys(w http.ResponseWriter, r *http.Request) { + // Check method again implicitly via mux, but extra check doesn't hurt if mux misconfigures. + // Actually mux handles it if using "GET /v1/keys". + if r.Method != http.MethodGet { + http.Error(w, "method not allowed", http.StatusMethodNotAllowed) + return + } + + keys, err := s.kemEnum.EnumerateKEMKeys() + if err != nil { + http.Error(w, fmt.Sprintf("failed to enumerate keys: %v", err), http.StatusInternalServerError) + return + } + + keyInfos := make([]KeyInfo, 0, len(keys)) + for _, k := range keys { + info := KeyInfo{ + BoundKemInfo: &BoundKEMInfo{ + KeyHandle: KeyHandle{Handle: k.ID.String()}, + KemPubKey: KemPublicKey{ + Algorithm: KemAlgorithm(k.KemAlgorithm), + PublicKey: base64.StdEncoding.EncodeToString(k.KEMPubKey), + }, + BindingPubKey: HpkePublicKey{ + Algorithm: HpkeAlgorithm{ + Kem: KemAlgorithm(k.KemAlgorithm), + Kdf: KdfAlgorithm(k.KdfAlgorithm), + Aead: AeadAlgorithm(k.AeadAlgorithm), + }, + PublicKey: base64.StdEncoding.EncodeToString(k.BindingPubKey), + }, + RemainingLifespan: ProtoDuration{Seconds: k.RemainingLifespanSecs}, + }, + } + keyInfos = append(keyInfos, info) + } + + resp := EnumerateKeysResponse{KeyInfos: keyInfos} + w.Header().Set("Content-Type", "application/json") + json.NewEncoder(w).Encode(resp) +} + func (s *Server) handleGenerateKem(w http.ResponseWriter, r *http.Request) { var req GenerateKemRequest if err := json.NewDecoder(r.Body).Decode(&req); err != nil { @@ -191,5 +280,3 @@ func writeError(w http.ResponseWriter, message string, code int) { w.WriteHeader(code) json.NewEncoder(w).Encode(map[string]string{"error": message}) } - - diff --git a/keymanager/workload_service/server_test.go b/keymanager/workload_service/server_test.go index fc9e35c9d..a1eadf445 100644 --- a/keymanager/workload_service/server_test.go +++ b/keymanager/workload_service/server_test.go @@ -2,6 +2,7 @@ package workload_service import ( "bytes" + "encoding/base64" "encoding/json" "fmt" "net/http" @@ -9,6 +10,7 @@ import ( "strings" "testing" + kpskcc "github.com/google/go-tpm-tools/keymanager/key_protection_service/key_custody_core" "github.com/google/uuid" algorithms "github.com/google/go-tpm-tools/keymanager/km_common/proto" @@ -40,6 +42,16 @@ func (m *mockKEMKeyGen) GenerateKEMKeypair(algo *algorithms.HpkeAlgorithm, bindi return m.uuid, m.pubKey, m.err } +// mockKEMKeyEnumerator implements KEMKeyEnumerator for testing. +type mockKEMKeyEnumerator struct { + keys []kpskcc.KEMKeyInfo + err error +} + +func (m *mockKEMKeyEnumerator) EnumerateKEMKeys() ([]kpskcc.KEMKeyInfo, error) { + return m.keys, m.err +} + func validGenerateBody() []byte { body, _ := json.Marshal(GenerateKemRequest{ Algorithm: KemAlgorithmDHKEMX25519HKDFSHA256, @@ -65,6 +77,7 @@ func TestHandleGenerateKemSuccess(t *testing.T) { srv := NewServer( &mockBindingKeyGen{uuid: bindingUUID, pubKey: bindingPubKey}, kemGen, + &mockKEMKeyEnumerator{}, ) req := httptest.NewRequest(http.MethodPost, "/v1/keys:generate_kem", bytes.NewReader(validGenerateBody())) @@ -113,6 +126,7 @@ func TestHandleGenerateKemInvalidMethod(t *testing.T) { srv := NewServer( &mockBindingKeyGen{pubKey: make([]byte, 32)}, &mockKEMKeyGen{pubKey: make([]byte, 32)}, + &mockKEMKeyEnumerator{}, ) req := httptest.NewRequest(http.MethodGet, "/v1/keys:generate_kem", nil) @@ -128,6 +142,7 @@ func TestHandleGenerateKemBadRequest(t *testing.T) { srv := NewServer( &mockBindingKeyGen{uuid: uuid.New(), pubKey: make([]byte, 32)}, &mockKEMKeyGen{uuid: uuid.New(), pubKey: make([]byte, 32)}, + &mockKEMKeyEnumerator{}, ) tests := []struct { @@ -182,6 +197,7 @@ func TestHandleGenerateKemBadJSON(t *testing.T) { srv := NewServer( &mockBindingKeyGen{pubKey: make([]byte, 32)}, &mockKEMKeyGen{pubKey: make([]byte, 32)}, + &mockKEMKeyEnumerator{}, ) badBodies := []struct { @@ -212,6 +228,7 @@ func TestHandleGenerateKemBindingGenError(t *testing.T) { srv := NewServer( &mockBindingKeyGen{err: fmt.Errorf("binding FFI error")}, &mockKEMKeyGen{pubKey: make([]byte, 32)}, + &mockKEMKeyEnumerator{}, ) req := httptest.NewRequest(http.MethodPost, "/v1/keys:generate_kem", bytes.NewReader(validGenerateBody())) @@ -228,6 +245,7 @@ func TestHandleGenerateKemFlexibleLifespan(t *testing.T) { srv := NewServer( &mockBindingKeyGen{uuid: uuid.New(), pubKey: make([]byte, 32)}, &mockKEMKeyGen{uuid: uuid.New(), pubKey: make([]byte, 32)}, + &mockKEMKeyEnumerator{}, ) tests := []struct { @@ -270,6 +288,7 @@ func TestHandleGenerateKemKEMGenError(t *testing.T) { srv := NewServer( &mockBindingKeyGen{uuid: uuid.New(), pubKey: make([]byte, 32)}, &mockKEMKeyGen{err: fmt.Errorf("KEM FFI error")}, + &mockKEMKeyEnumerator{}, ) req := httptest.NewRequest(http.MethodPost, "/v1/keys:generate_kem", bytes.NewReader(validGenerateBody())) @@ -282,6 +301,166 @@ func TestHandleGenerateKemKEMGenError(t *testing.T) { } } +func TestHandleEnumerateKeysEmpty(t *testing.T) { + srv := NewServer( + &mockBindingKeyGen{}, + &mockKEMKeyGen{}, + &mockKEMKeyEnumerator{keys: []kpskcc.KEMKeyInfo{}}, + ) + + req := httptest.NewRequest(http.MethodGet, "/v1/keys", nil) + w := httptest.NewRecorder() + srv.Handler().ServeHTTP(w, req) + + if w.Code != http.StatusOK { + t.Fatalf("expected status 200, got %d: %s", w.Code, w.Body.String()) + } + + var resp EnumerateKeysResponse + if err := json.NewDecoder(w.Body).Decode(&resp); err != nil { + t.Fatalf("failed to decode response: %v", err) + } + if len(resp.KeyInfos) != 0 { + t.Fatalf("expected 0 key infos, got %d", len(resp.KeyInfos)) + } +} + +func TestHandleEnumerateKeysWithKeys(t *testing.T) { + kem1 := uuid.New() + kem2 := uuid.New() + kemPubKey1 := make([]byte, 32) + kemPubKey2 := make([]byte, 32) + bindingPubKey1 := make([]byte, 32) + bindingPubKey2 := make([]byte, 32) + for i := range kemPubKey1 { + kemPubKey1[i] = byte(i) + kemPubKey2[i] = byte(i + 50) + bindingPubKey1[i] = byte(i + 100) + bindingPubKey2[i] = byte(i + 150) + } + + mockEnum := &mockKEMKeyEnumerator{ + keys: []kpskcc.KEMKeyInfo{ + { + ID: kem1, + KemAlgorithm: 1, + KdfAlgorithm: 1, + AeadAlgorithm: 1, + KEMPubKey: kemPubKey1, + BindingPubKey: bindingPubKey1, + RemainingLifespanSecs: 3500, + }, + { + ID: kem2, + KemAlgorithm: 1, + KdfAlgorithm: 1, + AeadAlgorithm: 1, + KEMPubKey: kemPubKey2, + BindingPubKey: bindingPubKey2, + RemainingLifespanSecs: 7100, + }, + }, + } + + srv := NewServer( + &mockBindingKeyGen{}, + &mockKEMKeyGen{}, + mockEnum, + ) + + req := httptest.NewRequest(http.MethodGet, "/v1/keys", nil) + w := httptest.NewRecorder() + srv.Handler().ServeHTTP(w, req) + + if w.Code != http.StatusOK { + t.Fatalf("expected status 200, got %d: %s", w.Code, w.Body.String()) + } + + var resp EnumerateKeysResponse + if err := json.NewDecoder(w.Body).Decode(&resp); err != nil { + t.Fatalf("failed to decode response: %v", err) + } + if len(resp.KeyInfos) != 2 { + t.Fatalf("expected 2 key infos, got %d", len(resp.KeyInfos)) + } + + // Verify both keys appear (order-independent). + found := make(map[string]*BoundKEMInfo) + for _, ki := range resp.KeyInfos { + if ki.BoundKemInfo == nil { + t.Fatal("expected non-nil boundKemInfo") + } + found[ki.BoundKemInfo.KeyHandle.Handle] = ki.BoundKemInfo + } + + // Verify key 1. + info1, ok := found[kem1.String()] + if !ok { + t.Fatalf("expected kem1 %s in response", kem1) + } + if info1.KemPubKey.Algorithm != KemAlgorithmDHKEMX25519HKDFSHA256 { + t.Fatalf("expected algorithm %v, got %v", KemAlgorithmDHKEMX25519HKDFSHA256, info1.KemPubKey.Algorithm) + } + if info1.KemPubKey.PublicKey != base64.StdEncoding.EncodeToString(kemPubKey1) { + t.Fatalf("KEM pub key mismatch for kem1") + } + if info1.BindingPubKey.PublicKey != base64.StdEncoding.EncodeToString(bindingPubKey1) { + t.Fatalf("binding pub key mismatch for kem1") + } + expectedHPKE := HpkeAlgorithm{ + Kem: KemAlgorithmDHKEMX25519HKDFSHA256, + Kdf: KdfAlgorithmHKDFSHA256, + Aead: AeadAlgorithmAES256GCM, + } + if info1.BindingPubKey.Algorithm != expectedHPKE { + t.Fatalf("HPKE algorithm mismatch for kem1: expected %v, got %v", expectedHPKE, info1.BindingPubKey.Algorithm) + } + if info1.RemainingLifespan.Seconds != 3500 { + t.Fatalf("expected remaining lifespan 3500, got %d", info1.RemainingLifespan.Seconds) + } + + // Verify key 2. + info2, ok := found[kem2.String()] + if !ok { + t.Fatalf("expected kem2 %s in response", kem2) + } + if info2.RemainingLifespan.Seconds != 7100 { + t.Fatalf("expected remaining lifespan 7100, got %d", info2.RemainingLifespan.Seconds) + } +} + +func TestHandleEnumerateKeysMethodNotAllowed(t *testing.T) { + srv := NewServer( + &mockBindingKeyGen{}, + &mockKEMKeyGen{}, + &mockKEMKeyEnumerator{}, + ) + + req := httptest.NewRequest(http.MethodPost, "/v1/keys", nil) + w := httptest.NewRecorder() + srv.Handler().ServeHTTP(w, req) + + if w.Code != http.StatusMethodNotAllowed { + t.Fatalf("expected status 405, got %d", w.Code) + } +} + +func TestHandleEnumerateKeysError(t *testing.T) { + srv := NewServer( + &mockBindingKeyGen{}, + &mockKEMKeyGen{}, + &mockKEMKeyEnumerator{err: fmt.Errorf("enumerate error")}, + ) + + req := httptest.NewRequest(http.MethodGet, "/v1/keys", nil) + w := httptest.NewRecorder() + srv.Handler().ServeHTTP(w, req) + + if w.Code != http.StatusInternalServerError { + t.Fatalf("expected status 500, got %d: %s", w.Code, w.Body.String()) + } +} + func TestHandleGenerateKemMapUniqueness(t *testing.T) { bindingPubKey := make([]byte, 32) @@ -294,7 +473,7 @@ func TestHandleGenerateKemMapUniqueness(t *testing.T) { bindingGen := &mockBindingKeyGen{} kemGen := &mockKEMKeyGen{} - srv := NewServer(bindingGen, kemGen) + srv := NewServer(bindingGen, kemGen, &mockKEMKeyEnumerator{}) // First call. bindingGen.uuid = bindingUUID1 From ceb576d2311c9456f6728417e18d35e58807a629 Mon Sep 17 00:00:00 2001 From: Atul Patil Date: Fri, 20 Feb 2026 10:03:40 +0000 Subject: [PATCH 09/10] Refactor Enumerate Keys API Go implementation --- .../kps_key_custody_core_cgo.go | 10 ++++--- .../key_custody_core/types.go | 10 ++++--- .../key_protection_service/service_test.go | 10 ++++--- keymanager/workload_service/server.go | 8 +++--- keymanager/workload_service/server_test.go | 26 +++++++++++-------- 5 files changed, 38 insertions(+), 26 deletions(-) diff --git a/keymanager/key_protection_service/key_custody_core/kps_key_custody_core_cgo.go b/keymanager/key_protection_service/key_custody_core/kps_key_custody_core_cgo.go index f852debaf..0d91ac030 100644 --- a/keymanager/key_protection_service/key_custody_core/kps_key_custody_core_cgo.go +++ b/keymanager/key_protection_service/key_custody_core/kps_key_custody_core_cgo.go @@ -91,11 +91,15 @@ func EnumerateKEMKeys() ([]KEMKeyInfo, error) { bindingPubKey := make([]byte, e.binding_pub_key_len) copy(bindingPubKey, C.GoBytes(unsafe.Pointer(&e.binding_pub_key[0]), C.int(e.binding_pub_key_len))) + algoBytes := C.GoBytes(unsafe.Pointer(&e.algorithm[0]), C.int(e.algorithm_len)) + algo := &algorithms.HpkeAlgorithm{} + if err := proto.Unmarshal(algoBytes, algo); err != nil { + return nil, fmt.Errorf("failed to unmarshal algorithm for key %d: %w", i, err) + } + result[i] = KEMKeyInfo{ ID: id, - KemAlgorithm: int32(e.algorithm.kem), - KdfAlgorithm: int32(e.algorithm.kdf), - AeadAlgorithm: int32(e.algorithm.aead), + Algorithm: algo, KEMPubKey: kemPubKey, BindingPubKey: bindingPubKey, RemainingLifespanSecs: uint64(e.remaining_lifespan_secs), diff --git a/keymanager/key_protection_service/key_custody_core/types.go b/keymanager/key_protection_service/key_custody_core/types.go index b40f0e8f1..f0cc40457 100644 --- a/keymanager/key_protection_service/key_custody_core/types.go +++ b/keymanager/key_protection_service/key_custody_core/types.go @@ -1,13 +1,15 @@ package kpskcc -import "github.com/google/uuid" +import ( + algorithms "github.com/google/go-tpm-tools/keymanager/km_common/proto" + "github.com/google/uuid" +) +// KEMKeyInfo holds metadata for a single KEM key returned by EnumerateKEMKeys. // KEMKeyInfo holds metadata for a single KEM key returned by EnumerateKEMKeys. type KEMKeyInfo struct { ID uuid.UUID - KemAlgorithm int32 - KdfAlgorithm int32 - AeadAlgorithm int32 + Algorithm *algorithms.HpkeAlgorithm KEMPubKey []byte BindingPubKey []byte RemainingLifespanSecs uint64 diff --git a/keymanager/key_protection_service/service_test.go b/keymanager/key_protection_service/service_test.go index b6b532277..58a7d3580 100644 --- a/keymanager/key_protection_service/service_test.go +++ b/keymanager/key_protection_service/service_test.go @@ -63,10 +63,12 @@ func TestServiceGenerateKEMKeypairError(t *testing.T) { func TestServiceEnumerateKEMKeysSuccess(t *testing.T) { expectedKeys := []kpskcc.KEMKeyInfo{ { - ID: uuid.New(), - KemAlgorithm: 1, - KdfAlgorithm: 1, - AeadAlgorithm: 1, + ID: uuid.New(), + Algorithm: &algorithms.HpkeAlgorithm{ + Kem: algorithms.KemAlgorithm_KEM_ALGORITHM_DHKEM_X25519_HKDF_SHA256, + Kdf: algorithms.KdfAlgorithm_KDF_ALGORITHM_HKDF_SHA256, + Aead: algorithms.AeadAlgorithm_AEAD_ALGORITHM_AES_256_GCM, + }, KEMPubKey: make([]byte, 32), BindingPubKey: make([]byte, 32), RemainingLifespanSecs: 3500, diff --git a/keymanager/workload_service/server.go b/keymanager/workload_service/server.go index 2661b4d3c..2834f4cb9 100644 --- a/keymanager/workload_service/server.go +++ b/keymanager/workload_service/server.go @@ -191,14 +191,14 @@ func (s *Server) handleEnumerateKeys(w http.ResponseWriter, r *http.Request) { BoundKemInfo: &BoundKEMInfo{ KeyHandle: KeyHandle{Handle: k.ID.String()}, KemPubKey: KemPublicKey{ - Algorithm: KemAlgorithm(k.KemAlgorithm), + Algorithm: KemAlgorithm(k.Algorithm.Kem), PublicKey: base64.StdEncoding.EncodeToString(k.KEMPubKey), }, BindingPubKey: HpkePublicKey{ Algorithm: HpkeAlgorithm{ - Kem: KemAlgorithm(k.KemAlgorithm), - Kdf: KdfAlgorithm(k.KdfAlgorithm), - Aead: AeadAlgorithm(k.AeadAlgorithm), + Kem: KemAlgorithm(k.Algorithm.Kem), + Kdf: KdfAlgorithm(k.Algorithm.Kdf), + Aead: AeadAlgorithm(k.Algorithm.Aead), }, PublicKey: base64.StdEncoding.EncodeToString(k.BindingPubKey), }, diff --git a/keymanager/workload_service/server_test.go b/keymanager/workload_service/server_test.go index a1eadf445..af4072dc1 100644 --- a/keymanager/workload_service/server_test.go +++ b/keymanager/workload_service/server_test.go @@ -205,9 +205,9 @@ func TestHandleGenerateKemBadJSON(t *testing.T) { body string }{ {"not json", "not json"}, - {"lifespan as string", `{"algorithm":1,"key_protection_mechanism":2,"lifespan":"3600"}`}, - {"lifespan as string with suffix", `{"algorithm":1,"key_protection_mechanism":2,"lifespan":"3600s"}`}, - {"lifespan negative", `{"algorithm":1,"key_protection_mechanism":2,"lifespan":-1}`}, + {"lifespan as integer", `{"algorithm":1,"key_protection_mechanism":2,"lifespan":3600}`}, + {"lifespan missing s suffix", `{"algorithm":1,"key_protection_mechanism":2,"lifespan":"3600"}`}, + {"lifespan negative", `{"algorithm":1,"key_protection_mechanism":2,"lifespan":"-1s"}`}, } for _, tc := range badBodies { @@ -342,19 +342,23 @@ func TestHandleEnumerateKeysWithKeys(t *testing.T) { mockEnum := &mockKEMKeyEnumerator{ keys: []kpskcc.KEMKeyInfo{ { - ID: kem1, - KemAlgorithm: 1, - KdfAlgorithm: 1, - AeadAlgorithm: 1, + ID: kem1, + Algorithm: &algorithms.HpkeAlgorithm{ + Kem: algorithms.KemAlgorithm_KEM_ALGORITHM_DHKEM_X25519_HKDF_SHA256, + Kdf: algorithms.KdfAlgorithm_KDF_ALGORITHM_HKDF_SHA256, + Aead: algorithms.AeadAlgorithm_AEAD_ALGORITHM_AES_256_GCM, + }, KEMPubKey: kemPubKey1, BindingPubKey: bindingPubKey1, RemainingLifespanSecs: 3500, }, { - ID: kem2, - KemAlgorithm: 1, - KdfAlgorithm: 1, - AeadAlgorithm: 1, + ID: kem2, + Algorithm: &algorithms.HpkeAlgorithm{ + Kem: algorithms.KemAlgorithm_KEM_ALGORITHM_DHKEM_X25519_HKDF_SHA256, + Kdf: algorithms.KdfAlgorithm_KDF_ALGORITHM_HKDF_SHA256, + Aead: algorithms.AeadAlgorithm_AEAD_ALGORITHM_AES_256_GCM, + }, KEMPubKey: kemPubKey2, BindingPubKey: bindingPubKey2, RemainingLifespanSecs: 7100, From 6fd86d5248113dc4802997041794cdf1a5550308 Mon Sep 17 00:00:00 2001 From: Atul Patil Date: Fri, 20 Feb 2026 15:00:31 +0000 Subject: [PATCH 10/10] Refactor Enumerate API: simplify FFI and update Go wrapper --- .../kps_key_custody_core_cgo.go | 23 +++++++++++-------- .../key_custody_core/types.go | 2 -- keymanager/key_protection_service/service.go | 10 ++++---- .../key_protection_service/service_test.go | 16 +++++++------ keymanager/workload_service/server.go | 22 +++++++++--------- keymanager/workload_service/server_test.go | 21 +++-------------- 6 files changed, 42 insertions(+), 52 deletions(-) diff --git a/keymanager/key_protection_service/key_custody_core/kps_key_custody_core_cgo.go b/keymanager/key_protection_service/key_custody_core/kps_key_custody_core_cgo.go index 0d91ac030..b44bc4dfe 100644 --- a/keymanager/key_protection_service/key_custody_core/kps_key_custody_core_cgo.go +++ b/keymanager/key_protection_service/key_custody_core/kps_key_custody_core_cgo.go @@ -61,15 +61,24 @@ func GenerateKEMKeypair(algo *algorithms.HpkeAlgorithm, bindingPubKey []byte, li return id, pubkey, nil } -// EnumerateKEMKeys retrieves all active KEM key entries from the Rust KCC registry. -func EnumerateKEMKeys() ([]KEMKeyInfo, error) { - const maxEntries = 256 - var entries [maxEntries]C.KpsKeyInfo +// EnumerateKEMKeys retrieves active KEM key entries from the Rust KCC registry with pagination. +func EnumerateKEMKeys(limit, offset int) ([]KEMKeyInfo, error) { + if limit <= 0 { + return nil, fmt.Errorf("limit must be positive") + } + if offset < 0 { + return nil, fmt.Errorf("offset must be non-negative") + } + + // Dynamic allocation might be better, but for now using a slice on the heap is safer than large stack usage. + // C.KpsKeyInfo is large (~2KB+), so even 256 entries is 500KB. + entries := make([]C.KpsKeyInfo, limit) var count C.size_t rc := C.key_manager_enumerate_kem_keys( &entries[0], - C.size_t(maxEntries), + C.size_t(limit), + C.size_t(offset), &count, ) if rc != 0 { @@ -88,9 +97,6 @@ func EnumerateKEMKeys() ([]KEMKeyInfo, error) { kemPubKey := make([]byte, e.kem_pub_key_len) copy(kemPubKey, C.GoBytes(unsafe.Pointer(&e.kem_pub_key[0]), C.int(e.kem_pub_key_len))) - bindingPubKey := make([]byte, e.binding_pub_key_len) - copy(bindingPubKey, C.GoBytes(unsafe.Pointer(&e.binding_pub_key[0]), C.int(e.binding_pub_key_len))) - algoBytes := C.GoBytes(unsafe.Pointer(&e.algorithm[0]), C.int(e.algorithm_len)) algo := &algorithms.HpkeAlgorithm{} if err := proto.Unmarshal(algoBytes, algo); err != nil { @@ -101,7 +107,6 @@ func EnumerateKEMKeys() ([]KEMKeyInfo, error) { ID: id, Algorithm: algo, KEMPubKey: kemPubKey, - BindingPubKey: bindingPubKey, RemainingLifespanSecs: uint64(e.remaining_lifespan_secs), } } diff --git a/keymanager/key_protection_service/key_custody_core/types.go b/keymanager/key_protection_service/key_custody_core/types.go index f0cc40457..c981eb2b7 100644 --- a/keymanager/key_protection_service/key_custody_core/types.go +++ b/keymanager/key_protection_service/key_custody_core/types.go @@ -5,12 +5,10 @@ import ( "github.com/google/uuid" ) -// KEMKeyInfo holds metadata for a single KEM key returned by EnumerateKEMKeys. // KEMKeyInfo holds metadata for a single KEM key returned by EnumerateKEMKeys. type KEMKeyInfo struct { ID uuid.UUID Algorithm *algorithms.HpkeAlgorithm KEMPubKey []byte - BindingPubKey []byte RemainingLifespanSecs uint64 } diff --git a/keymanager/key_protection_service/service.go b/keymanager/key_protection_service/service.go index 206444c49..6ff5d7c8b 100644 --- a/keymanager/key_protection_service/service.go +++ b/keymanager/key_protection_service/service.go @@ -17,19 +17,19 @@ type KEMKeyGenerator interface { // KEMKeyEnumerator enumerates active KEM keys in the KPS registry. type KEMKeyEnumerator interface { - EnumerateKEMKeys() ([]kpskcc.KEMKeyInfo, error) + EnumerateKEMKeys(limit, offset int) ([]kpskcc.KEMKeyInfo, error) } // Service implements KEMKeyGenerator and KEMKeyEnumerator by delegating to the KPS KCC FFI. type Service struct { generateKEMKeypairFn func(algo *algorithms.HpkeAlgorithm, bindingPubKey []byte, lifespanSecs uint64) (uuid.UUID, []byte, error) - enumerateKEMKeysFn func() ([]kpskcc.KEMKeyInfo, error) + enumerateKEMKeysFn func(limit, offset int) ([]kpskcc.KEMKeyInfo, error) } // NewService creates a new KPS KOL service with the given KCC functions. func NewService( generateKEMKeypairFn func(algo *algorithms.HpkeAlgorithm, bindingPubKey []byte, lifespanSecs uint64) (uuid.UUID, []byte, error), - enumerateKEMKeysFn func() ([]kpskcc.KEMKeyInfo, error), + enumerateKEMKeysFn func(limit, offset int) ([]kpskcc.KEMKeyInfo, error), ) *Service { return &Service{ generateKEMKeypairFn: generateKEMKeypairFn, @@ -44,6 +44,6 @@ func (s *Service) GenerateKEMKeypair(algo *algorithms.HpkeAlgorithm, bindingPubK } // EnumerateKEMKeys retrieves all active KEM key entries from the KPS KCC registry. -func (s *Service) EnumerateKEMKeys() ([]kpskcc.KEMKeyInfo, error) { - return s.enumerateKEMKeysFn() +func (s *Service) EnumerateKEMKeys(limit, offset int) ([]kpskcc.KEMKeyInfo, error) { + return s.enumerateKEMKeysFn(limit, offset) } diff --git a/keymanager/key_protection_service/service_test.go b/keymanager/key_protection_service/service_test.go index 58a7d3580..5c33beae0 100644 --- a/keymanager/key_protection_service/service_test.go +++ b/keymanager/key_protection_service/service_test.go @@ -27,7 +27,7 @@ func TestServiceGenerateKEMKeypairSuccess(t *testing.T) { } return expectedUUID, expectedPubKey, nil }, - func() ([]kpskcc.KEMKeyInfo, error) { + func(limit, offset int) ([]kpskcc.KEMKeyInfo, error) { return nil, nil }, ) @@ -49,7 +49,7 @@ func TestServiceGenerateKEMKeypairError(t *testing.T) { func(algo *algorithms.HpkeAlgorithm, bindingPubKey []byte, lifespanSecs uint64) (uuid.UUID, []byte, error) { return uuid.Nil, nil, fmt.Errorf("FFI error") }, - func() ([]kpskcc.KEMKeyInfo, error) { + func(limit, offset int) ([]kpskcc.KEMKeyInfo, error) { return nil, nil }, ) @@ -70,7 +70,6 @@ func TestServiceEnumerateKEMKeysSuccess(t *testing.T) { Aead: algorithms.AeadAlgorithm_AEAD_ALGORITHM_AES_256_GCM, }, KEMPubKey: make([]byte, 32), - BindingPubKey: make([]byte, 32), RemainingLifespanSecs: 3500, }, } @@ -79,12 +78,15 @@ func TestServiceEnumerateKEMKeysSuccess(t *testing.T) { func(algo *algorithms.HpkeAlgorithm, bindingPubKey []byte, lifespanSecs uint64) (uuid.UUID, []byte, error) { return uuid.Nil, nil, nil }, - func() ([]kpskcc.KEMKeyInfo, error) { + func(limit, offset int) ([]kpskcc.KEMKeyInfo, error) { + if limit != 100 || offset != 0 { + return nil, fmt.Errorf("unexpected limit/offset: %d/%d", limit, offset) + } return expectedKeys, nil }, ) - keys, err := svc.EnumerateKEMKeys() + keys, err := svc.EnumerateKEMKeys(100, 0) if err != nil { t.Fatalf("unexpected error: %v", err) } @@ -101,12 +103,12 @@ func TestServiceEnumerateKEMKeysError(t *testing.T) { func(algo *algorithms.HpkeAlgorithm, bindingPubKey []byte, lifespanSecs uint64) (uuid.UUID, []byte, error) { return uuid.Nil, nil, nil }, - func() ([]kpskcc.KEMKeyInfo, error) { + func(limit, offset int) ([]kpskcc.KEMKeyInfo, error) { return nil, fmt.Errorf("enumerate error") }, ) - _, err := svc.EnumerateKEMKeys() + _, err := svc.EnumerateKEMKeys(100, 0) if err == nil { t.Fatal("expected error, got nil") } diff --git a/keymanager/workload_service/server.go b/keymanager/workload_service/server.go index 2834f4cb9..437ccecc4 100644 --- a/keymanager/workload_service/server.go +++ b/keymanager/workload_service/server.go @@ -30,9 +30,10 @@ type KEMKeyGenerator interface { GenerateKEMKeypair(algo *algorithms.HpkeAlgorithm, bindingPubKey []byte, lifespanSecs uint64) (uuid.UUID, []byte, error) } +// KEMKeyEnumerator enumerates active KEM keys from the KPS registry. // KEMKeyEnumerator enumerates active KEM keys from the KPS registry. type KEMKeyEnumerator interface { - EnumerateKEMKeys() ([]kpskcc.KEMKeyInfo, error) + EnumerateKEMKeys(limit, offset int) ([]kpskcc.KEMKeyInfo, error) } // KeyHandle represents a key handle returned from the API. @@ -98,7 +99,7 @@ type HpkePublicKey struct { type BoundKEMInfo struct { KeyHandle KeyHandle `json:"key_handle"` KemPubKey KemPublicKey `json:"kem_pub_key"` - BindingPubKey HpkePublicKey `json:"binding_pub_key"` + // BindingPubKey removed as it is no longer returned by KCC FFI RemainingLifespan ProtoDuration `json:"remaining_lifespan"` } @@ -179,7 +180,14 @@ func (s *Server) handleEnumerateKeys(w http.ResponseWriter, r *http.Request) { return } - keys, err := s.kemEnum.EnumerateKEMKeys() + // Default limit/offset for now. + // TODO: Parse from query params? + const ( + defaultLimit = 100 + defaultOffset = 0 + ) + + keys, err := s.kemEnum.EnumerateKEMKeys(defaultLimit, defaultOffset) if err != nil { http.Error(w, fmt.Sprintf("failed to enumerate keys: %v", err), http.StatusInternalServerError) return @@ -194,14 +202,6 @@ func (s *Server) handleEnumerateKeys(w http.ResponseWriter, r *http.Request) { Algorithm: KemAlgorithm(k.Algorithm.Kem), PublicKey: base64.StdEncoding.EncodeToString(k.KEMPubKey), }, - BindingPubKey: HpkePublicKey{ - Algorithm: HpkeAlgorithm{ - Kem: KemAlgorithm(k.Algorithm.Kem), - Kdf: KdfAlgorithm(k.Algorithm.Kdf), - Aead: AeadAlgorithm(k.Algorithm.Aead), - }, - PublicKey: base64.StdEncoding.EncodeToString(k.BindingPubKey), - }, RemainingLifespan: ProtoDuration{Seconds: k.RemainingLifespanSecs}, }, } diff --git a/keymanager/workload_service/server_test.go b/keymanager/workload_service/server_test.go index af4072dc1..d8787318b 100644 --- a/keymanager/workload_service/server_test.go +++ b/keymanager/workload_service/server_test.go @@ -48,7 +48,7 @@ type mockKEMKeyEnumerator struct { err error } -func (m *mockKEMKeyEnumerator) EnumerateKEMKeys() ([]kpskcc.KEMKeyInfo, error) { +func (m *mockKEMKeyEnumerator) EnumerateKEMKeys(limit, offset int) ([]kpskcc.KEMKeyInfo, error) { return m.keys, m.err } @@ -330,13 +330,10 @@ func TestHandleEnumerateKeysWithKeys(t *testing.T) { kem2 := uuid.New() kemPubKey1 := make([]byte, 32) kemPubKey2 := make([]byte, 32) - bindingPubKey1 := make([]byte, 32) - bindingPubKey2 := make([]byte, 32) + // BindingPubKey no longer used in response for i := range kemPubKey1 { kemPubKey1[i] = byte(i) kemPubKey2[i] = byte(i + 50) - bindingPubKey1[i] = byte(i + 100) - bindingPubKey2[i] = byte(i + 150) } mockEnum := &mockKEMKeyEnumerator{ @@ -349,7 +346,6 @@ func TestHandleEnumerateKeysWithKeys(t *testing.T) { Aead: algorithms.AeadAlgorithm_AEAD_ALGORITHM_AES_256_GCM, }, KEMPubKey: kemPubKey1, - BindingPubKey: bindingPubKey1, RemainingLifespanSecs: 3500, }, { @@ -360,7 +356,6 @@ func TestHandleEnumerateKeysWithKeys(t *testing.T) { Aead: algorithms.AeadAlgorithm_AEAD_ALGORITHM_AES_256_GCM, }, KEMPubKey: kemPubKey2, - BindingPubKey: bindingPubKey2, RemainingLifespanSecs: 7100, }, }, @@ -408,17 +403,7 @@ func TestHandleEnumerateKeysWithKeys(t *testing.T) { if info1.KemPubKey.PublicKey != base64.StdEncoding.EncodeToString(kemPubKey1) { t.Fatalf("KEM pub key mismatch for kem1") } - if info1.BindingPubKey.PublicKey != base64.StdEncoding.EncodeToString(bindingPubKey1) { - t.Fatalf("binding pub key mismatch for kem1") - } - expectedHPKE := HpkeAlgorithm{ - Kem: KemAlgorithmDHKEMX25519HKDFSHA256, - Kdf: KdfAlgorithmHKDFSHA256, - Aead: AeadAlgorithmAES256GCM, - } - if info1.BindingPubKey.Algorithm != expectedHPKE { - t.Fatalf("HPKE algorithm mismatch for kem1: expected %v, got %v", expectedHPKE, info1.BindingPubKey.Algorithm) - } + // BindingPubKey check removed if info1.RemainingLifespan.Seconds != 3500 { t.Fatalf("expected remaining lifespan 3500, got %d", info1.RemainingLifespan.Seconds) }