luciano/amazon-kinesis-client
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Simplify StsAssumeRoleCredentialsProvider logic

Browse source
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
Show stats Download patch file Download diff file Expand all files Collapse all files

37
amazon-kinesis-client-multilang/src/main/java/software/amazon/kinesis/multilang/auth/KclStsAssumeRoleCredentialsProvider.java
View file

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

24
amazon-kinesis-client-multilang/src/main/java/software/amazon/kinesis/multilang/auth/StsAssumeRoleCredentialsProviderConfig.java
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;
}
}

40
amazon-kinesis-client-multilang/src/main/java/software/amazon/kinesis/multilang/auth/StsAssumeRoleCredentialsProviderFactory.java
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();
}
}