From 119ef4220607a3d15af5d82d41178390e320bea0 Mon Sep 17 00:00:00 2001 From: Ethan Katnic Date: Mon, 19 Aug 2024 17:41:06 -0700 Subject: [PATCH] Rewrite KCLSTSAssumeRoleCredentialsProvider to use AWS SDK v2 --- ...SAssumeRoleSessionCredentialsProvider.java | 86 ------------------- .../KclStsAssumeRoleCredentialsProvider.java | 53 ++++++++++++ ...tsAssumeRoleCredentialsProviderConfig.java | 24 ++++++ ...sAssumeRoleCredentialsProviderFactory.java | 40 +++++++++ 4 files changed, 117 insertions(+), 86 deletions(-) delete mode 100644 amazon-kinesis-client-multilang/src/main/java/software/amazon/kinesis/multilang/auth/KclSTSAssumeRoleSessionCredentialsProvider.java create mode 100644 amazon-kinesis-client-multilang/src/main/java/software/amazon/kinesis/multilang/auth/KclStsAssumeRoleCredentialsProvider.java create mode 100644 amazon-kinesis-client-multilang/src/main/java/software/amazon/kinesis/multilang/auth/StsAssumeRoleCredentialsProviderConfig.java create mode 100644 amazon-kinesis-client-multilang/src/main/java/software/amazon/kinesis/multilang/auth/StsAssumeRoleCredentialsProviderFactory.java diff --git a/amazon-kinesis-client-multilang/src/main/java/software/amazon/kinesis/multilang/auth/KclSTSAssumeRoleSessionCredentialsProvider.java b/amazon-kinesis-client-multilang/src/main/java/software/amazon/kinesis/multilang/auth/KclSTSAssumeRoleSessionCredentialsProvider.java deleted file mode 100644 index b5b9f924..00000000 --- a/amazon-kinesis-client-multilang/src/main/java/software/amazon/kinesis/multilang/auth/KclSTSAssumeRoleSessionCredentialsProvider.java +++ /dev/null @@ -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); - } -} diff --git a/amazon-kinesis-client-multilang/src/main/java/software/amazon/kinesis/multilang/auth/KclStsAssumeRoleCredentialsProvider.java b/amazon-kinesis-client-multilang/src/main/java/software/amazon/kinesis/multilang/auth/KclStsAssumeRoleCredentialsProvider.java new file mode 100644 index 00000000..d077e6f2 --- /dev/null +++ b/amazon-kinesis-client-multilang/src/main/java/software/amazon/kinesis/multilang/auth/KclStsAssumeRoleCredentialsProvider.java @@ -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; + } +} diff --git a/amazon-kinesis-client-multilang/src/main/java/software/amazon/kinesis/multilang/auth/StsAssumeRoleCredentialsProviderConfig.java b/amazon-kinesis-client-multilang/src/main/java/software/amazon/kinesis/multilang/auth/StsAssumeRoleCredentialsProviderConfig.java new file mode 100644 index 00000000..f4e41257 --- /dev/null +++ b/amazon-kinesis-client-multilang/src/main/java/software/amazon/kinesis/multilang/auth/StsAssumeRoleCredentialsProviderConfig.java @@ -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; + } +} diff --git a/amazon-kinesis-client-multilang/src/main/java/software/amazon/kinesis/multilang/auth/StsAssumeRoleCredentialsProviderFactory.java b/amazon-kinesis-client-multilang/src/main/java/software/amazon/kinesis/multilang/auth/StsAssumeRoleCredentialsProviderFactory.java new file mode 100644 index 00000000..b3b55531 --- /dev/null +++ b/amazon-kinesis-client-multilang/src/main/java/software/amazon/kinesis/multilang/auth/StsAssumeRoleCredentialsProviderFactory.java @@ -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(); + } +}