Add confused deputy protection for KMS calls

Author: Charan000

cmd/server/main.go

@@ -48,6 +48,7 @@ func main() {
 		burstLimit         = flag.Int("burst-limit", 0, "(deprecated) number of tokens that can be consumed in a single call, use --retry-token-capacity instead")
 		retryTokenCapacity = flag.Int("retry-token-capacity", 0, "number of tokens for client-side AWS rate-limiting on retries")
 		encryptionCtxsArr  = flag.StringArray("encryption-context", []string{}, "AWS KMS Encryption Context (e.g. 'a=b,c=d')")
+		sourceArn          = flag.String("source-arn", "", "AWS source ARN for confused deputy protection")
 		debug              = flag.Bool("debug", false, "Print debug level logs")
 	)
 	flag.Parse()
@@ -92,7 +93,7 @@ func main() {
 		zap.Int("burst-limit", *burstLimit),
 		zap.Int("retry-token-capacity", *retryTokenCapacity),
 	)
-	c, err := cloud.New(*region, *kmsEndpoint, *qpsLimit, *burstLimit, *retryTokenCapacity)
+	c, err := cloud.New(*region, *kmsEndpoint, *qpsLimit, *burstLimit, *retryTokenCapacity, *sourceArn)
 	if err != nil {
 		zap.L().Fatal("Failed to create new KMS service", zap.Error(err))
 	}

pkg/cloud/cloud.go

@@ -23,15 +23,23 @@ import (
 	"github.com/aws/aws-sdk-go-v2/config"
 	"github.com/aws/aws-sdk-go-v2/feature/ec2/imds"
 	"github.com/aws/aws-sdk-go-v2/service/kms"
+	"github.com/aws/aws-sdk-go-v2/aws/arn"
+	smithymiddleware "github.com/aws/smithy-go/middleware"
+	smithyhttp "github.com/aws/smithy-go/transport/http"
 	"go.uber.org/zap"
 )
 
+const (
+	headerSourceArn     = "x-amz-source-arn"
+	headerSourceAccount = "x-amz-source-account"
+)
+
 type AWSKMSv2 interface {
 	Encrypt(ctx context.Context, params *kms.EncryptInput, optFns ...func(*kms.Options)) (*kms.EncryptOutput, error)
 	Decrypt(ctx context.Context, params *kms.DecryptInput, optFns ...func(*kms.Options)) (*kms.DecryptOutput, error)
 }
 
