Migrate STS client to aws-sdk-go-v2

kubernetes · Mar 31, 2024 · 856b819 · 856b819
1 parent 2a1faaa
commit 856b819
Show file tree

Hide file tree

Showing 4 changed files with 15 additions and 20 deletions.
diff --git a/go.mod b/go.mod
@@ -31,6 +31,7 @@ require (
 	github.com/aws/aws-sdk-go-v2/service/s3 v1.53.0
 	github.com/aws/aws-sdk-go-v2/service/sqs v1.31.4
 	github.com/aws/aws-sdk-go-v2/service/ssm v1.49.5
+	github.com/aws/aws-sdk-go-v2/service/sts v1.28.5
 	github.com/aws/smithy-go v1.20.2
 	github.com/blang/semver/v4 v4.0.0
 	github.com/cert-manager/cert-manager v1.14.4
@@ -120,7 +121,6 @@ require (
 	github.com/aws/aws-sdk-go-v2/service/internal/s3shared v1.17.4 // indirect
 	github.com/aws/aws-sdk-go-v2/service/sso v1.20.3 // indirect
 	github.com/aws/aws-sdk-go-v2/service/ssooidc v1.23.3 // indirect
-	github.com/aws/aws-sdk-go-v2/service/sts v1.28.5 // indirect
 	github.com/beorn7/perks v1.0.1 // indirect
 	github.com/cespare/xxhash/v2 v2.2.0 // indirect
 	github.com/chai2010/gettext-go v1.0.2 // indirect

diff --git a/upup/pkg/fi/cloudup/apply_cluster.go b/upup/pkg/fi/cloudup/apply_cluster.go
@@ -427,7 +427,7 @@ func (c *ApplyClusterCmd) Run(ctx context.Context) error {
 		{
 			awsCloud := cloud.(awsup.AWSCloud)
 
-			accountID, partition, err := awsCloud.AccountInfo()
+			accountID, partition, err := awsCloud.AccountInfo(ctx)
 			if err != nil {
 				return err
 			}

diff --git a/upup/pkg/fi/cloudup/awsup/aws_cloud.go b/upup/pkg/fi/cloudup/awsup/aws_cloud.go
@@ -39,6 +39,7 @@ import (
 	elbv2 "github.com/aws/aws-sdk-go-v2/service/elasticloadbalancingv2"
 	elbv2types "github.com/aws/aws-sdk-go-v2/service/elasticloadbalancingv2/types"
 	"github.com/aws/aws-sdk-go-v2/service/iam"
+	"github.com/aws/aws-sdk-go-v2/service/sts"
 	"github.com/aws/aws-sdk-go/aws"
 	"github.com/aws/aws-sdk-go/aws/awserr"
 	"github.com/aws/aws-sdk-go/aws/credentials/stscreds"
@@ -50,7 +51,6 @@ import (
 	"github.com/aws/aws-sdk-go/service/ec2/ec2iface"
 	"github.com/aws/aws-sdk-go/service/route53"
 	"github.com/aws/aws-sdk-go/service/route53/route53iface"
-	"github.com/aws/aws-sdk-go/service/sts"
 	"k8s.io/klog/v2"
 
 	v1 "k8s.io/api/core/v1"
@@ -191,7 +191,7 @@ type AWSCloud interface {
 	DescribeInstanceType(instanceType string) (*ec2.InstanceTypeInfo, error)
 
 	// AccountInfo returns the AWS account ID and AWS partition that we are deploying into
-	AccountInfo() (string, string, error)
+	AccountInfo(ctx context.Context) (string, string, error)
 }
 
 type awsCloudImplementation struct {
@@ -202,7 +202,7 @@ type awsCloudImplementation struct {
 	autoscaling *autoscaling.AutoScaling
 	route53     *route53.Route53
 	spotinst    spotinst.Cloud
-	sts         *sts.STS
+	sts         *sts.Client
 	sqs         *sqs.Client
 	eventbridge *eventbridge.Client
 	ssm         *ssm.Client
@@ -332,20 +332,15 @@ func NewAWSCloud(region string, tags map[string]string) (AWSCloud, error) {
 		c.ec2.Handlers.Send.PushFront(requestLogger)
 		c.addHandlers(region, &c.ec2.Handlers)
 
+		cfgV2, err := awsconfig.LoadDefaultConfig(ctx, loadOptions...)
+		if err != nil {
+			return c, fmt.Errorf("failed to load default aws config: %w", err)
+		}
+
 		c.iam = iam.NewFromConfig(cfgV2)
 		c.elb = elb.NewFromConfig(cfgV2)
 		c.elbv2 = elbv2.NewFromConfig(cfgV2)
-
-		sess, err = session.NewSessionWithOptions(session.Options{
-			Config:            *config,
-			SharedConfigState: session.SharedConfigEnable,
-		})
-		if err != nil {
-			return c, err
-		}
-		c.sts = sts.New(sess, config)
-		c.sts.Handlers.Send.PushFront(requestLogger)
-		c.addHandlers(region, &c.sts.Handlers)
+		c.sts = sts.NewFromConfig(cfgV2)
 
 		sess, err = session.NewSessionWithOptions(session.Options{
 			Config:            *config,
@@ -2407,17 +2402,17 @@ func describeInstanceType(c AWSCloud, instanceType string) (*ec2.InstanceTypeInf
 }
 
 // AccountInfo returns the AWS account ID and AWS partition that we are deploying into
-func (c *awsCloudImplementation) AccountInfo() (string, string, error) {
+func (c *awsCloudImplementation) AccountInfo(ctx context.Context) (string, string, error) {
 	request := &sts.GetCallerIdentityInput{}
 
-	response, err := c.sts.GetCallerIdentity(request)
+	response, err := c.sts.GetCallerIdentity(ctx, request)
 	if err != nil {
 		return "", "", fmt.Errorf("error getting AWS account ID: %v", err)
 	}
 
 	arn, err := arn.Parse(aws.StringValue(response.Arn))
 	if err != nil {
-		return "", "", fmt.Errorf("Failed to parse GetCallerIdentity ARN")
+		return "", "", fmt.Errorf("failed to parse GetCallerIdentity ARN: %w", err)
 	}
 
 	if arn.AccountID == "" {

diff --git a/upup/pkg/fi/cloudup/awsup/mock_aws_cloud.go b/upup/pkg/fi/cloudup/awsup/mock_aws_cloud.go
@@ -387,7 +387,7 @@ func (c *MockAWSCloud) DescribeInstanceType(instanceType string) (*ec2.InstanceT
 }
 
 // AccountInfo returns the AWS account ID and AWS partition that we are deploying into
-func (c *MockAWSCloud) AccountInfo() (string, string, error) {
+func (c *MockAWSCloud) AccountInfo(ctx context.Context) (string, string, error) {
 	return "123456789012", "aws-test", nil
 }