FEAT: encryption extension integration with kafka (#5625)

* FEAT: encryption extension integration with kafka buffer

Signed-off-by: George Chen <qchea@amazon.com>

Author: chenqi0805

data-prepper-plugins/encryption-plugin/src/main/java/org/opensearch/dataprepper/plugins/encryption/EncryptionPlugin.java

@@ -22,7 +22,7 @@
 import java.util.concurrent.TimeUnit;
 import java.util.stream.Collectors;
 
-@DataPrepperExtensionPlugin(modelType = EncryptionPluginConfig.class, rootKeyJsonPath = "/encryption")
+@DataPrepperExtensionPlugin(modelType = EncryptionPluginConfig.class, rootKeyJsonPath = "/encryption", allowInPipelineConfigurations = true)
 public class EncryptionPlugin implements ExtensionPlugin {
     static final int PERIOD_IN_SECONDS = 60;
     private static final Logger LOG = LoggerFactory.getLogger(EncryptionPlugin.class);

data-prepper-plugins/kafka-plugins/build.gradle

@@ -29,6 +29,7 @@ dependencies {
     implementation project(':data-prepper-plugins:buffer-common')
     implementation project(':data-prepper-plugins:blocking-buffer')
     implementation project(':data-prepper-plugins:aws-plugin-api')
+    implementation project(':data-prepper-plugins:encryption-plugin')
     // bump io.confluent:* dependencies correspondingly when bumping org.apache.kafka.*
     // https://docs.confluent.io/platform/current/release-notes/index.html
     implementation 'org.apache.kafka:kafka-clients:3.6.1'

data-prepper-plugins/kafka-plugins/src/integrationTest/java/org/opensearch/dataprepper/plugins/kafka/buffer/KafkaBufferIT.java

@@ -26,6 +26,7 @@
 import org.opensearch.dataprepper.model.event.Event;
 import org.opensearch.dataprepper.model.event.JacksonEvent;
 import org.opensearch.dataprepper.model.record.Record;
+import org.opensearch.dataprepper.plugins.encryption.EncryptionSupplier;
 import org.opensearch.dataprepper.plugins.kafka.configuration.KafkaProducerProperties;
 import org.opensearch.dataprepper.plugins.kafka.util.TestConsumer;
 import org.opensearch.dataprepper.plugins.kafka.util.TestProducer;
@@ -76,6 +77,8 @@ public class KafkaBufferIT {
     private AcknowledgementSetManager acknowledgementSetManager;
     @Mock
     private AcknowledgementSet acknowledgementSet;
+    @Mock
+    private EncryptionSupplier encryptionSupplier;
 
     private Random random;
 
@@ -123,11 +126,11 @@ void setUp() {
     }
 
     private KafkaBuffer createObjectUnderTestWithJsonDecoder() {
-        return new KafkaBuffer(pluginSetting, kafkaBufferConfig, acknowledgementSetManager, new JsonDecoder(), null, null);
+        return new KafkaBuffer(pluginSetting, kafkaBufferConfig, acknowledgementSetManager, new JsonDecoder(), null, null, encryptionSupplier);
     }
 
     private KafkaBuffer createObjectUnderTest() {
-        return new KafkaBuffer(pluginSetting, kafkaBufferConfig, acknowledgementSetManager, null, null, null);
+        return new KafkaBuffer(pluginSetting, kafkaBufferConfig, acknowledgementSetManager, null, null, null, encryptionSupplier);
     }
 
     @Test

data-prepper-plugins/kafka-plugins/src/integrationTest/java/org/opensearch/dataprepper/plugins/kafka/buffer/KafkaBufferOTelIT.java

@@ -28,6 +28,7 @@
 import org.opensearch.dataprepper.model.metric.JacksonHistogram;
 import org.opensearch.dataprepper.model.log.OpenTelemetryLog;
 import org.opensearch.dataprepper.model.trace.Span;
+import org.opensearch.dataprepper.plugins.encryption.EncryptionSupplier;
 import org.opensearch.dataprepper.plugins.kafka.configuration.EncryptionConfig;
 import org.opensearch.dataprepper.model.buffer.DelegatingBuffer;
 import org.opensearch.dataprepper.model.buffer.Buffer;
@@ -93,6 +94,8 @@ public class KafkaBufferOTelIT {
     private AcknowledgementSetManager acknowledgementSetManager;
     @Mock
     private BufferTopicConfig topicConfig;
+    @Mock
+    private EncryptionSupplier encryptionSupplier;
 
     private DelegatingBuffer buffer;
 
@@ -294,7 +297,7 @@ private void validateMetric(Event event) {
 
     @Test
     void test_otel_metrics_with_kafka_buffer() throws Exception {
-        KafkaBuffer kafkaBuffer = new KafkaBuffer(pluginSetting, kafkaBufferConfig, acknowledgementSetManager, new OTelMetricDecoder(OTelOutputFormat.OPENSEARCH), null, null);
+        KafkaBuffer kafkaBuffer = new KafkaBuffer(pluginSetting, kafkaBufferConfig, acknowledgementSetManager, new OTelMetricDecoder(OTelOutputFormat.OPENSEARCH), null, null, encryptionSupplier);
         buffer = new KafkaDelegatingBuffer(kafkaBuffer);
         final ExportMetricsServiceRequest request = createExportMetricsServiceRequest();
         buffer.writeBytes(request.toByteArray(), null, 10_000);
@@ -366,7 +369,7 @@ private void validateLog(OpenTelemetryLog logRecord) throws Exception {
 
     @Test
     void test_otel_logs_with_kafka_buffer() throws Exception {
-        KafkaBuffer kafkaBuffer = new KafkaBuffer(pluginSetting, kafkaBufferConfig, acknowledgementSetManager, new OTelLogsDecoder(OTelOutputFormat.OPENSEARCH), null, null);
+        KafkaBuffer kafkaBuffer = new KafkaBuffer(pluginSetting, kafkaBufferConfig, acknowledgementSetManager, new OTelLogsDecoder(OTelOutputFormat.OPENSEARCH), null, null, encryptionSupplier);
         buffer = new KafkaDelegatingBuffer(kafkaBuffer);
         final ExportLogsServiceRequest request = createExportLogsRequest();
         buffer.writeBytes(request.toByteArray(), null, 10_000);
@@ -437,7 +440,7 @@ private void validateSpan(Span span) throws Exception {
 
     @Test
     void test_otel_traces_with_kafka_buffer() throws Exception {
-        KafkaBuffer kafkaBuffer = new KafkaBuffer(pluginSetting, kafkaBufferConfig, acknowledgementSetManager, new OTelTraceDecoder(OTelOutputFormat.OPENSEARCH), null, null);
+        KafkaBuffer kafkaBuffer = new KafkaBuffer(pluginSetting, kafkaBufferConfig, acknowledgementSetManager, new OTelTraceDecoder(OTelOutputFormat.OPENSEARCH), null, null, encryptionSupplier);
         buffer = new KafkaDelegatingBuffer(kafkaBuffer);
         final ExportTraceServiceRequest request = createExportTraceRequest();
         buffer.writeBytes(request.toByteArray(), null, 10_000);

data-prepper-plugins/kafka-plugins/src/integrationTest/java/org/opensearch/dataprepper/plugins/kafka/buffer/KafkaBuffer_KmsIT.java

@@ -26,6 +26,7 @@
 import org.opensearch.dataprepper.model.event.JacksonEvent;
 import org.opensearch.dataprepper.model.plugin.PluginFactory;
 import org.opensearch.dataprepper.model.record.Record;
+import org.opensearch.dataprepper.plugins.encryption.EncryptionSupplier;
 import org.opensearch.dataprepper.plugins.kafka.util.TestProducer;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -78,6 +79,9 @@ public class KafkaBuffer_KmsIT {
     @Mock
     private AwsCredentialsSupplier awsCredentialsSupplier;
 
+    @Mock
+    private EncryptionSupplier encryptionSupplier;
+
     private Random random;
 
     private BufferTopicConfig topicConfig;
@@ -132,7 +136,7 @@ void setUp() {
     }
 
     private KafkaBuffer createObjectUnderTest() {
-        return new KafkaBuffer(pluginSetting, kafkaBufferConfig, acknowledgementSetManager, null, awsCredentialsSupplier, null);
+        return new KafkaBuffer(pluginSetting, kafkaBufferConfig, acknowledgementSetManager, null, awsCredentialsSupplier, null, encryptionSupplier);
     }
 
     @Nested

data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/buffer/BufferTopicConfig.java

@@ -8,6 +8,7 @@
 import com.fasterxml.jackson.annotation.JsonIgnore;
 import com.fasterxml.jackson.annotation.JsonProperty;
 import jakarta.validation.Valid;
+import jakarta.validation.constraints.AssertTrue;
 import jakarta.validation.constraints.Size;
 import org.opensearch.dataprepper.model.types.ByteCount;
 import org.opensearch.dataprepper.plugins.kafka.configuration.CommonTopicConfig;
@@ -37,6 +38,9 @@ class BufferTopicConfig extends CommonTopicConfig implements TopicProducerConfig
     static final Integer DEFAULT_NUM_OF_WORKERS = 2;
     static final Duration DEFAULT_HEART_BEAT_INTERVAL_DURATION = Duration.ofSeconds(5);
 
+    @JsonProperty("encryption_id")
+    private String encryptionId;
+
     @JsonProperty("encryption_key")
     private String encryptionKey;
 
@@ -119,6 +123,11 @@ public MessageFormat getSerdeFormat() {
         return MessageFormat.BYTES;
     }
 
+    @Override
+    public String getEncryptionId() {
+        return encryptionId;
+    }
+
     @Override
     public String getEncryptionKey() {
         return encryptionKey;
@@ -239,4 +248,13 @@ public long getFetchMinBytes() {
     public long getMaxPartitionFetchBytes() {
         return maxPartitionFetchBytes.getBytes();
     }
+
+    @AssertTrue(message = "Either encryption_id or encryption_key together with kms_config must be specified, " +
+            "and only one of them can be specified")
+    public boolean IsEncryptionAtRestSettingValid() {
+        if (encryptionId != null && (encryptionKey != null || kmsConfig != null)) {
+            return false;
+        }
+        return true;
+    }
 }

data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/buffer/KafkaBuffer.java

@@ -22,6 +22,7 @@
 import org.opensearch.dataprepper.plugins.buffer.blockingbuffer.BlockingBuffer;
 import org.apache.kafka.common.errors.RecordTooLargeException;
 import org.apache.kafka.common.errors.RecordBatchTooLargeException;
+import org.opensearch.dataprepper.plugins.encryption.EncryptionSupplier;
 import org.opensearch.dataprepper.plugins.kafka.admin.KafkaAdminAccessor;
 import org.opensearch.dataprepper.plugins.kafka.buffer.serialization.BufferSerializationFactory;
 import org.opensearch.dataprepper.plugins.kafka.common.KafkaMdc;
@@ -73,9 +74,10 @@ public class KafkaBuffer extends AbstractBuffer<Record<Event>> {
     public KafkaBuffer(final PluginSetting pluginSetting, final KafkaBufferConfig kafkaBufferConfig,
                        final AcknowledgementSetManager acknowledgementSetManager,
                        final ByteDecoder byteDecoder, final AwsCredentialsSupplier awsCredentialsSupplier,
-                       final CircuitBreaker circuitBreaker) {
+                       final CircuitBreaker circuitBreaker,
+                       final EncryptionSupplier encryptionSupplier) {
         super(kafkaBufferConfig.getCustomMetricPrefix().orElse(pluginSetting.getName()+"buffer"), pluginSetting.getPipelineName());
-        final SerializationFactory serializationFactory = new BufferSerializationFactory(new CommonSerializationFactory());
+        final SerializationFactory serializationFactory = new BufferSerializationFactory(new CommonSerializationFactory(), encryptionSupplier);
         final KafkaCustomProducerFactory kafkaCustomProducerFactory = new KafkaCustomProducerFactory(serializationFactory, awsCredentialsSupplier, new TopicServiceFactory());
         this.byteDecoder = byteDecoder;
         final String metricPrefixName = kafkaBufferConfig.getCustomMetricPrefix().orElse(pluginSetting.getName());

data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/buffer/serialization/BufferMessageEncryptionDeserializer.java

@@ -0,0 +1,50 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.dataprepper.plugins.kafka.buffer.serialization;
+
+import com.google.protobuf.InvalidProtocolBufferException;
+import org.apache.kafka.common.serialization.Deserializer;
+import org.opensearch.dataprepper.model.encryption.EncryptionEngine;
+import org.opensearch.dataprepper.plugins.encryption.DefaultEncryptionEnvelope;
+import org.opensearch.dataprepper.plugins.kafka.buffer.KafkaBufferMessage;
+
+/**
+ * A Kafka {@link Deserializer} which deserializes data from a KafkaBuffer Protobuf message.
+ *
+ * @param <T> The output type
+ */
+class BufferMessageEncryptionDeserializer<T> implements Deserializer<T> {
+    private final Deserializer<T> dataDeserializer;
+    private final EncryptionEngine encryptionEngine;
+
+    public BufferMessageEncryptionDeserializer(final Deserializer<T> dataDeserializer,
+                                               final EncryptionEngine encryptionEngine) {
+        this.dataDeserializer = dataDeserializer;
+        this.encryptionEngine = encryptionEngine;
+    }
+
+    @Override
+    public T deserialize(final String topic, final byte[] data) {
+        final KafkaBufferMessage.BufferData bufferedData;
+        try {
+            bufferedData = KafkaBufferMessage.BufferData.parseFrom(data);
+        } catch (final InvalidProtocolBufferException e) {
+            throw new RuntimeException(e);
+        }
+
+        final byte[] dataBytes = encryptionEngine.decrypt(
+                DefaultEncryptionEnvelope.builder()
+                        .encryptedDataKey(bufferedData.getEncryptedDataKey().toStringUtf8())
+                        .encryptedData(bufferedData.getData().toByteArray())
+                        .build());
+
+        return dataDeserializer.deserialize(topic, dataBytes);
+    }
+
+    Deserializer<T> getDataDeserializer() {
+        return dataDeserializer;
+    }
+}

data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/buffer/serialization/BufferMessageEncryptionSerializer.java

@@ -0,0 +1,56 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.dataprepper.plugins.kafka.buffer.serialization;
+
+import com.google.protobuf.ByteString;
+import org.apache.kafka.common.serialization.Serializer;
+import org.opensearch.dataprepper.model.encryption.EncryptionEngine;
+import org.opensearch.dataprepper.model.encryption.EncryptionEnvelope;
+import org.opensearch.dataprepper.plugins.kafka.buffer.KafkaBufferMessage;
+
+import java.util.Objects;
+
+/**
+ * A Kafka {@link Serializer} which serializes data into a KafkaBuffer Protobuf message.
+ *
+ * @param <T> The input type
+ */
+class BufferMessageEncryptionSerializer<T> implements Serializer<T> {
+    private final Serializer<T> dataSerializer;
+    private final EncryptionEngine encryptionEngine;
+
+    public BufferMessageEncryptionSerializer(final Serializer<T> dataSerializer,
+                                             final EncryptionEngine encryptionEngine) {
+        this.dataSerializer = Objects.requireNonNull(dataSerializer);
+        this.encryptionEngine = Objects.requireNonNull(encryptionEngine);
+    }
+
+    @Override
+    public byte[] serialize(final String topic, final T data) {
+        if (data == null)
+            return null;
+
+        final byte[] serializedData = dataSerializer.serialize(topic, data);
+        final EncryptionEnvelope encryptionEnvelope = encryptionEngine.encrypt(serializedData);
+
+        final KafkaBufferMessage.BufferData bufferedData = buildProtobufMessage(encryptionEnvelope);
+
+        return bufferedData.toByteArray();
+    }
+
+    Serializer<T> getDataSerializer() {
+        return dataSerializer;
+    }
+
+    private KafkaBufferMessage.BufferData buildProtobufMessage(final EncryptionEnvelope encryptionEnvelope) {
+        return KafkaBufferMessage.BufferData.newBuilder()
+                .setEncryptedDataKey(ByteString.copyFromUtf8(encryptionEnvelope.getEncryptedDataKey()))
+                .setData(ByteString.copyFrom(encryptionEnvelope.getEncryptedData()))
+                .setEncrypted(true)
+                .setMessageFormat(KafkaBufferMessage.MessageFormat.MESSAGE_FORMAT_BYTES)
+                .build();
+    }
+}

data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/buffer/serialization/BufferSerializationFactory.java

@@ -7,33 +7,60 @@
 
 import org.apache.kafka.common.serialization.Deserializer;
 import org.apache.kafka.common.serialization.Serializer;
+import org.opensearch.dataprepper.model.encryption.EncryptionEngine;
+import org.opensearch.dataprepper.plugins.encryption.EncryptionSupplier;
 import org.opensearch.dataprepper.plugins.kafka.common.KafkaDataConfig;
 import org.opensearch.dataprepper.plugins.kafka.common.serialization.SerializationFactory;
 
+import java.util.Objects;
+
 /**
  * An implementation of {@link SerializationFactory} specifically for the Kafka buffer.
  * It makes use the Kafka buffer's Protobuf wrapper.
  */
 public class BufferSerializationFactory implements SerializationFactory {
     private final SerializationFactory innerSerializationFactory;
+    private final EncryptionSupplier encryptionSupplier;
 
     /**
      * Create a new instance using an inner {@link SerializationFactory}. You typically pass in
      * the {@link org.opensearch.dataprepper.plugins.kafka.common.serialization.CommonSerializationFactory}.
      *
      * @param innerSerializationFactory The serialization factory for the inner data.
      */
-    public BufferSerializationFactory(final SerializationFactory innerSerializationFactory) {
+    public BufferSerializationFactory(final SerializationFactory innerSerializationFactory,
+                                      final EncryptionSupplier encryptionSupplier) {
         this.innerSerializationFactory = innerSerializationFactory;
+        this.encryptionSupplier = encryptionSupplier;
     }
 
     @Override
     public Deserializer<?> getDeserializer(final KafkaDataConfig dataConfig) {
-        return new BufferMessageDeserializer<>(innerSerializationFactory.getDeserializer(dataConfig));
+        final String encryptionId = dataConfig.getEncryptionId();
+        final Deserializer<?> dataDeserializer = innerSerializationFactory.getDeserializer(dataConfig);
+        if (Objects.isNull(encryptionId)) {
+            return new BufferMessageDeserializer<>(dataDeserializer);
+        } else {
+            final EncryptionEngine encryptionEngine = encryptionSupplier.getEncryptionEngine(encryptionId);
+            if (encryptionEngine == null) {
+                throw new IllegalArgumentException("Encryption engine not found for encryption_id: " + encryptionId);
+            }
+            return new BufferMessageEncryptionDeserializer<>(dataDeserializer, encryptionEngine);
+        }
     }
 
     @Override
     public Serializer<?> getSerializer(final KafkaDataConfig dataConfig) {
-        return new BufferMessageSerializer<>(innerSerializationFactory.getSerializer(dataConfig), dataConfig);
+        final String encryptionId = dataConfig.getEncryptionId();
+        final Serializer<?> dataSerializer = innerSerializationFactory.getSerializer(dataConfig);
+        if (Objects.isNull(encryptionId)) {
+            return new BufferMessageSerializer<>(innerSerializationFactory.getSerializer(dataConfig), dataConfig);
+        } else {
+            final EncryptionEngine encryptionEngine = encryptionSupplier.getEncryptionEngine(encryptionId);
+            if (encryptionEngine == null) {
+                throw new IllegalArgumentException("Encryption engine not found for encryption_id: " + encryptionId);
+            }
+            return new BufferMessageEncryptionSerializer<>(dataSerializer, encryptionEngine);
+        }
     }
 }