Add sts_header_overrides to s3 dlq configuration (#3845)

Signed-off-by: Taylor Gray <[email protected]>
opensearch-project · Dec 12, 2023 · 882c8f4 · 882c8f4
1 parent f9d9e97
commit 882c8f4
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 0 deletions.
diff --git a/...res-common/src/main/java/org/opensearch/dataprepper/plugins/dlq/s3/S3DlqWriterConfig.java b/...res-common/src/main/java/org/opensearch/dataprepper/plugins/dlq/s3/S3DlqWriterConfig.java
@@ -19,6 +19,7 @@
 import software.amazon.awssdk.services.sts.auth.StsAssumeRoleCredentialsProvider;
 import software.amazon.awssdk.services.sts.model.AssumeRoleRequest;
 
+import java.util.Map;
 import java.util.Optional;
 import java.util.UUID;
 
@@ -56,6 +57,9 @@ public class S3DlqWriterConfig {
     @Size(min = 2, max = 1224, message = "sts_external_id length should be between 2 and 1224 characters")
     private String stsExternalId;
 
+    @JsonProperty("sts_header_overrides")
+    private Map<String, String> stsHeaderOverrides;
+
     public String getBucket() {
         if (bucket.startsWith(S3_PREFIX)) {
             return bucket.substring(S3_PREFIX.length());
@@ -91,6 +95,11 @@ private AwsCredentialsProvider getAwsCredentialsProvider() {
             assumeRoleRequestBuilder = assumeRoleRequestBuilder.externalId(stsExternalId);
         }
 
+        if(stsHeaderOverrides != null && !stsHeaderOverrides.isEmpty()) {
+            assumeRoleRequestBuilder = assumeRoleRequestBuilder
+                    .overrideConfiguration(configuration -> stsHeaderOverrides.forEach(configuration::putHeader));
+        }
+
         return StsAssumeRoleCredentialsProvider.builder()
             .stsClient(stsClient)
             .refreshRequest(assumeRoleRequestBuilder.build())

diff --git a/...common/src/test/java/org/opensearch/dataprepper/plugins/dlq/s3/S3DlqWriterConfigTest.java b/...common/src/test/java/org/opensearch/dataprepper/plugins/dlq/s3/S3DlqWriterConfigTest.java
@@ -14,6 +14,7 @@
 import software.amazon.awssdk.services.s3.S3Client;
 
 import java.lang.reflect.Field;
+import java.util.Map;
 import java.util.UUID;
 
 import static org.hamcrest.MatcherAssert.assertThat;
@@ -68,6 +69,7 @@ public void getS3ClientWithValidStsRoleArnAndExternalId(final String stsRoleArn)
         final S3DlqWriterConfig config = new S3DlqWriterConfig();
         reflectivelySetField(config, "stsRoleArn", stsRoleArn);
         reflectivelySetField(config, "stsExternalId", UUID.randomUUID().toString());
+        reflectivelySetField(config, "stsHeaderOverrides", Map.of(UUID.randomUUID().toString(), UUID.randomUUID().toString()));
         final S3Client s3Client = config.getS3Client();
         assertThat(s3Client, is(notNullValue()));
     }