From 9a46073843ce6ab64cdddb00e9614c583dd3484d Mon Sep 17 00:00:00 2001
From: Jesse Rittner <rittneje@gmail.com>
Date: Thu, 9 Dec 2021 16:14:15 -0500
Subject: [PATCH] add options for customizing credential providers

---
 CHANGELOG_PENDING.md            |  2 ++
 aws/session/credentials.go      | 31 +++++++++++++-------
 aws/session/credentials_test.go | 51 +++++++++++++++++++++++++++++++++
 aws/session/session.go          |  5 ++++
 4 files changed, 79 insertions(+), 10 deletions(-)

diff --git a/CHANGELOG_PENDING.md b/CHANGELOG_PENDING.md
index 8a1927a39ca..d229672aebe 100644
--- a/CHANGELOG_PENDING.md
+++ b/CHANGELOG_PENDING.md
@@ -2,4 +2,6 @@
 
 ### SDK Enhancements
 
+* `aws/session`: Add options for customizing the construction of credential providers. Currently only supported for `stscreds.WebIdentityRoleProvider`.
+
 ### SDK Bugs
diff --git a/aws/session/credentials.go b/aws/session/credentials.go
index 3efdac29ff4..972b1bea393 100644
--- a/aws/session/credentials.go
+++ b/aws/session/credentials.go
@@ -14,8 +14,17 @@ import (
 	"github.com/aws/aws-sdk-go/aws/defaults"
 	"github.com/aws/aws-sdk-go/aws/request"
 	"github.com/aws/aws-sdk-go/internal/shareddefaults"
+	"github.com/aws/aws-sdk-go/service/sts"
 )
 
