From 4bbc4f25e1911aea2274e227ac48f807e0a53320 Mon Sep 17 00:00:00 2001
From: Frank Yin <franky@ziprecruiter.com>
Date: Tue, 7 Jan 2025 14:35:54 -0800
Subject: [PATCH] allow setting role session name for
 WebIdentityTokenFileCredentialsProvider

---
 .../msk/auth/iam/internals/MSKCredentialProvider.java  | 10 +++++++++-
 1 file changed, 9 insertions(+), 1 deletion(-)

diff --git a/src/main/java/software/amazon/msk/auth/iam/internals/MSKCredentialProvider.java b/src/main/java/software/amazon/msk/auth/iam/internals/MSKCredentialProvider.java
index ee694bf..cf94dfb 100644
--- a/src/main/java/software/amazon/msk/auth/iam/internals/MSKCredentialProvider.java
+++ b/src/main/java/software/amazon/msk/auth/iam/internals/MSKCredentialProvider.java
@@ -155,7 +155,6 @@ protected AwsCredentialsProvider getDefaultProvider() {
         return AwsCredentialsProviderChain.of(
             EnvironmentVariableCredentialsProvider.create(),
             SystemPropertyCredentialsProvider.create(),
-            WebIdentityTokenFileCredentialsProvider.builder().asyncCredentialUpdateEnabled(true).build(),
             ProfileCredentialsProvider.builder().profileFile(ProfileFileSupplier.defaultSupplier()).build(),
             ContainerCredentialsProvider.builder().asyncCredentialUpdateEnabled(true).build(),
             InstanceProfileCredentialsProvider.builder().asyncCredentialUpdateEnabled(true).build()
@@ -257,6 +256,7 @@ public List<AwsCredentialsProvider> getProviders() {
             List<AwsCredentialsProvider> providers = new ArrayList<>();
             getProfileProvider().ifPresent(providers::add);
             getStsRoleProvider().ifPresent(providers::add);
+            providers.add(getWebIdentityTokenProvider());
             return providers;
         }
 
@@ -347,6 +347,14 @@ else if (externalId != null) {
             });
         }
 
+        private WebIdentityTokenFileCredentialsProvider getWebIdentityTokenProvider() {
+            Optional<String> sessionName = Optional.ofNullable((String) optionsMap.get(AWS_ROLE_SESSION_KEY));
+            if (sessionName.isPresent()) {
+                return WebIdentityTokenFileCredentialsProvider.builder().asyncCredentialUpdateEnabled(true).roleSessionName(sessionName.get()).build();
+            }
+            return WebIdentityTokenFileCredentialsProvider.builder().asyncCredentialUpdateEnabled(true).build();
+        }
+
         StsAssumeRoleCredentialsProvider createSTSRoleCredentialProvider(
             String roleArn,
             String sessionName,