Skip to content

Commit

Permalink
Merge pull request #169 from mobsuccess-devops/fix/ec2-provider
Browse files Browse the repository at this point in the history 
fix: add missing InstanceProfileCredentialsProvider
  • Loading branch information
@hhkkxxx133
hhkkxxx133 authored Apr 19, 2024
2 parents c490e81 + f4abcc5 commit 16240fc
Show file tree
Hide file tree
Showing 2 changed files with 124 additions and 11 deletions.
4 changes: 3 additions & 1 deletion src/main/java/software/amazon/msk/auth/iam/internals/MSKCredentialProvider.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -36,6 +36,7 @@
import software.amazon.awssdk.auth.credentials.AwsSessionCredentials;
import software.amazon.awssdk.auth.credentials.ContainerCredentialsProvider;
import software.amazon.awssdk.auth.credentials.EnvironmentVariableCredentialsProvider;
import software.amazon.awssdk.auth.credentials.InstanceProfileCredentialsProvider;
import software.amazon.awssdk.auth.credentials.ProfileCredentialsProvider;
import software.amazon.awssdk.auth.credentials.StaticCredentialsProvider;
import software.amazon.awssdk.auth.credentials.SystemPropertyCredentialsProvider;
Expand Down Expand Up @@ -157,7 +158,8 @@ protected AwsCredentialsProvider getDefaultProvider() {
SystemPropertyCredentialsProvider.create(),
WebIdentityTokenFileCredentialsProvider.create(),
ProfileCredentialsProvider.create(),
ContainerCredentialsProvider.builder().build()
ContainerCredentialsProvider.builder().build(),
InstanceProfileCredentialsProvider.create()
);
}

Expand Down
131 changes: 121 additions & 10 deletions src/test/java/software/amazon/msk/auth/iam/internals/MSKCredentialProviderTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -36,6 +36,7 @@
import software.amazon.awssdk.auth.credentials.AwsSessionCredentials;
import software.amazon.awssdk.auth.credentials.ContainerCredentialsProvider;
import software.amazon.awssdk.auth.credentials.EnvironmentVariableCredentialsProvider;
import software.amazon.awssdk.auth.credentials.InstanceProfileCredentialsProvider;
import software.amazon.awssdk.auth.credentials.ProfileCredentialsProvider;
import software.amazon.awssdk.auth.credentials.SystemPropertyCredentialsProvider;
import software.amazon.awssdk.auth.credentials.WebIdentityTokenFileCredentialsProvider;
Expand Down Expand Up @@ -132,7 +133,8 @@ protected AwsCredentialsProvider getDefaultProvider() {
SystemPropertyCredentialsProvider.create(),
WebIdentityTokenFileCredentialsProvider.create(),
ProfileCredentialsProvider.builder().profileFile(profileFile).build(),
ContainerCredentialsProvider.builder().build()
ContainerCredentialsProvider.builder().build(),
InstanceProfileCredentialsProvider.create()
);
}
};
Expand Down Expand Up @@ -245,19 +247,54 @@ StsClient getStsClientForDebuggingCreds(AwsCredentials credentials) {
}

