feat(keymanager): Implement Decap and Seal orchestration in Go

This change functions as the Go-side implementation for the Decap and Seal flow.
It is based on PR #652 (Key Gen) and PR #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)

Author: atulpatildbz

keymanager/key_protection_service/key_custody_core/include/kps_key_custody_core.h

@@ -18,6 +18,27 @@ int32_t key_manager_generate_kem_keypair(KmHpkeAlgorithm algo,
                                          uint8_t *out_uuid,
                                          uint8_t *out_pubkey,
                                          size_t out_pubkey_len);
+// key_manager_decap_and_seal decapsulates a shared secret using the stored KEM
+// key identified by uuid_bytes, then reseals the shared secret with the
+// associated binding public key via HPKE.
+//
+// uuid_bytes must point to a 16-byte KEM key UUID.
+// encapsulated_key is the client-provided encapsulated key.
+// aad is optional Additional Authenticated Data (may be NULL if aad_len == 0).
+// out_encapsulated_key_len and out_ciphertext_len are in/out params.
+//
+// Returns 0 on success, -1 on invalid args/key not found, -2 if output buffers
+// are too small, -3 if decapsulation fails, -4 if sealing fails.
+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);
 
 #ifdef __cplusplus
 }  // extern "C"

keymanager/key_protection_service/key_custody_core/kps_key_custody_core_cgo.go

