feat: add Azure Key Vault support for payload encryption

codec/codec.go

@@ -4,6 +4,8 @@ import (
 	gcpKms "cloud.google.com/go/kms/apiv1"
 	"context"
 	"fmt"
+	"github.com/Azure/azure-sdk-for-go/sdk/azidentity"
+	"github.com/Azure/azure-sdk-for-go/sdk/security/keyvault/azkeys"
 	"github.com/aws/aws-sdk-go/aws"
 	"github.com/aws/aws-sdk-go/aws/session"
 	awsKms "github.com/aws/aws-sdk-go/service/kms"
@@ -136,9 +138,86 @@ func newCodecFactoryProvider(configProvider config.ConfigProvider) (EncryptionCo
 		), nil
 	}
 
+	cf.providers["azure-keyvault"] = func(args EncryptionCodecOptions) (converter.PayloadCodec, error) {
+		rawKeyID, ok := args.LocalEncryptionConfig.Config["key-id"]
+		if !ok {
+			return nil, fmt.Errorf("key not found in config")
+		}
+		keyID, ok := rawKeyID.(string)
+		if !ok {
+			return nil, fmt.Errorf("key is not a string")
+		}
+
+		vaultURL := os.Getenv(config.AzureVaultURLEnvVar)
+		if vaultURL == "" {
+			return nil, fmt.Errorf("AZURE_VAULT_URL environment variable is required")
+		}
+
+		cred, err := azidentity.NewDefaultAzureCredential(nil)
+		if err != nil {
+			return nil, fmt.Errorf("failed to create Azure credential: %v", err)
+		}
+
+		azClient, err := azkeys.NewClient(vaultURL, cred, nil)
+		if err != nil {
+			return nil, fmt.Errorf("failed to create Azure Key Vault client: %v", err)
+		}
+
+		azAdapter := &azureKeysClientAdapter{client: azClient}
+
+		azureMaterialsManager := crypto.NewAzureKeyVaultProvider(azAdapter, crypto.AzureKeyVaultOptions{
+			KeyID:     keyID,
+			Algorithm: "RSA-OAEP-256",
+		})
+
+		args.MetricsHandler.AddAttributes(attribute.String("encryption_key", keyID))
+
+		return NewEncryptionCodecWithCaching(
+			azureMaterialsManager,
+			args.CodecContext,
+			keyID,
+			args.MetricsHandler,
+			cf.cachingConfig,
+		), nil
+	}
+
 	return cf, nil
 }
 
