Allow injecting custom client options for AWS, Azure and GCP

Signed-off-by: Matheus Pimenta <matheuscscp@gmail.com>

Author: matheuscscp

azkv/keysource.go

@@ -60,6 +60,8 @@ type MasterKey struct {
 	// using TokenCredential.ApplyToMasterKey.
 	// If nil, azidentity.NewDefaultAzureCredential is used.
 	tokenCredential azcore.TokenCredential
+	// clientOptions contains the azkeys.ClientOptions used by the Azure client.
+	clientOptions *azkeys.ClientOptions
 }
 
 // NewMasterKey creates a new MasterKey from a URL, key name and version,
@@ -118,6 +120,17 @@ func (t TokenCredential) ApplyToMasterKey(key *MasterKey) {
 	key.tokenCredential = t.token
 }
 
+// ClientOptions is a wrapper around azkeys.ClientOptions to allow
+// configuration of the Azure Key Vault client.
+type ClientOptions struct {
+	o *azkeys.ClientOptions
+}
+
+// ApplyToMasterKey configures the ClientOptions on the provided key.
+func (c ClientOptions) ApplyToMasterKey(key *MasterKey) {
+	key.clientOptions = c.o
+}
+
 // Encrypt takes a SOPS data key, encrypts it with Azure Key Vault, and stores
 // the result in the EncryptedKey field.
 func (key *MasterKey) Encrypt(dataKey []byte) error {
@@ -182,7 +195,7 @@ func (key *MasterKey) Decrypt() ([]byte, error) {
 		return nil, fmt.Errorf("failed to base64 decode Azure Key Vault encrypted key: %w", err)
 	}
 
-	c, err := azkeys.NewClient(key.VaultURL, token, nil)
+	c, err := azkeys.NewClient(key.VaultURL, token, key.clientOptions)
 	if err != nil {
 		log.WithFields(logrus.Fields{"key": key.Name, "version": key.Version}).Info("Decryption failed")
 		return nil, fmt.Errorf("failed to construct Azure Key Vault client to decrypt data: %w", err)

gcpkms/keysource.go

@@ -66,6 +66,8 @@ type MasterKey struct {
 	// Mostly useful for testing at present, to wire the client to a mock
 	// server.
 	grpcConn *grpc.ClientConn
+	// grpcDialOpts are the gRPC dial options used to create the gRPC connection.
+	grpcDialOpts []grpc.DialOption
 }
 
 // NewMasterKeyFromResourceID creates a new MasterKey with the provided resource
@@ -116,6 +118,14 @@ func (c CredentialJSON) ApplyToMasterKey(key *MasterKey) {
 	key.credentialJSON = c
 }
 
+// DialOptions are the gRPC dial options used to create the gRPC connection.
+type DialOptions []grpc.DialOption
+
+// ApplyToMasterKey configures the DialOptions on the provided key.
+func (d DialOptions) ApplyToMasterKey(key *MasterKey) {
+	key.grpcDialOpts = d
+}
+
 // Encrypt takes a SOPS data key, encrypts it with GCP KMS, and stores the
 // result in the EncryptedKey field.
 func (key *MasterKey) Encrypt(dataKey []byte) error {
@@ -261,8 +271,13 @@ func (key *MasterKey) newKMSClient() (*kms.KeyManagementClient, error) {
 		}
 	}
 
-	if key.grpcConn != nil {
+	switch {
+	case key.grpcConn != nil:
 		opts = append(opts, option.WithGRPCConn(key.grpcConn))
+	case len(key.grpcDialOpts) > 0:
+		for _, opt := range key.grpcDialOpts {
+			opts = append(opts, option.WithGRPCDialOption(opt))
+		}
 	}
 
 	ctx := context.Background()

kms/keysource.go

@@ -9,6 +9,7 @@ import (
 	"context"
 	"encoding/base64"
 	"fmt"
+	"net/http"
 	"os"
 	"regexp"
 	"sort"
@@ -79,6 +80,8 @@ type MasterKey struct {
 	// injected using e.g. an environment variable. The field is not publicly
 	// exposed, nor configurable.
 	baseEndpoint string
+	// httpClient is used to override the default HTTP client used by the AWS client.
+	httpClient *http.Client
 }
 
 // NewMasterKey creates a new MasterKey from an ARN, role and context, setting
@@ -233,6 +236,17 @@ func (c CredentialsProvider) ApplyToMasterKey(key *MasterKey) {
 	key.credentialsProvider = c.provider
 }
 
+// HTTPClient is a wrapper around http.Client used for configuring the
+// AWS KMS client.
+type HTTPClient struct {
+	hc *http.Client
+}
+
+// ApplyToMasterKey configures the HTTP client on the provided key.
+func (h HTTPClient) ApplyToMasterKey(key *MasterKey) {
+	key.httpClient = h.hc
+}
+
 // Encrypt takes a SOPS data key, encrypts it with KMS and stores the result
 // in the EncryptedKey field.
 func (key *MasterKey) Encrypt(dataKey []byte) error {
@@ -369,6 +383,9 @@ func (key MasterKey) createKMSConfig() (*aws.Config, error) {
 			lo.SharedConfigProfile = key.AwsProfile
 		}
 		lo.Region = region
+		if key.httpClient != nil {
+			lo.HTTPClient = key.httpClient
+		}
 		return nil
 	})
 	if err != nil {