feat(record-encryption): decrypt share fetch responses

Apache Kafka is about to promote Queues for Kafka to stable. We now
decrypt records fetched by Share Group Consumers from encrypted topics.

If we cannot discover the name for a topic id, we:
* mark all share fetch response partitions for that topic id with the
  mapping error code and message.
* empty the memory records for all share fetch response partitions for
  that topic id.

This means that if we can obtain topic name for some of the response
topics we will try to decode and return those topic partitions, and if
we cannot obtain topic name those partitions will be marked as errored.

Signed-off-by: Robert Young <robertyoungnz@gmail.com>

Author: robobario

kroxylicious-filters/kroxylicious-record-encryption/etc/module-layering.xml

@@ -24,6 +24,8 @@
     <allow pkg="org.apache.kafka.common.message" local-only="true"/>
     <allow pkg="org.apache.kafka.common.record" local-only="true"/>
     <allow pkg="org.apache.kafka.common.errors" local-only="true"/>
+    <allow pkg="org.apache.kafka.common" local-only="true"/>
+    <allow pkg="org.apache.kafka.common.utils" local-only="true"/>
     <allow pkg="org.apache.kafka.common.protocol" local-only="true"/>
     <allow pkg="io.kroxylicious.proxy.filter" local-only="true"/> <!-- the filter api -->
     <allow pkg="io.kroxylicious.proxy.plugin" local-only="true"/> <!-- the plugin api -->

kroxylicious-filters/kroxylicious-record-encryption/src/main/java/io/kroxylicious/filter/encryption/RecordEncryptionFilter.java

@@ -7,15 +7,19 @@
 package io.kroxylicious.filter.encryption;
 
 import java.util.ArrayList;
+import java.util.Collection;
 import java.util.EnumSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.CompletionStage;
+import java.util.function.BiFunction;
 import java.util.function.Function;
+import java.util.function.ToIntFunction;
 import java.util.stream.Collectors;
 
+import org.apache.kafka.common.Uuid;
 import org.apache.kafka.common.errors.ApiException;
 import org.apache.kafka.common.message.ApiVersionsResponseData;
 import org.apache.kafka.common.message.FetchResponseData;
@@ -25,9 +29,13 @@
 import org.apache.kafka.common.message.ProduceRequestData.TopicProduceData;
 import org.apache.kafka.common.message.RequestHeaderData;
 import org.apache.kafka.common.message.ResponseHeaderData;
+import org.apache.kafka.common.message.ShareFetchResponseData;
+import org.apache.kafka.common.message.ShareFetchResponseData.ShareFetchableTopicResponse;
 import org.apache.kafka.common.protocol.ApiKeys;
 import org.apache.kafka.common.protocol.Errors;
+import org.apache.kafka.common.record.BaseRecords;
 import org.apache.kafka.common.record.MemoryRecords;
+import org.apache.kafka.common.utils.ImplicitLinkedHashCollection;
 import org.slf4j.Logger;
 
 import io.micrometer.core.instrument.Counter;
@@ -52,6 +60,9 @@
 import io.kroxylicious.proxy.filter.ProduceRequestFilter;
 import io.kroxylicious.proxy.filter.RequestFilterResult;
 import io.kroxylicious.proxy.filter.ResponseFilterResult;
+import io.kroxylicious.proxy.filter.ShareFetchResponseFilter;
+import io.kroxylicious.proxy.filter.metadata.TopicNameMapping;
+import io.kroxylicious.proxy.filter.metadata.TopicNameMappingException;
 
 import edu.umd.cs.findbugs.annotations.NonNull;
 
@@ -62,7 +73,7 @@
  * @param <K> The type of KEK reference
  */
 public class RecordEncryptionFilter<K>