+// azureKeysClientAdapter adapts azkeys.Client to the crypto.AzureKeyVaultClient interface
+type azureKeysClientAdapter struct {
+	client *azkeys.Client
+}
+
+func (a *azureKeysClientAdapter) WrapKey(ctx context.Context, keyID string, algorithm string, plaintext []byte) ([]byte, error) {
+	params := azkeys.KeyOperationParameters{
+		Algorithm: toAzKeysAlgorithm(algorithm),
+		Value:     plaintext,
+	}
+	resp, err := a.client.WrapKey(ctx, keyID, "", params, nil)
+	if err != nil {
+		return nil, err
+	}
+	return resp.Result, nil
+}
+
+func (a *azureKeysClientAdapter) UnwrapKey(ctx context.Context, keyID string, algorithm string, ciphertext []byte) ([]byte, error) {
+	params := azkeys.KeyOperationParameters{
+		Algorithm: toAzKeysAlgorithm(algorithm),
+		Value:     ciphertext,
+	}
+	resp, err := a.client.UnwrapKey(ctx, keyID, "", params, nil)
+	if err != nil {
+		return nil, err
+	}
+	return resp.Result, nil
+}
+
+func toAzKeysAlgorithm(algorithm string) *azkeys.EncryptionAlgorithm {
+	alg := azkeys.EncryptionAlgorithm(algorithm)
+	return &alg
+}
+
 func (e *encryptionCodecFactory) NewEncryptionCodec(args EncryptionCodecOptions) (converter.PayloadCodec, error) {
 	encryptionCodec, ok := e.providers[args.LocalEncryptionConfig.Type]
 	if !ok {

config/config.go

@@ -18,6 +18,8 @@ const (
 
 	GcpRegionEnvVar  = "GCP_REGION"
 	DefaultGcpRegion = "us-central1"
+
+	AzureVaultURLEnvVar = "AZURE_VAULT_URL"
 )
 
 type (

crypto/azure_keyvault_provider.go

@@ -0,0 +1,79 @@
+package crypto
+
+import (
+	"context"
+	"crypto/rand"
+	"fmt"
+)
+
+// AzureKeyVaultOptions contains configuration options for AzureKeyVaultProvider
+type AzureKeyVaultOptions struct {
+	// KeyID is the Azure Key Vault key identifier URL
+	// Format: https://{vault-name}.vault.azure.net/keys/{key-name}/{key-version}
+	KeyID string
+
+	// Algorithm is the key wrap algorithm to use (defaults to RSA-OAEP-256 if empty)
+	Algorithm string
+}
+
+// AzureKeyVaultClient defines the interface for Azure Key Vault key operations
+type AzureKeyVaultClient interface {
+	WrapKey(ctx context.Context, keyID string, algorithm string, plaintext []byte) ([]byte, error)
+	UnwrapKey(ctx context.Context, keyID string, algorithm string, ciphertext []byte) ([]byte, error)
+}
+
+// AzureKeyVaultProvider implements MaterialsManager using Azure Key Vault
+type AzureKeyVaultProvider struct {
+	client    AzureKeyVaultClient
+	keyID     string
+	algorithm string
+}
+
+// NewAzureKeyVaultProvider creates a new Azure Key Vault-based materials manager
+func NewAzureKeyVaultProvider(client AzureKeyVaultClient, options AzureKeyVaultOptions) *AzureKeyVaultProvider {
+	// Set default algorithm if not provided
+	algorithm := options.Algorithm
+	if algorithm == "" {
+		algorithm = "RSA-OAEP-256"
+	}
+
+	return &AzureKeyVaultProvider{
+		client:    client,
+		keyID:     options.KeyID,
+		algorithm: algorithm,
+	}
+}
+
+// GetMaterial generates new encryption materials using Azure Key Vault
+func (a *AzureKeyVaultProvider) GetMaterial(ctx context.Context, cryptoCtx CryptoContext) (*Material, error) {
+	// Generate a random 256-bit data key locally
+	plaintextKey := make([]byte, 32)
+	if _, err := rand.Read(plaintextKey); err != nil {
+		return nil, fmt.Errorf("failed to generate random key: %v", err)
+	}
+
+	// Wrap the data key using Azure Key Vault
+	encryptedKey, err := a.client.WrapKey(ctx, a.keyID, a.algorithm, plaintextKey)
+	if err != nil {
+		return nil, fmt.Errorf("failed to wrap data key: %v", err)
+	}
+
+	return &Material{
+		PlaintextKey: plaintextKey,
+		EncryptedKey: encryptedKey,
+	}, nil
+}
+
+// DecryptMaterial decrypts the encrypted key using Azure Key Vault
+func (a *AzureKeyVaultProvider) DecryptMaterial(ctx context.Context, cryptoCtx CryptoContext, material *Material) (*Material, error) {
+	// Unwrap the data key using Azure Key Vault
+	plaintextKey, err := a.client.UnwrapKey(ctx, a.keyID, a.algorithm, material.EncryptedKey)
+	if err != nil {
+		return nil, fmt.Errorf("failed to unwrap data key: %v", err)
+	}
+
+	return &Material{
+		PlaintextKey: plaintextKey,
+		EncryptedKey: material.EncryptedKey,
+	}, nil
+}

crypto/azure_keyvault_provider_test.go

@@ -0,0 +1,214 @@
+package crypto
+
+import (
+	"context"
+	"errors"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+// MockAzureKeyVaultClient implements AzureKeyVaultClient for testing
+type MockAzureKeyVaultClient struct {
+	wrapKeyOutput []byte
+	wrapKeyError  error
+	unwrapOutput  []byte
+	unwrapError   error
+	lastKeyID     string
+	lastAlgorithm string
+	lastPlaintext []byte
+	lastCipher    []byte
+}
+
+func (m *MockAzureKeyVaultClient) WrapKey(ctx context.Context, keyID string, algorithm string, plaintext []byte) ([]byte, error) {
+	m.lastKeyID = keyID
+	m.lastAlgorithm = algorithm
+	m.lastPlaintext = plaintext
+	return m.wrapKeyOutput, m.wrapKeyError
+}
+
+func (m *MockAzureKeyVaultClient) UnwrapKey(ctx context.Context, keyID string, algorithm string, ciphertext []byte) ([]byte, error) {
+	m.lastKeyID = keyID
+	m.lastAlgorithm = algorithm
+	m.lastCipher = ciphertext
+	return m.unwrapOutput, m.unwrapError
+}
+
+func TestNewAzureKeyVaultProvider(t *testing.T) {
+	tests := []struct {
+		name              string
+		options           AzureKeyVaultOptions
+		expectedAlgorithm string
+	}{
+		{
+			name:              "Default Algorithm",
+			options:           AzureKeyVaultOptions{KeyID: "https://myvault.vault.azure.net/keys/mykey/version1"},
+			expectedAlgorithm: "RSA-OAEP-256",
+		},
+		{
+			name:              "Custom Algorithm",
+			options:           AzureKeyVaultOptions{KeyID: "https://myvault.vault.azure.net/keys/mykey/version1", Algorithm: "RSA-OAEP"},
+			expectedAlgorithm: "RSA-OAEP",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			mockClient := &MockAzureKeyVaultClient{
+				wrapKeyOutput: []byte("test-wrapped-key"),
+			}
+			provider := NewAzureKeyVaultProvider(mockClient, tt.options)
+
+			// Verify algorithm by making a call that uses it
+			cryptoCtx := CryptoContext{"purpose": "test"}
+			ctx := context.Background()
+			_, err := provider.GetMaterial(ctx, cryptoCtx)
+			require.NoError(t, err)
+
+			assert.Equal(t, tt.options.KeyID, mockClient.lastKeyID)
+			assert.Equal(t, tt.expectedAlgorithm, mockClient.lastAlgorithm)
+		})
+	}
+}
+
+func TestAzureKeyVaultProvider_GetMaterial(t *testing.T) {
+	tests := []struct {
+		name          string
+		context       CryptoContext
+		wrapOutput    []byte
+		wrapError     error
+		expectedError bool
+	}{
+		{
+			name:          "Success",
+			context:       CryptoContext{"purpose": "test"},
+			wrapOutput:    []byte("test-wrapped-key"),
+			wrapError:     nil,
+			expectedError: false,
+		},
+		{
+			name:          "Wrap Key Error",
+			context:       CryptoContext{"purpose": "test"},
+			wrapOutput:    nil,
+			wrapError:     errors.New("Key Vault error"),
+			expectedError: true,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			mockClient := &MockAzureKeyVaultClient{
+				wrapKeyOutput: tt.wrapOutput,
+				wrapKeyError:  tt.wrapError,
+			}
+
+			provider := NewAzureKeyVaultProvider(mockClient, AzureKeyVaultOptions{
+				KeyID: "https://myvault.vault.azure.net/keys/mykey/version1",
+			})
+			ctx := context.Background()
+			material, err := provider.GetMaterial(ctx, tt.context)
+
+			if tt.expectedError {
+				assert.Error(t, err)
+				assert.Nil(t, material)
+			} else {
+				require.NoError(t, err)
+				assert.NotNil(t, material)
+				assert.Len(t, material.PlaintextKey, 32) // 256-bit key
+				assert.Equal(t, tt.wrapOutput, material.EncryptedKey)
+
+				// Verify the plaintext key was sent to WrapKey
+				assert.Equal(t, material.PlaintextKey, mockClient.lastPlaintext)
+				assert.Equal(t, "https://myvault.vault.azure.net/keys/mykey/version1", mockClient.lastKeyID)
+				assert.Equal(t, "RSA-OAEP-256", mockClient.lastAlgorithm)
+			}
+		})
+	}
+}
+
+func TestAzureKeyVaultProvider_DecryptMaterial(t *testing.T) {
+	tests := []struct {
+		name              string
+		context           CryptoContext
+		encryptedKey      []byte
+		unwrapOutput      []byte
+		unwrapError       error
+		expectedPlaintext []byte
+		expectedError     bool
+	}{
+		{
+			name:              "Success",
+			context:           CryptoContext{"purpose": "test"},
+			encryptedKey:      []byte("test-wrapped-key"),
+			unwrapOutput:      []byte("test-plaintext-key"),
+			unwrapError:       nil,
+			expectedPlaintext: []byte("test-plaintext-key"),
+			expectedError:     false,
+		},
+		{
+			name:              "Unwrap Key Error",
+			context:           CryptoContext{"purpose": "test"},
+			encryptedKey:      []byte("test-wrapped-key"),
+			unwrapOutput:      nil,
+			unwrapError:       errors.New("Key Vault error"),
+			expectedPlaintext: nil,
+			expectedError:     true,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			mockClient := &MockAzureKeyVaultClient{
+				unwrapOutput: tt.unwrapOutput,
+				unwrapError:  tt.unwrapError,
+			}
+
+			provider := NewAzureKeyVaultProvider(mockClient, AzureKeyVaultOptions{
+				KeyID: "https://myvault.vault.azure.net/keys/mykey/version1",
+			})
+			inputMaterial := &Material{
+				EncryptedKey: tt.encryptedKey,
+			}
+			ctx := context.Background()
+			material, err := provider.DecryptMaterial(ctx, tt.context, inputMaterial)
+
+			if tt.expectedError {
+				assert.Error(t, err)
+				assert.Nil(t, material)
+			} else {
+				require.NoError(t, err)
+				assert.Equal(t, tt.expectedPlaintext, material.PlaintextKey)
+				assert.Equal(t, tt.encryptedKey, material.EncryptedKey)
+
+				// Verify the correct ciphertext was sent to UnwrapKey
+				assert.Equal(t, tt.encryptedKey, mockClient.lastCipher)
+				assert.Equal(t, "https://myvault.vault.azure.net/keys/mykey/version1", mockClient.lastKeyID)
+				assert.Equal(t, "RSA-OAEP-256", mockClient.lastAlgorithm)
+			}
+		})
+	}
+}
+
+func TestAzureKeyVaultProvider_GetMaterialGeneratesUniqueKeys(t *testing.T) {
+	mockClient := &MockAzureKeyVaultClient{
+		wrapKeyOutput: []byte("test-wrapped-key"),
+	}
+
+	provider := NewAzureKeyVaultProvider(mockClient, AzureKeyVaultOptions{
+		KeyID: "https://myvault.vault.azure.net/keys/mykey/version1",
+	})
+
+	ctx := context.Background()
+	cryptoCtx := CryptoContext{"purpose": "test"}
+
+	// Generate two materials and verify they have different plaintext keys
+	material1, err := provider.GetMaterial(ctx, cryptoCtx)
+	require.NoError(t, err)
+
+	material2, err := provider.GetMaterial(ctx, cryptoCtx)
+	require.NoError(t, err)
+
+	assert.NotEqual(t, material1.PlaintextKey, material2.PlaintextKey,
+		"each call should generate a unique random key")
+}