Repository: kroxylicious/kroxylicious (java)

Issue:
Topic-encryption: Handle DEK creation failures
Currently if the `KeyContext` CompletionStage fails (so maybe something went wrong obtaining the EDEK from the KMS), it remains cached and will generate failures in perpetuity as more produce requests are received. There is no logic to retry failed EDEK generation. 

We likely want to guarantee that we only have one attempt in flight to obtain an EDEK for a kekId so there is some threading concerns to think about too.

Pull request description:
Encryption Filter: Retry DEK creation. Ensure encryptions after a failed future can succeed
### Type of change

- Enhancement / new feature

### Description

Previously if the DEK future from the KMS failed it would have been cached for all time and led to all future encryption calls for that KEK to fail.

We also want to give DEK creation naive handling of some simple cases like a failure due to a one-off network issue. So it will now retry the DEK retrieval operation 3 times as part of loading a DEK into the cache.

As I considered how to implement these safely myself Sam said it sounded a lot like the caffeine async loading cache which knows not to cache failed futures. Kafka already depends on caffeine in 3.6.0 so it's a fairly safe dependency to add. Caffeine has already solved the gnarly async problems I was already worrying about.

We will also be able to take advantage of it's LRU and refreshing ability later to manage the cache size and to pre-emptively load a new DEK encrypter while an existing one continues to be used. We could minimise pauses by refreshing in a new DEK when X% of the encryptions have been used.

Note tests now have to supply non-zero timeouts when waiting on futures as the cache is leaning on the fork join pool to farm out the async loading operations.

This isn't tackling everything, we could also have future improvements introducing configurable timeouts to the InBandKeyManager so that oddities like futures from the KMS that never complete are handled.

Closes #765

### Checklist

_Please go through this checklist and make sure all applicable tasks have been done_

- [ ] Write tests
- [ ] Make sure all tests pass
- [ ] Review performance test results. Ensure that any degradations to performance numbers are understood and justified.
- [ ] Make sure all Sonar-Lint warnings are addressed or are justifiably ignored.
- [ ] Update documentation
- [ ] Reference relevant issue(s) and close them after merging
- [ ] For user facing changes, update CHANGELOG.md (remember to include changes affecting the API of the test artefacts too).


Patch:
diff --git a/kroxylicious-filters/kroxylicious-encryption/pom.xml b/kroxylicious-filters/kroxylicious-encryption/pom.xml
index 7e982bac..091bd495 100644
--- a/kroxylicious-filters/kroxylicious-encryption/pom.xml
+++ b/kroxylicious-filters/kroxylicious-encryption/pom.xml
@@ -42,6 +42,10 @@
             <groupId>com.fasterxml.jackson.core</groupId>
             <artifactId>jackson-databind</artifactId>
         </dependency>
+        <dependency>
+            <groupId>com.github.ben-manes.caffeine</groupId>
+            <artifactId>caffeine</artifactId>
+        </dependency>
 
         <dependency>
             <groupId>com.github.spotbugs</groupId>
diff --git a/kroxylicious-filters/kroxylicious-encryption/src/main/java/io/kroxylicious/filter/encryption/inband/EncryptorCreationException.java b/kroxylicious-filters/kroxylicious-encryption/src/main/java/io/kroxylicious/filter/encryption/inband/EncryptorCreationException.java
new file mode 100644
index 00000000..81a8da63
--- /dev/null
+++ b/kroxylicious-filters/kroxylicious-encryption/src/main/java/io/kroxylicious/filter/encryption/inband/EncryptorCreationException.java
@@ -0,0 +1,13 @@
+/*
+ * Copyright Kroxylicious Authors.
+ *
+ * Licensed under the Apache Software License version 2.0, available at http://www.apache.org/licenses/LICENSE-2.0
+ */
+
+package io.kroxylicious.filter.encryption.inband;
+
+public class EncryptorCreationException extends RuntimeException {
+    public EncryptorCreationException(String message) {
+        super(message);
+    }
+}
diff --git a/kroxylicious-filters/kroxylicious-encryption/src/main/java/io/kroxylicious/filter/encryption/inband/InBandKeyManager.java b/kroxylicious-filters/kroxylicious-encryption/src/main/java/io/kroxylicious/filter/encryption/inband/InBandKeyManager.java
index f491f9f5..d9feb9be 100644
--- a/kroxylicious-filters/kroxylicious-encryption/src/main/java/io/kroxylicious/filter/encryption/inband/InBandKeyManager.java
+++ b/kroxylicious-filters/kroxylicious-encryption/src/main/java/io/kroxylicious/filter/encryption/inband/InBandKeyManager.java
@@ -13,12 +13,16 @@ import java.util.List;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.CompletionStage;
 import java.util.concurrent.ConcurrentHashMap;