-        implements ProduceRequestFilter, FetchResponseFilter, ApiVersionsResponseFilter {
+        implements ProduceRequestFilter, FetchResponseFilter, ShareFetchResponseFilter, ApiVersionsResponseFilter {
     private static final Logger log = getLogger(RecordEncryptionFilter.class);
     private final TopicNameBasedKekSelector<K> kekSelector;
 
@@ -123,7 +134,7 @@ else if (isEncryptionException(throwable.getCause())) {
 
     private CompletionStage<ProduceRequestData> maybeEncodeProduce(ProduceRequestData request, FilterContext context) {
         var plainRecordsTotal = RecordEncryptionMetrics.plainRecordsCounter(context.getVirtualClusterName());
-        var encyptedRecordsTotal = RecordEncryptionMetrics.encryptedRecordsCounter(context.getVirtualClusterName());
+        var encryptedRecordsTotal = RecordEncryptionMetrics.encryptedRecordsCounter(context.getVirtualClusterName());
         var topicNameToData = request.topicData().stream().collect(Collectors.toMap(TopicProduceData::name, Function.identity()));
         CompletionStage<TopicNameKekSelection<K>> keks = filterThreadExecutor.completingOnFilterThread(kekSelector.selectKek(topicNameToData.keySet()));
         return keks // figure out what keks we need
@@ -149,7 +160,7 @@ private CompletionStage<ProduceRequestData> maybeEncodeProduce(ProduceRequestDat
                                     context::createByteBufferOutputStream)
                                     .thenApply(ppd::setRecords)
                                     .thenApply(produceData -> {
-                                        encyptedRecordsTotal
+                                        encryptedRecordsTotal
                                                 .withTags(RecordEncryptionMetrics.TOPIC_NAME, topicName)
                                                 .increment(RecordEncryptionUtil.totalRecordsInBatches((MemoryRecords) produceData.records()));
                                         return null;
@@ -176,6 +187,40 @@ private void generatePlainRecordsMetrics(Meter.MeterProvider<Counter> plainRecor
                                 .increment(RecordEncryptionUtil.totalRecordsInBatches((MemoryRecords) produceData.records()))));
     }
 
+    @Override
+    public CompletionStage<ResponseFilterResult> onShareFetchResponse(short apiVersion, ResponseHeaderData header, ShareFetchResponseData response,
+                                                                      FilterContext context) {
+        Set<Uuid> topicIds = response.responses().stream().map(ShareFetchableTopicResponse::topicId).collect(Collectors.toSet());
+        return context.topicNames(topicIds).thenCompose(topicNameMapping -> maybeDecodeShareFetch(topicNameMapping, response.responses(), context)
+                .thenCompose(topicResponses -> {
+                    ShareFetchResponseData.ShareFetchableTopicResponseCollection collection = new ShareFetchResponseData.ShareFetchableTopicResponseCollection();
+                    // danger, the share fetch response uses a collection where the `add` method can silently do nothing
+                    // if the element being added already has element link fields populated it will not be added.
+                    // we reset next and prev so the existing elements can be added without calling `duplicate()`
+                    // which would copy the memory records again.
+                    // this may leave the original collection in an inconsistent state, but we are finished using it
+                    topicResponses.forEach(topicResponse -> {
+                        topicResponse.setNext(ImplicitLinkedHashCollection.INVALID_INDEX);
+                        topicResponse.setPrev(ImplicitLinkedHashCollection.INVALID_INDEX);
+                        collection.mustAdd(topicResponse);
+                    });
+                    return context.forwardResponse(header, response.setResponses(collection));
+                })
+                .exceptionallyCompose(throwable -> {
+                    if (throwable.getCause() instanceof UnknownKeyException) {
+                        // #maybeDecodePartitions will have set the RESOURCE_NOT_FOUND error code on the partition(s) that failed to decrypt
+                        // and will have logged the affected topic-partitions.
+                        // Remove all the records from the whole fetch to avoid the possibility that the client processes an incomplete response.
+                        response.responses().forEach(topicResponse -> topicResponse.partitions().forEach(p -> p.setRecords(MemoryRecords.EMPTY)));
+                        return context.forwardResponse(header, response);
+                    }
+                    else {
+                        // returning a failed stage is effectively asking the runtime to kill the connection.
+                        return logAndCreateFailedStage(throwable);
+                    }
+                }));
+    }
+
     @Override
     public CompletionStage<ResponseFilterResult> onFetchResponse(short apiVersion, ResponseHeaderData header, FetchResponseData response, FilterContext context) {
         return maybeDecodeFetch(response.responses(), context)
@@ -189,36 +234,83 @@ public CompletionStage<ResponseFilterResult> onFetchResponse(short apiVersion, R
                         return context.forwardResponse(header, response);
                     }
                     else {
-                        log.atWarn().setMessage("Failed to process records, connection will be closed, cause message: {}")
-                                .addArgument(throwable.getMessage())
-                                .setCause(log.isDebugEnabled() ? throwable : null)
-                                .log();
                         // returning a failed stage is effectively asking the runtime to kill the connection.
-                        return CompletableFuture.failedStage(throwable);
+                        return logAndCreateFailedStage(throwable);
                     }
                 });
     }
 
+    private static CompletionStage<ResponseFilterResult> logAndCreateFailedStage(Throwable throwable) {
+        log.atWarn().setMessage("Failed to process records, connection will be closed, cause message: {}. Raise log level to DEBUG to see the stack.")
+                .addArgument(throwable.getMessage())
+                .setCause(log.isDebugEnabled() ? throwable : null)
+                .log();
+        return CompletableFuture.failedStage(throwable);
+    }
+
+    private CompletionStage<List<ShareFetchableTopicResponse>> maybeDecodeShareFetch(TopicNameMapping mapping, Collection<ShareFetchableTopicResponse> topics,
+                                                                                     FilterContext context) {
+        List<CompletionStage<ShareFetchableTopicResponse>> result = new ArrayList<>(topics.size());
+        for (ShareFetchableTopicResponse topicData : topics) {
+            TopicNameMappingException failure = mapping.failures().get(topicData.topicId());
+            String topicName = mapping.topicNames().get(topicData.topicId());
+            List<ShareFetchResponseData.PartitionData> partitions = topicData.partitions();
+            if (topicName != null) {
+                result.add(maybeDecodePartitions(topicName, partitions, context, ShareFetchResponseData.PartitionData::records,
+                        ShareFetchResponseData.PartitionData::partitionIndex,
+                        partitionData -> partitionData::setRecords, ShareFetchResponseData.PartitionData::setErrorCode).thenApply(kk -> {
+                            topicData.setPartitions(kk);
+                            return topicData;
+                        }));
+            }
+            else {
+                result.add(setAllPartitionsToError(topicData, failure));
+            }
+        }
+        return RecordEncryptionUtil.join(result);
+    }
+
+    private static CompletableFuture<ShareFetchableTopicResponse> setAllPartitionsToError(ShareFetchableTopicResponse topicData,
+                                                                                          TopicNameMappingException failure) {
+        for (ShareFetchResponseData.PartitionData partition : topicData.partitions()) {
+            partition.setRecords(MemoryRecords.EMPTY);
+            Errors error = failure != null ? failure.getError() : Errors.UNKNOWN_SERVER_ERROR;
+            partition.setErrorCode(error.code());
+            partition.setErrorMessage(error.message());
+        }
+        return CompletableFuture.completedFuture(topicData);
+    }
+
     private CompletionStage<List<FetchableTopicResponse>> maybeDecodeFetch(List<FetchableTopicResponse> topics, FilterContext context) {
         List<CompletionStage<FetchableTopicResponse>> result = new ArrayList<>(topics.size());
         for (FetchableTopicResponse topicData : topics) {
-            result.add(maybeDecodePartitions(topicData.topic(), topicData.partitions(), context).thenApply(kk -> {
-                topicData.setPartitions(kk);
-                return topicData;
-            }));
+            String topicName = topicData.topic();
+            List<PartitionData> partitions = topicData.partitions();
+            result.add(maybeDecodePartitions(topicName, partitions, context, PartitionData::records, PartitionData::partitionIndex,
+                    partitionData -> partitionData::setRecords,
+                    PartitionData::setErrorCode).thenApply(kk -> {
+                        topicData.setPartitions(kk);
+                        return topicData;
+                    }));
         }
         return RecordEncryptionUtil.join(result);
     }
 
-    private CompletionStage<List<PartitionData>> maybeDecodePartitions(String topicName,
-                                                                       List<PartitionData> partitions,
-                                                                       FilterContext context) {
-        List<CompletionStage<PartitionData>> result = new ArrayList<>(partitions.size());
-        for (PartitionData partitionData : partitions) {
-            if (!(partitionData.records() instanceof MemoryRecords)) {
+    private <T> CompletionStage<List<T>> maybeDecodePartitions(String topicName,
+                                                               List<T> partitions,
+                                                               FilterContext context,
+                                                               Function<T, BaseRecords> recordExtractor,
+                                                               ToIntFunction<T> partitionIndexExtractor,
+                                                               Function<T, Function<MemoryRecords, T>> setRecords,
+                                                               BiFunction<T, Short, T> errorsConsumer) {
+        List<CompletionStage<T>> result = new ArrayList<>(partitions.size());
+        for (T partitionData : partitions) {
+            BaseRecords records = recordExtractor.apply(partitionData);
+            if (!(records instanceof MemoryRecords)) {
                 throw new IllegalStateException();
             }
-            var stage = maybeDecodeRecords(topicName, partitionData, (MemoryRecords) partitionData.records(), context)
+            var stage = maybeDecodeRecords(topicName, (MemoryRecords) records, context, partitionIndexExtractor.applyAsInt(partitionData),
+                    setRecords.apply(partitionData))
                     .exceptionallyCompose(t -> {
                         var cause = t.getCause();
                         if (cause instanceof UnknownKeyException) {
@@ -229,11 +321,11 @@ private CompletionStage<List<PartitionData>> maybeDecodePartitions(String topicN
                                             + "Cause message: {}. "
                                             + "Raise log level to DEBUG to see the stack.")
                                     .addArgument(topicName)
-                                    .addArgument(partitionData.partitionIndex())
+                                    .addArgument(() -> partitionIndexExtractor.applyAsInt(partitionData))
                                     .addArgument(cause.getMessage())
                                     .setCause(log.isDebugEnabled() ? cause : null)
                                     .log();
-                            partitionData.setErrorCode(Errors.RESOURCE_NOT_FOUND.code());
+                            errorsConsumer.apply(partitionData, Errors.RESOURCE_NOT_FOUND.code());
                         }
                         return CompletableFuture.failedFuture(t);
                     });
@@ -242,16 +334,16 @@ private CompletionStage<List<PartitionData>> maybeDecodePartitions(String topicN
         return RecordEncryptionUtil.join(result);
     }
 
-    private CompletionStage<PartitionData> maybeDecodeRecords(String topicName,
-                                                              PartitionData fpr,
-                                                              MemoryRecords memoryRecords,
-                                                              FilterContext context) {
+    private <T> CompletionStage<T> maybeDecodeRecords(String topicName,
+                                                      MemoryRecords memoryRecords,
+                                                      FilterContext context, int partition,
+                                                      Function<MemoryRecords, T> setRecords) {
         return decryptionManager.decrypt(
                 topicName,
-                fpr.partitionIndex(),
+                partition,
                 memoryRecords,
                 context::createByteBufferOutputStream)
-                .thenApply(fpr::setRecords);
+                .thenApply(setRecords);
     }
 
     private static boolean isEncryptionException(Throwable throwable) {