feat: assume role support

Author: rxnew

config.go

@@ -0,0 +1,34 @@
+package main
+
+import (
+	"context"
+
+	"github.com/aws/aws-sdk-go-v2/aws"
+	"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/sts"
+)
+
+func loadConfig(ctx context.Context) (*aws.Config, error) {
+	cfg, err := config.LoadDefaultConfig(ctx)
+	if err != nil {
+		return nil, err
+	}
+
+	if e, ok := getEnvConfig(&cfg); ok && e.RoleARN != "" {
+		cfg.Credentials = stscreds.NewAssumeRoleProvider(sts.NewFromConfig(cfg), e.RoleARN, func(o *stscreds.AssumeRoleOptions) {
+			o.RoleSessionName = e.RoleSessionName
+		})
+	}
+
+	return &cfg, nil
+}
+
+func getEnvConfig(cfg *aws.Config) (*config.EnvConfig, bool) {
+	for _, s := range cfg.ConfigSources {
+		if c, ok := s.(config.EnvConfig); ok {
+			return &c, true
+		}
+	}
+	return nil, false
+}

go.mod

@@ -5,20 +5,20 @@ go 1.21
 require (
 	github.com/aws/aws-sdk-go-v2 v1.22.0
 	github.com/aws/aws-sdk-go-v2/config v1.20.0
+	github.com/aws/aws-sdk-go-v2/credentials v1.14.0
+	github.com/aws/aws-sdk-go-v2/service/sts v1.24.0
 	github.com/hashicorp/go-retryablehttp v0.7.4
 	github.com/spf13/cobra v1.7.0
 )
 
 require (
-	github.com/aws/aws-sdk-go-v2/credentials v1.14.0 // indirect
 	github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.14.0 // indirect
 	github.com/aws/aws-sdk-go-v2/internal/configsources v1.2.0 // indirect
 	github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.5.0 // indirect
 	github.com/aws/aws-sdk-go-v2/internal/ini v1.4.0 // indirect
 	github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.10.0 // indirect
 	github.com/aws/aws-sdk-go-v2/service/sso v1.16.0 // indirect
 	github.com/aws/aws-sdk-go-v2/service/ssooidc v1.18.0 // indirect
-	github.com/aws/aws-sdk-go-v2/service/sts v1.24.0 // indirect
 	github.com/aws/smithy-go v1.16.0 // indirect
 	github.com/hashicorp/go-cleanhttp v0.5.2 // indirect
 	github.com/inconshreveable/mousetrap v1.1.0 // indirect

main.go

@@ -12,7 +12,6 @@ import (
 	"strings"
 	"time"
 
-	"github.com/aws/aws-sdk-go-v2/config"
 	"github.com/hashicorp/go-retryablehttp"
 	"github.com/spf13/cobra"
 
@@ -52,7 +51,7 @@ func run(cmd *cobra.Command, args []string) {
 	ctx, stop := signal.NotifyContext(cmd.Context(), os.Interrupt)
 	defer stop()
 
-	cfg, err := config.LoadDefaultConfig(ctx)
+	cfg, err := loadConfig(ctx)
 	if err != nil {
 		log.Fatalf("failed to load configuration: %v", err)
 	}
@@ -62,7 +61,7 @@ func run(cmd *cobra.Command, args []string) {
 		log.Fatalf("failed to create HTTP request: %v", err)
 	}
 
-	resp, err := sigv4.NewHTTPClient(&cfg, opt.Service, newRetryableHTTPClient()).Do(req)
+	resp, err := sigv4.NewHTTPClient(cfg, opt.Service, newRetryableHTTPClient()).Do(req)
 	if err != nil {
 		log.Fatalf("failed to HTTP request: %v", err)
 	}