Skip to content

Commit

Permalink
WIP: Update to aws-sdk-go-v2
Browse files Browse the repository at this point in the history
  • Loading branch information
micahhausler committed Aug 27, 2024
1 parent b154c1d commit af2b5b0
Show file tree
Hide file tree
Showing 10 changed files with 381 additions and 279 deletions.
16 changes: 15 additions & 1 deletion go.mod
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,13 @@ go 1.22.5

require (
github.com/aws/aws-sdk-go v1.54.6
github.com/aws/aws-sdk-go-v2 v1.30.4
github.com/aws/aws-sdk-go-v2/config v1.27.30
github.com/aws/aws-sdk-go-v2/credentials v1.17.29
github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.16.12
github.com/aws/aws-sdk-go-v2/service/ec2 v1.176.0
github.com/aws/aws-sdk-go-v2/service/sts v1.30.5
github.com/aws/smithy-go v1.20.4
github.com/fsnotify/fsnotify v1.7.0
github.com/gofrs/flock v0.8.1
github.com/google/go-cmp v0.6.0
Expand All @@ -24,12 +31,19 @@ require (
)

require (
github.com/aws/aws-sdk-go-v2/internal/configsources v1.3.16 // indirect
github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.6.16 // indirect
github.com/aws/aws-sdk-go-v2/internal/ini v1.8.1 // indirect
github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.11.4 // indirect
github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.11.18 // indirect
github.com/aws/aws-sdk-go-v2/service/sso v1.22.5 // indirect
github.com/aws/aws-sdk-go-v2/service/ssooidc v1.26.5 // indirect
github.com/beorn7/perks v1.0.1 // indirect
github.com/blang/semver/v4 v4.0.0 // indirect
github.com/cespare/xxhash/v2 v2.3.0 // indirect
github.com/chzyer/readline v0.0.0-20180603132655-2972be24d48e // indirect
github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc // indirect
github.com/emicklei/go-restful/v3 v3.11.1 // indirect
github.com/emicklei/go-restful/v3 v3.11.3 // indirect
github.com/fxamacker/cbor/v2 v2.7.0 // indirect
github.com/go-logr/logr v1.4.2 // indirect
github.com/go-openapi/jsonpointer v0.20.2 // indirect
Expand Down
32 changes: 30 additions & 2 deletions go.sum
Original file line number Diff line number Diff line change
@@ -1,5 +1,33 @@
github.com/aws/aws-sdk-go v1.54.6 h1:HEYUib3yTt8E6vxjMWM3yAq5b+qjj/6aKA62mkgux9g=
github.com/aws/aws-sdk-go v1.54.6/go.mod h1:eRwEWoyTWFMVYVQzKMNHWP5/RV4xIUGMQfXQHfHkpNU=
github.com/aws/aws-sdk-go-v2 v1.30.4 h1:frhcagrVNrzmT95RJImMHgabt99vkXGslubDaDagTk8=
github.com/aws/aws-sdk-go-v2 v1.30.4/go.mod h1:CT+ZPWXbYrci8chcARI3OmI/qgd+f6WtuLOoaIA8PR0=
github.com/aws/aws-sdk-go-v2/config v1.27.30 h1:AQF3/+rOgeJBQP3iI4vojlPib5X6eeOYoa/af7OxAYg=
github.com/aws/aws-sdk-go-v2/config v1.27.30/go.mod h1:yxqvuubha9Vw8stEgNiStO+yZpP68Wm9hLmcm+R/Qk4=
github.com/aws/aws-sdk-go-v2/credentials v1.17.29 h1:CwGsupsXIlAFYuDVHv1nnK0wnxO0wZ/g1L8DSK/xiIw=
github.com/aws/aws-sdk-go-v2/credentials v1.17.29/go.mod h1:BPJ/yXV92ZVq6G8uYvbU0gSl8q94UB63nMT5ctNO38g=
github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.16.12 h1:yjwoSyDZF8Jth+mUk5lSPJCkMC0lMy6FaCD51jm6ayE=
github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.16.12/go.mod h1:fuR57fAgMk7ot3WcNQfb6rSEn+SUffl7ri+aa8uKysI=
github.com/aws/aws-sdk-go-v2/internal/configsources v1.3.16 h1:TNyt/+X43KJ9IJJMjKfa3bNTiZbUP7DeCxfbTROESwY=
github.com/aws/aws-sdk-go-v2/internal/configsources v1.3.16/go.mod h1:2DwJF39FlNAUiX5pAc0UNeiz16lK2t7IaFcm0LFHEgc=
github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.6.16 h1:jYfy8UPmd+6kJW5YhY0L1/KftReOGxI/4NtVSTh9O/I=
github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.6.16/go.mod h1:7ZfEPZxkW42Afq4uQB8H2E2e6ebh6mXTueEpYzjCzcs=
github.com/aws/aws-sdk-go-v2/internal/ini v1.8.1 h1:VaRN3TlFdd6KxX1x3ILT5ynH6HvKgqdiXoTxAF4HQcQ=
github.com/aws/aws-sdk-go-v2/internal/ini v1.8.1/go.mod h1:FbtygfRFze9usAadmnGJNc8KsP346kEe+y2/oyhGAGc=
github.com/aws/aws-sdk-go-v2/service/ec2 v1.176.0 h1:fWhkSvaQqa5eWiRwBw10FUnk1YatAQ9We4GdGxKiCtg=
github.com/aws/aws-sdk-go-v2/service/ec2 v1.176.0/go.mod h1:ISODge3zgdwOEa4Ou6WM9PKbxJWJ15DYKnr2bfmCAIA=
github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.11.4 h1:KypMCbLPPHEmf9DgMGw51jMj77VfGPAN2Kv4cfhlfgI=
github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.11.4/go.mod h1:Vz1JQXliGcQktFTN/LN6uGppAIRoLBR2bMvIMP0gOjc=
github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.11.18 h1:tJ5RnkHCiSH0jyd6gROjlJtNwov0eGYNz8s8nFcR0jQ=
github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.11.18/go.mod h1:++NHzT+nAF7ZPrHPsA+ENvsXkOO8wEu+C6RXltAG4/c=
github.com/aws/aws-sdk-go-v2/service/sso v1.22.5 h1:zCsFCKvbj25i7p1u94imVoO447I/sFv8qq+lGJhRN0c=
github.com/aws/aws-sdk-go-v2/service/sso v1.22.5/go.mod h1:ZeDX1SnKsVlejeuz41GiajjZpRSWR7/42q/EyA/QEiM=
github.com/aws/aws-sdk-go-v2/service/ssooidc v1.26.5 h1:SKvPgvdvmiTWoi0GAJ7AsJfOz3ngVkD/ERbs5pUnHNI=
github.com/aws/aws-sdk-go-v2/service/ssooidc v1.26.5/go.mod h1:20sz31hv/WsPa3HhU3hfrIet2kxM4Pe0r20eBZ20Tac=
github.com/aws/aws-sdk-go-v2/service/sts v1.30.5 h1:OMsEmCyz2i89XwRwPouAJvhj81wINh+4UK+k/0Yo/q8=
github.com/aws/aws-sdk-go-v2/service/sts v1.30.5/go.mod h1:vmSqFK+BVIwVpDAGZB3CoCXHzurt4qBE8lf+I/kRTh0=
github.com/aws/smithy-go v1.20.4 h1:2HK1zBdPgRbjFOHlfeQZfpC4r72MOb9bZkiFwggKO+4=
github.com/aws/smithy-go v1.20.4/go.mod h1:irrKGvNn1InZwb2d7fkIRNucdfwR8R+Ts3wxYa/cJHg=
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/blang/semver/v4 v4.0.0 h1:1PFHFE6yCCTv8C1TeyNNarDzntLi7wMI5i/pzqYIsAM=
Expand All @@ -17,8 +45,8 @@ github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSs
github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc h1:U9qPSI2PIWSS1VwoXQT9A3Wy9MM3WgvqSxFWenqJduM=
github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
github.com/emicklei/go-restful/v3 v3.11.1 h1:S+9bSbua1z3FgCnV0KKOSSZ3mDthb5NyEPL5gEpCvyk=
github.com/emicklei/go-restful/v3 v3.11.1/go.mod h1:6n3XBCmQQb25CM2LCACGz8ukIrRry+4bhvbpWn3mrbc=
github.com/emicklei/go-restful/v3 v3.11.3 h1:yagOQz/38xJmcNeZJtrUcKjkHRltIaIFXKWeG1SkWGE=
github.com/emicklei/go-restful/v3 v3.11.3/go.mod h1:6n3XBCmQQb25CM2LCACGz8ukIrRry+4bhvbpWn3mrbc=
github.com/frankban/quicktest v1.14.6 h1:7Xjx+VpznH+oBnejlPUj8oUpdxnVs4f8XU8WnHkI4W8=
github.com/frankban/quicktest v1.14.6/go.mod h1:4ptaffx2x8+WTWXmUCuVU6aPUX1/Mz7zb5vbUoiM6w0=
github.com/fsnotify/fsnotify v1.7.0 h1:8JEhPFa5W2WU7YfeZzPNqzMP6Lwt7L2715Ggo0nosvA=
Expand Down
2 changes: 1 addition & 1 deletion pkg/arn/arn.go
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@ import (
"fmt"
"strings"

awsarn "github.com/aws/aws-sdk-go/aws/arn"
awsarn "github.com/aws/aws-sdk-go-v2/aws/arn"
"github.com/aws/aws-sdk-go/aws/endpoints"
)

Expand Down
143 changes: 53 additions & 90 deletions pkg/ec2provider/ec2provider.go
Original file line number Diff line number Diff line change
@@ -1,25 +1,24 @@
package ec2provider

import (
"context"
"errors"
"fmt"
"sync"
"time"

"github.com/aws/aws-sdk-go/aws"
"github.com/aws/aws-sdk-go/aws/arn"
"github.com/aws/aws-sdk-go/aws/credentials"
"github.com/aws/aws-sdk-go/aws/credentials/stscreds"
"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/ec2"
"github.com/aws/aws-sdk-go/service/ec2/ec2iface"
"github.com/aws/aws-sdk-go/service/sts"
"github.com/sirupsen/logrus"
"sigs.k8s.io/aws-iam-authenticator/pkg"
"sigs.k8s.io/aws-iam-authenticator/pkg/httputil"
"sigs.k8s.io/aws-iam-authenticator/pkg/metrics"

"github.com/aws/aws-sdk-go-v2/aws"
"github.com/aws/aws-sdk-go-v2/aws/middleware"
"github.com/aws/aws-sdk-go-v2/config"
"github.com/aws/aws-sdk-go-v2/credentials/stscreds"
"github.com/aws/aws-sdk-go-v2/service/ec2"
"github.com/aws/aws-sdk-go-v2/service/sts"
smithymiddleware "github.com/aws/smithy-go/middleware"
"github.com/sirupsen/logrus"
)

const (
Expand All @@ -36,11 +35,6 @@ const (
// Maximum time in Milliseconds to wait for a new batch call this also depends on if the instance size has
// already become 100 then it will not respect this limit
maxWaitIntervalForBatch = 200

// Headers for STS request for source ARN
headerSourceArn = "x-amz-source-arn"
// Headers for STS request for source account
headerSourceAccount = "x-amz-source-account"
)

// Get a node name from instance ID
Expand All @@ -60,13 +54,13 @@ type ec2Requests struct {
}

type ec2ProviderImpl struct {
ec2 ec2iface.EC2API
ec2 ec2.DescribeInstancesAPIClient
privateDNSCache ec2PrivateDNSCache
ec2Requests ec2Requests
instanceIdsChannel chan string
}

func New(roleARN, sourceARN, region string, qps int, burst int) EC2Provider {
func New(roleARN, sourceARN, region string, qps int, burst int) (EC2Provider, error) {
dnsCache := ec2PrivateDNSCache{
cache: make(map[string]string),
lock: sync.RWMutex{},
Expand All @@ -75,50 +69,56 @@ func New(roleARN, sourceARN, region string, qps int, burst int) EC2Provider {
set: make(map[string]bool),
lock: sync.RWMutex{},
}
cfg, err := newConfig(roleARN, sourceARN, region, qps, burst)
if err != nil {
return nil, err
}

return &ec2ProviderImpl{
ec2: ec2.New(newSession(roleARN, sourceARN, region, qps, burst)),
ec2: ec2.NewFromConfig(cfg),
privateDNSCache: dnsCache,
ec2Requests: ec2Requests,
instanceIdsChannel: make(chan string, maxChannelSize),
}
}, nil
}

// Initial credentials loaded from SDK's default credential chain, such as
// the environment, shared credentials (~/.aws/credentials), or EC2 Instance
// Role.

func newSession(roleARN, sourceARN, region string, qps int, burst int) *session.Session {
sess := session.Must(session.NewSession())
sess.Handlers.Build.PushFrontNamed(request.NamedHandler{
Name: "authenticatorUserAgent",
Fn: request.MakeAddToUserAgentHandler(
"aws-iam-authenticator", pkg.Version),
})
if aws.StringValue(sess.Config.Region) == "" {
sess.Config.Region = aws.String(region)
func newConfig(roleARN, sourceArn, region string, qps, burst int) (aws.Config, error) {
rateLimitedClient, err := httputil.NewRateLimitedClient(qps, burst)
if err != nil {
logrus.Errorf("error creating rate limited client %s", err)
return aws.Config{}, err
}
loadOpts := []func(*config.LoadOptions) error{
config.WithRegion(region),
config.WithAPIOptions(
[]func(*smithymiddleware.Stack) error{
middleware.AddUserAgentKeyValue("aws-iam-authenticator", pkg.Version),
}),
config.WithHTTPClient(rateLimitedClient),
}

if roleARN != "" {
logrus.WithFields(logrus.Fields{
"roleARN": roleARN,
}).Infof("Using assumed role for EC2 API")

rateLimitedClient, err := httputil.NewRateLimitedClient(qps, burst)

cfg, err := config.LoadDefaultConfig(context.Background(), loadOpts...)
if err != nil {
logrus.Errorf("Getting error = %s while creating rate limited client ", err)
logrus.Errorf("error loading AWS config %s", err)
return aws.Config{}, err
}

stsClient := applySTSRequestHeaders(sts.New(sess, aws.NewConfig().WithHTTPClient(rateLimitedClient).WithSTSRegionalEndpoint(endpoints.RegionalSTSEndpoint)), sourceARN)
ap := &stscreds.AssumeRoleProvider{
Client: stsClient,
RoleARN: roleARN,
Duration: time.Duration(60) * time.Minute,
stsOpts := []func(*sts.Options){}
if sourceArn != "" {
stsOpts = append(stsOpts, WithSourceHeaders(sourceArn))
}

sess.Config.Credentials = credentials.NewCredentials(ap)
stsCli := sts.NewFromConfig(cfg, stsOpts...)
creds := stscreds.NewAssumeRoleProvider(stsCli, roleARN,
func(o *stscreds.AssumeRoleOptions) {
o.Duration = time.Duration(60) * time.Minute
})
loadOpts = append(loadOpts, config.WithCredentialsProvider(creds))
}
return sess
return config.LoadDefaultConfig(context.Background(), loadOpts...)
}

func (p *ec2ProviderImpl) setPrivateDNSNameCache(id string, privateDNSName string) {
Expand Down Expand Up @@ -197,17 +197,17 @@ func (p *ec2ProviderImpl) GetPrivateDNSName(id string) (string, error) {
logrus.Infof("Calling ec2:DescribeInstances for the InstanceId = %s ", id)
metrics.Get().EC2DescribeInstanceCallCount.Inc()
// Look up instance from EC2 API
output, err := p.ec2.DescribeInstances(&ec2.DescribeInstancesInput{
InstanceIds: aws.StringSlice([]string{id}),
output, err := p.ec2.DescribeInstances(context.Background(), &ec2.DescribeInstancesInput{
InstanceIds: []string{id},
})
if err != nil {
p.unsetRequestInFlightForInstanceId(id)
return "", fmt.Errorf("failed querying private DNS from EC2 API for node %s: %s ", id, err.Error())
}
for _, reservation := range output.Reservations {
for _, instance := range reservation.Instances {
if aws.StringValue(instance.InstanceId) == id {
privateDNSName = aws.StringValue(instance.PrivateDnsName)
if aws.ToString(instance.InstanceId) == id {
privateDNSName = aws.ToString(instance.PrivateDnsName)
p.setPrivateDNSNameCache(id, privateDNSName)
p.unsetRequestInFlightForInstanceId(id)
}
Expand Down Expand Up @@ -258,8 +258,8 @@ func (p *ec2ProviderImpl) getPrivateDnsAndPublishToCache(instanceIdList []string
// Look up instance from EC2 API
logrus.Infof("Making Batch Query to DescribeInstances for %v instances ", len(instanceIdList))
metrics.Get().EC2DescribeInstanceCallCount.Inc()
output, err := p.ec2.DescribeInstances(&ec2.DescribeInstancesInput{
InstanceIds: aws.StringSlice(instanceIdList),
output, err := p.ec2.DescribeInstances(context.Background(), &ec2.DescribeInstancesInput{
InstanceIds: instanceIdList,
})
if err != nil {
logrus.Errorf("Batch call failed querying private DNS from EC2 API for nodes [%s] : with error = []%s ", instanceIdList, err.Error())
Expand All @@ -272,8 +272,8 @@ func (p *ec2ProviderImpl) getPrivateDnsAndPublishToCache(instanceIdList []string
// Adding the result to privateDNSChache as well as removing from the requestQueueMap.
for _, reservation := range output.Reservations {
for _, instance := range reservation.Instances {
id := aws.StringValue(instance.InstanceId)
privateDNSName := aws.StringValue(instance.PrivateDnsName)
id := aws.ToString(instance.InstanceId)
privateDNSName := aws.ToString(instance.PrivateDnsName)
p.setPrivateDNSNameCache(id, privateDNSName)
}
}
Expand All @@ -284,40 +284,3 @@ func (p *ec2ProviderImpl) getPrivateDnsAndPublishToCache(instanceIdList []string
p.unsetRequestInFlightForInstanceId(id)
}
}

func applySTSRequestHeaders(stsClient *sts.STS, sourceARN string) *sts.STS {
// parse both source account and source arn from the sourceARN, and add them as headers to the STS client
if sourceARN != "" {
sourceAcct, err := getSourceAccount(sourceARN)
if err != nil {
panic(fmt.Sprintf("%s is not a valid arn, err: %v", sourceARN, err))
}
reqHeaders := map[string]string{
headerSourceAccount: sourceAcct,
headerSourceArn: sourceARN,
}
stsClient.Handlers.Sign.PushFront(func(s *request.Request) {
s.ApplyOptions(request.WithSetRequestHeaders(reqHeaders))
})
logrus.Infof("configuring STS client with extra headers, %v", reqHeaders)
}
return stsClient
}

// getSourceAccount constructs source acct and return them for use
func getSourceAccount(roleARN 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)
}

parsedArn, err := arn.Parse(roleARN)
if err != nil {
return "", err
}

return parsedArn.AccountID, nil
}
Loading

0 comments on commit af2b5b0

Please sign in to comment.