feat(keymanager): implement enumerate KEM keys API in Go

atulpatildbz · atulpatildbz · commit 02da1b85778a · 2026-02-27T12:45:11.000Z
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
@@ -7,6 +7,7 @@ package kpskcc
 #cgo LDFLAGS: -L${SRCDIR}/../../target/release -L${SRCDIR}/../../target/debug -lkps_key_custody_core
 #cgo LDFLAGS: -lcrypto -lssl
 #cgo LDFLAGS: -lpthread -ldl -lm -lstdc++
+#include <stdbool.h>
 #include "include/kps_key_custody_core.h"
 */
 import "C"
@@ -61,3 +62,53 @@ 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.
+// Returns a list of keys and a boolean indicating if there are more keys to fetch.
+func EnumerateKEMKeys(limit, offset int) ([]KEMKeyInfo, bool, error) {
+	if limit <= 0 {
+		return nil, false, fmt.Errorf("limit must be positive")
+	}
+	if offset < 0 {
+		return nil, false, fmt.Errorf("offset must be non-negative")
+	}
+
+	entries := make([]C.KpsKeyInfo, limit)
+	var hasMore C.bool
+
+	rc := C.key_manager_enumerate_kem_keys(
+		&entries[0],
+		C.size_t(limit),
+		C.size_t(offset),
+		&hasMore,
+	)
+	if rc < 0 {
+		return nil, false, fmt.Errorf("key_manager_enumerate_kem_keys failed with code %d", rc)
+	}
+
+	count := int(rc)
+	result := make([]KEMKeyInfo, count)
+	for i, e := range entries[:count] {
+		id, err := uuid.FromBytes(C.GoBytes(unsafe.Pointer(&e.uuid[0]), 16))
+		if err != nil {
+			return nil, false, fmt.Errorf("invalid UUID at index %d: %w", i, err)
+		}
+
+		kemPubKey := C.GoBytes(unsafe.Pointer(&e.pub_key[0]), C.int(e.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, false, 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, bool(hasMore), nil
+}
diff --git a/keymanager/key_protection_service/key_custody_core/types.go b/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
+}
diff --git a/keymanager/key_protection_service/service.go b/keymanager/key_protection_service/service.go
@@ -1,9 +1,10 @@
 // Package keyprotectionservice 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 keyprotectionservice
 
 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, bool, 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, bool, 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, bool, 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, bool, error) {
+	return s.enumerateKEMKeysFn(limit, offset)
+}
diff --git 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(_ *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(limit, offset int) ([]kpskcc.KEMKeyInfo, bool, error) {
+			return nil, false, nil
+		},
+	)
 
 	id, pubKey, err := svc.GenerateKEMKeypair(&algorithms.HpkeAlgorithm{}, make([]byte, 32), 7200)
 	if err != nil {
@@ -39,12 +45,71 @@ func TestServiceGenerateKEMKeypairSuccess(t *testing.T) {
 }
 
 func TestServiceGenerateKEMKeypairError(t *testing.T) {
-	svc := NewService(func(_ *algorithms.HpkeAlgorithm, _ []byte, _ 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(limit, offset int) ([]kpskcc.KEMKeyInfo, bool, error) {
+			return nil, false, 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(),
+			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),
+			RemainingLifespanSecs: 3500,
+		},
+	}
+
+	svc := NewService(
+		func(algo *algorithms.HpkeAlgorithm, bindingPubKey []byte, lifespanSecs uint64) (uuid.UUID, []byte, error) {
+			return uuid.Nil, nil, nil
+		},
+		func(limit, offset int) ([]kpskcc.KEMKeyInfo, bool, error) {
+			if limit != 100 || offset != 0 {
+				return nil, false, fmt.Errorf("unexpected limit/offset: %d/%d", limit, offset)
+			}
+			return expectedKeys, false, nil
+		},
+	)
+
+	keys, _, err := svc.EnumerateKEMKeys(100, 0)
+	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(limit, offset int) ([]kpskcc.KEMKeyInfo, bool, error) {
+			return nil, false, fmt.Errorf("enumerate error")
+		},
+	)
+
+	_, _, err := svc.EnumerateKEMKeys(100, 0)
+	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
@@ -97,3 +97,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
diff --git a/keymanager/workload_service/server_test.go b/keymanager/workload_service/server_test.go