-import java.util.function.Supplier;
 
 import org.apache.kafka.common.header.Header;
 import org.apache.kafka.common.header.internals.RecordHeader;
 import org.apache.kafka.common.record.Record;
 import org.apache.kafka.common.utils.ByteUtils;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import com.github.benmanes.caffeine.cache.AsyncLoadingCache;
+import com.github.benmanes.caffeine.cache.Caffeine;
 
 import io.kroxylicious.filter.encryption.AadSpec;
 import io.kroxylicious.filter.encryption.CipherCode;
@@ -61,11 +65,12 @@ public class InBandKeyManager<K, E> implements KeyManager<K> {
     private final BufferPool bufferPool;
     private final Serde<E> edekSerde;
     // TODO cache expiry, with key descruction
-    private final ConcurrentHashMap<K, CompletionStage<KeyContext>> keyContextCache;
+    private final AsyncLoadingCache<K, KeyContext> keyContextCache;
     private final ConcurrentHashMap<E, CompletionStage<AesGcmEncryptor>> decryptorCache;
     private final long dekTtlNanos;
     private final int maxEncryptionsPerDek;
     private final Header[] encryptionHeader;
+    private static final Logger LOGGER = LoggerFactory.getLogger(InBandKeyManager.class);
 
     public InBandKeyManager(Kms<K, E> kms,
                             BufferPool bufferPool,
@@ -76,32 +81,27 @@ public class InBandKeyManager<K, E> implements KeyManager<K> {
         this.dekTtlNanos = 5_000_000_000L;
         this.maxEncryptionsPerDek = maxEncryptionsPerDek;
         // TODO This ^^ must be > the maximum size of a batch to avoid an infinite loop
-        this.keyContextCache = new ConcurrentHashMap<>();
+        this.keyContextCache = Caffeine.newBuilder()
+                .buildAsync((key, executor) -> makeKeyContext(key));
         this.decryptorCache = new ConcurrentHashMap<>();
         this.encryptionVersion = EncryptionVersion.V1; // TODO read from config
         this.encryptionHeader = new Header[]{ new RecordHeader(ENCRYPTION_HEADER_NAME, new byte[]{ encryptionVersion.code() }) };
     }
 
-    private CompletionStage<KeyContext> getKeyContext(K key,
-                                                      Supplier<CompletionStage<KeyContext>> valueSupplier) {
-        return keyContextCache.compute(key, (k, v) -> {
-            if (v == null) {
-                return valueSupplier.get();
-                // TODO what happens if the CS complete exceptionally
-                // TODO what happens if the CS doesn't complete at all in a reasonably time frame?
-            }
-            else {
-                return v;
-            }
-        });
+    private CompletionStage<KeyContext> currentDekContext(@NonNull K kekId) {
+        // todo should we add some scheduled timeout as well? or should we rely on the KMS to timeout appropriately.
+        return keyContextCache.get(kekId);
     }
 
-    private CompletionStage<KeyContext> currentDekContext(@NonNull K kekId) {
-        return getKeyContext(kekId, makeKeyContext(kekId));
+    private CompletableFuture<KeyContext> makeKeyContext(@NonNull K kekId) {
+        return attemptMakeDekContext(kekId, 0);
     }
 
-    private Supplier<CompletionStage<KeyContext>> makeKeyContext(@NonNull K kekId) {
-        return () -> kms.generateDekPair(kekId)
+    private CompletableFuture<KeyContext> attemptMakeDekContext(@NonNull K kekId, int attempt) {
+        if (attempt >= MAX_ATTEMPTS) {
+            return CompletableFuture.failedFuture(new EncryptorCreationException("failed to create encryptor after " + attempt + " attempts"));
+        }
+        return kms.generateDekPair(kekId)
                 .thenApply(dekPair -> {
                     E edek = dekPair.edek();
                     short edekSize = (short) edekSerde.sizeOf(edek);
@@ -116,6 +116,11 @@ public class InBandKeyManager<K, E> implements KeyManager<K> {
                             // or we need mutex
                             // or we externalize the state
                             AesGcmEncryptor.forEncrypt(new AesGcmIvGenerator(new SecureRandom()), dekPair.dek()));
+                }).toCompletableFuture()
+                // todo wire in a scheduler so we can delay/jitter DEK creation attempts
+                .exceptionallyComposeAsync(throwable -> {
+                    LOGGER.error("failed to create DEK encryption context for {} on attempt {}", kekId, attempt, throwable);
+                    return attemptMakeDekContext(kekId, attempt + 1);
                 });
     }
 
@@ -142,13 +147,14 @@ public class InBandKeyManager<K, E> implements KeyManager<K> {
         }
         return currentDekContext(encryptionScheme.kekId()).thenCompose(keyContext -> {
             synchronized (keyContext) {
-                // if it's not alive we know a previous encrypt call has replaced the stage in the cache and fall through to retry encrypt
+                // if it's not alive we know a previous encrypt call has removed this stage from the cache and fall through to retry encrypt
                 if (!keyContext.isDestroyed()) {
                     if (!keyContext.hasAtLeastRemainingEncryptions(records.size())) {
-                        // replace the key context stage in the cache, then call encrypt again
+                        // remove the key context from the cache, then call encrypt again to drive caffeine to recreate it
                         rotateKeyContext(encryptionScheme, keyContext);
                     }
                     else {
+                        // todo ensure that a failure during encryption terminates the entire operation with a failed future
                         return encrypt(encryptionScheme, records, receiver, keyContext);
                     }
                 }
@@ -194,7 +200,7 @@ public class InBandKeyManager<K, E> implements KeyManager<K> {
     private void rotateKeyContext(@NonNull EncryptionScheme<K> encryptionScheme, KeyContext keyContext) {
         keyContext.destroy();
         K kekId = encryptionScheme.kekId();
-        keyContextCache.put(kekId, makeKeyContext(kekId).get());
+        keyContextCache.synchronous().invalidate(kekId);
     }
 
     private void encryptRecords(@NonNull EncryptionScheme<K> encryptionScheme,
diff --git a/kroxylicious-filters/kroxylicious-encryption/src/main/java/io/kroxylicious/filter/encryption/inband/KeyContext.java b/kroxylicious-filters/kroxylicious-encryption/src/main/java/io/kroxylicious/filter/encryption/inband/KeyContext.java
index 9e12c6ad..295dc74f 100644
--- a/kroxylicious-filters/kroxylicious-encryption/src/main/java/io/kroxylicious/filter/encryption/inband/KeyContext.java
+++ b/kroxylicious-filters/kroxylicious-encryption/src/main/java/io/kroxylicious/filter/encryption/inband/KeyContext.java
@@ -33,7 +33,7 @@ final class KeyContext implements Destroyable {
     private int remainingEncryptions;
     private boolean destroyed = false;
 
-    private static final Logger LOGGER = LoggerFactory.getLogger(InBandKeyManager.class);
+    private static final Logger LOGGER = LoggerFactory.getLogger(KeyContext.class);
     private static final Map<Class<? extends Destroyable>, Boolean> LOGGED_DESTROY_FAILED = new ConcurrentHashMap<>();
 
     KeyContext(@NonNull ByteBuffer serializedEdek,
diff --git a/pom.xml b/pom.xml
index c66e5d2b..88ce66d9 100644
--- a/pom.xml
+++ b/pom.xml
@@ -42,6 +42,7 @@
         <kafka.version>3.6.1</kafka.version>
         <kroxy.extension.version>0.8.0</kroxy.extension.version>
         <log4j.version>2.22.0</log4j.version>
+        <caffeine.version>3.1.8</caffeine.version>
         <picocli.version>4.7.5</picocli.version>
         <netty.version>4.1.101.Final</netty.version>
         <netty.io_uring.version>0.0.24.Final</netty.io_uring.version>
@@ -163,6 +164,12 @@
                 <scope>import</scope>
             </dependency>
 
+            <dependency>
+                <groupId>com.github.ben-manes.caffeine</groupId>
+                <artifactId>caffeine</artifactId>
+                <version>${caffeine.version}</version>
+            </dependency>
+
             <dependency>
                 <groupId>com.github.spotbugs</groupId>
                 <artifactId>spotbugs-annotations</artifactId>


Test patch:
diff --git a/kroxylicious-filters/kroxylicious-encryption/src/test/java/io/kroxylicious/filter/encryption/inband/InBandKeyManagerTest.java b/kroxylicious-filters/kroxylicious-encryption/src/test/java/io/kroxylicious/filter/encryption/inband/InBandKeyManagerTest.java
index 81c4035c..ed4e52af 100644
--- a/kroxylicious-filters/kroxylicious-encryption/src/test/java/io/kroxylicious/filter/encryption/inband/InBandKeyManagerTest.java
+++ b/kroxylicious-filters/kroxylicious-encryption/src/test/java/io/kroxylicious/filter/encryption/inband/InBandKeyManagerTest.java
@@ -12,6 +12,7 @@ import java.util.ArrayList;
 import java.util.Collection;
 import java.util.EnumSet;
 import java.util.List;
+import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.CompletionStage;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.TimeUnit;
@@ -22,6 +23,7 @@ import org.apache.kafka.common.header.Header;
 import org.apache.kafka.common.header.internals.RecordHeader;
 import org.apache.kafka.common.utils.ByteUtils;
 import org.junit.jupiter.api.Test;
+import org.mockito.Mockito;
 
 import io.kroxylicious.filter.encryption.EncryptionScheme;
 import io.kroxylicious.filter.encryption.Receiver;
@@ -29,6 +31,7 @@ import io.kroxylicious.filter.encryption.RecordField;
 import io.kroxylicious.kms.provider.kroxylicious.inmemory.InMemoryKms;
 import io.kroxylicious.kms.provider.kroxylicious.inmemory.UnitTestingKmsService;
 import io.kroxylicious.kms.service.DekPair;
+import io.kroxylicious.kms.service.KmsException;
 import io.kroxylicious.kms.service.Serde;
 
 import edu.umd.cs.findbugs.annotations.NonNull;
@@ -37,6 +40,7 @@ import static org.assertj.core.api.Assertions.assertThat;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertFalse;
 import static org.junit.jupiter.api.Assertions.assertNotEquals;
+import static org.mockito.Mockito.when;
 
 class InBandKeyManagerTest {
 
@@ -181,8 +185,89 @@ class InBandKeyManagerTest {
         CompletionStage encrypt = km.encrypt("topic", 1, new EncryptionScheme<>(kekId, EnumSet.of(RecordField.RECORD_VALUE)),
                 initial,
                 recordReceivedRecord(encrypted));
-        String expectedMessage = "failed to encrypt records for topic topic partition 1 after 3 attempts";
-        assertThat(encrypt).failsWithin(Duration.ZERO).withThrowableThat().withMessageContaining(expectedMessage);
+        assertThat(encrypt).failsWithin(Duration.ofSeconds(5)).withThrowableThat()
+                .withMessageMatching(".*failed to encrypt records for topic topic partition 1 after [0-9]+ attempts");
+    }
+
+    @Test
+    void dekCreationRetryFailurePropagatedToEncryptCompletionStage() {
+        var kmsService = UnitTestingKmsService.newInstance();
+        InMemoryKms kms = kmsService.buildKms(new UnitTestingKmsService.Config());
+        var kekId = kms.generateKey();
+        InMemoryKms spyKms = Mockito.spy(kms);
+        when(spyKms.generateDekPair(kekId)).thenReturn(CompletableFuture.failedFuture(new EncryptorCreationException("failed to create that DEK")));
+        var km = new InBandKeyManager<>(spyKms, BufferPool.allocating(), 500000);
+
+        var value = ByteBuffer.wrap(new byte[]{ 1, 2, 3 });
+        var value2 = ByteBuffer.wrap(new byte[]{ 4, 5, 6 });
+        TestingRecord record = new TestingRecord(value);
+        TestingRecord record2 = new TestingRecord(value2);
+
+        List<TestingRecord> encrypted = new ArrayList<>();
+        List<TestingRecord> initial = List.of(record, record2);
+        CompletionStage encrypt = km.encrypt("topic", 1, new EncryptionScheme<>(kekId, EnumSet.of(RecordField.RECORD_VALUE)),
+                initial,
+                recordReceivedRecord(encrypted));
+        assertThat(encrypt).failsWithin(Duration.ofSeconds(5)).withThrowableThat().withMessageMatching(".*failed to create encryptor after [0-9]+ attempts");
+    }
+
+    @Test
+    void dekCreationRetrySuccess() {
+        var kmsService = UnitTestingKmsService.newInstance();
+        InMemoryKms kms = kmsService.buildKms(new UnitTestingKmsService.Config());
+        var kekId = kms.generateKey();
+        InMemoryKms spyKms = Mockito.spy(kms);
+        when(spyKms.generateDekPair(kekId)).thenReturn(CompletableFuture.failedFuture(new EncryptorCreationException("failed to create that DEK")),
+                CompletableFuture.failedFuture(new EncryptorCreationException("failed to create that DEK"))).thenCallRealMethod();
+        var km = new InBandKeyManager<>(spyKms, BufferPool.allocating(), 50000);
+
+        var value = ByteBuffer.wrap(new byte[]{ 1, 2, 3 });
+        var value2 = ByteBuffer.wrap(new byte[]{ 4, 5, 6 });
+        TestingRecord record = new TestingRecord(value);
+        TestingRecord record2 = new TestingRecord(value2);
+
+        List<TestingRecord> encrypted = new ArrayList<>();
+        List<TestingRecord> initial = List.of(record, record2);
+        CompletionStage encrypt = km.encrypt("topic", 1, new EncryptionScheme<>(kekId, EnumSet.of(RecordField.RECORD_VALUE)),
+                initial,
+                recordReceivedRecord(encrypted));
+        assertThat(encrypt).succeedsWithin(Duration.ofSeconds(5));
+        assertThat(encrypted).hasSize(2);
+    }
+
+    @Test
+    void afterWeFailToLoadADekTheNextEncryptionAttemptCanSucceed() {
+        var kmsService = UnitTestingKmsService.newInstance();
+        InMemoryKms kms = kmsService.buildKms(new UnitTestingKmsService.Config());
+        var kekId = kms.generateKey();
+        InMemoryKms spyKms = Mockito.spy(kms);
+        when(spyKms.generateDekPair(kekId)).thenReturn(CompletableFuture.failedFuture(new KmsException("failed to create that DEK")));
+
+        var km = new InBandKeyManager<>(spyKms, BufferPool.allocating(), 50000);
+
+        var value = ByteBuffer.wrap(new byte[]{ 1, 2, 3 });
+        var value2 = ByteBuffer.wrap(new byte[]{ 4, 5, 6 });
+        TestingRecord record = new TestingRecord(value);
+        TestingRecord record2 = new TestingRecord(value2);
+
+        List<TestingRecord> encrypted = new ArrayList<>();
+        List<TestingRecord> initial = List.of(record, record2);
+        CompletionStage encrypt = km.encrypt("topic", 1, new EncryptionScheme<>(kekId, EnumSet.of(RecordField.RECORD_VALUE)),
+                initial,
+                recordReceivedRecord(encrypted));
+        assertThat(encrypt).failsWithin(Duration.ofSeconds(5)).withThrowableThat().withMessageMatching(".*failed to create encryptor after [0-9]+ attempts");
+
+        // given KMS is no longer generating failed futures
+        when(spyKms.generateDekPair(kekId)).thenCallRealMethod();
+
+        // when
+        CompletionStage encrypt2 = km.encrypt("topic", 1, new EncryptionScheme<>(kekId, EnumSet.of(RecordField.RECORD_VALUE)),
+                initial,
+                recordReceivedRecord(encrypted));
+
+        // then
+        assertThat(encrypt2).succeedsWithin(Duration.ofSeconds(5));
+        assertThat(encrypted).hasSize(2);
     }
 
     @NonNull