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

keymanager/go.mod

@@ -1,5 +1,10 @@
 module github.com/google/go-tpm-tools/keymanager
 
-go 1.22
+go 1.23
 
-require github.com/google/uuid v1.6.0
+toolchain go1.24.13
+
+require (
+	github.com/google/uuid v1.6.0
+	google.golang.org/protobuf v1.36.11
+)

keymanager/go.sum

@@ -1,2 +1,6 @@
+github.com/google/go-cmp v0.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8=
+github.com/google/go-cmp v0.7.0/go.mod h1:pXiqmnSA92OHEEa9HXL2W4E7lf9JzCmGVUdgjX3N/iU=
 github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0=
 github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
+google.golang.org/protobuf v1.36.11 h1:fV6ZwhNocDyBLK0dj+fg8ektcVegBBuEolpbTQyBNVE=
+google.golang.org/protobuf v1.36.11/go.mod h1:HTf+CrKn2C3g5S8VImy6tdcUvCska2kB7j23XfzDpco=

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"
@@ -68,6 +69,56 @@ func GenerateKEMKeypair(algo *keymanager.HpkeAlgorithm, bindingPubKey []byte, li
 	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 := &keymanager.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
+}
+
 // DecapAndSeal decapsulates a shared secret using the stored KEM key and
 // reseals it with the associated binding public key via Rust FFI.
 // Returns the new encapsulated key and sealed ciphertext.

keymanager/key_protection_service/key_custody_core/kps_key_custody_core_stub.go

@@ -15,6 +15,11 @@ func GenerateKEMKeypair(_ *algorithms.HpkeAlgorithm, _ []byte, _ uint64) (uuid.U
 	return uuid.Nil, nil, fmt.Errorf("GenerateKEMKeypair is not supported on this architecture")
 }
 
