[keymanager/wsd] Add Go orchestration layer (KOL) for WSD key generation by atulpatildbz · Pull Request #652 · google/go-tpm-tools

atulpatildbz · 2026-02-09T10:06:44Z

Summary

Implement the Go Key Orchestration Layer (KOL) for the Workload Service Daemon (WSD)
Expose POST /v1/keys:generate_kem endpoint that orchestrates the full key generation flow:

WSD KOL calls WSD KCC (Rust FFI) to generate a binding keypair
WSD KOL passes the binding public key to KPS KOL (direct Go method call)
KPS KOL calls KPS KCC (Rust FFI) to generate a KEM keypair
WSD KOL stores the KEM UUID → Binding UUID mapping
Returns the KEM UUID to the workload

Extend Go CGO wrappers to match updated Rust FFI signatures (KmHpkeAlgorithm, usize)
Align request/response format with API spec (JSON body with algorithm, key_protection_mechanism, lifespan)

Dependencies

This PR is created on top of #655

Commits to Review

a8545e1 onwards ("feat(keymanager): orchestrate Key Generation in Go"), which contains all orchestration logic and CGO adaptations.

Details

CGO bridges (wskcc, kpskcc): Go wrappers calling Rust FFI via CGO, returning (uuid.UUID, []byte, error)
KPS KOL (key_protection_service/service.go): KEMKeyGenerator interface wrapping KPS KCC FFI
WSD KOL (workload_service/server.go): HTTP server with /v1/keys:generate_kem, kemToBindingMap, dependency-injected interfaces
Request validation: algorithm (DHKEM_X25519 only), key_protection_mechanism (KEY_PROTECTION_VM only), lifespan > 0

Verification

Component Integration Tests:

go test -tags integration ./... passed.

Verifies full flow from Go Service layer -> CGO -> Rust Core FFI.
Covers:

Binding Key Gen (WSD)
KEM Key Gen (KPS)
End-to-End Orchestration (WSD -> KPS)

Manual Verification :

1647b5e : Created a manual test binary that initializes the WSD server with Rust FFI bindings and exposed the generate_kem endpoint for verification via curl.

> go run keymanager/cmd/wsd_manual_test/main.go

> curl --unix-socket /tmp/wsd.sock http://unix/v1/keys:generate_kem \
  -H "Content-Type: application/json" \
  -d '{"algorithm":"DHKEM_X25519_HKDF_SHA256", "key_protection_mechanism":"KEY_PROTECTION_VM", "lifespan":"3600s"}'

{"key_handle":{"handle":"4cbe1d70-d637-44f9-aa53-a08d5ba445e0"}}

NilanjanDaw

Thanks for the PR. Some initial comments.

keymanager/go.mod

keymanager/key_protection_service/key_custody_core/include/kps_key_custody_core.h

keymanager/key_protection_service/key_custody_core/kps_key_custody_core_cgo.go

keymanager/workload_service/key_custody_core/include/ws_key_custody_core.h

keymanager/workload_service/server.go

keymanager/key_protection_service/key_custody_core/kps_key_custody_core_cgo.go

Generate FFI headers for km_common, workload_service key custody core, and key_protection_service key custody core using cbindgen. Headers were generated from source with: CBINDGEN_BIN=cbindgen ./keymanager/generate_ffi_headers.sh Generation output was then compared against the headers in PR google#652 for parity.

This change functions as the Go-side implementation for the Decap and Seal flow. It is based on PR google#652 (Key Gen) and PR google#649 (FFI Decap). Base Commit: 2030fa6 (Align Key Generation API with contract) Squashed Commits: - Fix compilation and tests for wsd_decaps_go (API alignment to /v1/keys:decap) - keymanager/wsd: align decaps payloads with proto messages - keymanager/wsd: define decaps API proto schema - [keymanager/wsd] Add /keys:decaps endpoint with DecapAndSeal + Open orchestration Key Features: - endpoint: POST /v1/keys:decap - Request: DecapsRequest (snake_case) - Response: DecapsResponse (snake_case) - Flows: DecapAndSeal (KPS) -> Open (WSD)

keymanager/go.sum

NilanjanDaw · 2026-02-10T14:20:17Z