@Test
public void testEc2CredsWithDebugCredsNoAccessToSts_Succeed() {
public void testEcsCredsWithDebugCredsNoAccessToSts_Succeed() {
Map<String, String> optionsMap = new HashMap<>();
optionsMap.put(AWS_DEBUG_CREDS_NAME, "true");


ContainerCredentialsProvider mockEc2CredsProvider = Mockito.mock(ContainerCredentialsProvider.class);
Mockito.when(mockEc2CredsProvider.resolveIdentity())
ContainerCredentialsProvider mockEcsCredsProvider = Mockito.mock(ContainerCredentialsProvider.class);
Mockito.when(mockEcsCredsProvider.resolveIdentity())
.thenAnswer(i -> CompletableFuture.completedFuture(AwsBasicCredentials.create(ACCESS_KEY_VALUE_TWO, SECRET_KEY_VALUE_TWO)));

StsClient mockSts = Mockito.mock(StsClient.class);
Mockito.when(mockSts.getCallerIdentity())
.thenThrow(SdkClientException.create("TEST TEST"));

MSKCredentialProvider provider = new MSKCredentialProvider(optionsMap) {
protected AwsCredentialsProvider getDefaultProvider() {
return mockEcsCredsProvider;
}

StsClient getStsClientForDebuggingCreds(AwsCredentials credentials) {
return mockSts;
}
};
assertTrue(provider.getShouldDebugCreds());

AwsCredentials credentials = provider.resolveCredentials();

validateBasicCredentialsTwo(credentials);

provider.close();
Mockito.verify(mockSts, times(1)).getCallerIdentity();
Mockito.verify(mockEcsCredsProvider, times(1)).resolveIdentity();
Mockito.verifyNoMoreInteractions(mockEcsCredsProvider);
}

@Test
public void testEc2CredsWithDebugCredsNoAccessToSts_Succeed() {
Map<String, String> optionsMap = new HashMap<>();
optionsMap.put(AWS_DEBUG_CREDS_NAME, "true");


InstanceProfileCredentialsProvider mockEc2CredsProvider = Mockito.mock(InstanceProfileCredentialsProvider.class);
Mockito.when(mockEc2CredsProvider.resolveIdentity())
.thenAnswer(i -> CompletableFuture.completedFuture(AwsBasicCredentials.create(ACCESS_KEY_VALUE_TWO, SECRET_KEY_VALUE_TWO)));

StsClient mockSts = Mockito.mock(StsClient.class);
Mockito.when(mockSts.getCallerIdentity())
.thenThrow(SdkClientException.create("TEST TEST"));

MSKCredentialProvider provider = new MSKCredentialProvider(optionsMap) {
protected AwsCredentialsProvider getDefaultProvider() {
return mockEc2CredsProvider;
Expand Down Expand Up @@ -460,7 +497,7 @@ public void testEc2CredsWithSixRetriableErrorsCustomRetry_ThrowsException() {
Map<String, String> optionsMap = new HashMap<>();
optionsMap.put("awsMaxRetries", "5");

AwsCredentialsProvider mockEc2CredsProvider = setupMockDefaultProviderWithRetriableExceptions(numExceptions);
AwsCredentialsProvider mockEc2CredsProvider = setupMockEc2DefaultProviderWithRetriableExceptions(numExceptions);

MSKCredentialProvider provider = new MSKCredentialProvider(optionsMap) {
protected AwsCredentialsProvider getDefaultProvider() {
Expand All @@ -481,7 +518,7 @@ public void testEc2CredsWithOnrRetriableErrorsCustomZeroRetry_ThrowsException()
Map<String, String> optionsMap = new HashMap<>();
optionsMap.put("awsMaxRetries", "0");

AwsCredentialsProvider mockEc2CredsProvider = setupMockDefaultProviderWithRetriableExceptions(numExceptions);
AwsCredentialsProvider mockEc2CredsProvider = setupMockEc2DefaultProviderWithRetriableExceptions(numExceptions);

MSKCredentialProvider provider = new MSKCredentialProvider(optionsMap) {
protected AwsCredentialsProvider getDefaultProvider() {
Expand All @@ -500,7 +537,7 @@ private void testEc2CredsWithRetriableErrorsCustomRetry(int numExceptions) {
Map<String, String> optionsMap = new HashMap<>();
optionsMap.put("awsMaxRetries", "5");

AwsCredentialsProvider mockEc2CredsProvider = setupMockDefaultProviderWithRetriableExceptions(numExceptions);
AwsCredentialsProvider mockEc2CredsProvider = setupMockEc2DefaultProviderWithRetriableExceptions(numExceptions);

MSKCredentialProvider provider = new MSKCredentialProvider(optionsMap) {
protected AwsCredentialsProvider getDefaultProvider() {
Expand All @@ -518,6 +555,70 @@ protected AwsCredentialsProvider getDefaultProvider() {
Mockito.verifyNoMoreInteractions(mockEc2CredsProvider);
}

@Test
public void testEcsCredsWithSixRetriableErrorsCustomRetry_ThrowsException() {
int numExceptions = 6;
Map<String, String> optionsMap = new HashMap<>();
optionsMap.put("awsMaxRetries", "5");

AwsCredentialsProvider mockEcsCredsProvider = setupMockEcsDefaultProviderWithRetriableExceptions(numExceptions);

MSKCredentialProvider provider = new MSKCredentialProvider(optionsMap) {
protected AwsCredentialsProvider getDefaultProvider() {
return mockEcsCredsProvider;
}
};
assertFalse(provider.getShouldDebugCreds());

assertThrows(SdkClientException.class, () -> provider.resolveCredentials());

Mockito.verify(mockEcsCredsProvider, times(numExceptions)).resolveIdentity();
Mockito.verifyNoMoreInteractions(mockEcsCredsProvider);
}

@Test
public void testEcsCredsWithOnrRetriableErrorsCustomZeroRetry_ThrowsException() {
int numExceptions = 1;
Map<String, String> optionsMap = new HashMap<>();
optionsMap.put("awsMaxRetries", "0");

AwsCredentialsProvider mockEcsCredsProvider = setupMockEcsDefaultProviderWithRetriableExceptions(numExceptions);

MSKCredentialProvider provider = new MSKCredentialProvider(optionsMap) {
protected AwsCredentialsProvider getDefaultProvider() {
return mockEcsCredsProvider;
}
};
assertFalse(provider.getShouldDebugCreds());

assertThrows(SdkClientException.class, () -> provider.resolveCredentials());

Mockito.verify(mockEcsCredsProvider, times(numExceptions)).resolveIdentity();
Mockito.verifyNoMoreInteractions(mockEcsCredsProvider);
}

private void testEcsCredsWithRetriableErrorsCustomRetry(int numExceptions) {
Map<String, String> optionsMap = new HashMap<>();
optionsMap.put("awsMaxRetries", "5");

AwsCredentialsProvider mockEcsCredsProvider = setupMockEcsDefaultProviderWithRetriableExceptions(numExceptions);

MSKCredentialProvider provider = new MSKCredentialProvider(optionsMap) {
protected AwsCredentialsProvider getDefaultProvider() {
return mockEcsCredsProvider;
}
};
assertFalse(provider.getShouldDebugCreds());

AwsCredentials credentials = provider.resolveCredentials();

validateBasicCredentialsTwo(credentials);

provider.close();
Mockito.verify(mockEcsCredsProvider, times(numExceptions + 1)).resolveIdentity();
Mockito.verifyNoMoreInteractions(mockEcsCredsProvider);
}

private void testRoleCredsWithRetriableErrors(int numExceptions) {
StsAssumeRoleCredentialsProvider mockStsRoleProvider = setupMockStsRoleCredentialsProviderWithRetriableExceptions(
numExceptions);
Expand Down Expand Up @@ -608,13 +709,23 @@ private SdkException[] getSdkBaseExceptions(int numErrors) {
.collect(Collectors.toList()).toArray(new SdkException[numErrors]);
}

private AwsCredentialsProvider setupMockDefaultProviderWithRetriableExceptions(int numErrors) {
private AwsCredentialsProvider setupMockEcsDefaultProviderWithRetriableExceptions(int numErrors) {
SdkException[] exceptionsToThrow = getSdkBaseExceptions(numErrors);
ContainerCredentialsProvider mockEc2Provider = Mockito.mock(ContainerCredentialsProvider.class);
ContainerCredentialsProvider mockEcsProvider = Mockito.mock(ContainerCredentialsProvider.class);

Mockito.when(mockEc2Provider.resolveIdentity())
Mockito.when(mockEcsProvider.resolveIdentity())
.thenThrow(exceptionsToThrow)
.thenAnswer(i -> CompletableFuture.completedFuture(AwsBasicCredentials.create(ACCESS_KEY_VALUE_TWO, SECRET_KEY_VALUE_TWO)));
return mockEcsProvider;
}

private AwsCredentialsProvider setupMockEc2DefaultProviderWithRetriableExceptions(int numErrors) {
SdkException[] exceptionsToThrow = getSdkBaseExceptions(numErrors);
InstanceProfileCredentialsProvider mockEc2Provider = Mockito.mock(InstanceProfileCredentialsProvider.class);

Mockito.when(mockEc2Provider.resolveIdentity())
.thenThrow(exceptionsToThrow)
.thenAnswer(i -> CompletableFuture.completedFuture(AwsBasicCredentials.create(ACCESS_KEY_VALUE_TWO, SECRET_KEY_VALUE_TWO)));
return mockEc2Provider;
}

Expand Down

0 comments on commit 16240fc

Please sign in to comment.