Skip to content

Commit f26280f

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

File tree

4 files changed

+29
-109
lines changed

4 files changed

+29
-109
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(kpsSvc, &realWorkloadService{})
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 & 67 deletions
Original file line numberDiff line numberDiff line change
@@ -57,69 +57,12 @@ func (k *KemAlgorithm) UnmarshalJSON(data []byte) error {
5757
return fmt.Errorf("unknown KemAlgorithm: %q", s)
5858
}
5959

60-
// KeyProtectionMechanism represents the requested key protection backend.
61-
type KeyProtectionMechanism int32
62-
63-
const (
64-
KeyProtectionMechanismUnspecified KeyProtectionMechanism = 0
65-
// KeyProtectionMechanismDefault is the default but invalid value.
66-
KeyProtectionMechanismDefault KeyProtectionMechanism = 1
67-
// KeyProtectionMechanismVM specifies that the key is protected by the VM.
68-
KeyProtectionMechanismVM KeyProtectionMechanism = 2
69-
KeyProtectionMechanismVMEmulated KeyProtectionMechanism = 3
70-
)
71-
72-
var (
73-
keyProtectionMechanismToString = map[KeyProtectionMechanism]string{
74-
KeyProtectionMechanismUnspecified: "KEY_PROTECTION_UNSPECIFIED",
75-
KeyProtectionMechanismDefault: "DEFAULT",
76-
KeyProtectionMechanismVM: "KEY_PROTECTION_VM",
77-
KeyProtectionMechanismVMEmulated: "KEY_PROTECTION_VM_EMULATED",
78-
}
79-
stringToKeyProtectionMechanism = map[string]KeyProtectionMechanism{
80-
"KEY_PROTECTION_UNSPECIFIED": KeyProtectionMechanismUnspecified,
81-
"DEFAULT": KeyProtectionMechanismDefault,
82-
"KEY_PROTECTION_VM": KeyProtectionMechanismVM,
83-
"KEY_PROTECTION_VM_EMULATED": KeyProtectionMechanismVMEmulated,
84-
}
85-
)
86-
87-
func (k KeyProtectionMechanism) String() string {
88-
if s, ok := keyProtectionMechanismToString[k]; ok {
89-
return s
90-
}
91-
return fmt.Sprintf("KEY_PROTECTION_MECHANISM_UNKNOWN(%d)", k)
92-
}
93-
94-
// MarshalJSON converts a KeyProtectionMechanism enum value to its JSON string representation.
95-
func (k KeyProtectionMechanism) MarshalJSON() ([]byte, error) {
96-
return json.Marshal(k.String())
97-
}
98-
99-
// UnmarshalJSON parses a JSON string into a KeyProtectionMechanism enum value.
100-
func (k *KeyProtectionMechanism) UnmarshalJSON(data []byte) error {
101-
var s string
102-
if err := json.Unmarshal(data, &s); err != nil {
103-
return fmt.Errorf("KeyProtectionMechanism must be a string")
104-
}
105-
if v, ok := stringToKeyProtectionMechanism[s]; ok {
106-
*k = v
107-
return nil
108-
}
109-
return fmt.Errorf("unknown KeyProtectionMechanism: %q", s)
110-
}
111-
11260
// Supported algorithms and mechanisms.
11361
var (
11462
// SupportedKemAlgorithms is the source of truth for supported algorithms.
11563
SupportedKemAlgorithms = []KemAlgorithm{
11664
KemAlgorithmDHKEMX25519HKDFSHA256,
11765
}
118-
119-
// SupportedKeyProtectionMechanisms is the source of truth for supported mechanisms.
120-
SupportedKeyProtectionMechanisms = []KeyProtectionMechanism{
121-
KeyProtectionMechanismVMEmulated,
122-
}
12366
)
12467

12568
// IsSupported returns true if the KEM algorithm is supported.
@@ -132,16 +75,6 @@ func (k KemAlgorithm) IsSupported() bool {
13275
return false
13376
}
13477