-func New(region, kmsEndpoint string, qps, burst, retryTokenCapacity int) (AWSKMSv2, error) {
+func New(region, kmsEndpoint string, qps, burst, retryTokenCapacity int, sourceArn string) (AWSKMSv2, error) {
 	var optFns []func(*config.LoadOptions) error
 	if region != "" {
 		optFns = append(optFns, config.WithRegion(region))
@@ -69,6 +77,11 @@ func New(region, kmsEndpoint string, qps, burst, retryTokenCapacity int) (AWSKMS
 		return nil, fmt.Errorf("failed to create AWS config: %w", err)
 	}
 
+	err = addConfusedDeputyHeaders(&cfg, sourceArn)
+	if err != nil {
+		return nil, err
+	}
+
 	if cfg.Region == "" {
 		ec2 := imds.NewFromConfig(cfg)
 		region, err := ec2.GetRegion(context.Background(), &imds.GetRegionInput{})
@@ -88,3 +101,48 @@ func New(region, kmsEndpoint string, qps, burst, retryTokenCapacity int) (AWSKMS
 	client := kms.NewFromConfig(cfg, kmsOptFns...)
 	return client, nil
 }
+
+func addConfusedDeputyHeaders(cfg *aws.Config, sourceArn string) error {
+	if sourceArn != "" {
+		sourceAccount, err := getSourceAccount(sourceArn)
+		if err != nil {
+			return err
+		}
+
+		cfg.APIOptions = append(cfg.APIOptions, func(stack *smithymiddleware.Stack) error {
+			return stack.Build.Add(smithymiddleware.BuildMiddlewareFunc("KMSConfusedDeputyHeaders", func(
+				ctx context.Context, in smithymiddleware.BuildInput, next smithymiddleware.BuildHandler,
+			) (smithymiddleware.BuildOutput, smithymiddleware.Metadata, error) {
+				req, ok := in.Request.(*smithyhttp.Request)
+				if ok {
+					req.Header.Set(headerSourceAccount, sourceAccount)
+					req.Header.Set(headerSourceArn, sourceArn)
+				}
+				return next.HandleBuild(ctx, in)
+			}), smithymiddleware.Before)
+		})
+
+		zap.L().Info("configuring KMS client with confused deputy headers",
+			zap.String("sourceArn", sourceArn),
+			zap.String("sourceAccount", sourceAccount),
+		)
+	}
+	return nil
+}
+
+// getSourceAccount constructs source account and return them for use
+func getSourceAccount(sourceArn string) (string, error) {
+	// ARN format (https://docs.aws.amazon.com/IAM/latest/UserGuide/reference-arns.html)
+	// arn:partition:service:region:account-id:resource-type/resource-id
+	// arn:aws:eks:region:account:cluster/cluster-name
+	if !arn.IsARN(sourceArn) {
+		return "", fmt.Errorf("incorrect ARN format for source arn: %s", sourceArn)
+	}
+
+	parsedArn, err := arn.Parse(sourceArn)
+	if err != nil {
+		return "", err
+	}
+
+	return parsedArn.AccountID, nil
+}

pkg/cloud/cloud_test.go

@@ -25,7 +25,7 @@ zssmrkdYYvn9aUhjc3XK3tjAoDpsPpeBeTBamuUKDHoH/dNRXxerZ8vu6uPR3Pgs
 `)
 
 func TestNewSessionClientWithoutEnv(t *testing.T) {
-	kmsObjet, err := New("us-west-2", "https://kms.us-west-2.amazonaws.com", 0, 0, 500)
+	kmsObjet, err := New("us-west-2", "https://kms.us-west-2.amazonaws.com", 0, 0, 500, "")
 	assert.NoError(t, err, "Failed to create object with error (%v)", err)
 	assert.NotNil(t, kmsObjet, "Failed to create object with error (%v)", err)
 }
@@ -36,7 +36,7 @@ func TestNewSessionClientWithEnv(t *testing.T) {
 	defer os.Remove(tempFile)            //nolint:errcheck
 	os.Setenv("AWS_CA_BUNDLE", tempFile) //nolint:errcheck
 	defer os.Unsetenv("AWS_CA_BUNDLE")   //nolint:errcheck
-	kmsObjet, err := New("us-west-2", "https://kms.us-west-2.amazonaws.com", 0, 0, 500)
+	kmsObjet, err := New("us-west-2", "https://kms.us-west-2.amazonaws.com", 0, 0, 500, "")
 	assert.NoError(t, err, "Failed to create object with error (%v)", err)
 	assert.NotNil(t, kmsObjet, "Failed to create object with error (%v)", err)
 }
@@ -95,7 +95,7 @@ func TestNewConfig(t *testing.T) {
 
 	for _, test := range tests {
 		t.Run(test.name, func(t *testing.T) {
-			_, err := New(test.region, test.endpoint, test.qps, test.burst, test.retryTokenCapacity)
+			_, err := New(test.region, test.endpoint, test.qps, test.burst, test.retryTokenCapacity, "")
 			if test.expectErr {
 				assert.Error(t, err)
 			} else {
@@ -104,3 +104,27 @@ func TestNewConfig(t *testing.T) {
 		})
 	}
 }
+
+func TestNewWithSourceArn(t *testing.T) {
+	client, err := New("us-east-1", "", 0, 0, 0, "arn:aws:eks:us-east-1:123456789012:cluster/test")
+	assert.NoError(t, err)
+	assert.NotNil(t, client)
+}
+
+func TestNewWithEmptySourceArn(t *testing.T) {
+	client, err := New("us-east-1", "", 0, 0, 0, "")
+	assert.NoError(t, err)
+	assert.NotNil(t, client)
+}
+
+func TestNewWithMalformedSourceArn(t *testing.T) {
+	cfg, err := New("us-east-1", "", 0, 0, 0, "invalid-arn-format")
+	assert.Nil(t, cfg)
+	assert.Error(t, err)
+}
+
+func TestGetSourceAccount(t *testing.T) {
+	account, err := getSourceAccount("arn:aws:eks:us-east-1:123456789012:cluster/test")
+	assert.NoError(t, err)
+	assert.Equal(t, "123456789012", account)
+}