keymanager/km_common/build.rs

    }

    let mut config = prost_build::Config::new();
-    config.type_attribute("HpkeAlgorithm", "#[repr(C)]");


I see you removed the type_attribute injection in build.rs. Is this because cbindgen couldn't pick up the generated struct from OUT_DIR?

NilanjanDaw · 2026-02-10T14:22:59Z

keymanager/key_protection_service/key_custody_core/include/kps_key_custody_core.h

+extern "C" {
+#endif // __cplusplus
+
+int32_t key_manager_generate_kem_keypair(KmHpkeAlgorithm algo,


Can we remove the km prefixes from the algorithm definitions, seems a bit weird to have them

i'm trying to remove KmHpkeAlgorithm altogether and reuse HpkeAlgorithm. as discussed here

Keeping this thread open until that is done

keymanager/workload_service/server.go

NilanjanDaw · 2026-02-17T09:59:26Z

keymanager/workload_service/server.go

+	if r.Method != http.MethodPost {
+		http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
+		return
+	}


Does Golang support function decorators? I am wondering if there is a way to handle the list of supported HTTP methods in a cleaner way.

Something like https://go.dev/blog/routing-enhancements maybe?

i don't actually see a problem with this. i had referred to the teeserver implementation https://github.com/google/go-tpm-tools/blob/main/launcher/teeserver/tee_server.go#L142

But makes sense to follow Go 1.22's routing. handled here

NilanjanDaw · 2026-02-17T10:04:31Z

keymanager/workload_service/server.go

+
+	// Validate algorithm: only DHKEM_X25519_HKDF_SHA256 supported.
+	if req.Algorithm != KemAlgorithmDHKEMX25519HKDFSHA256 {
+		http.Error(


Nit: Should we declare a global error handler? Probably would make it easier to handle and return standard HTTP error codes.

handled it using a helper function. here

NilanjanDaw · 2026-02-17T10:05:19Z

keymanager/workload_service/server.go

+			w,
+			fmt.Sprintf(
+				"unsupported algorithm: only DHKEM_X25519_HKDF_SHA256 (%d) is supported",
+				KemAlgorithmDHKEMX25519HKDFSHA256,


Should be a list of supported algorithms coming from algorithms.proto

makes sense. handled here

NilanjanDaw · 2026-02-17T10:16:18Z

keymanager/workload_service/server.go

+	}
+
+	// Validate keyProtectionMechanism: only KEY_PROTECTION_VM supported.
+	if req.KeyProtectionMechanism != KeyProtectionMechanismVM {


Wondering if there is a better way to sanitize the request params. Can something like https://github.com/bufbuild/protovalidate-go for this?

Since GenerateKemRequest is a manually defined Go struct (not a generated Protobuf message), protovalidate-go cannot be used directly on it as it requires proto.Message instances

keymanager/key_protection_service/key_custody_core/kps_key_custody_core_cgo.go

keymanager/workload_service/key_custody_core/ws_key_custody_core_cgo.go

Update the /v1/keys:generate_kem API to accept strictly numeric values (seconds) for the 'lifespan' field. Previously it accepted string values like '300s', which is now removed to simplify the API contract. as per comment: google#652 (comment)

Evaluate support for KEM algorithms and Key Protection Mechanisms using IsSupported() methods on the enums instead of hardcoded checks. This simplifies adding new supported algorithms in the future. comment: google#652 (comment)

This change introduces a helper function 'writeError' in server.go to ensure all error responses are returned as JSON, consistent with the API's Content-Type. This addresses the review comment: google#652 (comment)

Address review comment: google#652 (comment)

This change introduces a helper function 'writeError' in server.go to ensure all error responses are returned as JSON, consistent with the API's Content-Type. This addresses the review comment: google#652 (comment)

Address review comment: google#652 (comment)

keymanager/key_protection_service/key_custody_core/kps_key_custody_core_cgo.go

keymanager/workload_service/key_custody_core/ws_key_custody_core_cgo.go

keymanager/workload_service/server.go

keymanager/key_protection_service/key_custody_core/kps_key_custody_core_cgo.go

keymanager/workload_service/server.go

…afety - keymanager/go.mod: Removed explicit toolchain pin and aligned module to go 1.22. - keymanager/ws/include/ws_key_custody_core.h: Updated WSD header to include shared generated km_algorithms.h. - keymanager/ws/server.go: Replaced hardcoded algorithm checks with typed proto_enums validation. - keymanager/ws/server.go: Refactored request fields to use explicitly typed enum definitions. - keymanager/ws/server.go: Implemented parsing of duration strings in JSON requests. - keymanager/ws/server.go: Updated routing to follow Go 1.22 patterns. - keymanager/ws/ws_key_custody_core_cgo.go: Parameterized algorithm selection to use request values instead of hardcoding. - keymanager/kps/kps_key_custody_core_cgo.go: Updated CGO to use shared generated FFI constants from km_common. - keymanager/kps/kps_key_custody_core_cgo.go: Updated CGO LDFLAGS to link against both debug and release Rust targets. - keymanager/go.sum: Checked in go.sum for reproducible builds. - keymanager/km_common/build.rs: Verified type_attribute injection for HpkeAlgorithm. - keymanager/ws/lib.rs: Updated Rust FFI to accept serialized HpkeAlgorithm bytes explicitly.

Update the /v1/keys:generate_kem API to accept strictly numeric values (seconds) for the 'lifespan' field. Previously it accepted string values like '300s', which is now removed to simplify the API contract. as per comment: google#652 (comment)

Evaluate support for KEM algorithms and Key Protection Mechanisms using IsSupported() methods on the enums instead of hardcoded checks. This simplifies adding new supported algorithms in the future. comment: google#652 (comment)

This change introduces a helper function 'writeError' in server.go to ensure all error responses are returned as JSON, consistent with the API's Content-Type. This addresses the review comment: google#652 (comment)

Address review comment: google#652 (comment)

…afety - keymanager/go.mod: Removed explicit toolchain pin and aligned module to go 1.22. - keymanager/ws/include/ws_key_custody_core.h: Updated WSD header to include shared generated km_algorithms.h. - keymanager/ws/server.go: Replaced hardcoded algorithm checks with typed proto_enums validation. - keymanager/ws/server.go: Refactored request fields to use explicitly typed enum definitions. - keymanager/ws/server.go: Implemented parsing of duration strings in JSON requests. - keymanager/ws/server.go: Updated routing to follow Go 1.22 patterns. - keymanager/ws/ws_key_custody_core_cgo.go: Parameterized algorithm selection to use request values instead of hardcoding. - keymanager/kps/kps_key_custody_core_cgo.go: Updated CGO to use shared generated FFI constants from km_common. - keymanager/kps/kps_key_custody_core_cgo.go: Updated CGO LDFLAGS to link against both debug and release Rust targets. - keymanager/go.sum: Checked in go.sum for reproducible builds. - keymanager/km_common/build.rs: Verified type_attribute injection for HpkeAlgorithm. - keymanager/ws/lib.rs: Updated Rust FFI to accept serialized HpkeAlgorithm bytes explicitly.

This change introduces automated C header generation using binding and fixes a critical ABI mismatch in the key custody FFI layer. Changes: 1. FFI Header Generation: - Added keymanager/generate_ffi_headers.sh script to generate C headers for all key custody components. - Added cbindgen.toml configurations for km_common, kps_key_custody_core, and ws_key_custody_core. - Exposed km_common::ffi with stable C-ABI structs and constants to ensure compatibility. 2. FFI ABI Safety Fix: - Problem: Previous FFI functions accepted HpkeAlgorithm (Protobuf-generated repr(Rust) struct) by value. This caused undefined behavior as C callers expected a C-compatible layout. - Fix: Updated key_manager_generate_kem_keypair and key_manager_generate_binding_keypair to accept KmHpkeAlgorithm, a dedicated repr(C) struct. - Mechanism: Implemented safe conversion (impl From<KmHpkeAlgorithm> for HpkeAlgorithm) to bridge the FFI boundary safely. - Verification: Added FFI-specific tests covering KmHpkeAlgorithm usage.

This commit implements the Go orchestration layer changes for key generation. It aligns the Key Generation API with the contract.

Update the /v1/keys:generate_kem API to accept strictly numeric values (seconds) for the 'lifespan' field. Previously it accepted string values like '300s', which is now removed to simplify the API contract. as per comment: google#652 (comment)

Use http.ServeMux features to strictly enforce POST method for /v1/keys:generate_kem, removing manual checks in the handler.

Evaluate support for KEM algorithms and Key Protection Mechanisms using IsSupported() methods on the enums instead of hardcoded checks. This simplifies adding new supported algorithms in the future. comment: google#652 (comment)

This change introduces a helper function 'writeError' in server.go to ensure all error responses are returned as JSON, consistent with the API's Content-Type. This addresses the review comment: google#652 (comment)

- Replaced KmHpkeAlgorithm struct with serialized HpkeAlgorithm protobuf bytes in FFI. - Added prost dependency to kps and ws crates for decoding. - Removed keymanager/km_common/src/ffi.rs and include/km_algorithms.h to avoid duplication. - Regenerated C headers for kps and ws components. - Verified with cargo test.

This change updates the Go FFI bindings to match the recent Rust changes where KmHpkeAlgorithm C-struct was replaced with a serialized HpkeAlgorithm proto.

Address review comment: google#652 (comment)

…afety - keymanager/go.mod: Removed explicit toolchain pin and aligned module to go 1.22. - keymanager/ws/include/ws_key_custody_core.h: Updated WSD header to include shared generated km_algorithms.h. - keymanager/ws/server.go: Replaced hardcoded algorithm checks with typed proto_enums validation. - keymanager/ws/server.go: Refactored request fields to use explicitly typed enum definitions. - keymanager/ws/server.go: Implemented parsing of duration strings in JSON requests. - keymanager/ws/server.go: Updated routing to follow Go 1.22 patterns. - keymanager/ws/ws_key_custody_core_cgo.go: Parameterized algorithm selection to use request values instead of hardcoding. - keymanager/kps/kps_key_custody_core_cgo.go: Updated CGO to use shared generated FFI constants from km_common. - keymanager/kps/kps_key_custody_core_cgo.go: Updated CGO LDFLAGS to link against both debug and release Rust targets. - keymanager/go.sum: Checked in go.sum for reproducible builds. - keymanager/km_common/build.rs: Verified type_attribute injection for HpkeAlgorithm. - keymanager/ws/lib.rs: Updated Rust FFI to accept serialized HpkeAlgorithm bytes explicitly.

Centralizes the algorithm mapping logic by moving it from a private helper in server.go to a public method ToHpkeAlgorithm on the KemAlgorithm type in proto_enums.go. This change simplifies future introspection APIs (e.g., getCapabilities) by providing a single source of truth for algorithm configuration without needing to instantiate backend services or duplicate logic. It ensures that the API layer can validate and describe supported algorithms consistently.

keymanager/workload_service/server.go

atulpatildbz · 2026-02-19T11:46:38Z

/gcbrun

This was referenced Feb 9, 2026

[keymanager/wsd] GET /v1/keys enumerate keys #653

Open

[keymanager/wsd] Add /keys:decaps KOL API with DecapAndSeal + Open orchestration #650

Open

[keymanager/wsd] Add /keys:destroy Go KOL handler with KEM + binding key destruction #651

Open

atulpatildbz requested a review from NilanjanDaw February 9, 2026 10:13

NilanjanDaw reviewed Feb 9, 2026

View reviewed changes

atulpatildbz added a commit to atulpatildbz/go-tpm-tools that referenced this pull request Feb 9, 2026

keymanager: address PR google#652 review feedback

ed213d4

atulpatildbz added a commit to atulpatildbz/go-tpm-tools that referenced this pull request Feb 10, 2026

keymanager: address PR google#652 review feedback

ad03588

atulpatildbz force-pushed the wsd_generate_key_go branch 2 times, most recently from e333851 to 4d0831d Compare February 10, 2026 17:24

atulpatildbz marked this pull request as draft February 10, 2026 18:02

atulpatildbz force-pushed the wsd_generate_key_go branch from 4d0831d to 91d3bd7 Compare February 13, 2026 14:45

atulpatildbz added a commit to atulpatildbz/go-tpm-tools that referenced this pull request Feb 14, 2026

keymanager: address PR google#652 review feedback

991f87e

atulpatildbz force-pushed the wsd_generate_key_go branch 3 times, most recently from bba524b to 426856a Compare February 15, 2026 05:15

atulpatildbz marked this pull request as ready for review February 15, 2026 05:19

atulpatildbz requested a review from NilanjanDaw February 15, 2026 06:53

atulpatildbz mentioned this pull request Feb 15, 2026

[KeyManager] FFI for Enumerate keys #665

Open

atulpatildbz mentioned this pull request Feb 15, 2026

feat(workload_service): implement GetCapabilities API #666

Open

atulpatildbz force-pushed the wsd_generate_key_go branch 2 times, most recently from 696b9db to a8545e1 Compare February 15, 2026 11:25

atulpatildbz requested review from alexmwu and jkl73 February 16, 2026 16:50

NilanjanDaw reviewed Feb 17, 2026

View reviewed changes

atulpatildbz added a commit to atulpatildbz/go-tpm-tools that referenced this pull request Feb 18, 2026

keymanager: list supported algorithms in error message

0af243e

Address review comment: google#652 (comment)

atulpatildbz added a commit to atulpatildbz/go-tpm-tools that referenced this pull request Feb 18, 2026

keymanager: list supported algorithms in error message

d757900

Address review comment: google#652 (comment)

atulpatildbz force-pushed the wsd_generate_key_go branch from 0af243e to d757900 Compare February 18, 2026 14:19

NilanjanDaw requested changes Feb 18, 2026

View reviewed changes

atulpatildbz added a commit to atulpatildbz/go-tpm-tools that referenced this pull request Feb 19, 2026

keymanager: list supported algorithms in error message

05505d2

Address review comment: google#652 (comment)

atulpatildbz force-pushed the wsd_generate_key_go branch from 28b1e04 to 039ca49 Compare February 19, 2026 04:17

atulpatildbz added 11 commits February 19, 2026 04:21

feat(keymanager): orchestrate Key Generation in Go

4e8fdfb

This commit implements the Go orchestration layer changes for key generation. It aligns the Key Generation API with the contract.

refactor(workload_service): use Go 1.22 routing for method enforcement

7d70a3c

Use http.ServeMux features to strictly enforce POST method for /v1/keys:generate_kem, removing manual checks in the handler.

Refactor error handling to return JSON

e03894a

This change introduces a helper function 'writeError' in server.go to ensure all error responses are returned as JSON, consistent with the API's Content-Type. This addresses the review comment: google#652 (comment)

keymanager: update Go FFI to use serialized HpkeAlgorithm proto

45c49e8

This change updates the Go FFI bindings to match the recent Rust changes where KmHpkeAlgorithm C-struct was replaced with a serialized HpkeAlgorithm proto.

keymanager: list supported algorithms in error message

d0f4b34

Address review comment: google#652 (comment)

formatting fixes

f8fe23d

atulpatildbz force-pushed the wsd_generate_key_go branch from 039ca49 to f8fe23d Compare February 19, 2026 04:22

atulpatildbz added 2 commits February 19, 2026 04:42

keymanager: update KEMKeyGenerator interface to accept HpkeAlgorithm

98252d7

atulpatildbz requested a review from NilanjanDaw February 19, 2026 09:00

NilanjanDaw approved these changes Feb 19, 2026

View reviewed changes

keymanager/workload_service/server.go Show resolved Hide resolved

workload_service: cleanup socket file before listening

90cd4c6

atulpatildbz merged commit 8c93731 into google:main Feb 19, 2026
11 checks passed

Comments

Conversation

atulpatildbz commented Feb 9, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Dependencies

Commits to Review

Verification

Component Integration Tests:

Manual Verification :

Uh oh!

NilanjanDaw left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

atulpatildbz Feb 17, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

atulpatildbz Feb 17, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

atulpatildbz Feb 18, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

atulpatildbz commented Feb 19, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

atulpatildbz commented Feb 9, 2026 •

edited

Loading

atulpatildbz Feb 17, 2026 •

edited

Loading

atulpatildbz Feb 17, 2026 •

edited

Loading

atulpatildbz Feb 18, 2026 •

edited

Loading