Migrate ecs-init to aws-sdk-go-v2 (#4372)

ecs-init/cache/cache.go

@@ -17,6 +17,7 @@ package cache
 
 import (
 	"bufio"
+	"context"
 	"crypto/md5"
 	"fmt"
 	"io"
@@ -25,8 +26,8 @@ import (
 
 	"github.com/aws/amazon-ecs-agent/ecs-init/config"
 
-	"github.com/aws/aws-sdk-go/aws/ec2metadata"
-	"github.com/aws/aws-sdk-go/aws/session"
+	awsconfig "github.com/aws/aws-sdk-go-v2/config"
+	"github.com/aws/aws-sdk-go-v2/feature/ec2/imds"
 	log "github.com/cihub/seelog"
 	"github.com/pkg/errors"
 )
@@ -76,7 +77,7 @@ func NewDownloader() (*Downloader, error) {
 	if config.RunningInExternal() {
 		downloader.metadata = &blackholeInstanceMetadata{}
 	} else {
-		sessionInstance, err := session.NewSession()
+		cfg, err := awsconfig.LoadDefaultConfig(context.TODO())
 		if err != nil {
 			// metadata client is only used for retrieving the user's region.
 			// If it cannot be initialized, the region field is populated with the default value to prevent future
@@ -85,7 +86,7 @@ func NewDownloader() (*Downloader, error) {
 				err, config.DefaultRegionName)
 			downloader.region = config.DefaultRegionName
 		} else {
-			downloader.metadata = ec2metadata.New(sessionInstance)
+			downloader.metadata = imds.NewFromConfig(cfg)
 		}
 	}
 
@@ -181,13 +182,14 @@ func (d *Downloader) getRegion() string {
 		return d.region
 	}
 
-	region, err := d.metadata.Region()
+	output, err := d.metadata.GetRegion(context.TODO(), &imds.GetRegionInput{})
 	if err != nil {
 		log.Warnf("Could not retrieve the region from EC2 Instance Metadata. Error: %s", err.Error())
-		region = defaultRegion
+		d.region = defaultRegion
+		return d.region
 	}
-	d.region = region
 
+	d.region = output.Region
 	return d.region
 }
 

ecs-init/cache/dependencies.go

@@ -21,22 +21,23 @@ package cache
 //go:generate mockgen.sh cache $GOFILE
 
 import (
+	"context"
 	"io"
 	"os"
 	"path/filepath"
 
-	"github.com/aws/aws-sdk-go/aws"
-	"github.com/aws/aws-sdk-go/aws/credentials"
-	"github.com/aws/aws-sdk-go/aws/session"
-	"github.com/aws/aws-sdk-go/service/s3"
-	"github.com/aws/aws-sdk-go/service/s3/s3manager"
+	"github.com/aws/aws-sdk-go-v2/aws"
+	"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/feature/s3/manager"
+	"github.com/aws/aws-sdk-go-v2/service/s3"
 	log "github.com/cihub/seelog"
 	"github.com/pkg/errors"
 )
 
 // s3API captures the only method used from the s3 package
 type s3API interface {
-	Download(w io.WriterAt, input *s3.GetObjectInput, options ...func(*s3manager.Downloader)) (n int64, err error)
+	Download(ctx context.Context, w io.WriterAt, input *s3.GetObjectInput, options ...func(*manager.Downloader)) (int64, error)
 }
 
 // s3BucketDownloader wraps a bucket together with a downloader that can download from it
@@ -47,16 +48,17 @@ type s3BucketDownloader struct {
 }
 
 func newS3BucketDownloader(region, bucketName string) (*s3BucketDownloader, error) {
-	session, err := session.NewSession(&aws.Config{
-		Credentials: credentials.AnonymousCredentials,
-		Region:      aws.String(region),
-	})
+	cfg, err := config.LoadDefaultConfig(
+		context.TODO(),
+		config.WithCredentialsProvider((aws.AnonymousCredentials{})),
+		config.WithRegion((region)),
+	)
 	if err != nil {
 		return nil, errors.Wrapf(err, "failed to initialize downloader in region %s", region)
 	}
 
 	s3BucketDownloader := &s3BucketDownloader{
-		client: s3manager.NewDownloader(session),
+		client: manager.NewDownloader(s3.NewFromConfig(cfg)),
 		bucket: bucketName,
 		region: region,
 	}
@@ -77,7 +79,7 @@ func (bd *s3BucketDownloader) download(fileName, cacheDir string, fs fileSystem)
 		}
 	}()
 
-	_, err = bd.client.Download(file, &s3.GetObjectInput{
+	_, err = bd.client.Download(context.TODO(), file, &s3.GetObjectInput{
 		Bucket: aws.String(bd.bucket),
 		Key:    aws.String(fileName),
 	})
@@ -137,14 +139,14 @@ type fileSizeInfo interface {
 }
 
 type instanceMetadata interface {
-	Region() (string, error)
+	GetRegion(ctx context.Context, input *imds.GetRegionInput, opts ...func(*imds.Options)) (*imds.GetRegionOutput, error)
 }
 
 type blackholeInstanceMetadata struct {
 }
 
-func (b *blackholeInstanceMetadata) Region() (string, error) {
-	return "", errors.New("blackholed")
+func (b *blackholeInstanceMetadata) GetRegion(ctx context.Context, input *imds.GetRegionInput, opts ...func(*imds.Options)) (*imds.GetRegionOutput, error) {
+	return nil, errors.New("blackholed")
 }
 
 // standardFS delegates to the package-level functions

ecs-init/config/awsrulesfn/partition.go

@@ -0,0 +1,76 @@
+package awsrulesfn
+
+import "regexp"
+
+// Partition provides the metadata describing an AWS partition.
+type Partition struct {
+	ID            string                     `json:"id"`
+	Regions       map[string]RegionOverrides `json:"regions"`
+	RegionRegex   string                     `json:"regionRegex"`
+	DefaultConfig PartitionConfig            `json:"outputs"`
+}
+
+// PartitionConfig provides the endpoint metadata for an AWS region or partition.
+type PartitionConfig struct {
+	Name                 string `json:"name"`
+	DnsSuffix            string `json:"dnsSuffix"`
+	DualStackDnsSuffix   string `json:"dualStackDnsSuffix"`
+	SupportsFIPS         bool   `json:"supportsFIPS"`
+	SupportsDualStack    bool   `json:"supportsDualStack"`
+	ImplicitGlobalRegion string `json:"implicitGlobalRegion"`
+}
+
+type RegionOverrides struct {
+	Name               *string `json:"name"`
+	DnsSuffix          *string `json:"dnsSuffix"`
+	DualStackDnsSuffix *string `json:"dualStackDnsSuffix"`
+	SupportsFIPS       *bool   `json:"supportsFIPS"`
+	SupportsDualStack  *bool   `json:"supportsDualStack"`
+}
+
+const defaultPartition = "aws"
+
+func getPartition(partitions []Partition, region string) *PartitionConfig {
+	for _, partition := range partitions {
+		if v, ok := partition.Regions[region]; ok {
+			p := mergeOverrides(partition.DefaultConfig, v)
+			return &p
+		}
+	}
+
+	for _, partition := range partitions {
+		regionRegex := regexp.MustCompile(partition.RegionRegex)
+		if regionRegex.MatchString(region) {
+			v := partition.DefaultConfig
+			return &v
+		}
+	}
+
+	for _, partition := range partitions {
+		if partition.ID == defaultPartition {
+			v := partition.DefaultConfig
+			return &v
+		}
+	}
+
+	return nil
+}
+
+func mergeOverrides(into PartitionConfig, from RegionOverrides) PartitionConfig {
+	if from.Name != nil {
+		into.Name = *from.Name
+	}
+	if from.DnsSuffix != nil {
+		into.DnsSuffix = *from.DnsSuffix
+	}
+	if from.DualStackDnsSuffix != nil {
+		into.DualStackDnsSuffix = *from.DualStackDnsSuffix
+	}
+	if from.SupportsFIPS != nil {
+		into.SupportsFIPS = *from.SupportsFIPS
+	}
+	if from.SupportsDualStack != nil {
+		into.SupportsDualStack = *from.SupportsDualStack
+	}
+	return into
+}

ecs-init/config/awsrulesfn/partitions.gen.go

@@ -0,0 +1,26 @@
+package awsrulesfn
+
+import "regexp"
+
+// aws-sdk-go-v2 does not export partition metadata, so copy the files from vendor to make it accessible.
+
+//go:generate cp ../../vendor/github.com/aws/aws-sdk-go-v2/internal/endpoints/awsrulesfn/partition.go ../../vendor/github.com/aws/aws-sdk-go-v2/internal/endpoints/awsrulesfn/partitions.go .
+
+// GetPartitionForRegion returns an AWS partition for the region provided.
+// Unlike GetPartition, this function
+// 1. returns a Partition instead of a PartitionConfig
+// 2. returns nil instead of falling back to the default partition (aws) if no match is found
+func GetPartitionForRegion(region string) *Partition {
+	for _, partition := range partitions {
+		if _, ok := partition.Regions[region]; ok {
+			return &partition
+		}
+
+		regionRegex := regexp.MustCompile(partition.RegionRegex)
+		if regionRegex.MatchString(region) {
+			return &partition
+		}
+	}
+
+	return nil
+}