[KeyManager] End-to-end FFI error handling standardization (google#722)

This commit performs a comprehensive refactoring of the `KeyManager` component to utilize the standardized `keymanager.Status` Protobuf enum across all layers of the system. It replaces legacy negative `i32` error codes and ad-hoc `error` enums with a unified, type-safe `Status` reporting mechanism that spans Rust, C, and Go.

  Key Changes
   - Rust FFI & Core Updates:
       - Refactored `KeyCustodyCore` in both `key_protection_service` and `workload_service` to return the Status enum directly.
       - Simplified FFI implementation by utilizing the `ffi_call` and `ffi_call_i32` helpers introduced in Part 1.
       - Updated internal logic to map crate-specific errors to the standardized `Status` proto.
   - C-Binding & Header Standardization:
       - Modified C function signatures in `kps_key_custody_core.h` and `ws_key_custody_core.h` to return `Status` instead of `int32_t`.
   - Go CGO & Service Layer Integration:
       - Updated CGO wrappers to interpret the new Status return codes and convert them into idiomatic Go errors using the`ToStatus()` helper.
       - HTTP Status Mapping: Implemented `httpStatusFromError` in the Go server to automatically map standardized FFI errors to appropriate HTTP status codes (e.g., `ERROR_NOT_FOUND` → `404 Not Found`, `ERROR_INVALID_ARGUMENT` → `400 Bad` Request).
       - Removed redundant error aliases in `types.go` to maintain a single source of truth.
   - Test Alignment: Updated integration tests in Go and unit tests in Rust to verify the new error reporting behavior and ensure consistent error messages (e.g., `"FFI error: ERROR_NOT_FOUND"`).

Author: NilanjanDaw

keymanager/generate_ffi_headers.sh

@@ -4,6 +4,11 @@ set -euo pipefail
 ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
 CBINDGEN_BIN="${CBINDGEN_BIN:-cbindgen}"
 
+"${CBINDGEN_BIN}" --quiet \
+  "${ROOT_DIR}/km_common" \
+  --crate km_common \
+  --config "${ROOT_DIR}/km_common/cbindgen.toml" \
+  --output "${ROOT_DIR}/km_common/include/km_common_ffi.h"
 
 "${CBINDGEN_BIN}" --quiet \
   "${ROOT_DIR}/workload_service/key_custody_core" \
@@ -16,3 +21,4 @@ CBINDGEN_BIN="${CBINDGEN_BIN:-cbindgen}"
   --crate kps_key_custody_core \
   --config "${ROOT_DIR}/key_protection_service/key_custody_core/cbindgen.toml" \
   --output "${ROOT_DIR}/key_protection_service/key_custody_core/include/kps_key_custody_core.h"
+

keymanager/key_protection_service/key_custody_core/cbindgen.toml

@@ -8,14 +8,18 @@ style = "type"
 documentation = false
 usize_is_size_t = true
 sys_includes = ["stdbool.h", "stddef.h", "stdint.h"]
+includes = ["km_common_ffi.h"]
 
 
 [parse]
-parse_deps = false
+parse_deps = true
+include = ["km_common"]
 clean = true
 
 [export]
 item_types = ["functions", "structs", "constants", "enums"]
-include = ["key_manager_generate_kem_keypair", "key_manager_get_kem_key", "key_manager_destroy_kem_key", "key_manager_enumerate_kem_keys", "KpsKeyInfo", "MAX_ALGORITHM_LEN", "MAX_PUBLIC_KEY_LEN"]
-
+include = ["key_manager_generate_kem_keypair", "key_manager_get_kem_key", "key_manager_destroy_kem_key", "key_manager_enumerate_kem_keys", "KpsKeyInfo"]
+exclude = ["Status", "MAX_ALGORITHM_LEN", "MAX_PUBLIC_KEY_LEN"]
 
+[enum]
+prefix_with_name = true

keymanager/key_protection_service/key_custody_core/include/kps_key_custody_core.h

@@ -6,10 +6,7 @@
 #include <stdbool.h>
 #include <stddef.h>
 #include <stdint.h>
-
-#define MAX_ALGORITHM_LEN 128
-
-#define MAX_PUBLIC_KEY_LEN 2048
+#include "km_common_ffi.h"
 
 typedef struct {
   uint8_t uuid[16];
@@ -26,40 +23,40 @@ typedef struct {
 extern "C" {
 #endif // __cplusplus
 
-int32_t key_manager_generate_kem_keypair(const uint8_t *algo_ptr,
-                                         size_t algo_len,
-                                         const uint8_t *binding_pubkey,
-                                         size_t binding_pubkey_len,
-                                         uint64_t expiry_secs,
-                                         uint8_t *out_uuid,
-                                         uint8_t *out_pubkey,
-                                         size_t out_pubkey_len);
+Status key_manager_generate_kem_keypair(const uint8_t *algo_ptr,
+                                        size_t algo_len,
+                                        const uint8_t *binding_pubkey,
+                                        size_t binding_pubkey_len,
+                                        uint64_t expiry_secs,
+                                        uint8_t *out_uuid,
+                                        uint8_t *out_pubkey,
+                                        size_t out_pubkey_len);
 
-int32_t key_manager_destroy_kem_key(const uint8_t *uuid_bytes);
+Status key_manager_destroy_kem_key(const uint8_t *uuid_bytes);
 
 int32_t key_manager_enumerate_kem_keys(KpsKeyInfo *out_entries,
                                        size_t max_entries,
                                        size_t offset,
                                        bool *out_has_more);
 
-int32_t key_manager_decap_and_seal(const uint8_t *uuid_bytes,
-                                   const uint8_t *encapsulated_key,
-                                   size_t encapsulated_key_len,
-                                   const uint8_t *aad,
-                                   size_t aad_len,
-                                   uint8_t *out_encapsulated_key,
-                                   size_t out_encapsulated_key_len,
-                                   uint8_t *out_ciphertext,
-                                   size_t out_ciphertext_len);
+Status key_manager_decap_and_seal(const uint8_t *uuid_bytes,
+                                  const uint8_t *encapsulated_key,
+                                  size_t encapsulated_key_len,
+                                  const uint8_t *aad,
+                                  size_t aad_len,
+                                  uint8_t *out_encapsulated_key,
+                                  size_t out_encapsulated_key_len,
+                                  uint8_t *out_ciphertext,
+                                  size_t out_ciphertext_len);
 
-int32_t key_manager_get_kem_key(const uint8_t *uuid_bytes,
-                                uint8_t *out_kem_pubkey,
-                                size_t out_kem_pubkey_len,
-                                uint8_t *out_binding_pubkey,
-                                size_t out_binding_pubkey_len,
-                                uint8_t *out_algo,
-                                size_t *out_algo_len,
-                                uint64_t *out_delete_after);
+Status key_manager_get_kem_key(const uint8_t *uuid_bytes,
+                               uint8_t *out_kem_pubkey,
+                               size_t out_kem_pubkey_len,
+                               uint8_t *out_binding_pubkey,
+                               size_t out_binding_pubkey_len,
+                               uint8_t *out_algo,
+                               size_t *out_algo_len,
+                               uint64_t *out_remaining_lifespan_secs);
 
 #ifdef __cplusplus
 }  // extern "C"

keymanager/key_protection_service/key_custody_core/integration_test.go

@@ -106,7 +106,7 @@ func TestIntegrationGetKEMKeyNotFound(t *testing.T) {
 		t.Fatal("expected error for non-existent UUID")
 	}
 
-	expectedErrMsg := "key_manager_get_kem_key failed with code -1"
+	expectedErrMsg := "FFI status: STATUS_NOT_FOUND"
 	if !strings.Contains(err.Error(), expectedErrMsg) {
 		t.Fatalf("expected error containing %q, got: %v", expectedErrMsg, err)
 	}

keymanager/key_protection_service/key_custody_core/kps_key_custody_core_cgo.go

@@ -6,6 +6,7 @@
 package kpskcc
 
 /*
+#cgo CFLAGS: -I${SRCDIR}/../../km_common/include
 #cgo LDFLAGS: -L${SRCDIR}/../../target/release -L${SRCDIR}/../../target/debug -lkps_key_custody_core
 #cgo LDFLAGS: -lcrypto -lssl
 #cgo LDFLAGS: -lpthread -ldl -lm -lstdc++
@@ -57,8 +58,8 @@ func GenerateKEMKeypair(algo *keymanager.HpkeAlgorithm, bindingPubKey []byte, li
 		(*C.uint8_t)(unsafe.Pointer(&uuidBytes[0])),
 		(*C.uint8_t)(unsafe.Pointer(&pubkeyBuf[0])),
 		pubkeyLen,
-	); rc != 0 {
-		return uuid.Nil, nil, fmt.Errorf("key_manager_generate_kem_keypair failed with code %d", rc)
+	); keymanager.Status(rc) != keymanager.Status_STATUS_SUCCESS {
+		return uuid.Nil, nil, keymanager.Status(rc).ToStatus()
 	}
 
 	id, err := uuid.FromBytes(uuidBytes[:])
@@ -91,7 +92,7 @@ func EnumerateKEMKeys(limit, offset int) ([]KEMKeyInfo, bool, error) {
 		&hasMore,
 	)
 	if rc < 0 {
-		return nil, false, fmt.Errorf("key_manager_enumerate_kem_keys failed with code %d", rc)
+		return nil, false, keymanager.Status(-rc).ToStatus()
 	}
 
 	count := int(rc)
@@ -124,12 +125,10 @@ func EnumerateKEMKeys(limit, offset int) ([]KEMKeyInfo, bool, error) {
 // DestroyKEMKey destroys the KEM key identified by kemUUID via Rust FFI.
 func DestroyKEMKey(kemUUID uuid.UUID) error {
 	uuidBytes := kemUUID[:]
-	if rc := C.key_manager_destroy_kem_key(
+	rc := C.key_manager_destroy_kem_key(
 		(*C.uint8_t)(unsafe.Pointer(&uuidBytes[0])),
-	); rc != 0 {
-		return fmt.Errorf("key_manager_destroy_kem_key failed with code %d", rc)
-	}
-	return nil
+	)
+	return keymanager.Status(rc).ToStatus()
 }
 
 // GetKEMKey retrieves KEM and binding public keys, HpkeAlgorithm and remaining lifespan via Rust FFI.
@@ -153,8 +152,8 @@ func GetKEMKey(id uuid.UUID) ([]byte, []byte, *keymanager.HpkeAlgorithm, uint64,
 		&algoLenC,
 		&remainingLifespanSecs,
 	)
-	if rc != 0 {
-		return nil, nil, nil, 0, fmt.Errorf("key_manager_get_kem_key failed with code %d", rc)
+	if keymanager.Status(rc) != keymanager.Status_STATUS_SUCCESS {
+		return nil, nil, nil, 0, keymanager.Status(rc).ToStatus()
 	}
 
 	kemPubkey := make([]byte, len(kemPubkeyBuf))
@@ -201,8 +200,8 @@ func DecapAndSeal(kemUUID uuid.UUID, encapsulatedKey, aad []byte) ([]byte, []byt
 		outEncKeyLen,
 		(*C.uint8_t)(unsafe.Pointer(&outCT[0])),
 		outCTLen,
-	); rc != 0 {
-		return nil, nil, fmt.Errorf("key_manager_decap_and_seal failed with code %d", rc)
+	); keymanager.Status(rc) != keymanager.Status_STATUS_SUCCESS {
+		return nil, nil, keymanager.Status(rc).ToStatus()
 	}
 
 	sealEnc := make([]byte, outEncKeyLen)