@@ -57,3 +57,47 @@ func GenerateKEMKeypair(bindingPubKey []byte, lifespanSecs uint64) (uuid.UUID, [
 	copy(pubkey, pubkeyBuf[:pubkeyLen])
 	return id, pubkey, nil
 }
+
+// DecapAndSeal decapsulates a shared secret using the stored KEM key and
+// reseals it with the associated binding public key via Rust FFI.
+// Returns the new encapsulated key and sealed ciphertext.
+func DecapAndSeal(kemUUID uuid.UUID, encapsulatedKey, aad []byte) ([]byte, []byte, error) {
+	if len(encapsulatedKey) == 0 {
+		return nil, nil, fmt.Errorf("encapsulated key must not be empty")
+	}
+
+	uuidBytes := kemUUID[:]
+
+	var outEncKey [32]byte
+	outEncKeyLen := C.size_t(len(outEncKey))
+	var outCT [48]byte // 32-byte secret + 16-byte GCM tag
+	outCTLen := C.size_t(len(outCT))
+
+	var aadPtr *C.uint8_t
+	aadLen := C.size_t(0)
+	if len(aad) > 0 {
+		aadPtr = (*C.uint8_t)(unsafe.Pointer(&aad[0]))
+		aadLen = C.size_t(len(aad))
+	}
+
+	rc := C.key_manager_decap_and_seal(
+		(*C.uint8_t)(unsafe.Pointer(&uuidBytes[0])),
+		(*C.uint8_t)(unsafe.Pointer(&encapsulatedKey[0])),
+		C.size_t(len(encapsulatedKey)),
+		aadPtr,
+		aadLen,
+		(*C.uint8_t)(unsafe.Pointer(&outEncKey[0])),
+		&outEncKeyLen,
+		(*C.uint8_t)(unsafe.Pointer(&outCT[0])),
+		&outCTLen,
+	)
+	if rc != 0 {
+		return nil, nil, fmt.Errorf("key_manager_decap_and_seal failed with code %d", rc)
+	}
+
+	sealEnc := make([]byte, outEncKeyLen)
+	copy(sealEnc, outEncKey[:outEncKeyLen])
+	sealedCT := make([]byte, outCTLen)
+	copy(sealedCT, outCT[:outCTLen])
+	return sealEnc, sealedCT, nil
+}

keymanager/key_protection_service/key_custody_core/src/lib.rs

@@ -163,8 +163,8 @@ pub unsafe extern "C" fn key_manager_decap_and_seal(
     };
 
     let key_record = match KEY_REGISTRY.get_key(&uuid) {
-        Some(record) => record,
-        None => return -1,
+        Ok(record) => record,
+        Err(_) => return -1,
     };
 
     let (hpke_algo, binding_public_key) = match &key_record.meta.spec {
@@ -184,9 +184,14 @@ pub unsafe extern "C" fn key_manager_decap_and_seal(
     let enc_key_slice = unsafe { slice::from_raw_parts(encapsulated_key, encapsulated_key_len) };
 
     // Decapsulate
-    let mut shared_secret = match km_common::crypto::decaps(&key_record.private_key, enc_key_slice)
+    let private_key = match km_common::crypto::PrivateKey::try_from(key_record.private_key.as_bytes().to_vec()) {
+        Ok(pk) => pk,
+        Err(_) => return -3,
+    };
+
+    let mut shared_secret = match km_common::crypto::decaps(&private_key, enc_key_slice)
     {
-        Ok(s) => s,
+        Ok(secret) => secret,
         Err(_) => return -3,
     };
 
@@ -198,10 +203,10 @@ pub unsafe extern "C" fn key_manager_decap_and_seal(
 
     // Seal
     let (new_enc_key, sealed_ciphertext) = match km_common::crypto::hpke_seal(
-        binding_public_key,
+        &binding_public_key,
         &shared_secret,
         aad_slice,
-        hpke_algo,
+        &hpke_algo,
     ) {
         Ok(res) => res,
         Err(_) => {
@@ -427,21 +432,22 @@ mod tests {
     fn test_destroy_kem_key_success() {
         let binding_pubkey = [1u8; 32];
         let mut uuid_bytes = [0u8; 16];
-        let algo = HpkeAlgorithm {
+        let algo = KmHpkeAlgorithm {
             kem: KemAlgorithm::DhkemX25519HkdfSha256 as i32,
             kdf: KdfAlgorithm::HkdfSha256 as i32,
             aead: AeadAlgorithm::Aes256Gcm as i32,
         };
 
         unsafe {
+            let mut pubkey_bytes = [0u8; 32];
             key_manager_generate_kem_keypair(
                 algo,
                 binding_pubkey.as_ptr(),
                 binding_pubkey.len(),
                 3600,
                 uuid_bytes.as_mut_ptr(),
-                std::ptr::null_mut(),
-                std::ptr::null_mut(),
+                pubkey_bytes.as_mut_ptr(),
+                32,
             );
         }
 
@@ -466,6 +472,38 @@ mod tests {
         assert_eq!(result, -1);
     }
 
+    #[test]
+    fn test_key_manager_generate_kem_keypair() {
+        // 1. Setup binding key
+        let binding_kem_algo = KemAlgorithm::DhkemX25519HkdfSha256;
+        let (binding_pk, _) = km_common::crypto::generate_x25519_keypair(binding_kem_algo).unwrap();
+
+        // 2. Call generate_kem_keypair
+        let mut uuid_bytes = [0u8; 16];
+        let mut kem_pubkey_bytes = [0u8; 32];
+        let kem_pubkey_len = 32;
+        let algo = KmHpkeAlgorithm {
+            kem: KemAlgorithm::DhkemX25519HkdfSha256 as i32,
+            kdf: KdfAlgorithm::HkdfSha256 as i32,
+            aead: AeadAlgorithm::Aes256Gcm as i32,
+        };
+        let result = unsafe {
+            key_manager_generate_kem_keypair(
+                algo,
+                binding_pk.as_ptr(),
+                binding_pk.len(),
+                3600,
+                uuid_bytes.as_mut_ptr(),
+                kem_pubkey_bytes.as_mut_ptr(),
+                kem_pubkey_len,
+            )
+        };
+
+        assert_eq!(result, 0);
+        assert_ne!(uuid_bytes, [0u8; 16]);
+        assert_ne!(kem_pubkey_bytes, [0u8; 32]);
+    }
+
     #[test]
     fn test_decap_and_seal_success() {
         // 1. Setup binding key (receiver for seal)
@@ -476,12 +514,13 @@ mod tests {
         // 2. Generate KEM key in registry
         let mut uuid_bytes = [0u8; 16];
         let mut kem_pubkey_bytes = [0u8; 32];
-        let mut kem_pubkey_len = 32;
-        let algo = HpkeAlgorithm {
+        let kem_pubkey_len = 32;
+        let algo = KmHpkeAlgorithm {
             kem: KemAlgorithm::DhkemX25519HkdfSha256 as i32,
             kdf: KdfAlgorithm::HkdfSha256 as i32,
             aead: AeadAlgorithm::Aes256Gcm as i32,
         };
+        let hpke_algo: HpkeAlgorithm = algo.into();
         unsafe {
             key_manager_generate_kem_keypair(
                 algo,
@@ -490,7 +529,7 @@ mod tests {
                 3600,
                 uuid_bytes.as_mut_ptr(),
                 kem_pubkey_bytes.as_mut_ptr(),
-                &mut kem_pubkey_len,
+                kem_pubkey_len,
             );
         }
 
@@ -499,7 +538,7 @@ mod tests {
         let aad = b"test_aad";
         // We use `hpke_seal_raw` to act as the client to generate a valid encapsulation
         let (client_enc, client_ct) =
-            km_common::crypto::hpke_seal_raw(&kem_pubkey_bytes, pt, aad, &algo).unwrap();
+            km_common::crypto::hpke_seal_raw(&kem_pubkey_bytes, pt, aad, &hpke_algo).unwrap();
 
         // Step 3: Call `decap_and_seal`.
         let mut out_enc_key = [0u8; 32];
@@ -524,32 +563,33 @@ mod tests {
         assert_eq!(result, 0);
 
         // 4. Verify we can decrypt the result using binding_sk
+        let binding_sk_priv = km_common::crypto::PrivateKey::try_from(binding_sk.clone()).unwrap();
         let recovered_shared_secret =
-            km_common::crypto::hpke_open(&binding_sk, &out_enc_key, &out_ct, aad, &algo)
+            km_common::crypto::hpke_open(&binding_sk_priv, &out_enc_key, &out_ct, aad, &hpke_algo)
                 .expect("Failed to decrypt the resealed secret");
 
-        assert_eq!(recovered_shared_secret.len(), 32);
+        assert_eq!(recovered_shared_secret.as_slice().len(), 32);
 
         // 5. Verify the recovered secret matches what decaps would produce
         let key_record = KEY_REGISTRY.get_key(&Uuid::from_bytes(uuid_bytes)).unwrap();
+        let private_key = km_common::crypto::PrivateKey::try_from(key_record.private_key.as_bytes().to_vec()).unwrap();
         let expected_shared_secret = km_common::crypto::decaps(
-            key_record.private_key.as_bytes(),
+            &private_key,
             &client_enc,
-            KemAlgorithm::DhkemX25519HkdfSha256,
         )
         .expect("decaps failed");
         assert_eq!(
-            recovered_shared_secret, expected_shared_secret,
+            recovered_shared_secret.as_slice(), expected_shared_secret.as_slice(),
             "Recovered secret mismatch"
         );
 
         // 6. Verify that this secret correctly decrypts the original client ciphertext
         // using the shared secret directly instead of the private key.
         let decrypted_pt = km_common::crypto::hpke_open_with_shared_secret(
-            &recovered_shared_secret,
+            recovered_shared_secret.as_slice(),
             &client_ct,
             aad,
-            &algo,
+            &hpke_algo,
         )
         .expect("Failed to decrypt client message with shared secret");
 
@@ -610,20 +650,21 @@ mod tests {
 
         // 2. Generate KEM key
         let mut uuid_bytes = [0u8; 16];
-        let algo = HpkeAlgorithm {
+        let algo = KmHpkeAlgorithm {
             kem: KemAlgorithm::DhkemX25519HkdfSha256 as i32,
             kdf: KdfAlgorithm::HkdfSha256 as i32,
             aead: AeadAlgorithm::Aes256Gcm as i32,
         };
+        let mut kem_pubkey_bytes = [0u8; 32];
         unsafe {
             key_manager_generate_kem_keypair(
                 algo,
                 binding_pk.as_ptr(),
                 binding_pk.len(),
                 3600,
                 uuid_bytes.as_mut_ptr(),
-                std::ptr::null_mut(),
-                std::ptr::null_mut(),
+                kem_pubkey_bytes.as_mut_ptr(),
+                32,
             );
         }
 
@@ -659,8 +700,14 @@ mod tests {
         // 2. Generate KEM key
         let mut uuid_bytes = [0u8; 16];
         let mut kem_pubkey_bytes = [0u8; 32];
+        let mut kem_pubkey_bytes = [0u8; 32];
         let mut kem_pubkey_len = 32;
-        let algo = HpkeAlgorithm {
+        let algo = KmHpkeAlgorithm {
+            kem: KemAlgorithm::DhkemX25519HkdfSha256 as i32,
+            kdf: KdfAlgorithm::HkdfSha256 as i32,
+            aead: AeadAlgorithm::Aes256Gcm as i32,
+        };
+        let client_algo = HpkeAlgorithm {
             kem: KemAlgorithm::DhkemX25519HkdfSha256 as i32,
             kdf: KdfAlgorithm::HkdfSha256 as i32,
             aead: AeadAlgorithm::Aes256Gcm as i32,
@@ -673,13 +720,13 @@ mod tests {
                 3600,
                 uuid_bytes.as_mut_ptr(),
                 kem_pubkey_bytes.as_mut_ptr(),
-                &mut kem_pubkey_len,
+                kem_pubkey_len,
             );
         }
 
         // 3. Generate valid client encapsulation
         let (client_enc, _) =
-            km_common::crypto::hpke_seal_raw(&kem_pubkey_bytes, b"secret", b"", &algo).unwrap();
+            km_common::crypto::hpke_seal_raw(&kem_pubkey_bytes, b"secret", b"", &client_algo).unwrap();
 
         // 4. Call with small output buffers
         let mut out_enc_key = [0u8; 31]; // Small

keymanager/key_protection_service/service.go

@@ -1,6 +1,6 @@
 // Package key_protection_service implements the Key Orchestration Layer (KOL)
 // for the Key Protection Service. It wraps the KPS Key Custody Core (KCC) FFI
-// to provide a Go-native interface for KEM key generation.
+// to provide a Go-native interface for KEM key operations.
 package key_protection_service
 
 import "github.com/google/uuid"
@@ -10,18 +10,36 @@ type KEMKeyGenerator interface {
 	GenerateKEMKeypair(bindingPubKey []byte, lifespanSecs uint64) (uuid.UUID, []byte, error)
 }
 
-// Service implements KEMKeyGenerator by delegating to the KPS KCC FFI.
+// DecapSealer decapsulates a shared secret and reseals it with the binding key.
+type DecapSealer interface {
+	DecapAndSeal(kemUUID uuid.UUID, encapsulatedKey, aad []byte) (sealEnc []byte, sealedCT []byte, err error)
+}
+
+// Service implements KEMKeyGenerator and DecapSealer by delegating to the KPS KCC FFI.
 type Service struct {
 	generateKEMKeypairFn func(bindingPubKey []byte, lifespanSecs uint64) (uuid.UUID, []byte, error)
+	decapAndSealFn       func(kemUUID uuid.UUID, encapsulatedKey, aad []byte) ([]byte, []byte, error)
 }
 
-// NewService creates a new KPS KOL service with the given KCC function.
-func NewService(generateKEMKeypairFn func(bindingPubKey []byte, lifespanSecs uint64) (uuid.UUID, []byte, error)) *Service {
-	return &Service{generateKEMKeypairFn: generateKEMKeypairFn}
+// NewService creates a new KPS KOL service with the given KCC functions.
+func NewService(
+	generateKEMKeypairFn func(bindingPubKey []byte, lifespanSecs uint64) (uuid.UUID, []byte, error),
+	decapAndSealFn func(kemUUID uuid.UUID, encapsulatedKey, aad []byte) ([]byte, []byte, error),
+) *Service {
+	return &Service{
+		generateKEMKeypairFn: generateKEMKeypairFn,
+		decapAndSealFn:       decapAndSealFn,
+	}
 }
 
 // GenerateKEMKeypair generates a KEM keypair linked to the provided binding
 // public key by calling the KPS KCC FFI.
 func (s *Service) GenerateKEMKeypair(bindingPubKey []byte, lifespanSecs uint64) (uuid.UUID, []byte, error) {
 	return s.generateKEMKeypairFn(bindingPubKey, lifespanSecs)
 }
+
+// DecapAndSeal decapsulates a shared secret using the stored KEM key and
+// reseals it with the associated binding public key by calling the KPS KCC FFI.
+func (s *Service) DecapAndSeal(kemUUID uuid.UUID, encapsulatedKey, aad []byte) ([]byte, []byte, error) {
+	return s.decapAndSealFn(kemUUID, encapsulatedKey, aad)
+}