[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"
@@ -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
+}

keymanager/key_protection_service/key_custody_core/kps_key_custody_core_stub.go

@@ -14,3 +14,8 @@ import (
 func GenerateKEMKeypair(_ *algorithms.HpkeAlgorithm, _ []byte, _ uint64) (uuid.UUID, []byte, error) {
 	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")
+}

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 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 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"
 
 	algorithms "github.com/google/go-tpm-tools/keymanager/km_common/proto"
@@ -14,18 +15,36 @@ 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)
+}
+
+// KeyCustodyCore defines the required FFI interactions for KPS.
+type KeyCustodyCore interface {
+	GenerateKEMKeypair(algo *algorithms.HpkeAlgorithm, bindingPubKey []byte, lifespanSecs uint64) (uuid.UUID, []byte, error)
+	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)
+	kcc KeyCustodyCore
 }
 
-// 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 implementation.
+func NewService(kcc KeyCustodyCore) *Service {
+	return &Service{
+		kcc: kcc,
+	}
 }
 
 // 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)
+	return s.kcc.GenerateKEMKeypair(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.kcc.EnumerateKEMKeys(limit, offset)
 }

keymanager/key_protection_service/service_test.go

@@ -4,26 +4,48 @@ 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"
 )
 
+type mockKCC struct {
+	generateFn  func(algo *algorithms.HpkeAlgorithm, bindingPubKey []byte, lifespanSecs uint64) (uuid.UUID, []byte, error)
+	enumerateFn func(limit, offset int) ([]kpskcc.KEMKeyInfo, bool, error)
+}
+
+func (m *mockKCC) GenerateKEMKeypair(algo *algorithms.HpkeAlgorithm, bindingPubKey []byte, lifespanSecs uint64) (uuid.UUID, []byte, error) {
+	if m.generateFn != nil {
+		return m.generateFn(algo, bindingPubKey, lifespanSecs)
+	}
+	return uuid.Nil, nil, nil
+}
+
+func (m *mockKCC) EnumerateKEMKeys(limit, offset int) ([]kpskcc.KEMKeyInfo, bool, error) {
+	if m.enumerateFn != nil {
+		return m.enumerateFn(limit, offset)
+	}
+	return nil, false, nil
+}
+
 func TestServiceGenerateKEMKeypairSuccess(t *testing.T) {
 	expectedUUID := uuid.New()
 	expectedPubKey := make([]byte, 32)
 	for i := range expectedPubKey {
 		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(&mockKCC{
+		generateFn: 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
+		},
 	})
 
 	id, pubKey, err := svc.GenerateKEMKeypair(&algorithms.HpkeAlgorithm{}, make([]byte, 32), 7200)
@@ -39,12 +61,62 @@ 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(&mockKCC{
+		generateFn: func(_ *algorithms.HpkeAlgorithm, _ []byte, _ uint64) (uuid.UUID, []byte, error) {
+			return uuid.Nil, nil, fmt.Errorf("FFI error")
+		},
 	})
 
 	_, _, 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(&mockKCC{
+		enumerateFn: 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(&mockKCC{
+		enumerateFn: func(_, _ 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")
+	}
+}

keymanager/workload_service/integration_test.go

@@ -29,15 +29,14 @@ func TestIntegrationGenerateKeysEndToEnd(t *testing.T) {
 	kpsSvc := kps.NewService(kpskcc.GenerateKEMKeypair)
 	srv := NewServer(kpsSvc, &realWorkloadService{})
 
-	reqBody, err := json.Marshal(GenerateKemRequest{
-		Algorithm:              KemAlgorithmDHKEMX25519HKDFSHA256,
-		KeyProtectionMechanism: KeyProtectionMechanismVM,
-		Lifespan:               ProtoDuration{Seconds: 3600},
+	reqBody, err := json.Marshal(GenerateKeyRequest{
+		Algorithm: AlgorithmDetails{Type: "kem", Params: AlgorithmParams{KemID: KemAlgorithmDHKEMX25519HKDFSHA256}},
+		Lifespan:  ProtoDuration{Seconds: 3600},
 	})
 	if err != nil {
 		t.Fatalf("failed to marshal request: %v", err)
 	}
-	req := httptest.NewRequest(http.MethodPost, "/v1/keys:generate_kem", bytes.NewReader(reqBody))
+	req := httptest.NewRequest(http.MethodPost, "/v1/keys:generate_key", bytes.NewReader(reqBody))
 	req.Header.Set("Content-Type", "application/json")
 	w := httptest.NewRecorder()
 	srv.Handler().ServeHTTP(w, req)
@@ -46,7 +45,7 @@ func TestIntegrationGenerateKeysEndToEnd(t *testing.T) {
 		t.Fatalf("expected status 200, got %d: %s", w.Code, w.Body.String())
 	}
 
-	var resp GenerateKemResponse
+	var resp GenerateKeyResponse
 	if err := json.NewDecoder(w.Body).Decode(&resp); err != nil {
 		t.Fatalf("failed to decode response: %v", err)
 	}
@@ -78,15 +77,14 @@ func TestIntegrationGenerateKeysUniqueMappings(t *testing.T) {
 	// Generate two key sets.
 	var kemUUIDs [2]uuid.UUID
 	for i := 0; i < 2; i++ {
-		reqBody, err := json.Marshal(GenerateKemRequest{
-			Algorithm:              KemAlgorithmDHKEMX25519HKDFSHA256,
-			KeyProtectionMechanism: KeyProtectionMechanismVM,
-			Lifespan:               ProtoDuration{Seconds: 3600},
+		reqBody, err := json.Marshal(GenerateKeyRequest{
+			Algorithm: AlgorithmDetails{Type: "kem", Params: AlgorithmParams{KemID: KemAlgorithmDHKEMX25519HKDFSHA256}},
+			Lifespan:  ProtoDuration{Seconds: 3600},
 		})
 		if err != nil {
 			t.Fatalf("call %d: failed to marshal request: %v", i+1, err)
 		}
-		req := httptest.NewRequest(http.MethodPost, "/v1/keys:generate_kem", bytes.NewReader(reqBody))
+		req := httptest.NewRequest(http.MethodPost, "/v1/keys:generate_key", bytes.NewReader(reqBody))
 		req.Header.Set("Content-Type", "application/json")
 		w := httptest.NewRecorder()
 		srv.Handler().ServeHTTP(w, req)
@@ -95,7 +93,7 @@ func TestIntegrationGenerateKeysUniqueMappings(t *testing.T) {
 			t.Fatalf("call %d: expected status 200, got %d: %s", i+1, w.Code, w.Body.String())
 		}
 
-		var resp GenerateKemResponse
+		var resp GenerateKeyResponse
 		if err := json.NewDecoder(w.Body).Decode(&resp); err != nil {
 			t.Fatalf("call %d: failed to decode response: %v", i+1, err)
 		}