135-
// IsSupported returns true if the key protection mechanism is supported.
136-
func (k KeyProtectionMechanism) IsSupported() bool {
137-
for _, supported := range SupportedKeyProtectionMechanisms {
138-
if k == supported {
139-
return true
140-
}
141-
}
142-
return false
143-
}
144-
14578
// SupportedKemAlgorithmsString returns a comma-separated list of supported KEM algorithms.
14679
func SupportedKemAlgorithmsString() string {
14780
var names []string

‎keymanager/workload_service/server.go‎

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

7171
// GenerateKeyRequest is the JSON body for POST /v1/keys:generate_key.
7272
type GenerateKeyRequest struct {
73-
Algorithm AlgorithmDetails `json:"algorithm"`
74-
KeyProtectionMechanism KeyProtectionMechanism `json:"key_protection_mechanism,omitempty"`
75-
Lifespan ProtoDuration `json:"lifespan"`
73+
Algorithm AlgorithmDetails `json:"algorithm"`
74+
Lifespan ProtoDuration `json:"lifespan"`
7675
}
7776

7877
// GenerateKeyResponse is returned by POST /v1/keys:generate_key.
@@ -173,32 +172,27 @@ func (s *Server) handleGenerateKey(w http.ResponseWriter, r *http.Request) {
173172
return
174173
}
175174

176-
if req.Algorithm.Type != "kem" {
177-
writeError(w, fmt.Sprintf("unsupported algorithm type: %q. Only 'kem' is supported.", req.Algorithm.Type), http.StatusBadRequest)
175+
// Validate lifespan is positive.
176+
if req.Lifespan.Seconds == 0 {
177+
writeError(w, "lifespan must be greater than 0s", http.StatusBadRequest)
178178
return
179179
}
180180

181+
switch req.Algorithm.Type {
182+
case "kem":
183+
s.generateKEMKey(w, req)
184+
default:
185+
writeError(w, fmt.Sprintf("unsupported algorithm type: %q. Only 'kem' is supported.", req.Algorithm.Type), http.StatusBadRequest)
186+
}
187+
}
188+
189+
func (s *Server) generateKEMKey(w http.ResponseWriter, req GenerateKeyRequest) {
181190
// Validate algorithm.
182191
if !req.Algorithm.Params.KemID.IsSupported() {
183192
writeError(w, fmt.Sprintf("unsupported algorithm: %s. Supported algorithms: %s", req.Algorithm.Params.KemID, SupportedKemAlgorithmsString()), http.StatusBadRequest)
184193
return
185194
}
186195

187-
// Validate keyProtectionMechanism.
188-
if req.KeyProtectionMechanism == KeyProtectionMechanismUnspecified {
189-
req.KeyProtectionMechanism = KeyProtectionMechanismVMEmulated
190-
}
191-
if !req.KeyProtectionMechanism.IsSupported() {
192-
writeError(w, fmt.Sprintf("unsupported keyProtectionMechanism: %s", req.KeyProtectionMechanism), http.StatusBadRequest)
193-
return
194-
}
195-
196-
// Validate lifespan is positive.
197-
if req.Lifespan.Seconds == 0 {
198-
writeError(w, "lifespan must be greater than 0s", http.StatusBadRequest)
199-
return
200-
}
201-
202196
// Construct the full HPKE algorithm suite based on the requested KEM.
203197
// We currently only support one suite.
204198
algo, err := req.Algorithm.Params.KemID.ToHpkeAlgorithm()

‎keymanager/workload_service/server_test.go‎

Lines changed: 11 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -57,8 +57,7 @@ func validGenerateBody() []byte {
5757
KemID: KemAlgorithmDHKEMX25519HKDFSHA256,
5858
},
5959
},
60-
KeyProtectionMechanism: KeyProtectionMechanismVMEmulated,
61-
Lifespan: ProtoDuration{Seconds: 3600},
60+
Lifespan: ProtoDuration{Seconds: 3600},
6261
})
6362
return body
6463
}
@@ -150,23 +149,19 @@ func TestHandleGenerateKeyBadRequest(t *testing.T) {
150149
}{
151150
{
152151
name: "unsupported algorithm type",
153-
body: GenerateKeyRequest{Algorithm: AlgorithmDetails{Type: "mac", Params: AlgorithmParams{KemID: KemAlgorithmDHKEMX25519HKDFSHA256}}, KeyProtectionMechanism: KeyProtectionMechanismVMEmulated, Lifespan: ProtoDuration{Seconds: 3600}},
152+
body: GenerateKeyRequest{Algorithm: AlgorithmDetails{Type: "mac", Params: AlgorithmParams{KemID: KemAlgorithmDHKEMX25519HKDFSHA256}}, Lifespan: ProtoDuration{Seconds: 3600}},
154153
},
155154
{
156155
name: "unsupported algorithm",
157-
body: GenerateKeyRequest{Algorithm: AlgorithmDetails{Type: "kem", Params: AlgorithmParams{KemID: KemAlgorithmUnspecified}}, KeyProtectionMechanism: KeyProtectionMechanismVMEmulated, Lifespan: ProtoDuration{Seconds: 3600}},
158-
},
159-
{
160-
name: "unsupported key protection mechanism",
161-
body: GenerateKeyRequest{Algorithm: AlgorithmDetails{Type: "kem", Params: AlgorithmParams{KemID: KemAlgorithmDHKEMX25519HKDFSHA256}}, KeyProtectionMechanism: KeyProtectionMechanism(99), Lifespan: ProtoDuration{Seconds: 3600}},
156+
body: GenerateKeyRequest{Algorithm: AlgorithmDetails{Type: "kem", Params: AlgorithmParams{KemID: KemAlgorithmUnspecified}}, Lifespan: ProtoDuration{Seconds: 3600}},
162157
},
163158
{
164159
name: "zero lifespan",
165-
body: GenerateKeyRequest{Algorithm: AlgorithmDetails{Type: "kem", Params: AlgorithmParams{KemID: KemAlgorithmDHKEMX25519HKDFSHA256}}, KeyProtectionMechanism: KeyProtectionMechanismVMEmulated, Lifespan: ProtoDuration{Seconds: 0}},
160+
body: GenerateKeyRequest{Algorithm: AlgorithmDetails{Type: "kem", Params: AlgorithmParams{KemID: KemAlgorithmDHKEMX25519HKDFSHA256}}, Lifespan: ProtoDuration{Seconds: 0}},
166161
},
167162
{
168163
name: "missing algorithm (defaults to 0, type empty)",
169-
body: GenerateKeyRequest{KeyProtectionMechanism: KeyProtectionMechanismVMEmulated, Lifespan: ProtoDuration{Seconds: 3600}},
164+
body: GenerateKeyRequest{Lifespan: ProtoDuration{Seconds: 3600}},
170165
},
171166
}
172167

