From a3ab249eca14d2b0a4c2e7ac262c811e1935ba3f Mon Sep 17 00:00:00 2001
From: Sami S <25616506+ThaSami@users.noreply.github.com>
Date: Fri, 10 Nov 2023 18:33:44 +0000
Subject: [PATCH] assume roles using webIdentity

---
 pkg/scalers/aws_common.go | 21 +++++++++++++++++++--
 1 file changed, 19 insertions(+), 2 deletions(-)

diff --git a/pkg/scalers/aws_common.go b/pkg/scalers/aws_common.go
index a4f64c60783..1e296e503d6 100644
--- a/pkg/scalers/aws_common.go
+++ b/pkg/scalers/aws_common.go
@@ -4,6 +4,7 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"os"
 
 	"github.com/aws/aws-sdk-go-v2/aws"
 	"github.com/aws/aws-sdk-go-v2/config"
@@ -52,8 +53,24 @@ func getAwsConfig(ctx context.Context, awsRegion string, awsAuthorization awsAut
 
 	if metadata.awsAuthorization.awsRoleArn != "" {
 		stsSvc := sts.NewFromConfig(cfg)
-		stsCredentialProvider := stscreds.NewAssumeRoleProvider(stsSvc, metadata.awsAuthorization.awsRoleArn, func(options *stscreds.AssumeRoleOptions) {})
-		cfg.Credentials = aws.NewCredentialsCache(stsCredentialProvider)
+
+		// Create the web identity role provider
+		stsCredentialProvider := stscreds.NewWebIdentityRoleProvider(
+			stsSvc,
+			metadata.awsAuthorization.awsRoleArn,
+			stscreds.IdentityTokenFile(os.Getenv("AWS_WEB_IDENTITY_TOKEN_FILE")),
+		)
+
+		// Attempt to retrieve credentials
+		_, err := stsCredentialProvider.Retrieve(context.Background())
+		if err != nil {
+			// Setup AssumeRoleProvider as a fallback
+			assumeRoleCredentialProvider := stscreds.NewAssumeRoleProvider(stsSvc, metadata.awsAuthorization.awsRoleArn, func(options *stscreds.AssumeRoleOptions) {})
+			cfg.Credentials = aws.NewCredentialsCache(assumeRoleCredentialProvider)
+		} else {
+			// If the retrieval is successful, use the web identity credentials
+			cfg.Credentials = aws.NewCredentialsCache(stsCredentialProvider)
+		}
 	}
 
 	return &cfg, err