From 027748f7932c0c630b787e7579edd0093e5ef7c6 Mon Sep 17 00:00:00 2001
From: Jesse Haka <haka.jesse@gmail.com>
Date: Wed, 15 Mar 2023 19:20:42 +0200
Subject: [PATCH] make it possible to use assume role

---
 upup/pkg/fi/cloudup/awsup/aws_cloud.go | 24 +++++++++++++++++++-----
 1 file changed, 19 insertions(+), 5 deletions(-)

diff --git a/upup/pkg/fi/cloudup/awsup/aws_cloud.go b/upup/pkg/fi/cloudup/awsup/aws_cloud.go
index 1616eab583144..779b75631e308 100644
--- a/upup/pkg/fi/cloudup/awsup/aws_cloud.go
+++ b/upup/pkg/fi/cloudup/awsup/aws_cloud.go
@@ -19,6 +19,7 @@ package awsup
 import (
 	"context"
 	"fmt"
+	"os"
 	"strconv"
 	"strings"
 	"sync"
@@ -33,6 +34,7 @@ import (
 	"github.com/aws/aws-sdk-go/aws"
 	"github.com/aws/aws-sdk-go/aws/arn"
 	"github.com/aws/aws-sdk-go/aws/awserr"
+	"github.com/aws/aws-sdk-go/aws/credentials/stscreds"
 	"github.com/aws/aws-sdk-go/aws/request"
 	"github.com/aws/aws-sdk-go/aws/session"
 	"github.com/aws/aws-sdk-go/service/autoscaling"
@@ -236,6 +238,13 @@ func (c *awsCloudImplementation) Region() string {
 
 var awsCloudInstances map[string]AWSCloud = make(map[string]AWSCloud)
 
+func setConfig(config *aws.Config) *aws.Config {
+	// This avoids a confusing error message when we fail to get credentials
+	// e.g. https://github.com/kubernetes/kops/issues/605
+	config = config.WithCredentialsChainVerboseErrors(true)
+	return request.WithRetryer(config, newLoggingRetryer(ClientMaxRetries))
+}
+
 func NewAWSCloud(region string, tags map[string]string) (AWSCloud, error) {
 	raw := awsCloudInstances[region]
 	if raw == nil {
@@ -250,11 +259,7 @@ func NewAWSCloud(region string, tags map[string]string) (AWSCloud, error) {
 		}
 
 		config := aws.NewConfig().WithRegion(region)
-
-		// This avoids a confusing error message when we fail to get credentials
-		// e.g. https://github.com/kubernetes/kops/issues/605
-		config = config.WithCredentialsChainVerboseErrors(true)
-		config = request.WithRetryer(config, newLoggingRetryer(ClientMaxRetries))
+		config = setConfig(config)
 
 		requestLogger := newRequestLogger(2)
 
@@ -265,6 +270,15 @@ func NewAWSCloud(region string, tags map[string]string) (AWSCloud, error) {
 		if err != nil {
 			return c, err
 		}
+
+		// assumes the role before executing commands
+		roleARN := os.Getenv("AWS_ASSUME_ROLE_ARN")
+		if roleARN != "" {
+			creds := stscreds.NewCredentials(sess, roleARN)
+			config = &aws.Config{Credentials: creds}
+			config = setConfig(config).WithRegion(region)
+		}
+
 		c.ec2 = ec2.New(sess, config)
 		c.ec2.Handlers.Send.PushFront(requestLogger)
 		c.addHandlers(region, &c.ec2.Handlers)