This is an automated email from the ASF dual-hosted git repository.
dannycranmer pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/flink.git
The following commit(s) were added to refs/heads/master by this push:
new 79712e16adc [FLINK-29496][Connector/Kinesis] Add configuration option
for setting STS endpoint for AWS connectors
79712e16adc is described below
commit 79712e16adc13a9ccac016178fd3dcb551afbab3
Author: Aleksandr Pilipenko <[email protected]>
AuthorDate: Mon Oct 17 17:17:48 2022 +0000
[FLINK-29496][Connector/Kinesis] Add configuration option for setting STS
endpoint for AWS connectors
---
docs/content.zh/docs/connectors/table/kinesis.md | 8 +++
docs/content/docs/connectors/table/kinesis.md | 8 +++
.../connector/aws/config/AWSConfigConstants.java | 10 +++
.../flink/connector/aws/util/AWSGeneralUtil.java | 36 ++++++++---
.../connector/aws/util/AWSGeneralUtilTest.java | 31 +++++++++
.../streaming/connectors/kinesis/util/AWSUtil.java | 34 +++++++---
.../connectors/kinesis/testutils/TestUtils.java | 14 +++++
.../connectors/kinesis/util/AWSUtilTest.java | 73 ++++++++++++++++++++++
8 files changed, 195 insertions(+), 19 deletions(-)
diff --git a/docs/content.zh/docs/connectors/table/kinesis.md
b/docs/content.zh/docs/connectors/table/kinesis.md
index bc4d7dd51da..6e82cf96d7f 100644
--- a/docs/content.zh/docs/connectors/table/kinesis.md
+++ b/docs/content.zh/docs/connectors/table/kinesis.md
@@ -240,6 +240,14 @@ Connector Options
<td>String</td>
<td>The external ID to use when credential provider type is set to
ASSUME_ROLE.</td>
</tr>
+ <tr>
+ <td><h5>aws.credentials.role.stsEndpoint</h5></td>
+ <td>optional</td>
+ <td>no</td>
+ <td style="word-wrap: break-word;">(none)</td>
+ <td>String</td>
+ <td>The AWS endpoint for STS (derived from the AWS region setting if
not set) to use when credential provider type is set to ASSUME_ROLE.</td>
+ </tr>
<tr>
<td><h5>aws.credentials.role.provider</h5></td>
<td>optional</td>
diff --git a/docs/content/docs/connectors/table/kinesis.md
b/docs/content/docs/connectors/table/kinesis.md
index e68f9d67b06..d37617decb1 100644
--- a/docs/content/docs/connectors/table/kinesis.md
+++ b/docs/content/docs/connectors/table/kinesis.md
@@ -255,6 +255,14 @@ Connector Options
<td>String</td>
<td>The external ID to use when credential provider type is set to
ASSUME_ROLE.</td>
</tr>
+ <tr>
+ <td><h5>aws.credentials.role.stsEndpoint</h5></td>
+ <td>optional</td>
+ <td>no</td>
+ <td style="word-wrap: break-word;">(none)</td>
+ <td>String</td>
+ <td>The AWS endpoint for STS (derived from the AWS region setting if
not set) to use when credential provider type is set to ASSUME_ROLE.</td>
+ </tr>
<tr>
<td><h5>aws.credentials.role.provider</h5></td>
<td>optional</td>
diff --git
a/flink-connectors/flink-connector-aws-base/src/main/java/org/apache/flink/connector/aws/config/AWSConfigConstants.java
b/flink-connectors/flink-connector-aws-base/src/main/java/org/apache/flink/connector/aws/config/AWSConfigConstants.java
index b244b2c56a9..f67f8ad852e 100644
---
a/flink-connectors/flink-connector-aws-base/src/main/java/org/apache/flink/connector/aws/config/AWSConfigConstants.java
+++
b/flink-connectors/flink-connector-aws-base/src/main/java/org/apache/flink/connector/aws/config/AWSConfigConstants.java
@@ -93,6 +93,12 @@ public class AWSConfigConstants {
/** Optional configuration for profile name if credential provider type is
set to be PROFILE. */
public static final String AWS_PROFILE_NAME =
profileName(AWS_CREDENTIALS_PROVIDER);
+ /**
+ * The AWS endpoint for the STS (derived from the AWS region setting if
not set) to use if
+ * credential provider type is set to be ASSUME_ROLE.
+ */
+ public static final String AWS_ROLE_STS_ENDPOINT =
roleStsEndpoint(AWS_CREDENTIALS_PROVIDER);
+
/**
* The role ARN to use when credential provider type is set to ASSUME_ROLE
or
* WEB_IDENTITY_TOKEN.
@@ -170,6 +176,10 @@ public class AWSConfigConstants {
return prefix + ".role.provider";
}
+ public static String roleStsEndpoint(String prefix) {
+ return prefix + ".role.stsEndpoint";
+ }
+
public static String webIdentityTokenFile(String prefix) {
return prefix + ".webIdentityToken.file";
}
diff --git
a/flink-connectors/flink-connector-aws-base/src/main/java/org/apache/flink/connector/aws/util/AWSGeneralUtil.java
b/flink-connectors/flink-connector-aws-base/src/main/java/org/apache/flink/connector/aws/util/AWSGeneralUtil.java
index e1d7f982835..7b37e0bd954 100644
---
a/flink-connectors/flink-connector-aws-base/src/main/java/org/apache/flink/connector/aws/util/AWSGeneralUtil.java
+++
b/flink-connectors/flink-connector-aws-base/src/main/java/org/apache/flink/connector/aws/util/AWSGeneralUtil.java
@@ -38,11 +38,13 @@ import
software.amazon.awssdk.http.nio.netty.NettyNioAsyncHttpClient;
import software.amazon.awssdk.profiles.ProfileFile;
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.utils.AttributeMap;
import software.amazon.awssdk.utils.SdkAutoCloseable;
+import java.net.URI;
import java.nio.file.Paths;
import java.time.Duration;
import java.util.Map;
@@ -188,6 +190,16 @@ public class AWSGeneralUtil {
private static AwsCredentialsProvider getAssumeRoleCredentialProvider(
final Properties configProps, final String configPrefix) {
+ final StsClientBuilder stsClientBuilder =
+ StsClient.builder()
+ .credentialsProvider(
+ getCredentialsProvider(
+ configProps,
+
AWSConfigConstants.roleCredentialsProvider(configPrefix)))
+ .region(getRegion(configProps));
+ Optional.ofNullable(getStsEndpoint(configProps))
+ .ifPresent(stsClientBuilder::endpointOverride);
+
return StsAssumeRoleCredentialsProvider.builder()
.refreshRequest(
AssumeRoleRequest.builder()
@@ -201,15 +213,7 @@ public class AWSGeneralUtil {
configProps.getProperty(
AWSConfigConstants.externalId(configPrefix)))
.build())
- .stsClient(
- StsClient.builder()
- .credentialsProvider(
- getCredentialsProvider(
- configProps,
-
AWSConfigConstants.roleCredentialsProvider(
- configPrefix)))
- .region(getRegion(configProps))
- .build())
+ .stsClient(stsClientBuilder.build())
.build();
}
@@ -307,6 +311,20 @@ public class AWSGeneralUtil {
return
Region.of(configProps.getProperty(AWSConfigConstants.AWS_REGION));
}
+ /**
+ * Creates STS endpoint URI object from the given Properties.
+ *
+ * @param configProps the properties containing the endpoint
+ * @return STS endpoint URI specified by the properties, or {@code null}
if endpoint not
+ * specified
+ */
+ public static URI getStsEndpoint(final Properties configProps) {
+ return Optional.ofNullable(
+
configProps.getProperty(AWSConfigConstants.AWS_ROLE_STS_ENDPOINT))
+ .map(URI::create)
+ .orElse(null);
+ }
+
/**
* Checks whether or not a region is valid.
*
diff --git
a/flink-connectors/flink-connector-aws-base/src/test/java/org/apache/flink/connector/aws/util/AWSGeneralUtilTest.java
b/flink-connectors/flink-connector-aws-base/src/test/java/org/apache/flink/connector/aws/util/AWSGeneralUtilTest.java
index ca5613afc60..7f4a0777bf7 100644
---
a/flink-connectors/flink-connector-aws-base/src/test/java/org/apache/flink/connector/aws/util/AWSGeneralUtilTest.java
+++
b/flink-connectors/flink-connector-aws-base/src/test/java/org/apache/flink/connector/aws/util/AWSGeneralUtilTest.java
@@ -39,6 +39,7 @@ import
software.amazon.awssdk.services.sts.auth.StsAssumeRoleCredentialsProvider
import software.amazon.awssdk.utils.AttributeMap;
import software.amazon.awssdk.utils.ImmutableMap;
+import java.net.URI;
import java.nio.file.Paths;
import java.time.Duration;
import java.util.Map;
@@ -46,6 +47,7 @@ import java.util.Properties;
import static
org.apache.flink.connector.aws.config.AWSConfigConstants.AWS_CREDENTIALS_PROVIDER;
import static
org.apache.flink.connector.aws.config.AWSConfigConstants.AWS_REGION;
+import static
org.apache.flink.connector.aws.config.AWSConfigConstants.AWS_ROLE_STS_ENDPOINT;
import static
org.apache.flink.connector.aws.config.AWSConfigConstants.CredentialProvider.ASSUME_ROLE;
import static
org.apache.flink.connector.aws.config.AWSConfigConstants.CredentialProvider.AUTO;
import static
org.apache.flink.connector.aws.config.AWSConfigConstants.CredentialProvider.BASIC;
@@ -201,6 +203,25 @@ class AWSGeneralUtilTest {
verify(properties).getProperty(AWS_REGION);
}
+ @Test
+ void testGetCredentialsProviderAssumeRoleWithPropertyAwsStsEndpoint() {
+ Properties properties =
spy(TestUtil.properties(AWS_CREDENTIALS_PROVIDER, "ASSUME_ROLE"));
+ properties.setProperty(AWS_REGION, "eu-west-2");
+ properties.setProperty(AWS_ROLE_STS_ENDPOINT,
"https://sts.eu-west-2.amazonaws.com";);
+
+ AwsCredentialsProvider credentialsProvider =
+ AWSGeneralUtil.getCredentialsProvider(properties);
+
+
assertThat(credentialsProvider).isInstanceOf(StsAssumeRoleCredentialsProvider.class);
+
+
verify(properties).getProperty(AWSConfigConstants.roleArn(AWS_CREDENTIALS_PROVIDER));
+ verify(properties)
+
.getProperty(AWSConfigConstants.roleSessionName(AWS_CREDENTIALS_PROVIDER));
+
verify(properties).getProperty(AWSConfigConstants.externalId(AWS_CREDENTIALS_PROVIDER));
+ verify(properties).getProperty(AWS_REGION);
+ verify(properties).getProperty(AWS_ROLE_STS_ENDPOINT);
+ }
+
@Test
void testGetCredentialsProviderBasic() {
Properties properties = TestUtil.properties(AWS_CREDENTIALS_PROVIDER,
"BASIC");
@@ -696,6 +717,16 @@ class AWSGeneralUtilTest {
assertThat(AWSGeneralUtil.isValidRegion(Region.of("unstructured-string"))).isFalse();
}
+ @Test
+ void testGetStsEndpoint() {
+ String endpointValue = "https://sts.us-east-1.amazonaws.com";;
+ URI stsEndpoint =
+ AWSGeneralUtil.getStsEndpoint(
+ TestUtil.properties(AWS_ROLE_STS_ENDPOINT,
endpointValue));
+
+ assertThat(stsEndpoint).isEqualTo(URI.create(endpointValue));
+ }
+
@Test
void testUnrecognizableAwsRegionInConfig() {
diff --git
a/flink-connectors/flink-connector-kinesis/src/main/java/org/apache/flink/streaming/connectors/kinesis/util/AWSUtil.java
b/flink-connectors/flink-connector-kinesis/src/main/java/org/apache/flink/streaming/connectors/kinesis/util/AWSUtil.java
index 40f6e64e19b..04dd35bde0e 100644
---
a/flink-connectors/flink-connector-kinesis/src/main/java/org/apache/flink/streaming/connectors/kinesis/util/AWSUtil.java
+++
b/flink-connectors/flink-connector-kinesis/src/main/java/org/apache/flink/streaming/connectors/kinesis/util/AWSUtil.java
@@ -167,15 +167,6 @@ public class AWSUtil {
};
case ASSUME_ROLE:
- final AWSSecurityTokenService baseCredentials =
- AWSSecurityTokenServiceClientBuilder.standard()
- .withCredentials(
- getCredentialsProvider(
- configProps,
-
AWSConfigConstants.roleCredentialsProvider(
- configPrefix)))
-
.withRegion(configProps.getProperty(AWSConfigConstants.AWS_REGION))
- .build();
return new STSAssumeRoleSessionCredentialsProvider.Builder(
configProps.getProperty(AWSConfigConstants.roleArn(configPrefix)),
configProps.getProperty(
@@ -183,7 +174,7 @@ public class AWSUtil {
.withExternalId(
configProps.getProperty(
AWSConfigConstants.externalId(configPrefix)))
- .withStsClient(baseCredentials)
+ .withStsClient(createStsClient(configProps,
configPrefix))
.build();
case WEB_IDENTITY_TOKEN:
@@ -209,6 +200,29 @@ public class AWSUtil {
}
}
+ private static AWSSecurityTokenService createStsClient(
+ final Properties configProps, final String configPrefix) {
+ final String region =
configProps.getProperty(AWSConfigConstants.AWS_REGION);
+ final AWSSecurityTokenServiceClientBuilder stsClientBuilder =
+ AWSSecurityTokenServiceClientBuilder.standard()
+ .withCredentials(
+ getCredentialsProvider(
+ configProps,
+
AWSConfigConstants.roleCredentialsProvider(configPrefix)));
+
+ if (configProps.containsKey(AWSConfigConstants.AWS_ROLE_STS_ENDPOINT))
{
+ AwsClientBuilder.EndpointConfiguration endpointConfiguration =
+ new AwsClientBuilder.EndpointConfiguration(
+
configProps.getProperty(AWSConfigConstants.AWS_ROLE_STS_ENDPOINT),
+ region);
+ stsClientBuilder.withEndpointConfiguration(endpointConfiguration);
+ } else {
+ stsClientBuilder.withRegion(region);
+ }
+
+ return stsClientBuilder.build();
+ }
+
/**
* Checks whether or not a region ID is valid.
*
diff --git
a/flink-connectors/flink-connector-kinesis/src/test/java/org/apache/flink/streaming/connectors/kinesis/testutils/TestUtils.java
b/flink-connectors/flink-connector-kinesis/src/test/java/org/apache/flink/streaming/connectors/kinesis/testutils/TestUtils.java
index 7139d82e43f..c02fdb40f65 100644
---
a/flink-connectors/flink-connector-kinesis/src/test/java/org/apache/flink/streaming/connectors/kinesis/testutils/TestUtils.java
+++
b/flink-connectors/flink-connector-kinesis/src/test/java/org/apache/flink/streaming/connectors/kinesis/testutils/TestUtils.java
@@ -36,6 +36,7 @@ import org.apache.commons.lang3.RandomStringUtils;
import org.apache.commons.lang3.StringUtils;
import org.mockito.Mockito;
+import java.lang.reflect.Field;
import java.math.BigInteger;
import java.nio.ByteBuffer;
import java.util.ArrayList;
@@ -175,6 +176,19 @@ public class TestUtils {
return mockedRuntimeContext;
}
+ public static <T> T getField(String fieldName, Object obj) throws
Exception {
+ Field field = obj.getClass().getDeclaredField(fieldName);
+ field.setAccessible(true);
+ return (T) field.get(obj);
+ }
+
+ public static <T> T getField(String fieldName, Class<?> superclass, Object
obj)
+ throws Exception {
+ Field field = superclass.getDeclaredField(fieldName);
+ field.setAccessible(true);
+ return (T) field.get(obj);
+ }
+
/** A test record consumer used to capture messages from kinesis. */
public static class TestConsumer implements
RecordPublisher.RecordBatchConsumer {
private final List<RecordBatch> recordBatches = new ArrayList<>();
diff --git
a/flink-connectors/flink-connector-kinesis/src/test/java/org/apache/flink/streaming/connectors/kinesis/util/AWSUtilTest.java
b/flink-connectors/flink-connector-kinesis/src/test/java/org/apache/flink/streaming/connectors/kinesis/util/AWSUtilTest.java
index b9d5e6647c2..a8b1ba7c423 100644
---
a/flink-connectors/flink-connector-kinesis/src/test/java/org/apache/flink/streaming/connectors/kinesis/util/AWSUtilTest.java
+++
b/flink-connectors/flink-connector-kinesis/src/test/java/org/apache/flink/streaming/connectors/kinesis/util/AWSUtilTest.java
@@ -19,27 +19,38 @@ package org.apache.flink.streaming.connectors.kinesis.util;
import org.apache.flink.streaming.connectors.kinesis.config.AWSConfigConstants;
import org.apache.flink.streaming.connectors.kinesis.model.StartingPosition;
+import org.apache.flink.streaming.connectors.kinesis.testutils.TestUtils;
+import com.amazonaws.AmazonWebServiceClient;
import com.amazonaws.auth.AWSCredentials;
import com.amazonaws.auth.AWSCredentialsProvider;
import com.amazonaws.auth.DefaultAWSCredentialsProviderChain;
import com.amazonaws.auth.EnvironmentVariableCredentialsProvider;
+import com.amazonaws.auth.STSAssumeRoleSessionCredentialsProvider;
import com.amazonaws.auth.SystemPropertiesCredentialsProvider;
import com.amazonaws.auth.WebIdentityTokenCredentialsProvider;
import com.amazonaws.auth.profile.ProfileCredentialsProvider;
+import com.amazonaws.services.securitytoken.AWSSecurityTokenServiceClient;
import org.junit.Rule;
import org.junit.Test;
import org.junit.rules.ExpectedException;
import org.junit.runner.RunWith;
+import org.powermock.core.classloader.annotations.PowerMockIgnore;
import org.powermock.core.classloader.annotations.PrepareForTest;
import org.powermock.modules.junit4.PowerMockRunner;
+import java.net.URI;
import java.text.SimpleDateFormat;
import java.util.Date;
import java.util.Properties;
import static
com.amazonaws.services.kinesis.model.ShardIteratorType.AT_TIMESTAMP;
import static
org.apache.flink.connector.aws.config.AWSConfigConstants.AWS_CREDENTIALS_PROVIDER;
+import static
org.apache.flink.connector.aws.config.AWSConfigConstants.AWS_REGION;
+import static
org.apache.flink.connector.aws.config.AWSConfigConstants.AWS_ROLE_ARN;
+import static
org.apache.flink.connector.aws.config.AWSConfigConstants.AWS_ROLE_EXTERNAL_ID;
+import static
org.apache.flink.connector.aws.config.AWSConfigConstants.AWS_ROLE_SESSION_NAME;
+import static
org.apache.flink.connector.aws.config.AWSConfigConstants.AWS_ROLE_STS_ENDPOINT;
import static
org.apache.flink.streaming.connectors.kinesis.config.ConsumerConfigConstants.DEFAULT_STREAM_TIMESTAMP_DATE_FORMAT;
import static
org.apache.flink.streaming.connectors.kinesis.config.ConsumerConfigConstants.STREAM_INITIAL_TIMESTAMP;
import static
org.apache.flink.streaming.connectors.kinesis.model.SentinelSequenceNumber.SENTINEL_AT_TIMESTAMP_SEQUENCE_NUM;
@@ -49,6 +60,7 @@ import static org.assertj.core.api.Assertions.assertThat;
/** Tests for AWSUtil. */
@RunWith(PowerMockRunner.class)
@PrepareForTest(AWSUtil.class)
+@PowerMockIgnore({"javax.net.ssl.*", "javax.security.*"})
public class AWSUtilTest {
@Rule private final ExpectedException exception = ExpectedException.none();
@@ -163,6 +175,67 @@ public class AWSUtilTest {
.isEqualTo("wJalrXUtnFEMI/K7MDENG/bPxRfiCY2222222222");
}
+ @Test
+ public void testGetCredentialsProviderAssumeRole() throws Exception {
+ Properties testConfig = new Properties();
+ testConfig.setProperty(AWS_CREDENTIALS_PROVIDER, "ASSUME_ROLE");
+ testConfig.setProperty(AWS_REGION, "us-east-1");
+ testConfig.setProperty(AWS_ROLE_ARN, "arn");
+ testConfig.setProperty(AWS_ROLE_EXTERNAL_ID, "external_id");
+ testConfig.setProperty(AWS_ROLE_SESSION_NAME, "session_name");
+
+ AWSCredentialsProvider credentialsProvider =
AWSUtil.getCredentialsProvider(testConfig);
+
+
assertThat(credentialsProvider).isInstanceOf(STSAssumeRoleSessionCredentialsProvider.class);
+
+ STSAssumeRoleSessionCredentialsProvider assumeRoleCredentialsProvider =
+ (STSAssumeRoleSessionCredentialsProvider) credentialsProvider;
+
+ assertThat(TestUtils.<String>getField("roleArn",
assumeRoleCredentialsProvider))
+ .isEqualTo("arn");
+ assertThat(TestUtils.<String>getField("roleSessionName",
assumeRoleCredentialsProvider))
+ .isEqualTo("session_name");
+ assertThat(TestUtils.<String>getField("roleExternalId",
assumeRoleCredentialsProvider))
+ .isEqualTo("external_id");
+
+ AWSSecurityTokenServiceClient stsService =
+ TestUtils.getField("securityTokenService",
assumeRoleCredentialsProvider);
+
+ boolean isEndpointOverridden =
+ TestUtils.getField(
+ "isEndpointOverridden", AmazonWebServiceClient.class,
stsService);
+ assertThat(isEndpointOverridden).isEqualTo(false);
+ }
+
+ @Test
+ public void testGetCredentialsProviderAssumeRoleWithCustomStsEndpoint()
throws Exception {
+ Properties testConfig = new Properties();
+ testConfig.setProperty(AWS_CREDENTIALS_PROVIDER, "ASSUME_ROLE");
+ testConfig.setProperty(AWS_REGION, "us-east-1");
+ testConfig.setProperty(AWS_ROLE_ARN, "arn");
+ testConfig.setProperty(AWS_ROLE_EXTERNAL_ID, "external_id");
+ testConfig.setProperty(AWS_ROLE_SESSION_NAME, "session_name");
+ testConfig.setProperty(AWS_ROLE_STS_ENDPOINT,
"https://sts.us-east-1.amazonaws.com";);
+
+ AWSCredentialsProvider credentialsProvider =
AWSUtil.getCredentialsProvider(testConfig);
+
+
assertThat(credentialsProvider).isInstanceOf(STSAssumeRoleSessionCredentialsProvider.class);
+
+ STSAssumeRoleSessionCredentialsProvider assumeRoleCredentialsProvider =
+ (STSAssumeRoleSessionCredentialsProvider) credentialsProvider;
+
+ AWSSecurityTokenServiceClient stsService =
+ TestUtils.getField("securityTokenService",
assumeRoleCredentialsProvider);
+
+ boolean isEndpointOverridden =
+ TestUtils.getField(
+ "isEndpointOverridden", AmazonWebServiceClient.class,
stsService);
+ URI endpoint = TestUtils.getField("endpoint",
AmazonWebServiceClient.class, stsService);
+
+ assertThat(isEndpointOverridden).isEqualTo(true);
+
assertThat(endpoint).isEqualTo(URI.create("https://sts.us-east-1.amazonaws.com";));
+ }
+
@Test
public void testValidRegion() {
assertThat(AWSUtil.isValidRegion("us-east-1")).isTrue();
