From d265ea71645898b87ca98bab1a88aa2f64b2e20a Mon Sep 17 00:00:00 2001
From: Jo Rhett <jrhett@tubi.tv>
Date: Wed, 9 Mar 2022 17:04:21 -0800
Subject: [PATCH] Add roleprovider for role assumption by ec2 instance

If WebIdentityTokenFile is empty, try to assume the role using pre-existing creds
  https://github.com/aws/aws-sdk-go/blob/main/aws/credentials/stscreds/assume_role_provider.go#L17-L24
---
 awsutil/generate_credentials.go | 56 +++++++++++++++++++++++----------
 1 file changed, 39 insertions(+), 17 deletions(-)

diff --git a/awsutil/generate_credentials.go b/awsutil/generate_credentials.go
index ac0fb00..23c9dd4 100644
--- a/awsutil/generate_credentials.go
+++ b/awsutil/generate_credentials.go
@@ -50,10 +50,10 @@ type CredentialsConfig struct {
 	// The profile for the shared credentials provider, if being used
 	Profile string
 
-	// The role ARN to use if using the web identity token provider
+	// The role ARN to assume
 	RoleARN string
 
-	// The role session name to use if using the web identity token provider
+	// The role session name to use when assuming another role
 	RoleSessionName string
 
 	// The web identity token file to use if using the web identity token provider
@@ -160,24 +160,46 @@ func (c *CredentialsConfig) GenerateCredentialChain(opt ...Option) (*credentials
 	if roleSessionName == "" {
 		roleSessionName = os.Getenv("AWS_ROLE_SESSION_NAME")
 	}
-	if roleARN != "" && tokenPath != "" {
-		// this session is only created to create the WebIdentityRoleProvider, as the env variables are already there
-		// this automatically assumes the role, but the provider needs to be added to the chain
-		c.log(hclog.Debug, "adding web identity provider", "roleARN", roleARN)
-		sess, err := session.NewSession()
-		if err != nil {
-			return nil, errors.Wrap(err, "error creating a new session to create a WebIdentityRoleProvider")
-		}
-		webIdentityProvider := stscreds.NewWebIdentityRoleProvider(sts.New(sess), roleARN, roleSessionName, tokenPath)
+	if roleARN != "" {
+		if  tokenPath != "" {
+			// this session is only created to create the WebIdentityRoleProvider, as the env variables are already there
+			// this automatically assumes the role, but the provider needs to be added to the chain
+			c.log(hclog.Debug, "adding web identity provider", "roleARN", roleARN)
+			sess, err := session.NewSession()
+			if err != nil {
+				return nil, errors.Wrap(err, "error creating a new session to create a WebIdentityRoleProvider")
+			}
+			webIdentityProvider := stscreds.NewWebIdentityRoleProvider(sts.New(sess), roleARN, roleSessionName, tokenPath)
+
+			// Check if the webIdentityProvider can successfully retrieve
+			// credentials (via sts:AssumeRole), and warn if there's a problem.
+			if _, err := webIdentityProvider.Retrieve(); err != nil {
+				c.log(hclog.Warn, "error assuming role", "roleARN", roleARN, "tokenPath", tokenPath, "sessionName", roleSessionName, "err", err)
+			}
 
-		// Check if the webIdentityProvider can successfully retrieve
-		// credentials (via sts:AssumeRole), and warn if there's a problem.
-		if _, err := webIdentityProvider.Retrieve(); err != nil {
-			c.log(hclog.Warn, "error assuming role", "roleARN", roleARN, "tokenPath", tokenPath, "sessionName", roleSessionName, "err", err)
+			// Add the web identity role credential provider
+			providers = append(providers, webIdentityProvider)
 		}
+		else {
+			// this session is only created to create the instanceRoleProvider, as the env variables are already there
+			// this automatically assumes the role, but the provider needs to be added to the chain
+			c.log(hclog.Debug, "adding ec2-instance role provider", "roleARN", roleARN)
+			sess, err := session.NewSession()
+			if err != nil {
+				return nil, errors.Wrap(err, "error creating a new session for ec2 instance role credentials")
+			}
+			instanceRoleProvider := stscreds.NewCredentials(sess, roleARN, func(p *stscreds.AssumeRoleProvider) {
+				p.RoleSessionName = roleSessionName
+			})
+
+			// Check if the provider can successfully retrieve credentials (via sts:AssumeRole), and warn if there's a problem.
+			if _, err := instanceRoleProvider.Retrieve(); err != nil {
+				c.log(hclog.Warn, "error assuming role", "roleARN", roleARN, "sessionName", roleSessionName, "err", err)
+			}
 
-		// Add the web identity role credential provider
-		providers = append(providers, webIdentityProvider)
+			// Add the instance role credential provider
+			providers = append(providers, instanceRoleProvider)
+		}
 	}
 
 	if opts.withEnvironmentCredentials {