Update aws-sdk-go and change way to get regional sts endpoint

Author: jaydeokar

go.mod

@@ -4,7 +4,7 @@ go 1.22.5
 
 require (
 	github.com/aws/amazon-vpc-cni-k8s v1.18.1
-	github.com/aws/aws-sdk-go v1.51.32
+	github.com/aws/aws-sdk-go v1.55.5
 	github.com/go-logr/logr v1.4.2
 	github.com/go-logr/zapr v1.3.0
 	github.com/golang/mock v1.6.0

go.sum

@@ -4,6 +4,8 @@ github.com/aws/amazon-vpc-cni-k8s v1.18.1 h1:u/OeBgnUUX6f3PCEOpA4dbG0+iZ71CnY6tE
 github.com/aws/amazon-vpc-cni-k8s v1.18.1/go.mod h1:m/J5GsxF0Th2iQTOE3ww4W9LFvwdC0tGyA9dIL4h6iQ=
 github.com/aws/aws-sdk-go v1.51.32 h1:A6mPui7QP4mwmovyzgtdedbRbNur1Iu0/El7hBWNHms=
 github.com/aws/aws-sdk-go v1.51.32/go.mod h1:LF8svs817+Nz+DmiMQKTO3ubZ/6IaTpq3TjupRn3Eqk=
+github.com/aws/aws-sdk-go v1.55.5 h1:KKUZBfBoyqy5d3swXyiC7Q76ic40rYcbqH7qjh59kzU=
+github.com/aws/aws-sdk-go v1.55.5/go.mod h1:eRwEWoyTWFMVYVQzKMNHWP5/RV4xIUGMQfXQHfHkpNU=
 github.com/beorn7/perks v1.0.1 h1:VlbKKnNfV8bJzeqoa4cOKqO6bYr3WgKZxO8Z16+hsOM=
 github.com/beorn7/perks v1.0.1/go.mod h1:G2ZrVWU2WbWT9wwq4/hrbKbnv/1ERSJQ0ibhJ6rlkpw=
 github.com/cespare/xxhash/v2 v2.2.0 h1:DC2CZ1Ep5Y4k3ZQ899DldepgrayRUGE6BBZ/cd9Cj44=

pkg/aws/ec2/api/wrapper.go

@@ -510,21 +510,21 @@ func (e *ec2Wrapper) getClientUsingAssumedRole(instanceRegion, roleARN, clusterN
 	}
 	e.log.Info("created rate limited http client", "qps", qps, "burst", burst)
 
-	// Get the regional sts end point
-	regionalSTSEndpoint, err := endpoints.DefaultResolver().
-		EndpointFor("sts", aws.StringValue(userStsSession.Config.Region), endpoints.STSRegionalEndpointOption)
-	if err != nil {
-		return nil, fmt.Errorf("failed to get the regional sts endoint for region %s: %v",
-			*userStsSession.Config.Region, err)
-	}
-
+	// GetPartition ID, SourceAccount and SourceARN
 	roleARN = strings.Trim(roleARN, "\"")
 
-	sourceAcct, sourceArn, err := utils.GetSourceAcctAndArn(roleARN, region, clusterName)
+	sourceAcct, partitionID, sourceArn, err := utils.GetSourceAcctAndArn(roleARN, region, clusterName)
 	if err != nil {
 		return nil, err
 	}
 
+	// Get the regional sts end point
+	regionalSTSEndpoint, err := e.getRegionalStsEndpoint(partitionID, region)
+	if err != nil {
+		return nil, fmt.Errorf("failed to get the regional sts endpoint for region %s: %v %v",
+			*userStsSession.Config.Region, err, partitionID)
+	}
+
 	regionalProvider := &stscreds.AssumeRoleProvider{
 		Client:          e.createSTSClient(userStsSession, client, regionalSTSEndpoint, sourceAcct, sourceArn),
 		RoleARN:         roleARN,
@@ -547,7 +547,7 @@ func (e *ec2Wrapper) getClientUsingAssumedRole(instanceRegion, roleARN, clusterN
 	// If the regional STS endpoint is different than the global STS endpoint then add the global sts endpoint
 	if regionalSTSEndpoint.URL != globalSTSEndpoint.URL {
 		globalProvider := &stscreds.AssumeRoleProvider{
-			Client:   e.createSTSClient(userStsSession, client, regionalSTSEndpoint, sourceAcct, sourceArn),
+			Client:   e.createSTSClient(userStsSession, client, globalSTSEndpoint, sourceAcct, sourceArn),
 			RoleARN:  roleARN,
 			Duration: time.Minute * 60,
 		}
@@ -892,3 +892,35 @@ func (e *ec2Wrapper) DisassociateTrunkInterface(input *ec2.DisassociateTrunkInte
 	}
 	return err
 }
+
+func (e *ec2Wrapper) getRegionalStsEndpoint(partitionID, region string) (endpoints.ResolvedEndpoint, error) {
+	var partition *endpoints.Partition
+	var stsServiceID = "sts"
+	for _, p := range endpoints.DefaultPartitions() {
+		if partitionID == p.ID() {
+			partition = &p
+			break
+		}
+	}
+	if partition == nil {
+		return endpoints.ResolvedEndpoint{}, fmt.Errorf("partition %s not valid", partitionID)
+	}
+
+	stsSvc, ok := partition.Services()[stsServiceID]
+	if !ok {
+		e.log.Info("STS service not found in partition, generating default endpoint.", "Partition:", partitionID)
+		// Add the host of the current instances region if the service doesn't already exists in the partition
+		// so we don't fail if the service is not present in the go sdk but matches the instances region.
+		res, err := partition.EndpointFor(stsServiceID, region, endpoints.STSRegionalEndpointOption, endpoints.ResolveUnknownServiceOption)
+		if err != nil {
+			return endpoints.ResolvedEndpoint{}, fmt.Errorf("error resolving endpoint for %s in partition %s. err: %v", region, partition.ID(), err)
+		}
+		return res, nil
+	}
+
+	res, err := stsSvc.ResolveEndpoint(region, endpoints.STSRegionalEndpointOption)
+	if err != nil {
+		return endpoints.ResolvedEndpoint{}, fmt.Errorf("error resolving endpoint for %s in partition %s. err: %v", region, partition.ID(), err)
+	}
+	return res, nil
+}

pkg/utils/helper.go

@@ -213,22 +213,22 @@ func IsNitroInstance(instanceType string) (bool, error) {
 }
 
 // GetSourceAcctAndArn constructs source acct and arn and return them for use
-func GetSourceAcctAndArn(roleARN, region, clusterName string) (string, string, error) {
+func GetSourceAcctAndArn(roleARN, region, clusterName string) (string, 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
 	// IAM format, region is always blank
 	// arn:aws:iam::account:role/role-name-with-path
 	if !arn.IsARN(roleARN) {
-		return "", "", fmt.Errorf("incorrect ARN format for role %s", roleARN)
+		return "", "", "", fmt.Errorf("incorrect ARN format for role %s", roleARN)
 	} else if region == "" {
-		return "", "", nil
+		return "", "", "", nil
 	}
 
 	parsedArn, err := arn.Parse(roleARN)
 	if err != nil {
-		return "", "", err
+		return "", "", "", err
 	}
 
 	sourceArn := fmt.Sprintf("arn:%s:eks:%s:%s:cluster/%s", parsedArn.Partition, region, parsedArn.AccountID, clusterName)
-	return parsedArn.AccountID, sourceArn, nil
+	return parsedArn.AccountID, parsedArn.Partition, sourceArn, nil
 }