Refactor Enumerate Keys API Go implementation

Author: atulpatildbz

keymanager/key_protection_service/key_custody_core/kps_key_custody_core_cgo.go

@@ -91,11 +91,15 @@ func EnumerateKEMKeys() ([]KEMKeyInfo, error) {
 		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)))
 
+		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,
-			KemAlgorithm:          int32(e.algorithm.kem),
-			KdfAlgorithm:          int32(e.algorithm.kdf),
-			AeadAlgorithm:         int32(e.algorithm.aead),
+			Algorithm:             algo,
 			KEMPubKey:             kemPubKey,
 			BindingPubKey:         bindingPubKey,
 			RemainingLifespanSecs: uint64(e.remaining_lifespan_secs),

keymanager/key_protection_service/key_custody_core/types.go

@@ -1,13 +1,15 @@
 package kpskcc
 
-import "github.com/google/uuid"
+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.
 // KEMKeyInfo holds metadata for a single KEM key returned by EnumerateKEMKeys.
 type KEMKeyInfo struct {
 	ID                    uuid.UUID
-	KemAlgorithm          int32
-	KdfAlgorithm          int32
-	AeadAlgorithm         int32
+	Algorithm             *algorithms.HpkeAlgorithm
 	KEMPubKey             []byte
 	BindingPubKey         []byte
 	RemainingLifespanSecs uint64

keymanager/key_protection_service/service_test.go

@@ -63,10 +63,12 @@ func TestServiceGenerateKEMKeypairError(t *testing.T) {
 func TestServiceEnumerateKEMKeysSuccess(t *testing.T) {
 	expectedKeys := []kpskcc.KEMKeyInfo{
 		{
-			ID:                    uuid.New(),
-			KemAlgorithm:          1,
-			KdfAlgorithm:          1,
-			AeadAlgorithm:         1,
+			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),
 			BindingPubKey:         make([]byte, 32),
 			RemainingLifespanSecs: 3500,

keymanager/workload_service/server.go

@@ -191,14 +191,14 @@ func (s *Server) handleEnumerateKeys(w http.ResponseWriter, r *http.Request) {
 			BoundKemInfo: &BoundKEMInfo{
 				KeyHandle: KeyHandle{Handle: k.ID.String()},
 				KemPubKey: KemPublicKey{
-					Algorithm: KemAlgorithm(k.KemAlgorithm),
+					Algorithm: KemAlgorithm(k.Algorithm.Kem),
 					PublicKey: base64.StdEncoding.EncodeToString(k.KEMPubKey),
 				},
 				BindingPubKey: HpkePublicKey{
 					Algorithm: HpkeAlgorithm{
-						Kem:  KemAlgorithm(k.KemAlgorithm),
-						Kdf:  KdfAlgorithm(k.KdfAlgorithm),
-						Aead: AeadAlgorithm(k.AeadAlgorithm),
+						Kem:  KemAlgorithm(k.Algorithm.Kem),
+						Kdf:  KdfAlgorithm(k.Algorithm.Kdf),
+						Aead: AeadAlgorithm(k.Algorithm.Aead),
 					},
 					PublicKey: base64.StdEncoding.EncodeToString(k.BindingPubKey),
 				},

keymanager/workload_service/server_test.go

@@ -205,9 +205,9 @@ func TestHandleGenerateKemBadJSON(t *testing.T) {
 		body string
 	}{
 		{"not json", "not json"},
-		{"lifespan as string", `{"algorithm":1,"key_protection_mechanism":2,"lifespan":"3600"}`},
-		{"lifespan as string with suffix", `{"algorithm":1,"key_protection_mechanism":2,"lifespan":"3600s"}`},
-		{"lifespan negative", `{"algorithm":1,"key_protection_mechanism":2,"lifespan":-1}`},
+		{"lifespan as integer", `{"algorithm":1,"key_protection_mechanism":2,"lifespan":3600}`},
+		{"lifespan missing s suffix", `{"algorithm":1,"key_protection_mechanism":2,"lifespan":"3600"}`},
+		{"lifespan negative", `{"algorithm":1,"key_protection_mechanism":2,"lifespan":"-1s"}`},
 	}
 
 	for _, tc := range badBodies {
@@ -342,19 +342,23 @@ func TestHandleEnumerateKeysWithKeys(t *testing.T) {
 	mockEnum := &mockKEMKeyEnumerator{
 		keys: []kpskcc.KEMKeyInfo{
 			{
-				ID:                    kem1,
-				KemAlgorithm:          1,
-				KdfAlgorithm:          1,
-				AeadAlgorithm:         1,
+				ID: kem1,
+				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:             kemPubKey1,
 				BindingPubKey:         bindingPubKey1,
 				RemainingLifespanSecs: 3500,
 			},
 			{
-				ID:                    kem2,
-				KemAlgorithm:          1,
-				KdfAlgorithm:          1,
-				AeadAlgorithm:         1,
+				ID: kem2,
+				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:             kemPubKey2,
 				BindingPubKey:         bindingPubKey2,
 				RemainingLifespanSecs: 7100,