S3 sink server-side encryption with KMS (#6655)

S3 sink server-side encryption with KMS

Adds new configuration for encryption options in the S3 sink. Allow configuring a custom KMS key for S3 server-side encryption. Support SSE-KMS and DSSE-KMS. Supports multi-part and locally buffered options.

Resolves #6528.

Signed-off-by: David Venable <dlv@amazon.com>

Author: dlvenable

data-prepper-plugins/common/src/main/java/org/opensearch/dataprepper/plugins/accumulator/BufferFactory.java

@@ -1,6 +1,10 @@
 /*
  * Copyright OpenSearch Contributors
  * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
  */
 
 package org.opensearch.dataprepper.plugins.accumulator;

data-prepper-plugins/s3-sink/build.gradle

@@ -1,6 +1,10 @@
 /*
  * Copyright OpenSearch Contributors
  * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
  */
 
 dependencies {
@@ -77,6 +81,7 @@ task integrationTest(type: Test) {
     systemProperty 'log4j.configurationFile', 'src/test/resources/log4j2.properties'
     systemProperty 'tests.s3sink.bucket', System.getProperty('tests.s3sink.bucket')
     systemProperty 'tests.s3sink.region', System.getProperty('tests.s3sink.region')
+    systemProperty 'tests.s3sink.kms_key', System.getProperty('tests.s3sink.kms_key')
 
     filter {
         includeTestsMatching '*IT'

data-prepper-plugins/s3-sink/src/integrationTest/java/org/opensearch/dataprepper/plugins/sink/s3/S3SinkIT.java

@@ -1,6 +1,10 @@
 /*
  * Copyright OpenSearch Contributors
  * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
  */
 
 package org.opensearch.dataprepper.plugins.sink.s3;
@@ -36,6 +40,8 @@
 import org.opensearch.dataprepper.plugins.sink.s3.configuration.AggregateThresholdOptions;
 import org.opensearch.dataprepper.plugins.sink.s3.configuration.AwsAuthenticationOptions;
 import org.opensearch.dataprepper.plugins.sink.s3.configuration.ObjectKeyOptions;
+import org.opensearch.dataprepper.plugins.sink.s3.configuration.ServerSideEncryptionConfig;
+import org.opensearch.dataprepper.plugins.sink.s3.configuration.ServerSideEncryptionType;
 import org.opensearch.dataprepper.plugins.sink.s3.configuration.ThresholdOptions;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -44,14 +50,18 @@
 import software.amazon.awssdk.regions.Region;
 import software.amazon.awssdk.services.s3.S3Client;
 import software.amazon.awssdk.services.s3.model.GetObjectRequest;
+import software.amazon.awssdk.services.s3.model.HeadObjectRequest;
+import software.amazon.awssdk.services.s3.model.HeadObjectResponse;
 import software.amazon.awssdk.services.s3.model.ListObjectsV2Request;
 import software.amazon.awssdk.services.s3.model.ListObjectsV2Response;
 import software.amazon.awssdk.services.s3.model.S3Object;
+import software.amazon.awssdk.services.s3.model.ServerSideEncryption;
 
 import java.io.File;
 import java.io.FileInputStream;
 import java.io.FileOutputStream;
 import java.io.IOException;
+import java.lang.reflect.Field;
 import java.time.Duration;
 import java.time.LocalDate;
 import java.time.LocalDateTime;
@@ -424,6 +434,57 @@ void testWithDynamicGroupsAndAggregateThreshold() throws IOException {
         }
     }
 
+    @ParameterizedTest
+    @ArgumentsSource(EncryptionArguments.class)
+    void test_server_side_encryption(final BufferTypeOptions bufferTypeOptions,
+                                     final ServerSideEncryptionConfig encryptionConfig,
+                                     final ServerSideEncryption expectedEncryption,
+                                     final String expectedKmsKeyId) {
+        final OutputScenario outputScenario = new NdjsonOutputScenario();
+        final String testRun = "encryption-" + bufferTypeOptions + "-" + expectedEncryption + "-" + UUID.randomUUID();
+        final String pathPrefix = pathPrefixForTestSuite + testRun;
+        when(objectKeyOptions.getPathPrefix()).thenReturn(pathPrefix + "/");
+
+        when(pluginFactory.loadPlugin(eq(OutputCodec.class), any())).thenReturn(outputScenario.getCodec());
+        when(s3SinkConfig.getBufferType()).thenReturn(bufferTypeOptions);
+        when(s3SinkConfig.getCompression()).thenReturn(new NoneCompressionScenario().getCompressionOption());
+        when(thresholdOptions.getEventCount()).thenReturn(1);
+
+        when(expressionEvaluator.extractDynamicExpressionsFromFormatExpression(anyString()))
+                .thenReturn(Collections.emptyList());
+        when(expressionEvaluator.extractDynamicKeysFromFormatExpression(anyString()))
+                .thenReturn(Collections.emptyList());
+
+        when(s3SinkConfig.getServerSideEncryptionConfig()).thenReturn(encryptionConfig);
+
+        final S3Sink objectUnderTest = createObjectUnderTest();
+
+        final List<Record<Event>> events = List.of(new Record<>(generateTestEvent(generateEventData(1))));
+        objectUnderTest.doOutput(events);
+
+        final ListObjectsV2Response listObjectsResponse = s3Client.listObjectsV2(ListObjectsV2Request.builder()
+                .bucket(bucketName)
+                .prefix(pathPrefix + "/")
+                .build());
+
+        assertThat(listObjectsResponse.contents().size(), equalTo(1));
+
+        final String objectKey = listObjectsResponse.contents().get(0).key();
+        final HeadObjectResponse headResponse = s3Client.headObject(HeadObjectRequest.builder()
+                .bucket(bucketName)
+                .key(objectKey)
+                .build());
+
+        assertThat(headResponse.serverSideEncryption(), equalTo(expectedEncryption));
+        assertThat(headResponse.ssekmsKeyId(), equalTo(expectedKmsKeyId));
+    }
+
+    private static void setField(final Object object, final String fieldName, final Object value) throws Exception {
+        final Field field = object.getClass().getDeclaredField(fieldName);
+        field.setAccessible(true);
+        field.set(object, value);
+    }
+
     private File decompressFileIfNecessary(OutputScenario outputScenario, CompressionScenario compressionScenario, String pathPrefix, File target) throws IOException {
 
         if (outputScenario.isCompressionInternal() || !compressionScenario.requiresDecompression())
@@ -552,6 +613,33 @@ public Stream<? extends Arguments> provideArguments(ExtensionContext extensionCo
         }
     }
 
+    static class EncryptionArguments implements ArgumentsProvider {
+        @Override
+        public Stream<? extends Arguments> provideArguments(final ExtensionContext context) throws Exception {
+            final String kmsKeyId = System.getProperty("tests.s3sink.kms_key");
+
+            final ServerSideEncryptionConfig s3Config = new ServerSideEncryptionConfig();
+            setField(s3Config, "type", ServerSideEncryptionType.S3);
+
+            final ServerSideEncryptionConfig kmsConfig = new ServerSideEncryptionConfig();
+            setField(kmsConfig, "type", ServerSideEncryptionType.KMS);
+            setField(kmsConfig, "kmsKeyId", kmsKeyId);
+
+            final ServerSideEncryptionConfig dsseConfig = new ServerSideEncryptionConfig();
+            setField(dsseConfig, "type", ServerSideEncryptionType.KMS_DSSE);
+            setField(dsseConfig, "kmsKeyId", kmsKeyId);
+
+            final List<BufferTypeOptions> bufferTypes = List.of(BufferTypeOptions.INMEMORY, BufferTypeOptions.MULTI_PART);
+
+            return bufferTypes.stream().flatMap(bufferType -> Stream.of(
+                    arguments(bufferType, null, ServerSideEncryption.AES256, null),
+                    arguments(bufferType, s3Config, ServerSideEncryption.AES256, null),
+                    arguments(bufferType, kmsConfig, ServerSideEncryption.AWS_KMS, kmsKeyId),
+                    arguments(bufferType, dsseConfig, ServerSideEncryption.AWS_KMS_DSSE, kmsKeyId)
+            ));
+        }
+    }
+
     private static Stream<? extends Arguments> generateCombinedArguments(
             final List<BufferScenario> bufferScenarios,
             final List<OutputScenario> outputScenarios,

data-prepper-plugins/s3-sink/src/main/java/org/opensearch/dataprepper/plugins/codec/parquet/S3OutputStream.java

@@ -1,11 +1,16 @@
 /*
  * Copyright OpenSearch Contributors
  * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
  */
 package org.opensearch.dataprepper.plugins.codec.parquet;
 
 
 import org.apache.parquet.io.PositionOutputStream;
+import org.opensearch.dataprepper.plugins.sink.s3.configuration.ServerSideEncryptionConfig;
 import org.opensearch.dataprepper.plugins.sink.s3.ownership.BucketOwnerProvider;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -87,6 +92,8 @@ public class S3OutputStream extends PositionOutputStream {
 
     private final ExecutorService executorService;
 
+    private final ServerSideEncryptionConfig serverSideEncryptionConfig;
+
     /**
      * Creates a new S3 OutputStream
      *
@@ -95,12 +102,14 @@ public class S3OutputStream extends PositionOutputStream {
      * @param keySupplier     path within the bucket
      * @param defaultBucket default bucket
      * @param bucketOwnerProvider bucket owner provider
+     * @param serverSideEncryptionConfig server-side encryption config
      */
     public S3OutputStream(final S3AsyncClient s3Client,
                           final Supplier<String> bucketSupplier,
                           final Supplier<String> keySupplier,
                           final String defaultBucket,
-                          final BucketOwnerProvider bucketOwnerProvider) {
+                          final BucketOwnerProvider bucketOwnerProvider,
+                          final ServerSideEncryptionConfig serverSideEncryptionConfig) {
         this.s3Client = s3Client;
         this.bucket = bucketSupplier.get();
         this.key = keySupplier.get();
@@ -111,6 +120,15 @@ public S3OutputStream(final S3AsyncClient s3Client,
         this.defaultBucket = defaultBucket;
         this.executorService = Executors.newSingleThreadExecutor();
         this.bucketOwnerProvider = bucketOwnerProvider;
+        this.serverSideEncryptionConfig = serverSideEncryptionConfig;
+    }
+
+    public S3OutputStream(final S3AsyncClient s3Client,
+                          final Supplier<String> bucketSupplier,
+                          final Supplier<String> keySupplier,
+                          final String defaultBucket,
+                          final BucketOwnerProvider bucketOwnerProvider) {
+        this(s3Client, bucketSupplier, keySupplier, defaultBucket, bucketOwnerProvider, null);
     }
 
     @Override
@@ -285,12 +303,14 @@ public long getPos() throws IOException {
     }
 
     private void createMultipartUpload() {
-        CreateMultipartUploadRequest uploadRequest = CreateMultipartUploadRequest.builder()
+        CreateMultipartUploadRequest.Builder builder = CreateMultipartUploadRequest.builder()
                 .bucket(bucket)
                 .key(key)
-                .expectedBucketOwner(bucketOwnerProvider.getBucketOwner(bucket).orElse(null))
-                .build();
-        CompletableFuture<CreateMultipartUploadResponse> multipartUpload = s3Client.createMultipartUpload(uploadRequest);
+                .expectedBucketOwner(bucketOwnerProvider.getBucketOwner(bucket).orElse(null));
+        if (serverSideEncryptionConfig != null) {
+            serverSideEncryptionConfig.applyTo(builder);
+        }
+        CompletableFuture<CreateMultipartUploadResponse> multipartUpload = s3Client.createMultipartUpload(builder.build());
 
         final CreateMultipartUploadResponse response = multipartUpload.join();
 

data-prepper-plugins/s3-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/s3/S3SinkConfig.java

@@ -1,6 +1,10 @@
 /*
  * Copyright OpenSearch Contributors
  * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
  */
 
 package org.opensearch.dataprepper.plugins.sink.s3;
@@ -18,6 +22,7 @@
 import org.opensearch.dataprepper.plugins.sink.s3.configuration.AwsAuthenticationOptions;
 import org.opensearch.dataprepper.plugins.sink.s3.configuration.ClientOptions;
 import org.opensearch.dataprepper.plugins.sink.s3.configuration.ObjectKeyOptions;
+import org.opensearch.dataprepper.plugins.sink.s3.configuration.ServerSideEncryptionConfig;
 import org.opensearch.dataprepper.plugins.sink.s3.configuration.ThresholdOptions;
 
 import java.util.Map;
@@ -108,6 +113,10 @@ private boolean isValidBucketConfig() {
     @JsonProperty("client")
     private ClientOptions clientOptions;
 
+    @JsonProperty("server_side_encryption")
+    @Valid
+    private ServerSideEncryptionConfig serverSideEncryptionConfig;
+
     /**
      * Aws Authentication configuration Options.
      * @return aws authentication options.
@@ -213,4 +222,8 @@ public String getDefaultBucketOwner() {
     public ClientOptions getClientOptions() {
         return clientOptions;
     }
+
+    public ServerSideEncryptionConfig getServerSideEncryptionConfig() {
+        return serverSideEncryptionConfig;
+    }
 }

data-prepper-plugins/s3-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/s3/accumulator/BufferFactory.java

@@ -1,10 +1,15 @@
 /*
  * Copyright OpenSearch Contributors
  * SPDX-License-Identifier: Apache-2.0
+ *
+ *  The OpenSearch Contributors require contributions made to
+ *  this file be licensed under the Apache-2.0 license or a
+ *  compatible open source license.
  */
 
 package org.opensearch.dataprepper.plugins.sink.s3.accumulator;
 
+import org.opensearch.dataprepper.plugins.sink.s3.configuration.ServerSideEncryptionConfig;
 import org.opensearch.dataprepper.plugins.sink.s3.ownership.BucketOwnerProvider;
 import software.amazon.awssdk.services.s3.S3AsyncClient;
 
@@ -13,5 +18,11 @@
 import java.util.function.Function;
 
 public interface BufferFactory {
-    Buffer getBuffer(S3AsyncClient s3Client, Supplier<String> bucketSupplier, Supplier<String> keySupplier, String defaultBucket, Function<Integer, Map<String, String>> metadataSupplier, BucketOwnerProvider bucketOwnerProvider);
+    Buffer getBuffer(S3AsyncClient s3Client,
+                     Supplier<String> bucketSupplier,
+                     Supplier<String> keySupplier,
+                     String defaultBucket,
+                     Function<Integer, Map<String, String>> metadataSupplier,
+                     BucketOwnerProvider bucketOwnerProvider,
+                     ServerSideEncryptionConfig serverSideEncryptionConfig);
 }

data-prepper-plugins/s3-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/s3/accumulator/BufferUtilities.java

@@ -1,10 +1,15 @@
 /*
  * Copyright OpenSearch Contributors
  * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
  */
 
 package org.opensearch.dataprepper.plugins.sink.s3.accumulator;
 
+import org.opensearch.dataprepper.plugins.sink.s3.configuration.ServerSideEncryptionConfig;
 import org.opensearch.dataprepper.plugins.sink.s3.ownership.BucketOwnerProvider;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -34,17 +39,12 @@ static CompletableFuture<PutObjectResponse> putObjectOrSendToDefaultBucket(final
                                                                                final String targetBucket,
                                                                                final String defaultBucket,
                                                                                final Map<String, String> objectMetadata,
-                                                                               final BucketOwnerProvider bucketOwnerProvider) {
+                                                                               final BucketOwnerProvider bucketOwnerProvider,
+                                                                               final ServerSideEncryptionConfig serverSideEncryptionConfig) {
 
         final boolean[] defaultBucketAttempted = new boolean[1];
-        PutObjectRequest.Builder builder =  PutObjectRequest.builder()
-                .bucket(targetBucket)
-                .key(objectKey)
-                .expectedBucketOwner(bucketOwnerProvider.getBucketOwner(targetBucket).orElse(null));
-        if (objectMetadata != null) {
-            builder = builder.metadata(objectMetadata);
-        }
-        return s3Client.putObject(builder.build(), requestBody)
+        final PutObjectRequest putObjectRequest = buildPutObjectRequest(targetBucket, objectKey, objectMetadata, bucketOwnerProvider, serverSideEncryptionConfig);
+        return s3Client.putObject(putObjectRequest, requestBody)
                 .handle((result, ex) -> {
                     if (ex != null) {
                         runOnFailure.accept(ex);
@@ -53,13 +53,8 @@ static CompletableFuture<PutObjectResponse> putObjectOrSendToDefaultBucket(final
                                 (ex instanceof NoSuchBucketException || ex.getCause() instanceof NoSuchBucketException || ex.getMessage().contains(ACCESS_DENIED) || ex.getMessage().contains(INVALID_BUCKET))) {
                             LOG.warn("Bucket {} could not be accessed, attempting to send to default_bucket {}", targetBucket, defaultBucket);
                             defaultBucketAttempted[0] = true;
-                            return s3Client.putObject(
-                                    PutObjectRequest.builder()
-                                            .bucket(defaultBucket)
-                                            .key(objectKey)
-                                            .expectedBucketOwner(bucketOwnerProvider.getBucketOwner(defaultBucket).orElse(null))
-                                            .build(),
-                                    requestBody);
+                            final PutObjectRequest defaultPutObjectRequest = buildPutObjectRequest(defaultBucket, objectKey, null, bucketOwnerProvider, serverSideEncryptionConfig);
+                            return s3Client.putObject(defaultPutObjectRequest, requestBody);
                         } else {
                             runOnCompletion.accept(false);
                             return CompletableFuture.completedFuture(result);
@@ -80,4 +75,22 @@ static CompletableFuture<PutObjectResponse> putObjectOrSendToDefaultBucket(final
                     }
                 });
     }
+
+    private static PutObjectRequest buildPutObjectRequest(final String bucket,
+                                                          final String objectKey,
+                                                          final Map<String, String> objectMetadata,
+                                                          final BucketOwnerProvider bucketOwnerProvider,
+                                                          final ServerSideEncryptionConfig serverSideEncryptionConfig) {
+        final PutObjectRequest.Builder builder = PutObjectRequest.builder()
+                .bucket(bucket)
+                .key(objectKey)
+                .expectedBucketOwner(bucketOwnerProvider.getBucketOwner(bucket).orElse(null));
+        if (objectMetadata != null) {
+            builder.metadata(objectMetadata);
+        }
+        if (serverSideEncryptionConfig != null) {
+            serverSideEncryptionConfig.applyTo(builder);
+        }
+        return builder.build();
+    }
 }