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

Author: atulpatildbz

keymanager/key_protection_service/key_custody_core/kps_key_custody_core_cgo.go

@@ -61,3 +61,47 @@ func GenerateKEMKeypair(algo *keymanager.HpkeAlgorithm, bindingPubKey []byte, li
 	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/kps_key_custody_core_stub.go

@@ -14,3 +14,8 @@ import (
 func GenerateKEMKeypair(_ *algorithms.HpkeAlgorithm, _ []byte, _ uint64) (uuid.UUID, []byte, error) {
 	return uuid.Nil, nil, fmt.Errorf("GenerateKEMKeypair is not supported on this architecture")
 }
+
+// DecapAndSeal is a stub for architectures where the Rust library is not supported.
+func DecapAndSeal(_ uuid.UUID, _, _ []byte) ([]byte, []byte, error) {
+	return nil, nil, fmt.Errorf("DecapAndSeal is not supported on this architecture")
+}

keymanager/key_protection_service/service.go

@@ -1,6 +1,6 @@
 // Package keyprotectionservice 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 keyprotectionservice
 
 import (
@@ -14,18 +14,36 @@ type KEMKeyGenerator interface {
 	GenerateKEMKeypair(algo *keymanager.HpkeAlgorithm, bindingPubKey []byte, lifespanSecs uint64) ([]byte, []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(algo *keymanager.HpkeAlgorithm, 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(algo *keymanager.HpkeAlgorithm, 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(algo *keymanager.HpkeAlgorithm, 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(algo *keymanager.HpkeAlgorithm, bindingPubKey []byte, lifespanSecs uint64) (uuid.UUID, []byte, error) {
 	return s.generateKEMKeypairFn(algo, 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)
+}

keymanager/key_protection_service/service_test.go

@@ -9,6 +9,11 @@ import (
 	keymanager "github.com/google/go-tpm-tools/keymanager/km_common/proto"
 )
 
+// noopDecapAndSeal is a placeholder for tests that don't exercise DecapAndSeal.
+func noopDecapAndSeal(_ uuid.UUID, _, _ []byte) ([]byte, []byte, error) {
+	return nil, nil, nil
+}
+
 func TestServiceGenerateKEMKeypairSuccess(t *testing.T) {
 	expectedUUID := uuid.New()
 	expectedPubKey := make([]byte, 32)
@@ -24,7 +29,7 @@ func TestServiceGenerateKEMKeypairSuccess(t *testing.T) {
 			t.Fatalf("expected lifespanSecs 7200, got %d", lifespanSecs)
 		}
 		return expectedUUID, expectedPubKey, nil
-	})
+	}, noopDecapAndSeal)
 
 	id, pubKey, err := svc.GenerateKEMKeypair(&keymanager.HpkeAlgorithm{}, make([]byte, 32), 7200)
 	if err != nil {
@@ -41,10 +46,45 @@ func TestServiceGenerateKEMKeypairSuccess(t *testing.T) {
 func TestServiceGenerateKEMKeypairError(t *testing.T) {
 	svc := NewService(func(_ *keymanager.HpkeAlgorithm, _ []byte, _ uint64) (uuid.UUID, []byte, error) {
 		return uuid.Nil, nil, fmt.Errorf("FFI error")
-	})
+	}, noopDecapAndSeal)
 
 	_, _, err := svc.GenerateKEMKeypair(&keymanager.HpkeAlgorithm{}, make([]byte, 32), 3600)
 	if err == nil {
 		t.Fatal("expected error, got nil")
 	}
 }
+
+func TestServiceDecapAndSealSuccess(t *testing.T) {
+	kemUUID := uuid.New()
+	expectedSealEnc := []byte("seal-enc-key")
+	expectedSealedCT := []byte("sealed-ciphertext")
+
+	svc := NewService(nil, func(id uuid.UUID, _, _ []byte) ([]byte, []byte, error) {
+		if id != kemUUID {
+			t.Fatalf("expected KEM UUID %s, got %s", kemUUID, id)
+		}
+		return expectedSealEnc, expectedSealedCT, nil
+	})
+
+	sealEnc, sealedCT, err := svc.DecapAndSeal(kemUUID, []byte("enc-key"), []byte("aad"))
+	if err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+	if string(sealEnc) != string(expectedSealEnc) {
+		t.Fatalf("expected seal enc %q, got %q", expectedSealEnc, sealEnc)
+	}
+	if string(sealedCT) != string(expectedSealedCT) {
+		t.Fatalf("expected sealed CT %q, got %q", expectedSealedCT, sealedCT)
+	}
+}
+
+func TestServiceDecapAndSealError(t *testing.T) {
+	svc := NewService(nil, func(_ uuid.UUID, _, _ []byte) ([]byte, []byte, error) {
+		return nil, nil, fmt.Errorf("decap FFI error")
+	})
+
+	_, _, err := svc.DecapAndSeal(uuid.New(), []byte("enc-key"), nil)
+	if err == nil {
+		t.Fatal("expected error, got nil")
+	}
+}

keymanager/workload_service/integration_test.go

@@ -24,10 +24,17 @@ func (r *realWorkloadService) GenerateBindingKeypair(algo *keymanager.HpkeAlgori
 	return wskcc.GenerateBindingKeypair(algo, lifespanSecs)
 }
 
+// realOpener wraps the actual WSD KCC FFI for Open.
+type realOpener struct{}
+
+func (r *realOpener) Open(bindingUUID uuid.UUID, enc, ciphertext, aad []byte) ([]byte, error) {
+	return wskcc.Open(bindingUUID, enc, ciphertext, aad)
+}
+
 func TestIntegrationGenerateKeysEndToEnd(t *testing.T) {
 	// Wire up real FFI calls: WSD KCC for binding, KPS KCC (via KPS KOL) for KEM.
-	kpsSvc := kps.NewService(kpskcc.GenerateKEMKeypair)
-	srv := NewServer(kpsSvc, &realWorkloadService{})
+	kpsSvc := kps.NewService(kpskcc.GenerateKEMKeypair, kpskcc.DecapAndSeal)
+	srv, _ := NewServer(kpsSvc, &realWorkloadService{}, "test.sock")
 
 	reqBody, err := json.Marshal(GenerateKemRequest{
 		Algorithm:              KemAlgorithmDHKEMX25519HKDFSHA256,
@@ -72,8 +79,8 @@ func TestIntegrationGenerateKeysEndToEnd(t *testing.T) {
 }
 
 func TestIntegrationGenerateKeysUniqueMappings(t *testing.T) {
-	kpsSvc := kps.NewService(kpskcc.GenerateKEMKeypair)
-	srv := NewServer(kpsSvc, &realWorkloadService{})
+	kpsSvc := kps.NewService(kpskcc.GenerateKEMKeypair, kpskcc.DecapAndSeal)
+	srv, _ := NewServer(kpsSvc, &realWorkloadService{}, "test.sock")
 
 	// Generate two key sets.
 	var kemUUIDs [2]uuid.UUID

keymanager/workload_service/key_custody_core/ws_key_custody_core_cgo.go

@@ -54,3 +54,49 @@ func GenerateBindingKeypair(algo *keymanager.HpkeAlgorithm, lifespanSecs uint64)
 	copy(pubkey, pubkeyBuf[:pubkeyLen])
 	return id, pubkey, nil
 }
+
+// Open decrypts a sealed ciphertext using the binding key identified by
+// bindingUUID via Rust FFI (HPKE Open).
+// Returns the decrypted plaintext (shared secret).
+func Open(bindingUUID uuid.UUID, enc, ciphertext, aad []byte) ([]byte, error) {
+	if len(enc) == 0 {
+		return nil, fmt.Errorf("enc must not be empty")
+	}
+	if len(ciphertext) == 0 {
+		return nil, fmt.Errorf("ciphertext must not be empty")
+	}
+
+	uuidBytes := bindingUUID[:]
+
+	var outPT [32]byte
+	outPTLen := C.size_t(len(outPT))
+
+	// Rust key_manager_open requires non-null aad pointer.
+	// Use a sentinel byte so the pointer is always valid.
+	var aadSentinel [1]byte
+	aadPtr := (*C.uint8_t)(unsafe.Pointer(&aadSentinel[0]))
+	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_open(
+		(*C.uint8_t)(unsafe.Pointer(&uuidBytes[0])),
+		(*C.uint8_t)(unsafe.Pointer(&enc[0])),
+		C.size_t(len(enc)),
+		(*C.uint8_t)(unsafe.Pointer(&ciphertext[0])),
+		C.size_t(len(ciphertext)),
+		aadPtr,
+		aadLen,
+		(*C.uint8_t)(unsafe.Pointer(&outPT[0])),
+		outPTLen,
+	)
+	if rc != 0 {
+		return nil, fmt.Errorf("key_manager_open failed with code %d", rc)
+	}
+
+	plaintext := make([]byte, outPTLen)
+	copy(plaintext, outPT[:outPTLen])
+	return plaintext, nil
+}

keymanager/workload_service/key_custody_core/ws_key_custody_core_stub.go

@@ -14,3 +14,8 @@ import (
 func GenerateBindingKeypair(_ *algorithms.HpkeAlgorithm, _ uint64) (uuid.UUID, []byte, error) {
 	return uuid.Nil, nil, fmt.Errorf("GenerateBindingKeypair is not supported on this architecture")
 }
+
+// Open is a stub for architectures where the Rust library is not supported.
+func Open(_ uuid.UUID, _, _, _ []byte) ([]byte, error) {
+	return nil, fmt.Errorf("Open is not supported on this architecture")
+}