Skip to content

Commit 5fa5451

Browse files
kkondakakkondaka
authored
Add support for named configurations in aws config block (opensearch-project#6905)
Signed-off-by: Kondaka <krishkdk@amazon.com>
1 parent f853bea commit 5fa5451

38 files changed

Lines changed: 399 additions & 52 deletions

File tree

‎data-prepper-plugins/aws-lambda/src/integrationTest/java/org/opensearch/dataprepper/plugins/lambda/LambdaProcessorSinkIT.java‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121
import org.opensearch.dataprepper.model.types.ByteCount;
2222

2323
import org.opensearch.dataprepper.aws.api.AwsCredentialsSupplier;
24+
import org.opensearch.dataprepper.aws.api.AwsCredentialsOptions;
2425
import org.opensearch.dataprepper.expression.ExpressionEvaluator;
2526
import org.opensearch.dataprepper.model.configuration.PluginSetting;
2627
import org.opensearch.dataprepper.plugins.lambda.common.config.InvocationType;
@@ -188,7 +189,7 @@ public void setup() {
188189
lambdaProcessorConfig = mock(LambdaProcessorConfig.class);
189190
expressionEvaluator = mock(ExpressionEvaluator.class);
190191
awsCredentialsProvider = DefaultCredentialsProvider.create();
191-
when(awsCredentialsSupplier.getProvider(any())).thenReturn(awsCredentialsProvider);
192+
when(awsCredentialsSupplier.getProvider(any(AwsCredentialsOptions.class))).thenReturn(awsCredentialsProvider);
192193
pluginFactory = mock(PluginFactory.class);
193194
JsonInputCodecConfig jsonInputCodecConfig = mock(JsonInputCodecConfig.class);
194195
when(jsonInputCodecConfig.getKeyName()).thenReturn(null);

‎data-prepper-plugins/aws-lambda/src/integrationTest/java/org/opensearch/dataprepper/plugins/lambda/processor/LambdaProcessorIT.java‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,7 @@
3838
import org.mockito.junit.jupiter.MockitoSettings;
3939
import org.mockito.quality.Strictness;
4040
import org.opensearch.dataprepper.aws.api.AwsCredentialsSupplier;
41+
import org.opensearch.dataprepper.aws.api.AwsCredentialsOptions;
4142
import org.opensearch.dataprepper.expression.ExpressionEvaluator;
4243
import org.opensearch.dataprepper.metrics.PluginMetrics;
4344
import org.opensearch.dataprepper.model.breaker.CircuitBreaker;
@@ -145,7 +146,7 @@ public void setup() {
145146
lambdaProcessorConfig = mock(LambdaProcessorConfig.class);
146147
expressionEvaluator = mock(ExpressionEvaluator.class);
147148
awsCredentialsProvider = DefaultCredentialsProvider.create();
148-
when(awsCredentialsSupplier.getProvider(any())).thenReturn(awsCredentialsProvider);
149+
when(awsCredentialsSupplier.getProvider(any(AwsCredentialsOptions.class))).thenReturn(awsCredentialsProvider);
149150
pluginFactory = mock(PluginFactory.class);
150151
JsonInputCodecConfig jsonInputCodecConfig = mock(JsonInputCodecConfig.class);
151152
when(jsonInputCodecConfig.getKeyName()).thenReturn(null);
@@ -476,7 +477,7 @@ def lambda_handler(event, context):
476477
when(lambdaProcessorConfig.getAwsAuthenticationOptions()).thenReturn(awsAuthenticationOptions);
477478

478479
// Setup the mock for getProvider
479-
when(awsCredentialsSupplier.getProvider(any())).thenReturn(awsCredentialsProvider);
480+
when(awsCredentialsSupplier.getProvider(any(AwsCredentialsOptions.class))).thenReturn(awsCredentialsProvider);
480481

481482
// Mock the factory to inject our CountingRetryCondition into the LambdaAsyncClient
482483
try (MockedStatic<LambdaClientFactory> mockedFactory = mockStatic(LambdaClientFactory.class)) {

‎data-prepper-plugins/aws-plugin-api/src/main/java/org/opensearch/dataprepper/aws/api/AwsConfig.java‎

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@
1111

1212
import com.fasterxml.jackson.annotation.JsonProperty;
1313
import jakarta.validation.constraints.Size;
14+
import jakarta.validation.constraints.AssertTrue;
1415
import software.amazon.awssdk.regions.Region;
1516

1617
import java.util.Map;
@@ -33,6 +34,10 @@ public class AwsConfig {
3334
@Size(max = 5, message = "sts_header_overrides supports a maximum of 5 headers to override")
3435
private Map<String, String> awsStsHeaderOverrides;
3536

37+
@JsonProperty("configuration")
38+
private String configuration;
39+
40+
3641
@JsonProperty("sts_external_id")
3742
@Size(min = 2, max = 1224, message = "awsStsExternalId length should be between 2 and 1224 characters")
3843
private String awsStsExternalId;
@@ -52,4 +57,18 @@ public String getAwsStsExternalId() {
5257
public Map<String, String> getAwsStsHeaderOverrides() {
5358
return awsStsHeaderOverrides;
5459
}
60+
61+
@AssertTrue(message = "'configuration' cannot be used together with 'region', 'sts_role_arn', 'sts_header_overrides', or 'sts_external_id'. " +
62+
"Use either a named configuration reference or inline credentials, not both.")
63+
public boolean isValidConfiguration() {
64+
if (configuration != null) {
65+
return awsRegion == null && awsStsRoleArn == null && awsStsHeaderOverrides == null && awsStsExternalId == null;
66+
}
67+
return true;
68+
}
69+
70+
public String getConfiguration() {
71+
return configuration;
72+
}
73+
5574
}

‎data-prepper-plugins/aws-plugin-api/src/main/java/org/opensearch/dataprepper/aws/api/AwsCredentialsSupplier.java‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,14 @@ public interface AwsCredentialsSupplier {
2727
*/
2828
AwsCredentialsProvider getProvider(AwsCredentialsOptions options);
2929

30+
/**
31+
* Gets an AWS SDK {@link AwsCredentialsProvider} resolved from a named configuration
32+
* defined in the extensions.aws.configurations block.
33+
* @param configurationName The name of the configuration (e.g. "ecs_task_role")
34+
* @return An {@link AwsCredentialsProvider} to use.
35+
*/
36+
AwsCredentialsProvider getProvider(String configurationName);
37+
3038
/**
3139
* Gets the default region if it is configured. Otherwise returns null
3240
* @return Default {@link Region}

‎data-prepper-plugins/aws-plugin-api/src/test/java/org/opensearch/dataprepper/aws/api/AwsConfigTest.java‎

Lines changed: 58 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -49,6 +49,64 @@ void TestConfigOptions_notNull() throws NoSuchFieldException, IllegalAccessExcep
4949
assertThat(awsConfig.getAwsRegion(), equalTo(Region.of(testRegion)));
5050
}
5151

52+
@Test
53+
void TestConfigOptions_configuration_returns_value() throws NoSuchFieldException, IllegalAccessException {
54+
final String configName = "ecs_task_role";
55+
reflectivelySetField(awsConfig, "configuration", configName);
56+
assertThat(awsConfig.getConfiguration(), equalTo(configName));
57+
}
58+
59+
@Test
60+
void TestConfigOptions_configuration_returns_null_by_default() {
61+
assertThat(awsConfig.getConfiguration(), equalTo(null));
62+
}
63+
64+
@Test
65+
void isValidConfiguration_returns_true_when_configuration_is_null() {
66+
assertThat(awsConfig.isValidConfiguration(), equalTo(true));
67+
}
68+
69+
@Test
70+
void isValidConfiguration_returns_true_when_only_configuration_is_set() throws NoSuchFieldException, IllegalAccessException {
71+
reflectivelySetField(awsConfig, "configuration", "my_config");
72+
assertThat(awsConfig.isValidConfiguration(), equalTo(true));
73+
}
74+
75+
@Test
76+
void isValidConfiguration_returns_false_when_configuration_and_region_are_set() throws NoSuchFieldException, IllegalAccessException {
77+
reflectivelySetField(awsConfig, "configuration", "my_config");
78+
reflectivelySetField(awsConfig, "awsRegion", "us-east-1");
79+
assertThat(awsConfig.isValidConfiguration(), equalTo(false));
80+
}
81+
82+
@Test
83+
void isValidConfiguration_returns_false_when_configuration_and_sts_role_arn_are_set() throws NoSuchFieldException, IllegalAccessException {
84+
reflectivelySetField(awsConfig, "configuration", "my_config");
85+
reflectivelySetField(awsConfig, "awsStsRoleArn", "arn:aws:iam::123456789012:role/TestRole");
86+
assertThat(awsConfig.isValidConfiguration(), equalTo(false));
87+
}
88+
89+
@Test
90+
void isValidConfiguration_returns_false_when_configuration_and_sts_header_overrides_are_set() throws NoSuchFieldException, IllegalAccessException {
91+
reflectivelySetField(awsConfig, "configuration", "my_config");
92+
reflectivelySetField(awsConfig, "awsStsHeaderOverrides", Map.of("key", "value"));
93+
assertThat(awsConfig.isValidConfiguration(), equalTo(false));
94+
}
95+
96+
@Test
97+
void isValidConfiguration_returns_false_when_configuration_and_sts_external_id_are_set() throws NoSuchFieldException, IllegalAccessException {
98+
reflectivelySetField(awsConfig, "configuration", "my_config");
99+
reflectivelySetField(awsConfig, "awsStsExternalId", "ext-id-123");
100+
assertThat(awsConfig.isValidConfiguration(), equalTo(false));
101+
}
102+
103+
@Test
104+
void isValidConfiguration_returns_true_when_inline_credentials_used_without_configuration() throws NoSuchFieldException, IllegalAccessException {
105+
reflectivelySetField(awsConfig, "awsRegion", "us-east-1");
106+
reflectivelySetField(awsConfig, "awsStsRoleArn", "arn:aws:iam::123456789012:role/TestRole");
107+
assertThat(awsConfig.isValidConfiguration(), equalTo(true));
108+
}
109+
52110
private void reflectivelySetField(final AwsConfig awsConfig, final String fieldName, final Object value) throws NoSuchFieldException, IllegalAccessException {
53111
final Field field = AwsConfig.class.getDeclaredField(fieldName);
54112
try {

‎data-prepper-plugins/aws-plugin/src/main/java/org/opensearch/dataprepper/plugins/aws/AwsPlugin.java‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -32,7 +32,9 @@ public AwsPlugin(final AwsPluginConfig awsPluginConfig) {
3232

3333
this.awsPluginConfig = awsPluginConfig;
3434

35-
final CredentialsProviderFactory credentialsProviderFactory = new CredentialsProviderFactory(awsPluginConfig != null ? awsPluginConfig.getDefaultStsConfiguration() : new AwsStsConfiguration());
35+
final CredentialsProviderFactory credentialsProviderFactory = new CredentialsProviderFactory(
36+
awsPluginConfig != null ? awsPluginConfig.getDefaultStsConfiguration() : new AwsStsConfiguration(),
37+
awsPluginConfig);
3638
final CredentialsCache credentialsCache = new CredentialsCache();
3739
defaultAwsCredentialsSupplier = new DefaultAwsCredentialsSupplier(credentialsProviderFactory, credentialsCache);
3840
}

‎data-prepper-plugins/aws-plugin/src/main/java/org/opensearch/dataprepper/plugins/aws/AwsPluginConfig.java‎

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,8 +9,14 @@
99

1010
package org.opensearch.dataprepper.plugins.aws;
1111

12+
import com.fasterxml.jackson.annotation.JsonAnyGetter;
13+
import com.fasterxml.jackson.annotation.JsonAnySetter;
1214
import com.fasterxml.jackson.annotation.JsonProperty;
1315

16+
import java.util.Collection;
17+
import java.util.HashMap;
18+
import java.util.Map;
19+
1420
public class AwsPluginConfig {
1521

1622
@JsonProperty("default")
@@ -19,4 +25,24 @@ public class AwsPluginConfig {
1925
public AwsStsConfiguration getDefaultStsConfiguration() {
2026
return defaultStsConfiguration;
2127
}
28+
29+
private Map<String, AwsStsConfiguration> allOtherConfigurations = new HashMap<>();
30+
31+
@JsonAnyGetter
32+
public Map<String, AwsStsConfiguration> getAllOtherConfigurations() {
33+
return allOtherConfigurations;
34+
}
35+
36+
@JsonAnySetter
37+
public void setAdditionalConfiguration(String name, AwsStsConfiguration configuration) {
38+
allOtherConfigurations.put(name, configuration);
39+
}
40+
41+
public Collection<String> listNonDefaultConfigurations() {
42+
return allOtherConfigurations.keySet();
43+
}
44+
45+
public AwsStsConfiguration getConfiguration(String name) {
46+
return allOtherConfigurations.get(name);
47+
}
2248
}

‎data-prepper-plugins/aws-plugin/src/main/java/org/opensearch/dataprepper/plugins/aws/AwsStsConfiguration.java‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,10 @@ public class AwsStsConfiguration {
2828
@Size(max = 5, message = "sts_header_overrides supports a maximum of 5 headers to override")
2929
private Map<String, String> awsStsHeaderOverrides;
3030

31+
@JsonProperty("use_aws_sdk_default")
32+
private boolean useAwsSdkDefault = false;
33+
34+
3135
public Region getAwsRegion() {
3236
return awsRegion != null ? Region.of(awsRegion) : null;
3337
}
@@ -39,4 +43,8 @@ public String getAwsStsRoleArn() {
3943
public Map<String, String> getStsHeaderOverrides() {
4044
return awsStsHeaderOverrides;
4145
}
46+
47+
public boolean isUseAwsSdkDefault() {
48+
return useAwsSdkDefault;
49+
}
4250
}

‎data-prepper-plugins/aws-plugin/src/main/java/org/opensearch/dataprepper/plugins/aws/CredentialsProviderFactory.java‎

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -41,10 +41,42 @@ class CredentialsProviderFactory {
4141
static final long STS_CLIENT_MAX_BACKOFF_MILLIS = 60000L;
4242

4343
private final AwsStsConfiguration defaultStsConfiguration;
44+
private final AwsPluginConfig awsPluginConfig;
4445

4546
public CredentialsProviderFactory(final AwsStsConfiguration defaultStsConfiguration) {
47+
this(defaultStsConfiguration, null);
48+
}
49+
50+
public CredentialsProviderFactory(final AwsStsConfiguration defaultStsConfiguration, final AwsPluginConfig awsPluginConfig) {
4651
Objects.requireNonNull(defaultStsConfiguration);
4752
this.defaultStsConfiguration = defaultStsConfiguration;
53+
this.awsPluginConfig = awsPluginConfig;
54+
}
55+
56+
AwsCredentialsOptions resolveNamedConfiguration(final String configurationName) {
57+
if (awsPluginConfig == null) {
58+
throw new IllegalArgumentException("Named configuration '" + configurationName +
59+
"' referenced but no aws extensions configured");
60+
}
61+
final AwsStsConfiguration namedConfig = awsPluginConfig.getConfiguration(configurationName);
62+
if (namedConfig == null) {
63+
throw new IllegalArgumentException("Named configuration '" + configurationName +
64+
"' not found in extensions.aws.configurations");
65+
}
66+
if (namedConfig.isUseAwsSdkDefault()) {
67+
return AwsCredentialsOptions.defaultOptionsWithDefaultCredentialsProvider();
68+
}
69+
final AwsCredentialsOptions.Builder builder = AwsCredentialsOptions.builder();
70+
if (namedConfig.getAwsRegion() != null) {
71+
builder.withRegion(namedConfig.getAwsRegion());
72+
}
73+
if (namedConfig.getAwsStsRoleArn() != null) {
74+
builder.withStsRoleArn(namedConfig.getAwsStsRoleArn());
75+
}
76+
if (namedConfig.getStsHeaderOverrides() != null) {
77+
builder.withStsHeaderOverrides(namedConfig.getStsHeaderOverrides());
78+
}
79+
return builder.build();
4880
}
4981

5082
Region getDefaultRegion() {

‎data-prepper-plugins/aws-plugin/src/main/java/org/opensearch/dataprepper/plugins/aws/DefaultAwsCredentialsSupplier.java‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,12 @@ public AwsCredentialsProvider getProvider(final AwsCredentialsOptions options) {
3131
return credentialsCache.getOrCreate(options, () -> credentialsProviderFactory.providerFromOptions(options));
3232
}
3333

34+
@Override
35+
public AwsCredentialsProvider getProvider(final String configurationName) {
36+
final AwsCredentialsOptions options = credentialsProviderFactory.resolveNamedConfiguration(configurationName);
37+
return getProvider(options);
38+
}
39+
3440
@Override
3541
public Optional<Region> getDefaultRegion() {
3642
return Optional.ofNullable(credentialsProviderFactory.getDefaultRegion());

0 commit comments

Comments
 (0)