Skip to content

Commit a7a2b09

Browse files
atulpatildbzatulpatildbz
committed
[keymanager] address review comments
* Refactor GenerateKey algorithm validation to support future key types * remove KeyProtectionMechanism
1 parent ff4220e commit a7a2b09

File tree

4 files changed

+30
-105
lines changed

4 files changed

+30
-105
lines changed

‎keymanager/workload_service/integration_test.go‎

Lines changed: 4 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -30,9 +30,8 @@ func TestIntegrationGenerateKeysEndToEnd(t *testing.T) {
3030
srv := NewServer(&realBindingKeyGen{}, kpsSvc)
3131

3232
reqBody, err := json.Marshal(GenerateKeyRequest{
33-
Algorithm: AlgorithmDetails{Type: "kem", Params: AlgorithmParams{KemID: KemAlgorithmDHKEMX25519HKDFSHA256}},
34-
KeyProtectionMechanism: KeyProtectionMechanismVMEmulated,
35-
Lifespan: ProtoDuration{Seconds: 3600},
33+
Algorithm: AlgorithmDetails{Type: "kem", Params: AlgorithmParams{KemID: KemAlgorithmDHKEMX25519HKDFSHA256}},
34+
Lifespan: ProtoDuration{Seconds: 3600},
3635
})
3736
if err != nil {
3837
t.Fatalf("failed to marshal request: %v", err)
@@ -79,9 +78,8 @@ func TestIntegrationGenerateKeysUniqueMappings(t *testing.T) {
7978
var kemUUIDs [2]uuid.UUID
8079
for i := 0; i < 2; i++ {
8180
reqBody, err := json.Marshal(GenerateKeyRequest{
82-
Algorithm: AlgorithmDetails{Type: "kem", Params: AlgorithmParams{KemID: KemAlgorithmDHKEMX25519HKDFSHA256}},
83-
KeyProtectionMechanism: KeyProtectionMechanismVMEmulated,
84-
Lifespan: ProtoDuration{Seconds: 3600},
81+
Algorithm: AlgorithmDetails{Type: "kem", Params: AlgorithmParams{KemID: KemAlgorithmDHKEMX25519HKDFSHA256}},
82+
Lifespan: ProtoDuration{Seconds: 3600},
8583
})
8684
if err != nil {
8785
t.Fatalf("call %d: failed to marshal request: %v", i+1, err)

‎keymanager/workload_service/proto_enums.go‎

Lines changed: 0 additions & 63 deletions
Original file line numberDiff line numberDiff line change
@@ -53,65 +53,12 @@ func (k *KemAlgorithm) UnmarshalJSON(data []byte) error {
5353
return fmt.Errorf("unknown KemAlgorithm: %q", s)
5454
}
5555

56-
// KeyProtectionMechanism represents the requested key protection backend.
57-
type KeyProtectionMechanism int32
58-
59-
const (
60-
KeyProtectionMechanismUnspecified KeyProtectionMechanism = 0
61-
KeyProtectionMechanismDefault KeyProtectionMechanism = 1
62-
KeyProtectionMechanismVM KeyProtectionMechanism = 2
63-
KeyProtectionMechanismVMEmulated KeyProtectionMechanism = 3
64-
)
65-
66-
var (
67-
keyProtectionMechanismToString = map[KeyProtectionMechanism]string{
68-
KeyProtectionMechanismUnspecified: "KEY_PROTECTION_UNSPECIFIED",
69-
KeyProtectionMechanismDefault: "DEFAULT",
70-
KeyProtectionMechanismVM: "KEY_PROTECTION_VM",
71-
KeyProtectionMechanismVMEmulated: "KEY_PROTECTION_VM_EMULATED",
72-
}
73-
stringToKeyProtectionMechanism = map[string]KeyProtectionMechanism{
74-
"KEY_PROTECTION_UNSPECIFIED": KeyProtectionMechanismUnspecified,
75-
"DEFAULT": KeyProtectionMechanismDefault,
76-
"KEY_PROTECTION_VM": KeyProtectionMechanismVM,
77-
"KEY_PROTECTION_VM_EMULATED": KeyProtectionMechanismVMEmulated,
78-
}
79-
)
80-
81-
func (k KeyProtectionMechanism) String() string {
82-
if s, ok := keyProtectionMechanismToString[k]; ok {
83-
return s
84-
}
85-
return fmt.Sprintf("KEY_PROTECTION_MECHANISM_UNKNOWN(%d)", k)
86-
}
87-
88-
func (k KeyProtectionMechanism) MarshalJSON() ([]byte, error) {
89-
return json.Marshal(k.String())
90-
}
91-
92-
func (k *KeyProtectionMechanism) UnmarshalJSON(data []byte) error {
93-
var s string
94-
if err := json.Unmarshal(data, &s); err != nil {
95-
return fmt.Errorf("KeyProtectionMechanism must be a string")
96-
}
97-
if v, ok := stringToKeyProtectionMechanism[s]; ok {
98-
*k = v
99-
return nil
100-
}
101-
return fmt.Errorf("unknown KeyProtectionMechanism: %q", s)
102-
}
103-
10456
// Supported algorithms and mechanisms.
10557
var (
10658
// SupportedKemAlgorithms is the source of truth for supported algorithms.
10759
SupportedKemAlgorithms = []KemAlgorithm{
10860
KemAlgorithmDHKEMX25519HKDFSHA256,
10961
}
110-
111-
// SupportedKeyProtectionMechanisms is the source of truth for supported mechanisms.
112-
SupportedKeyProtectionMechanisms = []KeyProtectionMechanism{
113-
KeyProtectionMechanismVMEmulated,
114-
}
11562
)
11663

11764
// IsSupported returns true if the KEM algorithm is supported.
@@ -124,16 +71,6 @@ func (k KemAlgorithm) IsSupported() bool {
12471
return false
12572
}
12673

127-
// IsSupported returns true if the key protection mechanism is supported.
128-
func (k KeyProtectionMechanism) IsSupported() bool {
129-
for _, supported := range SupportedKeyProtectionMechanisms {
130-
if k == supported {
131-
return true
132-
}
133-
}
134-
return false
135-
}
136-
13774
// SupportedKemAlgorithmsString returns a comma-separated list of supported KEM algorithms.
13875
func SupportedKemAlgorithmsString() string {
13976
var names []string

‎keymanager/workload_service/server.go‎

Lines changed: 15 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -58,9 +58,8 @@ func (d ProtoDuration) MarshalJSON() ([]byte, error) {
5858

5959
// GenerateKeyRequest is the JSON body for POST /v1/keys:generate_key.
6060
type GenerateKeyRequest struct {
61-
Algorithm AlgorithmDetails `json:"algorithm"`
62-
KeyProtectionMechanism KeyProtectionMechanism `json:"key_protection_mechanism,omitempty"`
63-
Lifespan ProtoDuration `json:"lifespan"`
61+
Algorithm AlgorithmDetails `json:"algorithm"`
62+
Lifespan ProtoDuration `json:"lifespan"`
6463
}
6564

6665
// GenerateKeyResponse is returned by POST /v1/keys:generate_key.
@@ -153,29 +152,25 @@ func (s *Server) handleGenerateKey(w http.ResponseWriter, r *http.Request) {
153152
return
154153
}
155154

156-
if req.Algorithm.Type != "kem" {
157-
writeError(w, fmt.Sprintf("unsupported algorithm type: %q. Only 'kem' is supported.", req.Algorithm.Type), http.StatusBadRequest)
158-
return
159-
}
160155

161-
// Validate algorithm.
162-
if !req.Algorithm.Params.KemID.IsSupported() {
163-
writeError(w, fmt.Sprintf("unsupported algorithm: %s. Supported algorithms: %s", req.Algorithm.Params.KemID, SupportedKemAlgorithmsString()), http.StatusBadRequest)
156+
// Validate lifespan is positive.
157+
if req.Lifespan.Seconds == 0 {
158+
writeError(w, "lifespan must be greater than 0s", http.StatusBadRequest)
164159
return
165160
}
166161

167-
// Validate keyProtectionMechanism.
168-
if req.KeyProtectionMechanism == KeyProtectionMechanismUnspecified {
169-
req.KeyProtectionMechanism = KeyProtectionMechanismVMEmulated
170-
}
171-
if !req.KeyProtectionMechanism.IsSupported() {
172-
writeError(w, fmt.Sprintf("unsupported keyProtectionMechanism: %s", req.KeyProtectionMechanism), http.StatusBadRequest)
173-
return
162+
switch req.Algorithm.Type {
163+
case "kem":
164+
s.generateKEMKey(w, req)
165+
default:
166+
writeError(w, fmt.Sprintf("unsupported algorithm type: %q. Only 'kem' is supported.", req.Algorithm.Type), http.StatusBadRequest)
174167
}
168+
}
175169

176-
// Validate lifespan is positive.
177-
if req.Lifespan.Seconds == 0 {
178-
writeError(w, "lifespan must be greater than 0s", http.StatusBadRequest)
170+
func (s *Server) generateKEMKey(w http.ResponseWriter, req GenerateKeyRequest) {
171+
// Validate algorithm.
172+
if !req.Algorithm.Params.KemID.IsSupported() {
173+
writeError(w, fmt.Sprintf("unsupported algorithm: %s. Supported algorithms: %s", req.Algorithm.Params.KemID, SupportedKemAlgorithmsString()), http.StatusBadRequest)
179174
return
180175
}
181176

‎keymanager/workload_service/server_test.go‎

Lines changed: 11 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -48,8 +48,7 @@ func validGenerateBody() []byte {
4848
KemID: KemAlgorithmDHKEMX25519HKDFSHA256,
4949
},
5050
},
51-
KeyProtectionMechanism: KeyProtectionMechanismVMEmulated,
52-
Lifespan: ProtoDuration{Seconds: 3600},
51+
Lifespan: ProtoDuration{Seconds: 3600},
5352
})
5453
return body
5554
}
@@ -141,23 +140,19 @@ func TestHandleGenerateKeyBadRequest(t *testing.T) {
141140
}{
142141
{
143142
name: "unsupported algorithm type",
144-
body: GenerateKeyRequest{Algorithm: AlgorithmDetails{Type: "mac", Params: AlgorithmParams{KemID: KemAlgorithmDHKEMX25519HKDFSHA256}}, KeyProtectionMechanism: KeyProtectionMechanismVMEmulated, Lifespan: ProtoDuration{Seconds: 3600}},
143+
body: GenerateKeyRequest{Algorithm: AlgorithmDetails{Type: "mac", Params: AlgorithmParams{KemID: KemAlgorithmDHKEMX25519HKDFSHA256}}, Lifespan: ProtoDuration{Seconds: 3600}},
145144
},
146145
{
147146
name: "unsupported algorithm",
148-
body: GenerateKeyRequest{Algorithm: AlgorithmDetails{Type: "kem", Params: AlgorithmParams{KemID: KemAlgorithmUnspecified}}, KeyProtectionMechanism: KeyProtectionMechanismVMEmulated, Lifespan: ProtoDuration{Seconds: 3600}},
149-
},
150-
{
151-
name: "unsupported key protection mechanism",
152-
body: GenerateKeyRequest{Algorithm: AlgorithmDetails{Type: "kem", Params: AlgorithmParams{KemID: KemAlgorithmDHKEMX25519HKDFSHA256}}, KeyProtectionMechanism: KeyProtectionMechanism(99), Lifespan: ProtoDuration{Seconds: 3600}},
147+
body: GenerateKeyRequest{Algorithm: AlgorithmDetails{Type: "kem", Params: AlgorithmParams{KemID: KemAlgorithmUnspecified}}, Lifespan: ProtoDuration{Seconds: 3600}},
153148
},
154149
{
155150
name: "zero lifespan",
156-
body: GenerateKeyRequest{Algorithm: AlgorithmDetails{Type: "kem", Params: AlgorithmParams{KemID: KemAlgorithmDHKEMX25519HKDFSHA256}}, KeyProtectionMechanism: KeyProtectionMechanismVMEmulated, Lifespan: ProtoDuration{Seconds: 0}},
151+
body: GenerateKeyRequest{Algorithm: AlgorithmDetails{Type: "kem", Params: AlgorithmParams{KemID: KemAlgorithmDHKEMX25519HKDFSHA256}}, Lifespan: ProtoDuration{Seconds: 0}},
157152
},
158153
{
159154
name: "missing algorithm (defaults to 0, type empty)",
160-
body: GenerateKeyRequest{KeyProtectionMechanism: KeyProtectionMechanismVMEmulated, Lifespan: ProtoDuration{Seconds: 3600}},
155+
body: GenerateKeyRequest{Lifespan: ProtoDuration{Seconds: 3600}},
161156
},
162157
}
163158

@@ -198,9 +193,9 @@ func TestHandleGenerateKeyBadJSON(t *testing.T) {
198193
body string
199194
}{
200195
{"not json", "not json"},
201-
{"lifespan as string", `{"algorithm":1,"key_protection_mechanism":2,"lifespan":"3600"}`},
202-
{"lifespan as string with suffix", `{"algorithm":1,"key_protection_mechanism":2,"lifespan":"3600s"}`},
203-
{"lifespan negative", `{"algorithm":1,"key_protection_mechanism":2,"lifespan":-1}`},
196+
{"lifespan as string", `{"algorithm":1,"lifespan":"3600"}`},
197+
{"lifespan as string with suffix", `{"algorithm":1,"lifespan":"3600s"}`},
198+
{"lifespan negative", `{"algorithm":1,"lifespan":-1}`},
204199
}
205200

206201
for _, tc := range badBodies {
@@ -246,17 +241,17 @@ func TestHandleGenerateKeyFlexibleLifespan(t *testing.T) {
246241
}{
247242
{
248243
name: "integer seconds",
249-
body: `{"algorithm":{"type":"kem","params":{"kem_id":"DHKEM_X25519_HKDF_SHA256"}},"key_protection_mechanism":"KEY_PROTECTION_VM_EMULATED","lifespan":3600}`,
244+
body: `{"algorithm":{"type":"kem","params":{"kem_id":"DHKEM_X25519_HKDF_SHA256"}},"lifespan":3600}`,
250245
expected: 3600,
251246
},
252247
{
253248
name: "float seconds",
254-
body: `{"algorithm":{"type":"kem","params":{"kem_id":"DHKEM_X25519_HKDF_SHA256"}},"key_protection_mechanism":"KEY_PROTECTION_VM_EMULATED","lifespan":1.5}`,
249+
body: `{"algorithm":{"type":"kem","params":{"kem_id":"DHKEM_X25519_HKDF_SHA256"}},"lifespan":1.5}`,
255250
expected: 1, // Truncated to 1
256251
},
257252
{
258253
name: "float seconds round down",
259-
body: `{"algorithm":{"type":"kem","params":{"kem_id":"DHKEM_X25519_HKDF_SHA256"}},"key_protection_mechanism":"KEY_PROTECTION_VM_EMULATED","lifespan":3600.9}`,
254+
body: `{"algorithm":{"type":"kem","params":{"kem_id":"DHKEM_X25519_HKDF_SHA256"}},"lifespan":3600.9}`,
260255
expected: 3600,
261256
},
262257
}

0 commit comments

Comments
 (0)