Add support for external id in auth (#109)

aws · Mar 1, 2023 · 648003c · 648003c
1 parent 87163dc
commit 648003c
Show file tree

Hide file tree

Showing 2 changed files with 58 additions and 2 deletions.
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
@@ -58,8 +58,8 @@
  * sasl.jaas.config = IAMLoginModule required awsProfileName={profile name};
  * The currently supported options are:
  * 1. A particular AWS Credential profile: awsProfileName={profile name}
- * 2. A particular AWS IAM Role, with optional access key id an secret key, and optionally AWS IAM role session name
- *    and AWS region for the STS endpoint:
+ * 2. A particular AWS IAM Role, with optional access key id and secret key OR optional external id,
+ *    and optionally AWS IAM role session name and AWS region for the STS endpoint:
  *     awsRoleArn={IAM Role ARN}, awsRoleAccessKeyId={access key id}, awsSecretAccessKey={secret access key},
  *     awsRoleSessionName={session name}, awsStsRegion={region name}
  * 3. Optional arguments to configure retries when we fail to load credentials:
@@ -74,6 +74,7 @@ public class MSKCredentialProvider implements AWSCredentialsProvider, AutoClosea
     private static final Logger log = LoggerFactory.getLogger(MSKCredentialProvider.class);
     private static final String AWS_PROFILE_NAME_KEY = "awsProfileName";
     private static final String AWS_ROLE_ARN_KEY = "awsRoleArn";
+    private static final String AWS_ROLE_EXTERNAL_ID = "awsRoleExternalId";
     private static final String AWS_ROLE_ACCESS_KEY_ID = "awsRoleAccessKeyId";
     private static final String AWS_ROLE_SECRET_ACCESS_KEY = "awsRoleSecretAccessKey";
     private static final String AWS_ROLE_SESSION_KEY = "awsRoleSessionName";
@@ -288,10 +289,14 @@ private Optional<STSAssumeRoleSessionCredentialsProvider> getStsRoleProvider() {
 
                 String accessKey = (String) optionsMap.getOrDefault(AWS_ROLE_ACCESS_KEY_ID, null);
                 String secretKey = (String) optionsMap.getOrDefault(AWS_ROLE_SECRET_ACCESS_KEY, null);
+                String externalId = (String) optionsMap.getOrDefault(AWS_ROLE_EXTERNAL_ID, null);
                 if (accessKey != null && secretKey != null) {
                     AWSCredentialsProvider credentials = new AWSStaticCredentialsProvider(new BasicAWSCredentials(accessKey, secretKey));
                     return createSTSRoleCredentialProvider((String) p, sessionName, stsRegion, credentials);
                 }
+                else if (externalId != null) {
+                    return createSTSRoleCredentialProvider((String) p, externalId, sessionName, stsRegion);
+                }
 
                 return createSTSRoleCredentialProvider((String) p, sessionName, stsRegion);
             });
@@ -319,6 +324,20 @@ STSAssumeRoleSessionCredentialsProvider createSTSRoleCredentialProvider(String r
                     .withStsClient(stsClient)
                     .build();
         }
+
+        STSAssumeRoleSessionCredentialsProvider createSTSRoleCredentialProvider(String roleArn,
+                                                                                String externalId,
+                                                                                String sessionName,
+                                                                                String stsRegion) {
+            AWSSecurityTokenService stsClient = AWSSecurityTokenServiceClientBuilder.standard()
+                    .withRegion(stsRegion)
+                    .build();
+
+            return new STSAssumeRoleSessionCredentialsProvider.Builder(roleArn, sessionName)
+                    .withStsClient(stsClient)
+                    .withExternalId(externalId)
+                    .build();
+        }
     }
 
 }
diff --git a/src/test/java/software/amazon/msk/auth/iam/internals/MSKCredentialProviderTest.java b/src/test/java/software/amazon/msk/auth/iam/internals/MSKCredentialProviderTest.java
@@ -61,9 +61,11 @@ public class MSKCredentialProviderTest {
     private static final String PROFILE_ACCESS_KEY_VALUE = "PROFILE_ACCESS_KEY";
     private static final String PROFILE_SECRET_KEY_VALUE = "PROFILE_SECRET_KEY";
     private static final String TEST_ROLE_ARN = "TEST_ROLE_ARN";
+    private static final String TEST_ROLE_EXTERNAL_ID = "TEST_EXTERNAL_ID";
     private static final String TEST_ROLE_SESSION_NAME = "TEST_ROLE_SESSION_NAME";
     private static final String SESSION_TOKEN = "SESSION_TOKEN";
     private static final String AWS_ROLE_ARN = "awsRoleArn";
+    private static final String AWS_ROLE_EXTERNAL_ID = "awsRoleExternalId";
     private static final String AWS_ROLE_ACCESS_KEY_ID = "awsRoleAccessKeyId";
     private static final String AWS_ROLE_SECRET_ACCESS_KEY = "awsRoleSecretAccessKey";
     private static final String AWS_PROFILE_NAME = "awsProfileName";
@@ -324,6 +326,41 @@ STSAssumeRoleSessionCredentialsProvider createSTSRoleCredentialProvider(String r
         Mockito.verify(mockStsRoleProvider, times(1)).close();
     }
 
+    @Test
+    public void testAwsRoleArnSessionNameStsRegionAndExternalId() {
+        STSAssumeRoleSessionCredentialsProvider mockStsRoleProvider = Mockito
+                .mock(STSAssumeRoleSessionCredentialsProvider.class);
+        Mockito.when(mockStsRoleProvider.getCredentials())
+                .thenReturn(new BasicSessionCredentials(ACCESS_KEY_VALUE, SECRET_KEY_VALUE, SESSION_TOKEN));
+
+        Map<String, String> optionsMap = new HashMap<>();
+        optionsMap.put(AWS_ROLE_ARN, TEST_ROLE_ARN);
+        optionsMap.put(AWS_ROLE_EXTERNAL_ID, TEST_ROLE_EXTERNAL_ID);
+        optionsMap.put("awsRoleSessionName", TEST_ROLE_SESSION_NAME);
+        optionsMap.put("awsStsRegion", "eu-west-1");
+
+        MSKCredentialProvider.ProviderBuilder providerBuilder = new MSKCredentialProvider.ProviderBuilder(optionsMap) {
+            STSAssumeRoleSessionCredentialsProvider createSTSRoleCredentialProvider(String roleArn,
+                                                                                    String externalId,
+                                                                                    String sessionName,
+                                                                                    String stsRegion) {
+                assertEquals(TEST_ROLE_ARN, roleArn);
+                assertEquals(TEST_ROLE_EXTERNAL_ID, externalId);
+                assertEquals(TEST_ROLE_SESSION_NAME, sessionName);
+                assertEquals("eu-west-1", stsRegion);
+                return mockStsRoleProvider;
+            }
+        };
+        MSKCredentialProvider provider = new MSKCredentialProvider(providerBuilder);
+        assertFalse(provider.getShouldDebugCreds());
+
+        AWSCredentials credentials = provider.getCredentials();
+        validateBasicSessionCredentials(credentials);
+
+        provider.close();
+        Mockito.verify(mockStsRoleProvider, times(1)).close();
+    }
+
     @Test
     public void testProfileNameAndRoleArn() {
         ProfileFile profileFile = getProfileFile();