[keymanager/wsd] GET /v1/keys enumerate keys

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"]
 
 

keymanager/key_protection_service/key_custody_core/include/kps_key_custody_core.h

@@ -6,6 +6,15 @@
 #include <stddef.h>
 #include <stdint.h>
 
+typedef struct {
+  uint8_t uuid[16];
+  uint8_t algorithm[128];
+  size_t algorithm_len;
+  uint8_t kem_pub_key[2048];
+  size_t kem_pub_key_len;
+  uint64_t remaining_lifespan_secs;
+} KpsKeyInfo;
+
 #ifdef __cplusplus
 extern "C" {
 #endif // __cplusplus
@@ -21,6 +30,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,
+                                       size_t *out_count);
+
 #ifdef __cplusplus
 }  // extern "C"
 #endif  // __cplusplus

keymanager/key_protection_service/key_custody_core/kps_key_custody_core_cgo.go

@@ -60,3 +60,56 @@ func GenerateKEMKeypair(algo *algorithms.HpkeAlgorithm, bindingPubKey []byte, li
 	copy(pubkey, pubkeyBuf[:pubkeyLen])
 	return id, pubkey, nil
 }
+
+// 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(limit),
+		C.size_t(offset),
+		&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)))
+
+		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,
+			Algorithm:             algo,
+			KEMPubKey:             kemPubKey,
+			RemainingLifespanSecs: uint64(e.remaining_lifespan_secs),
+		}
+	}
+
+	return result, nil
+}

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<KeyRegistry> = LazyLock::new(|| {
@@ -145,6 +145,77 @@ 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; 128],
+    pub algorithm_len: usize,
+    pub kem_pub_key: [u8; 2048],
+    pub kem_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,
+    offset: usize,
+    out_count: *mut usize,
+) -> i32 {
+    if out_entries.is_null() || out_count.is_null() {
+        return -1;
+    }
+
+    let metas = KEY_REGISTRY.list_kem_keys(offset, max_entries);
+    let now = Instant::now();
+    let mut count = 0usize;
+
+    for meta in &metas {
+        if let KeySpec::KemWithBindingPub {
+            algo,
+            kem_public_key,
+            binding_public_key: _,
+        } = &meta.spec
+        {
+            let algo_bytes = algo.encode_to_vec();
+            if kem_public_key.as_bytes().len() > 2048 || algo_bytes.len() > 128 {
+                eprintln!(
+                    "Skipping key {}: size exceeds buffer limits (algo={}, kem={})",
+                    meta.id,
+                    algo_bytes.len(),
+                    kem_public_key.as_bytes().len()
+                );
+            } else {
+                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 = [0u8; 128];
+            entry.algorithm[..algo_bytes.len()].copy_from_slice(&algo_bytes);
+            entry.algorithm_len = algo_bytes.len();
+
+            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.remaining_lifespan_secs = remaining;
+
+            count += 1;
+        }
+    }
+
+    unsafe {
+        *out_count = count;
+    }
+    0
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -326,10 +397,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 +433,79 @@ 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, 0, 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; 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; 128],
+            algorithm_len: 0,
+            kem_pub_key: [0; 2048],
+            kem_pub_key_len: 0,
+            remaining_lifespan_secs: 0,
+        });
+        let mut count: usize = 0;
+
+        let rc = unsafe {
+            // 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).
+        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;
+                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
+                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 (
+	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.
+type KEMKeyInfo struct {
+	ID                    uuid.UUID
+	Algorithm             *algorithms.HpkeAlgorithm
+	KEMPubKey             []byte
+	RemainingLifespanSecs uint64
+}

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,18 +15,35 @@ 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(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(limit, offset int) ([]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(limit, offset int) ([]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(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(limit, offset int) ([]kpskcc.KEMKeyInfo, error) {
+	return s.enumerateKEMKeysFn(limit, offset)
+}