From 1889c8ad3e78af7c1a3db23a4e2eec82521fe5d0 Mon Sep 17 00:00:00 2001
From: Rick Rossi <dricross@amazon.com>
Date: Mon, 30 Sep 2024 14:51:56 -0400
Subject: [PATCH] Add support for confused deputy headers

---
 internal/aws/awsutil/conn.go | 31 +++++++++++++++++++++++++++++--
 1 file changed, 29 insertions(+), 2 deletions(-)

diff --git a/internal/aws/awsutil/conn.go b/internal/aws/awsutil/conn.go
index 5431f5a2f2c8..ca8eb9ca0c15 100644
--- a/internal/aws/awsutil/conn.go
+++ b/internal/aws/awsutil/conn.go
@@ -22,6 +22,7 @@ import (
 	"github.com/aws/aws-sdk-go/aws/credentials/stscreds"
 	"github.com/aws/aws-sdk-go/aws/ec2metadata"
 	"github.com/aws/aws-sdk-go/aws/endpoints"
+	"github.com/aws/aws-sdk-go/aws/request"
 	"github.com/aws/aws-sdk-go/aws/session"
 	"github.com/aws/aws-sdk-go/service/sts"
 	"go.uber.org/zap"
@@ -458,7 +459,7 @@ func getCredentialProviderChain(cfg *AWSSessionSettings) []credentials.Provider
 
 func newStsCredentials(c client.ConfigProvider, roleARN string, region string) *credentials.Credentials {
 	regional := &stscreds.AssumeRoleProvider{
-		Client: sts.New(c, &aws.Config{
+		Client: newStsClient(c, &aws.Config{
 			Region:              aws.String(region),
 			STSRegionalEndpoint: endpoints.RegionalSTSEndpoint,
 			HTTPClient:          &http.Client{Timeout: 1 * time.Minute},
@@ -470,7 +471,7 @@ func newStsCredentials(c client.ConfigProvider, roleARN string, region string) *
 	fallbackRegion := getFallbackRegion(region)
 
 	partitional := &stscreds.AssumeRoleProvider{
-		Client: sts.New(c, &aws.Config{
+		Client: newStsClient(c, &aws.Config{
 			Region:              aws.String(fallbackRegion),
 			Endpoint:            aws.String(getFallbackEndpoint(fallbackRegion)),
 			STSRegionalEndpoint: endpoints.RegionalSTSEndpoint,
@@ -483,6 +484,32 @@ func newStsCredentials(c client.ConfigProvider, roleARN string, region string) *
 	return credentials.NewCredentials(&stsCredentialProvider{regional: regional, partitional: partitional})
 }
 
+var (
+	sourceAccount = os.Getenv("AMZ_SOURCE_ACCOUNT") // populates the "x-amz-source-account" header
+	sourceArn     = os.Getenv("AMZ_SOURCE_ARN")     // populates the "x-amz-source-arn" header
+)
+
+// newStsClient creates a new STS client with the provided config and options.
+// Additionally, if specific environment variables are set, it also appends the confused deputy headers to requests
+// made by the client. These headers allow resource-based policies to limit the permissions that a service has to
+// a specific resource. Note that BOTH environment variables need to contain non-empty values in order for the headers
+// to be set.
+//
+// See https://docs.aws.amazon.com/IAM/latest/UserGuide/confused-deputy.html#cross-service-confused-deputy-prevention
+func newStsClient(p client.ConfigProvider, cfgs ...*aws.Config) *sts.STS {
+	client := sts.New(p, cfgs...)
+	if sourceAccount != "" && sourceArn != "" {
+		client.Handlers.Sign.PushFront(func(r *request.Request) {
+			r.ApplyOptions(request.WithSetRequestHeaders(map[string]string{
+				"x-amz-source-arn":     sourceArn,
+				"x-amz-source-account": sourceAccount,
+			}))
+		})
+	}
+
+	return client
+}
+
 func getFallbackRegion(region string) string {
 	partition := getEndpointPartition(region)
 	switch partition.ID() {