@@ -207,9 +202,9 @@ func TestHandleGenerateKeyBadJSON(t *testing.T) {
207202
body string
208203
}{
209204
{"not json", "not json"},
210-
{"lifespan as string", `{"algorithm":1,"key_protection_mechanism":2,"lifespan":"3600"}`},
211-
{"lifespan as string with suffix", `{"algorithm":1,"key_protection_mechanism":2,"lifespan":"3600s"}`},
212-
{"lifespan negative", `{"algorithm":1,"key_protection_mechanism":2,"lifespan":-1}`},
205+
{"lifespan as string", `{"algorithm":1,"lifespan":"3600"}`},
206+
{"lifespan as string with suffix", `{"algorithm":1,"lifespan":"3600s"}`},
207+
{"lifespan negative", `{"algorithm":1,"lifespan":-1}`},
213208
}
214209

215210
for _, tc := range badBodies {
@@ -255,17 +250,17 @@ func TestHandleGenerateKeyFlexibleLifespan(t *testing.T) {
255250
}{
256251
{
257252
name: "integer seconds",
258-
body: `{"algorithm":{"type":"kem","params":{"kem_id":"DHKEM_X25519_HKDF_SHA256"}},"key_protection_mechanism":"KEY_PROTECTION_VM_EMULATED","lifespan":3600}`,
253+
body: `{"algorithm":{"type":"kem","params":{"kem_id":"DHKEM_X25519_HKDF_SHA256"}},"lifespan":3600}`,
259254
expected: 3600,
260255
},
261256
{
262257
name: "float seconds",
263-
body: `{"algorithm":{"type":"kem","params":{"kem_id":"DHKEM_X25519_HKDF_SHA256"}},"key_protection_mechanism":"KEY_PROTECTION_VM_EMULATED","lifespan":1.5}`,
258+
body: `{"algorithm":{"type":"kem","params":{"kem_id":"DHKEM_X25519_HKDF_SHA256"}},"lifespan":1.5}`,
264259
expected: 1, // Truncated to 1
265260
},
266261
{
267262
name: "float seconds round down",
268-
body: `{"algorithm":{"type":"kem","params":{"kem_id":"DHKEM_X25519_HKDF_SHA256"}},"key_protection_mechanism":"KEY_PROTECTION_VM_EMULATED","lifespan":3600.9}`,
263+
body: `{"algorithm":{"type":"kem","params":{"kem_id":"DHKEM_X25519_HKDF_SHA256"}},"lifespan":3600.9}`,
269264
expected: 3600,
270265
},
271266
}

0 commit comments

Comments
 (0)