feat(core): Encapsulate>Encrypt (#2676)

### Proposed Changes

- Prefer using the Encapsulate interface method, when possible
- Deprecate the Key's Export method, since we generally don't want to
export the keys and Encapsulate is more accurate for what we want.
- The goal here is to present an interface that supports a fully
delegated rewrap, e.g. using a [Luna Functionality
Module](https://thalesdocs.com/gphsm/luna/7/docs/network/Content/FM_SDK/preface.htm)
or other KAS. In this case, the call to Encapsulate would trigger the
external rewrap.

### Checklist

- [ ] I have added or updated unit tests
- [ ] I have added or updated integration tests (if appropriate)
- [ ] I have added or updated documentation

### Testing Instructions

Author: dmihalcik-virtru

lib/ocrypto/ec_key_pair.go

@@ -446,15 +446,15 @@ func ECPrivateKeyInPemFormat(privateKey ecdsa.PrivateKey) (string, error) {
 
 // ECPublicKeyInPemFormat Returns public key in pem format.
 func ECPublicKeyInPemFormat(publicKey ecdsa.PublicKey) (string, error) {
-	publicKeyBytes, err := x509.MarshalPKIXPublicKey(publicKey)
+	pkb, err := x509.MarshalPKIXPublicKey(publicKey)
 	if err != nil {
 		return "", fmt.Errorf("x509.MarshalPKIXPublicKey failed: %w", err)
 	}
 
 	publicKeyPem := pem.EncodeToMemory(
 		&pem.Block{
 			Type:  "PUBLIC KEY",
-			Bytes: publicKeyBytes,
+			Bytes: pkb,
 		},
 	)
 

lib/ocrypto/interfaces.go

@@ -27,6 +27,7 @@ type ProtectedKey interface {
 	VerifyBinding(ctx context.Context, policy, policyBinding []byte) error
 
 	// Export returns the raw key data, optionally encrypting it with the provided encapsulator
+	// Deprecated: Use the Encapsulator's Encapsulate method instead
 	Export(encapsulator Encapsulator) ([]byte, error)
 
 	// DecryptAESGCM decrypts encrypted policies and metadata

lib/ocrypto/protected_key.go

@@ -55,14 +55,15 @@ func (k *AESProtectedKey) DecryptAESGCM(iv []byte, body []byte, tagSize int) ([]
 }
 
 // Export returns the raw key data, optionally encrypting it with the provided Encapsulator
+// Deprecated: Use the Encapsulator's Encapsulate method instead
 func (k *AESProtectedKey) Export(encapsulator Encapsulator) ([]byte, error) {
 	if encapsulator == nil {
 		// Return error if encapsulator is nil
 		return nil, errors.New("encapsulator cannot be nil")
 	}
 
 	// Encrypt the key data before returning
-	encryptedKey, err := encapsulator.Encrypt(k.rawKey)
+	encryptedKey, err := encapsulator.Encapsulate(k)
 	if err != nil {
 		return nil, fmt.Errorf("failed to encrypt key data for export: %w", err)
 	}

lib/ocrypto/protected_key_test.go

@@ -75,14 +75,22 @@ func TestAESProtectedKey_Export_WithEncapsulator(t *testing.T) {
 	require.NoError(t, err)
 
 	// Mock encapsulator
+	mockEncryptFunc := func(data []byte) ([]byte, error) {
+		// Simple XOR encryption for testing
+		result := make([]byte, len(data))
+		for i, b := range data {
+			result[i] = b ^ 0xFF
+		}
+		return result, nil
+	}
 	mockEncapsulator := &mockEncapsulator{
-		encryptFunc: func(data []byte) ([]byte, error) {
-			// Simple XOR encryption for testing
-			result := make([]byte, len(data))
-			for i, b := range data {
-				result[i] = b ^ 0xFF
+		encryptFunc: mockEncryptFunc,
+		encapsulateFunc: func(k ProtectedKey) ([]byte, error) {
+			aeskey, ok := k.(*AESProtectedKey)
+			if !ok {
+				return nil, assert.AnError
 			}
-			return result, nil
+			return mockEncryptFunc(aeskey.rawKey)
 		},
 	}
 
@@ -105,7 +113,7 @@ func TestAESProtectedKey_Export_EncapsulatorError(t *testing.T) {
 	require.NoError(t, err)
 
 	mockEncapsulator := &mockEncapsulator{
-		encryptFunc: func(_ []byte) ([]byte, error) {
+		encapsulateFunc: func(_ ProtectedKey) ([]byte, error) {
 			return nil, assert.AnError
 		},
 	}
@@ -175,12 +183,13 @@ func TestAESProtectedKey_InterfaceCompliance(t *testing.T) {
 // Mock encapsulator for testing
 type mockEncapsulator struct {
 	encryptFunc      func([]byte) ([]byte, error)
+	encapsulateFunc  func(ProtectedKey) ([]byte, error)
 	publicKeyPEMFunc func() (string, error)
 	ephemeralKeyFunc func() []byte
 }
 
-func (m *mockEncapsulator) Encapsulate(_ ProtectedKey) ([]byte, error) {
-	return nil, nil
+func (m *mockEncapsulator) Encapsulate(pk ProtectedKey) ([]byte, error) {
+	return m.encapsulateFunc(pk)
 }
 
 func (m *mockEncapsulator) Encrypt(data []byte) ([]byte, error) {

service/internal/security/basic_manager.go

@@ -133,9 +133,24 @@ func (b *BasicManager) DeriveKey(ctx context.Context, keyDetails trust.KeyDetail
 	return NewInProcessAESKey(key), nil
 }
 
+type OCEncapsulator struct {
+	ocrypto.PublicKeyEncryptor
+}
+
+func (e *OCEncapsulator) Encapsulate(dek trust.ProtectedKey) ([]byte, error) {
+	ipk, ok := dek.(*InProcessAESKey)
+	if !ok {
+		return nil, errors.New("invalid DEK type for encapsulation")
+	}
+	return e.Encrypt(ipk.rawKey)
+}
+
 func (b *BasicManager) GenerateECSessionKey(_ context.Context, ephemeralPublicKey string) (trust.Encapsulator, error) {
-	// Implementation of GenerateECSessionKey method
-	return ocrypto.FromPublicPEMWithSalt(ephemeralPublicKey, NanoVersionSalt(), nil)
+	pke, err := ocrypto.FromPublicPEMWithSalt(ephemeralPublicKey, NanoVersionSalt(), nil)
+	if err != nil {
+		return nil, fmt.Errorf("failed to create public key encryptor: %w", err)
+	}
+	return &OCEncapsulator{PublicKeyEncryptor: pke}, nil
 }
 
 func (b *BasicManager) Close() {

service/internal/security/in_process_provider.go

@@ -57,7 +57,7 @@ func (k *InProcessAESKey) Export(encapsulator trust.Encapsulator) ([]byte, error
 	}
 
 	// If an encryptor is provided, encrypt the key data before returning
-	encryptedKey, err := encapsulator.Encrypt(k.rawKey)
+	encryptedKey, err := encapsulator.Encapsulate(k)
 	if err != nil {
 		if k.logger != nil {
 			k.logger.Warn("failed to encrypt key data for export", slog.Any("err", err))
@@ -349,7 +349,11 @@ func (a *InProcessProvider) DeriveKey(_ context.Context, keyDetails trust.KeyDet
 
 // GenerateECSessionKey generates a session key for NanoTDF
 func (a *InProcessProvider) GenerateECSessionKey(_ context.Context, ephemeralPublicKey string) (trust.Encapsulator, error) {
-	return ocrypto.FromPublicPEMWithSalt(ephemeralPublicKey, NanoVersionSalt(), nil)
+	pke, err := ocrypto.FromPublicPEMWithSalt(ephemeralPublicKey, NanoVersionSalt(), nil)
+	if err != nil {
+		return nil, fmt.Errorf("session key generation failed to create public key encryptor: %w", err)
+	}
+	return &OCEncapsulator{PublicKeyEncryptor: pke}, nil
 }
 
 // Close releases any resources held by the provider

service/kas/access/rewrap.go

@@ -712,6 +712,7 @@ func (p *Provider) tdf3Rewrap(ctx context.Context, requests []*kaspb.UnsignedRew
 		failAllKaos(requests, results, err400("invalid request"))
 		return "", results
 	}
+	encap := security.OCEncapsulator{PublicKeyEncryptor: asymEncrypt}
 
 	var sessionKey string
 	if e, ok := asymEncrypt.(ocrypto.ECEncryptor); ok {
@@ -771,7 +772,7 @@ func (p *Provider) tdf3Rewrap(ctx context.Context, requests []*kaspb.UnsignedRew
 			}
 
 			// Use the Export method with the asymEncrypt encryptor
-			encryptedKey, err := kaoRes.DEK.Export(asymEncrypt)
+			encryptedKey, err := encap.Encapsulate(kaoRes.DEK)
 			if err != nil {
 				//nolint:sloglint // reference to camelcase key is intentional
 				p.Logger.WarnContext(ctx, "rewrap: Export with encryptor failed", slog.String("clientPublicKey", clientPublicKey), slog.Any("error", err))

service/trust/delegating_key_service_test.go

@@ -172,6 +172,14 @@ type MockEncapsulator struct {
 	mock.Mock
 }
 
+func (m *MockEncapsulator) Encapsulate(dek ProtectedKey) ([]byte, error) {
+	args := m.Called(dek)
+	if a0, ok := args.Get(0).([]byte); ok {
+		return a0, args.Error(1)
+	}
+	return nil, args.Error(1)
+}
+
 func (m *MockEncapsulator) Encrypt(data []byte) ([]byte, error) {
 	args := m.Called(data)
 	if a0, ok := args.Get(0).([]byte); ok {

service/trust/key_manager.go

@@ -6,6 +6,9 @@ import (
 )
 
 type Encapsulator interface {
+	// Encapsulate wraps a secret key with the encapsulation key
+	Encapsulate(dek ProtectedKey) ([]byte, error)
+
 	// Encrypt wraps a secret key with the encapsulation key
 	Encrypt(data []byte) ([]byte, error)
 
@@ -23,6 +26,7 @@ type ProtectedKey interface {
 	VerifyBinding(ctx context.Context, policy, binding []byte) error
 
 	// Export returns the raw key data, optionally encrypting it with the provided encryptor
+	// Deprecated: Use the Encapsulator's Encapsulate method instead
 	Export(encryptor Encapsulator) ([]byte, error)
 
 	// Used to decrypt encrypted policies and metadata