Rewrite KCLSTSAssumeRoleCredentialsProvider to use AWS SDK v2

This commit is contained in:
Ethan Katnic 2024-08-19 17:41:06 -07:00
parent a8b70da3b3
commit 119ef42206
4 changed files with 117 additions and 86 deletions

View file

@ -1,86 +0,0 @@
/*
* Copyright 2023 Amazon.com, Inc. or its affiliates.
* Licensed under the Apache License, Version 2.0 (the
* "License"); you may not use this file except in compliance
* with the License. You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package software.amazon.kinesis.multilang.auth;
import java.util.Arrays;
import com.amazonaws.auth.AWSSessionCredentials;
import com.amazonaws.auth.AWSSessionCredentialsProvider;
import com.amazonaws.auth.STSAssumeRoleSessionCredentialsProvider;
import com.amazonaws.auth.STSAssumeRoleSessionCredentialsProvider.Builder;
import com.amazonaws.client.builder.AwsClientBuilder.EndpointConfiguration;
import com.amazonaws.regions.Regions;
import com.amazonaws.services.securitytoken.AWSSecurityTokenService;
import com.amazonaws.services.securitytoken.AWSSecurityTokenServiceClient;
import software.amazon.kinesis.multilang.NestedPropertyKey;
import software.amazon.kinesis.multilang.NestedPropertyProcessor;
/**
* An {@link AWSSessionCredentialsProvider} that is backed by STSAssumeRole.
*/
public class KclSTSAssumeRoleSessionCredentialsProvider
implements AWSSessionCredentialsProvider, NestedPropertyProcessor {
private final Builder builder;
private final STSAssumeRoleSessionCredentialsProvider provider;
/**
*
* @param params vararg parameters which must include roleArn at index=0,
* and roleSessionName at index=1
*/
public KclSTSAssumeRoleSessionCredentialsProvider(final String[] params) {
this(params[0], params[1], Arrays.copyOfRange(params, 2, params.length));
}
public KclSTSAssumeRoleSessionCredentialsProvider(
final String roleArn, final String roleSessionName, final String... params) {
builder = new Builder(roleArn, roleSessionName);
NestedPropertyKey.parse(this, params);
provider = builder.build();
}
@Override
public AWSSessionCredentials getCredentials() {
return provider.getCredentials();
}
@Override
public void refresh() {
// do nothing
}
@Override
public void acceptEndpoint(final String serviceEndpoint, final String signingRegion) {
final EndpointConfiguration endpoint = new EndpointConfiguration(serviceEndpoint, signingRegion);
final AWSSecurityTokenService stsClient = AWSSecurityTokenServiceClient.builder()
.withEndpointConfiguration(endpoint)
.build();
builder.withStsClient(stsClient);
}
@Override
public void acceptEndpointRegion(final Regions region) {
final AWSSecurityTokenService stsClient =
AWSSecurityTokenServiceClient.builder().withRegion(region).build();
builder.withStsClient(stsClient);
}
@Override
public void acceptExternalId(final String externalId) {
builder.withExternalId(externalId);
}
}

View file

@ -0,0 +1,53 @@
package software.amazon.kinesis.multilang.auth;
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.auth.StsAssumeRoleCredentialsProvider;
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;
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;
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();
}
@Override
public void acceptEndpoint(String serviceEndpoint, String signingRegion) {
this.serviceEndpoint = serviceEndpoint;
this.region = Region.of(signingRegion);
}
@Override
public void acceptEndpointRegion(Region region) {
this.region = region;
}
@Override
public void acceptExternalId(String externalId) {
this.externalId = externalId;
}
}

View file

@ -0,0 +1,24 @@
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

@ -0,0 +1,40 @@
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();
}
}