Simplify StsAssumeRoleCredentialsProvider logic

This commit is contained in:
Ethan Katnic 2024-08-27 09:06:07 -07:00
parent f8eafb457a
commit 6799c81f82
3 changed files with 21 additions and 80 deletions

View file

@ -1,53 +1,58 @@
package software.amazon.kinesis.multilang.auth;
import java.net.URI;
import java.util.Arrays;
import software.amazon.awssdk.auth.credentials.AwsCredentials;
import software.amazon.awssdk.auth.credentials.AwsCredentialsProvider;
import software.amazon.awssdk.regions.Region;
import software.amazon.awssdk.services.sts.StsClient;
import software.amazon.awssdk.services.sts.StsClientBuilder;
import software.amazon.awssdk.services.sts.auth.StsAssumeRoleCredentialsProvider;
import software.amazon.awssdk.services.sts.model.AssumeRoleRequest;
import software.amazon.awssdk.services.sts.model.AssumeRoleRequest.Builder;
import software.amazon.kinesis.multilang.NestedPropertyKey;
import software.amazon.kinesis.multilang.NestedPropertyProcessor;
public class KclStsAssumeRoleCredentialsProvider implements AwsCredentialsProvider, NestedPropertyProcessor {
private final String roleArn;
private final String roleSessionName;
private Region region;
private String serviceEndpoint;
private String externalId;
private final Builder assumeRoleRequestBuilder;
private final StsClientBuilder stsClientBuilder;
public KclStsAssumeRoleCredentialsProvider(String[] params) {
this(params[0], params[1], Arrays.copyOfRange(params, 2, params.length));
}
public KclStsAssumeRoleCredentialsProvider(String roleArn, String roleSessionName, String... params) {
this.roleArn = roleArn;
this.roleSessionName = roleSessionName;
this.assumeRoleRequestBuilder =
AssumeRoleRequest.builder().roleArn(roleArn).roleSessionName(roleSessionName);
this.stsClientBuilder = StsClient.builder();
NestedPropertyKey.parse(this, params);
}
@Override
public AwsCredentials resolveCredentials() {
StsAssumeRoleCredentialsProviderConfig config = new StsAssumeRoleCredentialsProviderConfig(
roleArn, roleSessionName, region, serviceEndpoint, externalId);
StsAssumeRoleCredentialsProvider stsAssumeRoleCredentialsProvider =
StsAssumeRoleCredentialsProviderFactory.createProvider(config);
return stsAssumeRoleCredentialsProvider.resolveCredentials();
StsClient stsClient = this.stsClientBuilder.build();
AssumeRoleRequest assumeRoleRequest = this.assumeRoleRequestBuilder.build();
StsAssumeRoleCredentialsProvider provider = StsAssumeRoleCredentialsProvider.builder()
.refreshRequest(assumeRoleRequest)
.stsClient(stsClient)
.build();
return provider.resolveCredentials();
}
@Override
public void acceptEndpoint(String serviceEndpoint, String signingRegion) {
this.serviceEndpoint = serviceEndpoint;
this.region = Region.of(signingRegion);
stsClientBuilder.endpointOverride(URI.create(serviceEndpoint));
stsClientBuilder.region(Region.of(signingRegion));
}
@Override
public void acceptEndpointRegion(Region region) {
this.region = region;
stsClientBuilder.region(region);
}
@Override
public void acceptExternalId(String externalId) {
this.externalId = externalId;
assumeRoleRequestBuilder.externalId(externalId);
}
}

View file

@ -1,24 +0,0 @@
package software.amazon.kinesis.multilang.auth;
import lombok.Getter;
import lombok.Setter;
import software.amazon.awssdk.regions.Region;
@Getter
@Setter
public class StsAssumeRoleCredentialsProviderConfig {
private final String roleArn;
private final String roleSessionName;
private final Region region;
private final String serviceEndpoint;
private final String externalId;
public StsAssumeRoleCredentialsProviderConfig(
String roleArn, String roleSessionName, Region region, String serviceEndpoint, String externalId) {
this.roleArn = roleArn;
this.roleSessionName = roleSessionName;
this.region = region;
this.serviceEndpoint = serviceEndpoint;
this.externalId = externalId;
}
}

View file

@ -1,40 +0,0 @@
package software.amazon.kinesis.multilang.auth;
import java.net.URI;
import java.net.URISyntaxException;
import software.amazon.awssdk.services.sts.StsClient;
import software.amazon.awssdk.services.sts.StsClientBuilder;
import software.amazon.awssdk.services.sts.auth.StsAssumeRoleCredentialsProvider;
import software.amazon.awssdk.services.sts.model.AssumeRoleRequest;
public class StsAssumeRoleCredentialsProviderFactory {
public static StsAssumeRoleCredentialsProvider createProvider(StsAssumeRoleCredentialsProviderConfig config) {
StsClientBuilder stsClientBuilder = StsClient.builder();
if (config.getRegion() != null) {
stsClientBuilder.region(config.getRegion());
}
if (config.getServiceEndpoint() != null) {
try {
stsClientBuilder.endpointOverride(new URI(config.getServiceEndpoint()));
} catch (URISyntaxException e) {
throw new IllegalArgumentException("Invalid service endpoint: " + config.getServiceEndpoint(), e);
}
}
StsClient stsClient = stsClientBuilder.build();
AssumeRoleRequest assumeRoleRequest = AssumeRoleRequest.builder()
.roleArn(config.getRoleArn())
.roleSessionName(config.getRoleSessionName())
.build();
return StsAssumeRoleCredentialsProvider.builder()
.refreshRequest(assumeRoleRequest)
.stsClient(stsClient)
.build();
}
}