+// CredentialsProviderOptions specifies additional options for configuring
+// credentials providers.
+type CredentialsProviderOptions struct {
+	// WebIdentityRoleProviderOptions configures a WebIdentityRoleProvider,
+	// such as setting its ExpiryWindow.
+	WebIdentityRoleProviderOptions func(*stscreds.WebIdentityRoleProvider)
+}
+
 func resolveCredentials(cfg *aws.Config,
 	envCfg envConfig, sharedCfg sharedConfig,
 	handlers request.Handlers,
@@ -40,6 +49,7 @@ func resolveCredentials(cfg *aws.Config,
 			envCfg.WebIdentityTokenFilePath,
 			envCfg.RoleARN,
 			envCfg.RoleSessionName,
+			sessOpts.CredentialsProviderOptions,
 		)
 
 	default:
@@ -59,6 +69,7 @@ var WebIdentityEmptyTokenFilePathErr = awserr.New(stscreds.ErrCodeWebIdentity, "
 func assumeWebIdentity(cfg *aws.Config, handlers request.Handlers,
 	filepath string,
 	roleARN, sessionName string,
+	options *CredentialsProviderOptions,
 ) (*credentials.Credentials, error) {
 
 	if len(filepath) == 0 {
@@ -69,17 +80,16 @@ func assumeWebIdentity(cfg *aws.Config, handlers request.Handlers,
 		return nil, WebIdentityEmptyRoleARNErr
 	}
 
-	creds := stscreds.NewWebIdentityCredentials(
-		&Session{
-			Config:   cfg,
-			Handlers: handlers.Copy(),
-		},
-		roleARN,
-		sessionName,
-		filepath,
-	)
+	svc := sts.New(&Session{
+		Config:   cfg,
+		Handlers: handlers.Copy(),
+	})
 
-	return creds, nil
+	p := stscreds.NewWebIdentityRoleProvider(svc, roleARN, sessionName, filepath)
+	if options != nil && options.WebIdentityRoleProviderOptions != nil {
+		options.WebIdentityRoleProviderOptions(p)
+	}
+	return credentials.NewCredentials(p), nil
 }
 
 func resolveCredsFromProfile(cfg *aws.Config,
@@ -114,6 +124,7 @@ func resolveCredsFromProfile(cfg *aws.Config,
 			sharedCfg.WebIdentityTokenFile,
 			sharedCfg.RoleARN,
 			sharedCfg.RoleSessionName,
+			sessOpts.CredentialsProviderOptions,
 		)
 
 	case sharedCfg.hasSSOConfiguration():
diff --git a/aws/session/credentials_test.go b/aws/session/credentials_test.go
index 914ee8a986e..f1a8f332f36 100644
--- a/aws/session/credentials_test.go
+++ b/aws/session/credentials_test.go
@@ -19,6 +19,7 @@ import (
 
 	"github.com/aws/aws-sdk-go/aws"
 	"github.com/aws/aws-sdk-go/aws/credentials"
+	"github.com/aws/aws-sdk-go/aws/credentials/stscreds"
 	"github.com/aws/aws-sdk-go/aws/defaults"
 	"github.com/aws/aws-sdk-go/aws/endpoints"
 	"github.com/aws/aws-sdk-go/aws/request"
@@ -814,6 +815,56 @@ func TestSessionAssumeRole_WithMFA_ExtendedDuration(t *testing.T) {
 	}
 }
 
+func TestSessionAssumeRoleWithWebIdentity_Options(t *testing.T) {
+	restoreEnvFn := initSessionTestEnv()
+	defer restoreEnvFn()
+
+	os.Setenv("AWS_REGION", "us-east-1")
+	os.Setenv("AWS_ROLE_ARN", "web_identity_role_arn")
+	os.Setenv("AWS_WEB_IDENTITY_TOKEN_FILE", "./testdata/wit.txt")
+
+	endpointResolver, teardown := setupCredentialsEndpoints(t)
+	defer teardown()
+
+	optionsCalled := false
+
+	sess, err := NewSessionWithOptions(Options{
+		Config: aws.Config{
+			EndpointResolver: endpointResolver,
+		},
+		CredentialsProviderOptions: &CredentialsProviderOptions{
+			WebIdentityRoleProviderOptions: func(*stscreds.WebIdentityRoleProvider) {
+				optionsCalled = true
+			},
+		},
+	})
+	if err != nil {
+		t.Fatalf("expect no error, got %v", err)
+	}
+
+	if !optionsCalled {
+		t.Errorf("expect options func to be called")
+	}
+
+	creds, err := sess.Config.Credentials.Get()
+	if err != nil {
+		t.Fatalf("expect no error, got %v", err)
+	}
+
+	if e, a := "WEB_IDENTITY_AKID", creds.AccessKeyID; e != a {
+		t.Errorf("expect %v, got %v", e, a)
+	}
+	if e, a := "WEB_IDENTITY_SECRET", creds.SecretAccessKey; e != a {
+		t.Errorf("expect %v, got %v", e, a)
+	}
+	if e, a := "WEB_IDENTITY_SESSION_TOKEN", creds.SessionToken; e != a {
+		t.Errorf("expect %v, got %v", e, a)
+	}
+	if e, a := stscreds.WebIdentityProviderName, creds.ProviderName; e != a {
+		t.Errorf("expect %v,got %v", e, a)
+	}
+}
+
 func ssoTestSetup() (func(), error) {
 	dir, err := ioutil.TempDir("", "sso-test")
 	if err != nil {
diff --git a/aws/session/session.go b/aws/session/session.go
index ebace4bb79d..4293dbe10bd 100644
--- a/aws/session/session.go
+++ b/aws/session/session.go
@@ -304,6 +304,11 @@ type Options struct {
 	//
 	// AWS_EC2_METADATA_SERVICE_ENDPOINT_MODE=IPv6
 	EC2IMDSEndpointMode endpoints.EC2IMDSEndpointModeState
+
+	// Specifies options for creating credential providers.
+	// These are only used if the aws.Config does not already
+	// include credentials.
+	CredentialsProviderOptions *CredentialsProviderOptions
 }
 
 // NewSessionWithOptions returns a new Session created from SDK defaults, config files,