Add confused deputy protection for KMS calls

- Add --source-arn flag to accept EKS cluster ARN for confused deputy protection
- Implement middleware to add x-amz-source-arn and x-amz-source-account headers to KMS API calls
- Add getSourceAccount function to parse account ID from ARN
- Update cloud.New function signature to accept sourceArn parameter
- Add comprehensive tests for new functionality
- Maintain backward compatibility when sourceArn is not provided

Fixes: KMS confused deputy protection requirement for EKS encryption provider

Author: Charan000

pkg/cloud/cloud.go

@@ -77,9 +77,7 @@ func New(region, kmsEndpoint string, qps, burst, retryTokenCapacity int, sourceA
 		return nil, fmt.Errorf("failed to create AWS config: %w", err)
 	}
 
-	if sourceArn != "" {
-		cfg.APIOptions = append(cfg.APIOptions, addConfusedDeputyHeaders(sourceArn))
-	}
+	addConfusedDeputyHeaders(&cfg, sourceArn)
 
 	if cfg.Region == "" {
 		ec2 := imds.NewFromConfig(cfg)
@@ -101,28 +99,46 @@ func New(region, kmsEndpoint string, qps, burst, retryTokenCapacity int, sourceA
 	return client, nil
 }
 
-func addConfusedDeputyHeaders(sourceArn string) func(*smithymiddleware.Stack) error {
-	return 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 {
-				sourceAccount, err := getSourceAccount(sourceArn)
-				if err == nil {
+func addConfusedDeputyHeaders(cfg *aws.Config, sourceArn string) {
+	if sourceArn != "" {
+		sourceAccount, err := getSourceAccount(sourceArn)
+		if err != nil {
+			panic(fmt.Sprintf("%s is not a valid arn, err: %v", sourceArn, 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)
+				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)
+	)
 }
 
+// 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(roleARN) {
+		return "", fmt.Errorf("incorrect ARN format for role %s", roleARN)
+	}
+
 	parsedArn, err := arn.Parse(sourceArn)
 	if err != nil {
 		return "", err
 	}
+
 	return parsedArn.AccountID, nil
 }

pkg/cloud/cloud_test.go

@@ -117,8 +117,14 @@ func TestNewWithEmptySourceArn(t *testing.T) {
 	assert.NotNil(t, client)
 }
 
+func TestNewWithMalformedSourceArn(t *testing.T) {
+	assert.Panics(t, func() {
+		New("us-east-1", "", 0, 0, 0, "invalid-arn-format")
+	})
+}
+
 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)
-}
+}