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.

Author: atulpatildbz

keymanager/key_protection_service/key_custody_core/include/kps_key_custody_core.h

@@ -19,8 +19,29 @@ int32_t key_manager_generate_kem_keypair(KmHpkeAlgorithm algo,
                                          uint8_t *out_pubkey,
                                          size_t out_pubkey_len);
 
+// 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;
+
+// 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_

keymanager/key_protection_service/key_custody_core/kps_key_custody_core_cgo.go

@@ -57,3 +57,47 @@ func GenerateKEMKeypair(bindingPubKey []byte, lifespanSecs uint64) (uuid.UUID, [
 	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
+}

keymanager/key_protection_service/key_custody_core/src/lib.rs

@@ -3,7 +3,7 @@ use km_common::crypto::PublicKey;
 use km_common::key_types::{KeyRecord, KeyRegistry, KeySpec};
 use std::slice;
 use std::sync::LazyLock;
-use std::time::Duration;
+use std::time::{Duration, Instant};
 
 static KEY_REGISTRY: LazyLock<KeyRegistry> = LazyLock::new(KeyRegistry::default);
 
@@ -99,6 +99,75 @@ pub unsafe extern "C" fn key_manager_generate_kem_keypair(
     }
 }
 
+#[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 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::*;
@@ -261,4 +330,82 @@ mod tests {
 
         assert_eq!(result, -1);
     }
+
+    #[test]
+    fn test_enumerate_kem_keys_null_pointers() {
+        let result = 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,
+        };
+
+        // Generate a key first.
+        let rc = unsafe {
+            key_manager_generate_kem_keypair(
+                KmHpkeAlgorithm {
+                    kem: algo.kem,
+                    kdf: algo.kdf,
+                    aead: algo.aead,
+                },
+                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);
+        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).
+        // Note: other tests may have added keys to the global registry too.
+        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");
+    }
 }

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
+}

keymanager/key_protection_service/service.go

@@ -1,27 +1,47 @@
 // 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 "github.com/google/uuid"
+import (
+	kpskcc "github.com/google/go-tpm-tools/keymanager/key_protection_service/key_custody_core"
+	"github.com/google/uuid"
+)
 
 // KEMKeyGenerator generates KEM keypairs linked to a binding public key.
 type KEMKeyGenerator interface {
 	GenerateKEMKeypair(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(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(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(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
 // public key by calling the KPS KCC FFI.
 func (s *Service) GenerateKEMKeypair(bindingPubKey []byte, lifespanSecs uint64) (uuid.UUID, []byte, error) {
 	return s.generateKEMKeypairFn(bindingPubKey, lifespanSecs)
 }
+
+// EnumerateKEMKeys retrieves all active KEM key entries from the KPS KCC registry.
+func (s *Service) EnumerateKEMKeys() ([]kpskcc.KEMKeyInfo, error) {
+	return s.enumerateKEMKeysFn()
+}