+// EnumerateKEMKeys is a stub for architectures where the Rust library is not supported.
+func EnumerateKEMKeys(limit, offset int) ([]KEMKeyInfo, bool, error) {
+	return nil, false, fmt.Errorf("EnumerateKEMKeys is not supported on this architecture")
+}
+
 // DecapAndSeal is a stub for architectures where the Rust library is not supported.
 func DecapAndSeal(_ uuid.UUID, _, _ []byte) ([]byte, []byte, error) {
 	return nil, nil, fmt.Errorf("DecapAndSeal is not supported on this architecture")

keymanager/key_protection_service/key_custody_core/types.go

@@ -0,0 +1,14 @@
+package kpskcc
+
+import (
+	keymanager "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             *keymanager.HpkeAlgorithm
+	KEMPubKey             []byte
+	RemainingLifespanSecs uint64
+}

keymanager/key_protection_service/service.go

@@ -1,30 +1,35 @@
 // 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 operations.
+// to provide a Go-native interface for cryptographic operations and key management.
 package keyprotectionservice
 
 import (
+	kpskcc "github.com/google/go-tpm-tools/keymanager/key_protection_service/key_custody_core"
 	"github.com/google/uuid"
 
-	kpscc "github.com/google/go-tpm-tools/keymanager/key_protection_service/key_custody_core"
 	keymanager "github.com/google/go-tpm-tools/keymanager/km_common/proto"
 )
 
 // KeyProtectionService provides the core key custody operations.
 type KeyProtectionService interface {
 	GenerateKEMKeypair(algo *keymanager.HpkeAlgorithm, bindingPubKey []byte, lifespanSecs uint64) (uuid.UUID, []byte, error)
 	DecapAndSeal(kemUUID uuid.UUID, encapsulatedKey, aad []byte) ([]byte, []byte, error)
+	EnumerateKEMKeys(limit, offset int) ([]kpskcc.KEMKeyInfo, bool, error)
 }
 
 // defaultKPS implements KeyProtectionService by delegating to the KPS KCC FFI.
 type defaultKPS struct{}
 
 func (d *defaultKPS) GenerateKEMKeypair(algo *keymanager.HpkeAlgorithm, bindingPubKey []byte, lifespanSecs uint64) (uuid.UUID, []byte, error) {
-	return kpscc.GenerateKEMKeypair(algo, bindingPubKey, lifespanSecs)
+	return kpskcc.GenerateKEMKeypair(algo, bindingPubKey, lifespanSecs)
 }
 
 func (d *defaultKPS) DecapAndSeal(kemUUID uuid.UUID, encapsulatedKey, aad []byte) ([]byte, []byte, error) {
-	return kpscc.DecapAndSeal(kemUUID, encapsulatedKey, aad)
+	return kpskcc.DecapAndSeal(kemUUID, encapsulatedKey, aad)
+}
+
+func (d *defaultKPS) EnumerateKEMKeys(limit, offset int) ([]kpskcc.KEMKeyInfo, bool, error) {
+	return kpskcc.EnumerateKEMKeys(limit, offset)
 }
 
 // Service implements KeyProtectionService by delegating to an underlying KeyProtectionService.
@@ -57,3 +62,8 @@ func (s *Service) GenerateKEMKeypair(algo *keymanager.HpkeAlgorithm, bindingPubK
 func (s *Service) DecapAndSeal(kemUUID uuid.UUID, encapsulatedKey, aad []byte) ([]byte, []byte, error) {
 	return s.kps.DecapAndSeal(kemUUID, encapsulatedKey, aad)
 }
+
+// EnumerateKEMKeys enumerates active KEM keys up to limit by offset calling the underlying KPS.
+func (s *Service) EnumerateKEMKeys(limit, offset int) ([]kpskcc.KEMKeyInfo, bool, error) {
+	return s.kps.EnumerateKEMKeys(limit, offset)
+}

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"
 
 	keymanager "github.com/google/go-tpm-tools/keymanager/km_common/proto"
@@ -12,6 +13,7 @@ import (
 type mockKPS struct {
 	generateKEMKeypairFn func(algo *keymanager.HpkeAlgorithm, bindingPubKey []byte, lifespanSecs uint64) (uuid.UUID, []byte, error)
 	decapAndSealFn       func(kemUUID uuid.UUID, encapsulatedKey, aad []byte) ([]byte, []byte, error)
+	enumerateKEMKeysFn   func(limit, offset int) ([]kpskcc.KEMKeyInfo, bool, error)
 }
 
 func (m *mockKPS) GenerateKEMKeypair(algo *keymanager.HpkeAlgorithm, bindingPubKey []byte, lifespanSecs uint64) (uuid.UUID, []byte, error) {
@@ -21,6 +23,13 @@ func (m *mockKPS) GenerateKEMKeypair(algo *keymanager.HpkeAlgorithm, bindingPubK
 	return uuid.Nil, nil, nil
 }
 
+func (m *mockKPS) EnumerateKEMKeys(limit, offset int) ([]kpskcc.KEMKeyInfo, bool, error) {
+	if m.enumerateKEMKeysFn != nil {
+		return m.enumerateKEMKeysFn(limit, offset)
+	}
+	return nil, false, nil
+}
+
 func (m *mockKPS) DecapAndSeal(kemUUID uuid.UUID, encapsulatedKey, aad []byte) ([]byte, []byte, error) {
 	if m.decapAndSealFn != nil {
 		return m.decapAndSealFn(kemUUID, encapsulatedKey, aad)
@@ -76,6 +85,56 @@ func TestServiceGenerateKEMKeypairError(t *testing.T) {
 	}
 }
 
+func TestServiceEnumerateKEMKeysSuccess(t *testing.T) {
+	expectedKeys := []kpskcc.KEMKeyInfo{
+		{
+			ID: uuid.New(),
+			Algorithm: &keymanager.HpkeAlgorithm{
+				Kem:  keymanager.KemAlgorithm_KEM_ALGORITHM_DHKEM_X25519_HKDF_SHA256,
+				Kdf:  keymanager.KdfAlgorithm_KDF_ALGORITHM_HKDF_SHA256,
+				Aead: keymanager.AeadAlgorithm_AEAD_ALGORITHM_AES_256_GCM,
+			},
+			KEMPubKey:             make([]byte, 32),
+			RemainingLifespanSecs: 3500,
+		},
+	}
+
+	mock := &mockKPS{
+		enumerateKEMKeysFn: 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
+		},
+	}
+	svc := newServiceWithKPS(mock)
+
+	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) {
+	mock := &mockKPS{
+		enumerateKEMKeysFn: func(_, _ int) ([]kpskcc.KEMKeyInfo, bool, error) {
+			return nil, false, fmt.Errorf("enumerate error")
+		},
+	}
+	svc := newServiceWithKPS(mock)
+
+	_, _, err := svc.EnumerateKEMKeys(100, 0)
+	if err == nil {
+		t.Fatal("expected error, got nil")
+	}
+}
+
 func TestServiceDecapAndSealSuccess(t *testing.T) {
 	kemUUID := uuid.New()
 	expectedSealEnc := []byte("seal-enc-key")

keymanager/workload_service/proto_enums.go

@@ -8,9 +8,6 @@ import (
 	keymanager "github.com/google/go-tpm-tools/keymanager/km_common/proto"
 )
 
-// These enum values mirror the proto definitions in PROTOS.md and are used by
-// the WSD JSON API contract.
-
 // KemAlgorithm represents the requested KEM algorithm.
 type KemAlgorithm int32
 
@@ -97,3 +94,95 @@ func (k KemAlgorithm) ToHpkeAlgorithm() (*keymanager.HpkeAlgorithm, error) {
 		return nil, fmt.Errorf("unsupported algorithm: %s", k)
 	}
 }
+
+// KdfAlgorithm represents the requested KDF algorithm.
+type KdfAlgorithm int32
+
+const (
+	// KdfAlgorithmUnspecified indicates an unspecified KDF algorithm.
+	KdfAlgorithmUnspecified KdfAlgorithm = 0
+	// KdfAlgorithmHKDFSHA256 specifies the HKDF-SHA256 KDF algorithm.
+	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)
+}
+
+// MarshalJSON converts a KdfAlgorithm enum value to its JSON string representation.
+func (k KdfAlgorithm) MarshalJSON() ([]byte, error) {
+	return json.Marshal(k.String())
+}
+
+// UnmarshalJSON converts a JSON string back into a KdfAlgorithm enum value.
+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 indicates an unspecified AEAD algorithm.
+	AeadAlgorithmUnspecified AeadAlgorithm = 0
+	// AeadAlgorithmAES256GCM specifies the AES-256-GCM AEAD algorithm.
+	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)
+}
+
+// MarshalJSON converts an AeadAlgorithm enum value to its JSON string representation.
+func (k AeadAlgorithm) MarshalJSON() ([]byte, error) {
+	return json.Marshal(k.String())
+}
+
+// UnmarshalJSON converts a JSON string back into an AeadAlgorithm enum value.
+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)
+}