Skip to content

Commit 8f0c3d6

Browse files
atulpatildbzatulpatildbz
authored
[KeyManager] Update API signatures for key generation and enumeration (google#702)
Modify the GenerateKeyResponse and KeyInfo structs to align with the latest API specifications. Changes: * Add PubKey field containing the algorithm details and public key. * Add KeyProtectionMechanism field set to KEY_PROTECTION_VM_EMULATED. * Replace RemainingLifespan with ExpirationTime (Unix timestamp). * Update server implementation and tests to populate and verify these new fields accordingly.
1 parent 49d04ab commit 8f0c3d6

File tree

3 files changed

+59
-58
lines changed

3 files changed

+59
-58
lines changed

‎keymanager/workload_service/integration_test.go‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -47,7 +47,7 @@ func TestIntegrationGenerateKeysEndToEnd(t *testing.T) {
4747

4848
reqBody, err := json.Marshal(GenerateKeyRequest{
4949
Algorithm: AlgorithmDetails{Type: "kem", Params: AlgorithmParams{KemID: KemAlgorithmDHKEMX25519HKDFSHA256}},
50-
Lifespan: ProtoDuration{Seconds: 3600},
50+
Lifespan: 3600,
5151
})
5252
if err != nil {
5353
t.Fatalf("failed to marshal request: %v", err)
@@ -98,7 +98,7 @@ func TestIntegrationGenerateKeysUniqueMappings(t *testing.T) {
9898
for i := 0; i < 2; i++ {
9999
reqBody, err := json.Marshal(GenerateKeyRequest{
100100
Algorithm: AlgorithmDetails{Type: "kem", Params: AlgorithmParams{KemID: KemAlgorithmDHKEMX25519HKDFSHA256}},
101-
Lifespan: ProtoDuration{Seconds: 3600},
101+
Lifespan: 3600,
102102
})
103103
if err != nil {
104104
t.Fatalf("call %d: failed to marshal request: %v", i+1, err)
@@ -149,7 +149,7 @@ func TestIntegrationDestroyKey(t *testing.T) {
149149
// 1. Generate a key first
150150
reqBody, _ := json.Marshal(GenerateKeyRequest{
151151
Algorithm: AlgorithmDetails{Type: "kem", Params: AlgorithmParams{KemID: KemAlgorithmDHKEMX25519HKDFSHA256}},
152-
Lifespan: ProtoDuration{Seconds: 3600},
152+
Lifespan: 3600,
153153
})
154154
reqGen := httptest.NewRequest(http.MethodPost, "/v1/keys:generate_key", bytes.NewReader(reqBody))
155155
reqGen.Header.Set("Content-Type", "application/json")
@@ -217,7 +217,7 @@ func TestIntegrationAutoDestroy(t *testing.T) {
217217
// 1. Generate a key with 1-second lifespan
218218
reqBody, _ := json.Marshal(GenerateKeyRequest{
219219
Algorithm: AlgorithmDetails{Type: "kem", Params: AlgorithmParams{KemID: KemAlgorithmDHKEMX25519HKDFSHA256}},
220-
Lifespan: ProtoDuration{Seconds: 1},
220+
Lifespan: 1,
221221
})
222222
reqGen := httptest.NewRequest(http.MethodPost, "/v1/keys:generate_key", bytes.NewReader(reqBody))
223223
reqGen.Header.Set("Content-Type", "application/json")
@@ -264,7 +264,7 @@ func TestIntegrationKeyClaims(t *testing.T) {
264264
// 1. Generate a KEM key
265265
reqBody, _ := json.Marshal(GenerateKeyRequest{
266266
Algorithm: AlgorithmDetails{Type: "kem", Params: AlgorithmParams{KemID: KemAlgorithmDHKEMX25519HKDFSHA256}},
267-
Lifespan: ProtoDuration{Seconds: 3600},
267+
Lifespan: 3600,
268268
})
269269
req := httptest.NewRequest(http.MethodPost, "/v1/keys:generate_key", bytes.NewReader(reqBody))
270270
w := httptest.NewRecorder()

‎keymanager/workload_service/server.go‎

Lines changed: 30 additions & 33 deletions
Original file line numberDiff line numberDiff line change
@@ -160,33 +160,10 @@ type KeyHandle struct {
160160
Handle string `json:"handle"`
161161
}
162162

163-
// ProtoDuration represents a google.protobuf.Duration in JSON (e.g. "3600s").
164-
type ProtoDuration struct {
165-
Seconds uint64
166-
}
167-
168-
// UnmarshalJSON parses a proto3 Duration JSON number (as seconds).
169-
func (d *ProtoDuration) UnmarshalJSON(data []byte) error {
170-
var v float64
171-
if err := json.Unmarshal(data, &v); err != nil {
172-
return fmt.Errorf("duration must be a numeric value (seconds): %w", err)
173-
}
174-
if v < 0 || v > math.MaxUint64 {
175-
return fmt.Errorf("duration %f out of range", v)
176-
}
177-
d.Seconds = uint64(v)
178-
return nil
179-
}
180-
181-
// MarshalJSON encodes as a proto3 Duration JSON number (as seconds).
182-
func (d ProtoDuration) MarshalJSON() ([]byte, error) {
183-
return json.Marshal(d.Seconds)
184-
}
185-
186163
// GenerateKeyRequest is the JSON body for POST /v1/keys:generate_key.
187164
type GenerateKeyRequest struct {
188165
Algorithm AlgorithmDetails `json:"algorithm"`
189-
Lifespan ProtoDuration `json:"lifespan"`
166+
Lifespan uint64 `json:"lifespan"`
190167
}
191168

192169
// DestroyRequest is the JSON body for POST /v1/keys:destroy.
@@ -196,7 +173,10 @@ type DestroyRequest struct {
196173

197174
// GenerateKeyResponse is returned by POST /v1/keys:generate_key.
198175
type GenerateKeyResponse struct {
199-
KeyHandle KeyHandle `json:"key_handle"`
176+
KeyHandle KeyHandle `json:"key_handle"`
177+
PubKey PubKeyInfo `json:"pub_key"`
178+
KeyProtectionMechanism string `json:"key_protection_mechanism"`
179+
ExpirationTime int64 `json:"expiration_time"`
200180
}
201181

202182
// AlgorithmParams represents the parameters for a specific algorithm type.
@@ -227,9 +207,10 @@ type EnumerateKeysResponse struct {
227207

228208
// KeyInfo contains information about a single key.
229209
type KeyInfo struct {
230-
KeyHandle KeyHandle `json:"key_handle"`
231-
PubKey PubKeyInfo `json:"pub_key"`
232-
RemainingLifespan ProtoDuration `json:"remaining_lifespan"`
210+
KeyHandle KeyHandle `json:"key_handle"`
211+
PubKey PubKeyInfo `json:"pub_key"`
212+
KeyProtectionMechanism string `json:"key_protection_mechanism"`
213+
ExpirationTime int64 `json:"expiration_time"`
233214
}
234215

235216
// PubKeyInfo contains the public key and its algorithm.
@@ -416,11 +397,15 @@ func (s *Server) handleGenerateKey(w http.ResponseWriter, r *http.Request) {
416397
return
417398
}
418399

419-
// Validate lifespan is positive.
420-
if req.Lifespan.Seconds == 0 {
400+
// Validate lifespan is positive and does not cause int64 overflow.
401+
if req.Lifespan == 0 {
421402
writeError(w, "lifespan must be greater than 0s", http.StatusBadRequest)
422403
return
423404
}
405+
if req.Lifespan > math.MaxInt64 {
406+
writeError(w, "lifespan exceeds maximum allowed value", http.StatusBadRequest)
407+
return
408+
}
424409

425410
switch req.Algorithm.Type {
426411
case "kem":
@@ -446,14 +431,14 @@ func (s *Server) generateKEMKey(w http.ResponseWriter, req GenerateKeyRequest) {
446431
}
447432

448433
// Generate binding keypair via WSD KCC FFI.
449-
bindingUUID, bindingPubKey, err := s.workloadService.GenerateBindingKeypair(algo, req.Lifespan.Seconds)
434+
bindingUUID, bindingPubKey, err := s.workloadService.GenerateBindingKeypair(algo, req.Lifespan)
450435
if err != nil {
451436
writeError(w, fmt.Sprintf("failed to generate binding keypair: %v", err), http.StatusInternalServerError)
452437
return
453438
}
454439

455440
// Generate KEM keypair via KPS KOL, passing the binding public key.
456-
kemUUID, _, err := s.keyProtectionService.GenerateKEMKeypair(algo, bindingPubKey, req.Lifespan.Seconds)
441+
kemUUID, kemPubKey, err := s.keyProtectionService.GenerateKEMKeypair(algo, bindingPubKey, req.Lifespan)
457442
if err != nil {
458443
writeError(w, fmt.Sprintf("failed to generate KEM keypair: %v", err), http.StatusInternalServerError)
459444
return
@@ -467,6 +452,17 @@ func (s *Server) generateKEMKey(w http.ResponseWriter, req GenerateKeyRequest) {
467452
// Return KEM UUID to workload.
468453
resp := GenerateKeyResponse{
469454
KeyHandle: KeyHandle{Handle: kemUUID.String()},
455+
PubKey: PubKeyInfo{
456+
Algorithm: AlgorithmDetails{
457+
Type: "kem",
458+
Params: AlgorithmParams{
459+
KemID: req.Algorithm.Params.KemID,
460+
},
461+
},
462+
PublicKey: base64.StdEncoding.EncodeToString(kemPubKey),
463+
},
464+
KeyProtectionMechanism: keymanager.KeyProtectionMechanism_KEY_PROTECTION_VM_EMULATED.String(),
465+
ExpirationTime: time.Now().Unix() + int64(req.Lifespan),
470466
}
471467
writeJSON(w, resp, http.StatusOK)
472468
}
@@ -520,7 +516,8 @@ func (s *Server) handleEnumerateKeys(w http.ResponseWriter, _ *http.Request) {
520516
},
521517
PublicKey: base64.StdEncoding.EncodeToString(key.KEMPubKey),
522518
},
523-
RemainingLifespan: ProtoDuration{Seconds: key.RemainingLifespanSecs},
519+
KeyProtectionMechanism: keymanager.KeyProtectionMechanism_KEY_PROTECTION_VM_EMULATED.String(),
520+
ExpirationTime: time.Now().Unix() + int64(key.RemainingLifespanSecs),
524521
})
525522
}
526523

‎keymanager/workload_service/server_test.go‎

Lines changed: 24 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -122,7 +122,7 @@ func validGenerateBody() []byte {
122122
KemID: KemAlgorithmDHKEMX25519HKDFSHA256,
123123
},
124124
},
125-
Lifespan: ProtoDuration{Seconds: 3600},
125+
Lifespan: 3600,
126126
})
127127
return body
128128
}
@@ -163,6 +163,15 @@ func TestHandleGenerateKeySuccess(t *testing.T) {
163163
if resp.KeyHandle.Handle != kemUUID.String() {
164164
t.Fatalf("expected KEM UUID %s, got %s", kemUUID, resp.KeyHandle.Handle)
165165
}
166+
if resp.PubKey.PublicKey != base64.StdEncoding.EncodeToString(kemPubKey) {
167+
t.Fatalf("expected KEM pub key %s, got %s", base64.StdEncoding.EncodeToString(kemPubKey), resp.PubKey.PublicKey)
168+
}
169+
if resp.KeyProtectionMechanism != keymanager.KeyProtectionMechanism_KEY_PROTECTION_VM_EMULATED.String() {
170+
t.Fatalf("expected %s, got %s", keymanager.KeyProtectionMechanism_KEY_PROTECTION_VM_EMULATED.String(), resp.KeyProtectionMechanism)
171+
}
172+
if resp.ExpirationTime <= time.Now().Unix() {
173+
t.Fatalf("expected expiration time in the future, got %d", resp.ExpirationTime)
174+
}
166175

167176
// Verify the binding public key was passed to KEM generator.
168177
if len(kemGen.receivedPubKey) != 32 {
@@ -216,19 +225,19 @@ func TestHandleGenerateKeyBadRequest(t *testing.T) {
216225
}{
217226
{
218227
name: "unsupported algorithm type",
219-
body: GenerateKeyRequest{Algorithm: AlgorithmDetails{Type: "mac", Params: AlgorithmParams{KemID: KemAlgorithmDHKEMX25519HKDFSHA256}}, Lifespan: ProtoDuration{Seconds: 3600}},
228+
body: GenerateKeyRequest{Algorithm: AlgorithmDetails{Type: "mac", Params: AlgorithmParams{KemID: KemAlgorithmDHKEMX25519HKDFSHA256}}, Lifespan: 3600},
220229
},
221230
{
222231
name: "unsupported algorithm",
223-
body: GenerateKeyRequest{Algorithm: AlgorithmDetails{Type: "kem", Params: AlgorithmParams{KemID: KemAlgorithmUnspecified}}, Lifespan: ProtoDuration{Seconds: 3600}},
232+
body: GenerateKeyRequest{Algorithm: AlgorithmDetails{Type: "kem", Params: AlgorithmParams{KemID: KemAlgorithmUnspecified}}, Lifespan: 3600},
224233
},
225234
{
226235
name: "zero lifespan",
227-
body: GenerateKeyRequest{Algorithm: AlgorithmDetails{Type: "kem", Params: AlgorithmParams{KemID: KemAlgorithmDHKEMX25519HKDFSHA256}}, Lifespan: ProtoDuration{Seconds: 0}},
236+
body: GenerateKeyRequest{Algorithm: AlgorithmDetails{Type: "kem", Params: AlgorithmParams{KemID: KemAlgorithmDHKEMX25519HKDFSHA256}}, Lifespan: 0},
228237
},
229238
{
230239
name: "missing algorithm (defaults to 0, type empty)",
231-
body: GenerateKeyRequest{Lifespan: ProtoDuration{Seconds: 3600}},
240+
body: GenerateKeyRequest{Lifespan: 3600},
232241
},
233242
}
234243

@@ -272,6 +281,7 @@ func TestHandleGenerateKeyBadJSON(t *testing.T) {
272281
{"lifespan as string", `{"algorithm":{"type":"kem","params":{"kem_id":"DHKEM_X25519_HKDF_SHA256"}},"lifespan":"3600"}`},
273282
{"lifespan as string with suffix", `{"algorithm":{"type":"kem","params":{"kem_id":"DHKEM_X25519_HKDF_SHA256"}},"lifespan":"3600s"}`},
274283
{"lifespan negative", `{"algorithm":{"type":"kem","params":{"kem_id":"DHKEM_X25519_HKDF_SHA256"}},"lifespan":-1}`},
284+
{"lifespan too large", `{"algorithm":{"type":"kem","params":{"kem_id":"DHKEM_X25519_HKDF_SHA256"}},"lifespan":9223372036854775808}`},
275285
}
276286

277287
for _, tc := range badBodies {
@@ -320,16 +330,6 @@ func TestHandleGenerateKeyFlexibleLifespan(t *testing.T) {
320330
body: `{"algorithm":{"type":"kem","params":{"kem_id":"DHKEM_X25519_HKDF_SHA256"}},"lifespan":3600}`,
321331
expected: 3600,
322332
},
323-
{
324-
name: "float seconds",
325-
body: `{"algorithm":{"type":"kem","params":{"kem_id":"DHKEM_X25519_HKDF_SHA256"}},"lifespan":1.5}`,
326-
expected: 1, // Truncated to 1
327-
},
328-
{
329-
name: "float seconds round down",
330-
body: `{"algorithm":{"type":"kem","params":{"kem_id":"DHKEM_X25519_HKDF_SHA256"}},"lifespan":3600.9}`,
331-
expected: 3600,
332-
},
333333
}
334334

335335
for _, tc := range tests {
@@ -522,18 +522,22 @@ func TestHandleEnumerateKeysWithKeys(t *testing.T) {
522522
if info1.PubKey.PublicKey != base64.StdEncoding.EncodeToString(kemPubKey1) {
523523
t.Fatalf("KEM pub key mismatch for kem1")
524524
}
525-
// BindingPubKey check removed
526-
if info1.RemainingLifespan.Seconds != 3500 {
527-
t.Fatalf("expected remaining lifespan 3500, got %d", info1.RemainingLifespan.Seconds)
525+
if info1.KeyProtectionMechanism != keymanager.KeyProtectionMechanism_KEY_PROTECTION_VM_EMULATED.String() {
526+
t.Fatalf("expected key protection mechanism %s, got %s", keymanager.KeyProtectionMechanism_KEY_PROTECTION_VM_EMULATED.String(), info1.KeyProtectionMechanism)
527+
}
528+
// Approximate check for expiration time
529+
if info1.ExpirationTime <= time.Now().Unix() {
530+
t.Fatalf("expected expiration time in the future, got %d", info1.ExpirationTime)
528531
}
529532

530533
// Verify key 2.
531534
info2, ok := found[kem2.String()]
532535
if !ok {
533536
t.Fatalf("expected kem2 %s in response", kem2)
534537
}
535-
if info2.RemainingLifespan.Seconds != 7100 {
536-
t.Fatalf("expected remaining lifespan 7100, got %d", info2.RemainingLifespan.Seconds)
538+
// Approximate check for expiration time
539+
if info2.ExpirationTime <= time.Now().Unix() {
540+
t.Fatalf("expected expiration time in the future, got %d", info2.ExpirationTime)
537541
}
538542
}
539543

0 commit comments

Comments
 (0)