Repository: kroxylicious/kroxylicious (java)

Issue:
Topic Encryption: Rotation of DEK after expiry or max encryptions not working
**Describe the bug**
If a DEK for kekId X is generated and used for some time, and it expires or is used to encrypt the maximum number of records allowed, then encryption starts failing with exception for new records for kekId X. We do not have code to handle obtaining a new DEK under these conditions.

Pull request description:
Encryption Filter: Rotate to a new DEK when the old one is exhausted
### Type of change

- Bugfix

### Description

A hard requirement we have is to stop using a DEK for encryption once it has been used to encrypt a certain amount of records.

Work done with the key context is now synchronized such that only one thread will encrypt OR update the future in the cache when rotation is required. If the key was destroyed while a thread is waiting for the lock it can retry the encrypt as the stage in the cache should have been atomically updated by a previous encrypt call.

Resolves #757 

### 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/CHANGELOG.md b/CHANGELOG.md
index 0e4ee5a4..eaa10ed3 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -6,6 +6,7 @@ Please enumerate **all user-facing** changes using format `<githib issue/pr numb
 
 ## 0.4.0
 
+* [#764](https://github.com/kroxylicious/kroxylicious/pull/764): Encryption Filter: Rotate to a new DEK when the old one is exhausted
 * [#752](https://github.com/kroxylicious/kroxylicious/issues/752): Remove redudant reinstallation of time-zone data in Dockerfile used for Kroxylicious container image
 * [#727](https://github.com/kroxylicious/kroxylicious/pull/727): Tease out simple transform filters into their own module
 * [#738](https://github.com/kroxylicious/kroxylicious/pull/738): Update to Kroxylicious Junit Ext 0.7.0
diff --git a/kroxylicious-filters/kroxylicious-encryption/src/main/java/io/kroxylicious/filter/encryption/EnvelopeEncryption.java b/kroxylicious-filters/kroxylicious-encryption/src/main/java/io/kroxylicious/filter/encryption/EnvelopeEncryption.java
index ea5d65fd..93684f1b 100644
--- a/kroxylicious-filters/kroxylicious-encryption/src/main/java/io/kroxylicious/filter/encryption/EnvelopeEncryption.java
+++ b/kroxylicious-filters/kroxylicious-encryption/src/main/java/io/kroxylicious/filter/encryption/EnvelopeEncryption.java
@@ -47,7 +47,7 @@ public class EnvelopeEncryption<K, E> implements FilterFactory<EnvelopeEncryptio
         KmsService<Object, K, E> kmsPlugin = context.pluginInstance(KmsService.class, configuration.kms());
         Kms<K, E> kms = kmsPlugin.buildKms(configuration.kmsConfig());
 
-        var keyManager = new InBandKeyManager<>(kms, BufferPool.allocating());
+        var keyManager = new InBandKeyManager<>(kms, BufferPool.allocating(), 500_000);
 
         KekSelectorService<Object, K> ksPlugin = context.pluginInstance(KekSelectorService.class, configuration.selector());
         TopicNameBasedKekSelector<K> kekSelector = ksPlugin.buildSelector(kms, configuration.selectorConfig());
diff --git a/kroxylicious-filters/kroxylicious-encryption/src/main/java/io/kroxylicious/filter/encryption/KeyManager.java b/kroxylicious-filters/kroxylicious-encryption/src/main/java/io/kroxylicious/filter/encryption/KeyManager.java
index f54ee0ac..a9cf4041 100644
--- a/kroxylicious-filters/kroxylicious-encryption/src/main/java/io/kroxylicious/filter/encryption/KeyManager.java
+++ b/kroxylicious-filters/kroxylicious-encryption/src/main/java/io/kroxylicious/filter/encryption/KeyManager.java
@@ -24,8 +24,8 @@ public interface KeyManager<K> {
      * Asynchronously encrypt the given {@code recordRequests} using the current DEK for the given KEK, calling the given receiver for each encrypted record
      * @param encryptionScheme The encryption scheme.
      * @param records The requests to be encrypted.
-     * @param receiver The receiver of the encrypted buffers.
-     * @return A completion stage that completes when all the records in the given {@code recordRequests} have been processed.
+     * @param receiver The receiver of the encrypted buffers. The receiver is guaranteed to receive the encrypted buffers sequentially, in the same order as {@code records}, with no possibility of parallel invocation.
+     * @return A completion stage that completes when all the records have been processed.
      */
     @NonNull
     CompletionStage<Void> encrypt(@NonNull EncryptionScheme<K> encryptionScheme,
@@ -35,8 +35,8 @@ public interface KeyManager<K> {
     /**
      * Asynchronously decrypt the given {@code kafkaRecords} (if they were, in fact, encrypted), calling the given {@code receiver} with the plaintext
      * @param records The records
-     * @param receiver The receiver of the plaintext buffers
-     * @return A completion stage that completes when the record has been processed.
+     * @param receiver The receiver of the plaintext buffers. The receiver is guaranteed to receive the decrypted buffers sequentially, in the same order as {@code records}, with no possibility of parallel invocation.
+     * @return A completion stage that completes when all the records have been processed.
      */
     @NonNull
     CompletionStage<Void> decrypt(@NonNull List<? extends Record> records,
diff --git a/kroxylicious-filters/kroxylicious-encryption/src/main/java/io/kroxylicious/filter/encryption/Receiver.java b/kroxylicious-filters/kroxylicious-encryption/src/main/java/io/kroxylicious/filter/encryption/Receiver.java
index 62ff4b90..389eee51 100644
--- a/kroxylicious-filters/kroxylicious-encryption/src/main/java/io/kroxylicious/filter/encryption/Receiver.java
+++ b/kroxylicious-filters/kroxylicious-encryption/src/main/java/io/kroxylicious/filter/encryption/Receiver.java
@@ -19,7 +19,8 @@ public interface Receiver {
      * Receive the ciphertext (encryption) or the plaintext (decryption) associated with the given record..
      *
      * @param kafkaRecord The record on which to base the revised record
-     * @param value The ciphertext or plaintext buffer.
+     * @param value The ciphertext or plaintext buffer. This buffer may be re-used, the implementor should extract all
+     * the bytes they need from the buffer before the end of the accept call.
      */
     void accept(Record kafkaRecord, ByteBuffer value, Header[] headers);
 }
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 825bdd39..f6f5d4f2 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
@@ -11,22 +11,17 @@ import java.security.SecureRandom;
 import java.util.ArrayList;
 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.concurrent.ConcurrentHashMap;
 import java.util.function.Supplier;
 
-import javax.security.auth.DestroyFailedException;
-import javax.security.auth.Destroyable;
-
 import org.apache.kafka.common.header.Header;
 import org.apache.kafka.common.header.internals.RecordHeader;
 import org.apache.kafka.common.record.Record;
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
 
+import io.kroxylicious.filter.encryption.EncryptionException;
 import io.kroxylicious.filter.encryption.EncryptionScheme;
 import io.kroxylicious.filter.encryption.KeyManager;
 import io.kroxylicious.filter.encryption.Receiver;
@@ -45,11 +40,9 @@ import edu.umd.cs.findbugs.annotations.Nullable;
  */
 public class InBandKeyManager<K, E> implements KeyManager<K> {
 
-    private static final Logger LOGGER = LoggerFactory.getLogger(InBandKeyManager.class);
-    private static final Map<Class<? extends Destroyable>, Boolean> LOGGED_DESTROY_FAILED = new ConcurrentHashMap<>();
-
     private static final String FIELDS_HEADER_NAME = "kroxylicious.io/encrypted";
     private static final String DEK_HEADER_NAME = "kroxylicious.io/dek";
+    private static final int MAX_ATTEMPTS = 3;
 
     private final Kms<K, E> kms;
     private final BufferPool bufferPool;
@@ -61,12 +54,13 @@ public class InBandKeyManager<K, E> implements KeyManager<K> {
     private final int maxEncryptionsPerDek;
 
     public InBandKeyManager(Kms<K, E> kms,
-                            BufferPool bufferPool) {
+                            BufferPool bufferPool,
+                            int maxEncryptionsPerDek) {
         this.kms = kms;
         this.bufferPool = bufferPool;
         this.edekSerde = kms.edekSerde();
         this.dekTtlNanos = 5_000_000_000L;
-        this.maxEncryptionsPerDek = 500_000;
+        this.maxEncryptionsPerDek = maxEncryptionsPerDek;
         // TODO This ^^ must be > the maximum size of a batch to avoid an infinite loop
         this.keyContextCache = new ConcurrentHashMap<>();
         this.decryptorCache = new ConcurrentHashMap<>();
@@ -86,37 +80,8 @@ public class InBandKeyManager<K, E> implements KeyManager<K> {
         });
     }
 
-    private CompletionStage<KeyContext> currentDekContext(@NonNull K kekId, int numRecords) {
-        return getKeyContext(kekId, makeKeyContext(kekId))
-                .thenCompose(cachedContext -> {
-                    if (!cachedContext.isExpiredForEncryption(System.nanoTime())
-                            && cachedContext.hasAtLeastRemainingEncryptions(numRecords)) {
-                        // TODO pre-populate decryptorCache?
-                        return CompletableFuture.completedFuture(cachedContext);
-                    }
-                    else {
-                        destroy(cachedContext);
-                        return currentDekContext(kekId, numRecords);
-                    }
-                });
-    }
-
-    static void destroy(Destroyable destroyable) {
-        try {
-            destroyable.destroy();
-        }
-        catch (DestroyFailedException e) {
-            var cls = destroyable.getClass();
-            LOGGED_DESTROY_FAILED.compute(cls, (k, logged) -> {
-                if (logged == null) {
-                    LOGGER.warn("Failed to destroy an instance of {}. "
-                            + "Note: this message is logged once per class even though there may be many occurrences of this event. "
-                            + "This event can happen because the JRE's SecretKeySpec class does not override destroy().",
-                            cls, e);
-                }
-                return Boolean.TRUE;
-            });
-        }
+    private CompletionStage<KeyContext> currentDekContext(@NonNull K kekId) {
+        return getKeyContext(kekId, makeKeyContext(kekId));
     }
 
     private Supplier<CompletionStage<KeyContext>> makeKeyContext(@NonNull K kekId) {
@@ -147,27 +112,60 @@ public class InBandKeyManager<K, E> implements KeyManager<K> {
     public CompletionStage<Void> encrypt(@NonNull EncryptionScheme<K> encryptionScheme,
                                          @NonNull List<? extends Record> records,
                                          @NonNull Receiver receiver) {
+        return attemptEncrypt(encryptionScheme, records, receiver, 0);
+    }
 
-        var fieldsHeader = createEncryptedFieldsHeader(encryptionScheme.recordFields());
-        return currentDekContext(encryptionScheme.kekId(), records.size()).thenAccept(keyContext -> {
-            var dekHeader = createEdekHeader(keyContext);
-            var maxValuePlaintextSize = encryptionScheme.recordFields().contains(RecordField.RECORD_VALUE)
-                    ? records.stream().mapToInt(Record::valueSize).max().orElse(-1)
-                    : -1;
+    @SuppressWarnings("java:S2445")
+    private CompletionStage<Void> attemptEncrypt(@NonNull EncryptionScheme<K> encryptionScheme, @NonNull List<? extends Record> records,
+                                                 @NonNull Receiver receiver, int attempt) {
+        if (attempt >= MAX_ATTEMPTS) {
+            return CompletableFuture.failedFuture(new EncryptionException("failed to encrypt records after " + attempt + " attempts"));
+        }
+        return currentDekContext(encryptionScheme.kekId()).thenCompose(keyContext -> {
             synchronized (keyContext) {
-                var valueCiphertext = maxValuePlaintextSize >= 0 ? bufferPool.acquire(keyContext.encodedSize(maxValuePlaintextSize)) : null;
-                try {
-                    encryptRecords(encryptionScheme, keyContext, records, fieldsHeader, dekHeader, valueCiphertext, receiver);
-                }
-                finally {
-                    if (maxValuePlaintextSize >= 0) {
-                        bufferPool.release(valueCiphertext);
+                // 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 (!keyContext.isDestroyed()) {
+                    if (!keyContext.hasAtLeastRemainingEncryptions(records.size())) {
+                        // replace the key context stage in the cache, then call encrypt again
+                        rotateKeyContext(encryptionScheme, keyContext);
+                    }
+                    else {
+                        return encrypt(encryptionScheme, records, receiver, keyContext);
                     }
                 }
             }
+            return attemptEncrypt(encryptionScheme, records, receiver, attempt + 1);
         });
     }
 
+    @NonNull
+    private CompletableFuture<Void> encrypt(@NonNull EncryptionScheme<K> encryptionScheme, @NonNull List<? extends Record> records,
+                                            @NonNull Receiver receiver, KeyContext keyContext) {
+        var fieldsHeader = createEncryptedFieldsHeader(encryptionScheme.recordFields());
+        var dekHeader = createEdekHeader(keyContext);
+        var maxValuePlaintextSize = encryptionScheme.recordFields().contains(RecordField.RECORD_VALUE)
+                ? records.stream().mapToInt(Record::valueSize).max().orElse(-1)
+                : -1;
+        var sharedValueCiphertextBuffer = maxValuePlaintextSize >= 0 ? bufferPool.acquire(keyContext.encodedSize(maxValuePlaintextSize)) : null;
+        try {
+            encryptRecords(encryptionScheme, keyContext, records, fieldsHeader, dekHeader, sharedValueCiphertextBuffer, receiver);
+        }
+        finally {
+            if (sharedValueCiphertextBuffer != null) {
+                bufferPool.release(sharedValueCiphertextBuffer);
+            }
+        }
+        keyContext.recordEncryptions(records.size());
+        return CompletableFuture.completedFuture(null);
+    }
+
+    // this must only be called while holding the lock on this keycontext
+    private void rotateKeyContext(@NonNull EncryptionScheme<K> encryptionScheme, KeyContext keyContext) {
+        keyContext.destroy();
+        K kekId = encryptionScheme.kekId();
+        keyContextCache.put(kekId, makeKeyContext(kekId).get());
+    }
+
     private void encryptRecords(@NonNull EncryptionScheme<K> encryptionScheme,
                                 @NonNull KeyContext keyContext,
                                 @NonNull List<? extends Record> records,
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 5e5e59b7..c4fcdb58 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
@@ -7,21 +7,34 @@
 package io.kroxylicious.filter.encryption.inband;
 
 import java.nio.ByteBuffer;
-import java.util.concurrent.atomic.AtomicInteger;
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
 
+import javax.annotation.concurrent.NotThreadSafe;
 import javax.security.auth.DestroyFailedException;
 import javax.security.auth.Destroyable;
 
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
 import edu.umd.cs.findbugs.annotations.NonNull;
 
 /**
- * A DekContext encapsulates an encryptor.
+ * A Context encapsulates an encryptor for a DEK.
+ * <p>
+ * KeyContext is not threadsafe and access should be externally co-ordinated.
+ * </p>
  */
+@NotThreadSafe
 final class KeyContext implements Destroyable {
     private final AesGcmEncryptor encryptor;
     private final byte[] prefix;
     private final long encryptionExpiryNanos;
-    private final AtomicInteger remainingEncryptions;
+    private int remainingEncryptions;
+    private boolean destroyed = false;
+
+    private static final Logger LOGGER = LoggerFactory.getLogger(InBandKeyManager.class);
+    private static final Map<Class<? extends Destroyable>, Boolean> LOGGED_DESTROY_FAILED = new ConcurrentHashMap<>();
 
     KeyContext(@NonNull ByteBuffer prefix,
                long encryptionExpiryNanos,
@@ -32,7 +45,7 @@ final class KeyContext implements Destroyable {
         }
         this.prefix = prefix.array();
         this.encryptionExpiryNanos = encryptionExpiryNanos;
-        this.remainingEncryptions = new AtomicInteger(maxEncryptions);
+        this.remainingEncryptions = maxEncryptions;
         this.encryptor = encryptor;
     }
 
@@ -48,7 +61,14 @@ final class KeyContext implements Destroyable {
         if (numEncryptions <= 0) {
             throw new IllegalArgumentException();
         }
-        return remainingEncryptions.getAndAdd(-numEncryptions) >= numEncryptions;
+        return remainingEncryptions >= numEncryptions;
+    }
+
+    public void recordEncryptions(int numEncryptions) {
+        if (numEncryptions <= 0) {
+            throw new IllegalArgumentException();
+        }
+        remainingEncryptions -= numEncryptions;
     }
 
     /**
@@ -67,14 +87,36 @@ final class KeyContext implements Destroyable {
      * @param output The output buffer
      */
     public void encode(@NonNull ByteBuffer plaintext, @NonNull ByteBuffer output) {
-        if (remainingEncryptions.get() < 0) {
+        if (remainingEncryptions <= 0) {
             throw new ExhaustedDekException("No more encryptions");
         }
         encryptor.encrypt(plaintext, output);
     }
 
     @Override
-    public void destroy() throws DestroyFailedException {
-        encryptor.destroy();
+    public void destroy() {
+        destroy(encryptor);
+        destroyed = true;
+    }
+
+    @Override
+    public boolean isDestroyed() {
+        return destroyed;
+    }
+
+    static void destroy(Destroyable destroyable) {
+        try {
+            destroyable.destroy();
+        }
+        catch (DestroyFailedException e) {
+            var cls = destroyable.getClass();
+            LOGGED_DESTROY_FAILED.computeIfAbsent(cls, (c) -> {
+                LOGGER.warn("Failed to destroy an instance of {}. "
+                        + "Note: this message is logged once per class even though there may be many occurrences of this event. "
+                        + "This event can happen because the JRE's SecretKeySpec class does not override destroy().",
+                        c, e);
+                return Boolean.TRUE;
+            });
+        }
     }
 }
diff --git a/kroxylicious-kms-provider-kroxylicious-inmemory/src/main/java/io/kroxylicious/kms/provider/kroxylicious/inmemory/InMemoryKms.java b/kroxylicious-kms-provider-kroxylicious-inmemory/src/main/java/io/kroxylicious/kms/provider/kroxylicious/inmemory/InMemoryKms.java
index 3ffc5a7c..748e9063 100644
--- a/kroxylicious-kms-provider-kroxylicious-inmemory/src/main/java/io/kroxylicious/kms/provider/kroxylicious/inmemory/InMemoryKms.java
+++ b/kroxylicious-kms-provider-kroxylicious-inmemory/src/main/java/io/kroxylicious/kms/provider/kroxylicious/inmemory/InMemoryKms.java
@@ -11,10 +11,11 @@ import java.security.InvalidKeyException;
 import java.security.NoSuchAlgorithmException;
 import java.security.SecureRandom;
 import java.util.HashMap;
+import java.util.List;
 import java.util.Map;
 import java.util.UUID;
 import java.util.concurrent.CompletableFuture;
-import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.CopyOnWriteArrayList;
 import java.util.function.Supplier;
 
 import javax.crypto.Cipher;
@@ -48,16 +49,13 @@ public class InMemoryKms implements
     private final int numAuthBits;
     private final SecureRandom secureRandom;
     private final Map<String, UUID> aliases;
-
-    private final AtomicInteger numDeksGenerated;
+    private final List<DekPair<InMemoryEdek>> edeksGenerated = new CopyOnWriteArrayList<>();
 
     InMemoryKms(int numIvBytes, int numAuthBits,
                 Map<UUID, SecretKey> keys,
-                Map<String, UUID> aliases,
-                AtomicInteger numDeksGenerated) {
+                Map<String, UUID> aliases) {
         this.keys = new HashMap<>(keys);
         this.aliases = new HashMap<>(aliases);
-        this.numDeksGenerated = numDeksGenerated;
         this.secureRandom = new SecureRandom();
         this.numIvBytes = numIvBytes;
         this.numAuthBits = numAuthBits;
@@ -97,7 +95,7 @@ public class InMemoryKms implements
      * @return the number of DEKs that have been generated.
      */
     public int numDeksGenerated() {
-        return numDeksGenerated.get();
+        return edeksGenerated.size();
     }
 
     private InMemoryEdek wrap(UUID kekRef, Supplier<SecretKey> generator) {
@@ -121,8 +119,9 @@ public class InMemoryKms implements
         try {
             var dek = this.aes.generateKey();
             var edek = wrap(kekRef, () -> dek);
-            numDeksGenerated.incrementAndGet();
-            return CompletableFuture.completedFuture(new DekPair<>(edek, dek));
+            DekPair<InMemoryEdek> dekPair = new DekPair<>(edek, dek);
+            edeksGenerated.add(dekPair);
+            return CompletableFuture.completedFuture(dekPair);
         }
         catch (KmsException e) {
             return CompletableFuture.failedFuture(e);
@@ -208,4 +207,8 @@ public class InMemoryKms implements
     public Serde<InMemoryEdek> edekSerde() {
         return InMemoryEdekSerde.instance();
     }
+
+    public DekPair<InMemoryEdek> getGeneratedEdek(int i) {
+        return edeksGenerated.get(i);
+    }
 }


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 d6d47f7a..228f6dbf 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
@@ -7,17 +7,16 @@
 package io.kroxylicious.filter.encryption.inband;
 
 import java.nio.ByteBuffer;
+import java.time.Duration;
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.Collection;
 import java.util.EnumSet;
 import java.util.List;
+import java.util.concurrent.CompletionStage;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.TimeoutException;
-import java.util.concurrent.atomic.AtomicBoolean;
-
-import javax.security.auth.DestroyFailedException;
-import javax.security.auth.Destroyable;
 
 import org.apache.kafka.common.header.Header;
 import org.apache.kafka.common.header.internals.RecordHeader;
@@ -28,12 +27,14 @@ import io.kroxylicious.filter.encryption.Receiver;
 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.Serde;
 
 import edu.umd.cs.findbugs.annotations.NonNull;
 
+import static org.assertj.core.api.Assertions.assertThat;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertNotEquals;
-import static org.junit.jupiter.api.Assertions.assertTrue;
 
 class InBandKeyManagerTest {
 
@@ -73,25 +74,11 @@ class InBandKeyManagerTest {
         assertEquals(myHeader, x1[0]);
     }
 
-    @Test
-    void testDestroy() {
-        AtomicBoolean called = new AtomicBoolean();
-        InBandKeyManager.destroy(new Destroyable() {
-            @Override
-            public void destroy() throws DestroyFailedException {
-                called.set(true);
-                throw new DestroyFailedException("Eeek!");
-            }
-        });
-
-        assertTrue(called.get());
-    }
-
     @Test
     void shouldEncryptRecordValue() {
         var kmsService = UnitTestingKmsService.newInstance();
         InMemoryKms kms = kmsService.buildKms(new UnitTestingKmsService.Config());
-        var km = new InBandKeyManager<>(kms, BufferPool.allocating());
+        var km = new InBandKeyManager<>(kms, BufferPool.allocating(), 500_000);
 
         var kekId = kms.generateKey();
 
@@ -114,6 +101,27 @@ class InBandKeyManagerTest {
         assertEquals(initial, decrypted);
     }
 
+    @Test
+    void encryptionRetry() {
+        var kmsService = UnitTestingKmsService.newInstance();
+        InMemoryKms kms = kmsService.buildKms(new UnitTestingKmsService.Config());
+        var kekId = kms.generateKey();
+        // configure 1 encryption per dek but then try to encrypt 2 records, will destroy and retry
+        var km = new InBandKeyManager<>(kms, BufferPool.allocating(), 1);
+
+        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(new EncryptionScheme<>(kekId, EnumSet.of(RecordField.RECORD_VALUE)),
+                initial,
+                recordReceivedRecord(encrypted));
+        assertThat(encrypt).failsWithin(Duration.ZERO).withThrowableThat().withMessageContaining("failed to encrypt records after 3 attempts");
+    }
+
     @NonNull
     private static Receiver recordReceivedRecord(Collection<TestingRecord> list) {
         return (r, v, h) -> {
@@ -125,7 +133,7 @@ class InBandKeyManagerTest {
     void shouldEncryptRecordValueForMultipleRecords() throws ExecutionException, InterruptedException, TimeoutException {
         var kmsService = UnitTestingKmsService.newInstance();
         InMemoryKms kms = kmsService.buildKms(new UnitTestingKmsService.Config());
-        var km = new InBandKeyManager<>(kms, BufferPool.allocating());
+        var km = new InBandKeyManager<>(kms, BufferPool.allocating(), 500_000);
 
         var kekId = kms.generateKey();
 
@@ -154,6 +162,120 @@ class InBandKeyManagerTest {
         assertEquals(initial, decrypted);
     }
 
+    @Test
+    void shouldGenerateNewDekIfOldDekHasNoRemainingEncryptions() throws ExecutionException, InterruptedException, TimeoutException {
+        var kmsService = UnitTestingKmsService.newInstance();
+        InMemoryKms kms = kmsService.buildKms(new UnitTestingKmsService.Config());
+        var km = new InBandKeyManager<>(kms, BufferPool.allocating(), 2);
+
+        var kekId = kms.generateKey();
+
+        var value = ByteBuffer.wrap(new byte[]{ 1, 2, 3 });
+        TestingRecord record = new TestingRecord(value);
+
+        var value2 = ByteBuffer.wrap(new byte[]{ 3, 4, 5 });
+        TestingRecord record2 = new TestingRecord(value2);
+
+        List<TestingRecord> encrypted = new ArrayList<>();
+        List<TestingRecord> initial = List.of(record, record2);
+        km.encrypt(new EncryptionScheme<>(kekId, EnumSet.of(RecordField.RECORD_VALUE)),
+                initial,
+                recordReceivedRecord(encrypted)).toCompletableFuture().get(10, TimeUnit.SECONDS);
+        record.value().rewind();
+        record2.value().rewind();
+
+        // at this point we have encrypted 2 records with the manager set to maximum 2 encryptions per dek
+
+        km.encrypt(new EncryptionScheme<>(kekId, EnumSet.of(RecordField.RECORD_VALUE)),
+                initial,
+                recordReceivedRecord(encrypted)).toCompletableFuture().get(10, TimeUnit.SECONDS);
+
+        record.value().rewind();
+        record2.value().rewind();
+
+        assertThat(kms.numDeksGenerated()).isEqualTo(2);
+        var edekOne = getSerializedGeneratedEdek(kms, 0);
+        var edekTwo = getSerializedGeneratedEdek(kms, 1);
+        assertThat(encrypted).hasSize(4);
+        List<TestingDek> deks = extractEdeks(encrypted);
+        assertThat(deks).containsExactly(edekOne, edekOne, edekTwo, edekTwo);
+    }
+
+    @Test
+    void shouldGenerateNewDekIfOldOneHasSomeRemainingEncryptionsButNotEnoughForWholeBatch() throws ExecutionException, InterruptedException, TimeoutException {
+        var kmsService = UnitTestingKmsService.newInstance();
+        InMemoryKms kms = kmsService.buildKms(new UnitTestingKmsService.Config());
+        var km = new InBandKeyManager<>(kms, BufferPool.allocating(), 3);
+
+        var kekId = kms.generateKey();
+
+        var value = ByteBuffer.wrap(new byte[]{ 1, 2, 3 });
+        TestingRecord record = new TestingRecord(value);
+
+        var value2 = ByteBuffer.wrap(new byte[]{ 3, 4, 5 });
+        TestingRecord record2 = new TestingRecord(value2);
+
+        List<TestingRecord> encrypted = new ArrayList<>();
+        List<TestingRecord> initial = List.of(record, record2);
+        km.encrypt(new EncryptionScheme<>(kekId, EnumSet.of(RecordField.RECORD_VALUE)),
+                initial,
+                recordReceivedRecord(encrypted)).toCompletableFuture().get(10, TimeUnit.SECONDS);
+        record.value().rewind();
+        record2.value().rewind();
+
+        // at this point we have encrypted 2 records with the manager set to maximum 3 encryptions per dek, so we need a new dek to encrypt 2 more records
+
+        km.encrypt(new EncryptionScheme<>(kekId, EnumSet.of(RecordField.RECORD_VALUE)),
+                initial,
+                recordReceivedRecord(encrypted)).toCompletableFuture().get(10, TimeUnit.SECONDS);
+
+        record.value().rewind();
+        record2.value().rewind();
+        assertThat(kms.numDeksGenerated()).isEqualTo(2);
+        var edekOne = getSerializedGeneratedEdek(kms, 0);
+        var edekTwo = getSerializedGeneratedEdek(kms, 1);
+        assertThat(encrypted).hasSize(4);
+        List<TestingDek> deks = extractEdeks(encrypted);
+        assertThat(deks).containsExactly(edekOne, edekOne, edekTwo, edekTwo);
+    }
+
+    @Test
+    void shouldUseSameDekForMultipleBatches() throws ExecutionException, InterruptedException, TimeoutException {
+        var kmsService = UnitTestingKmsService.newInstance();
+        InMemoryKms kms = kmsService.buildKms(new UnitTestingKmsService.Config());
+        var km = new InBandKeyManager<>(kms, BufferPool.allocating(), 4);
+
+        var kekId = kms.generateKey();
+
+        var value = ByteBuffer.wrap(new byte[]{ 1, 2, 3 });
+        TestingRecord record = new TestingRecord(value);
+
+        var value2 = ByteBuffer.wrap(new byte[]{ 3, 4, 5 });
+        TestingRecord record2 = new TestingRecord(value2);
+
+        List<TestingRecord> encrypted = new ArrayList<>();
+        List<TestingRecord> initial = List.of(record, record2);
+        km.encrypt(new EncryptionScheme<>(kekId, EnumSet.of(RecordField.RECORD_VALUE)),
+                initial,
+                recordReceivedRecord(encrypted)).toCompletableFuture().get(10, TimeUnit.SECONDS);
+        record.value().rewind();
+        record2.value().rewind();
+
+        // at this point we have encrypted 2 records with the manager set to maximum 4 encryptions per dek, so we do not need a new dek to encrypt 2 more records
+
+        km.encrypt(new EncryptionScheme<>(kekId, EnumSet.of(RecordField.RECORD_VALUE)),
+                initial,
+                recordReceivedRecord(encrypted)).toCompletableFuture().get(10, TimeUnit.SECONDS);
+
+        record.value().rewind();
+        record2.value().rewind();
+        assertThat(kms.numDeksGenerated()).isEqualTo(1);
+        var edekOne = getSerializedGeneratedEdek(kms, 0);
+        assertThat(encrypted).hasSize(4);
+        List<TestingDek> deks = extractEdeks(encrypted);
+        assertThat(deks).containsExactly(edekOne, edekOne, edekOne, edekOne);
+    }
+
     @NonNull
     private static ByteBuffer copyBytes(ByteBuffer v) {
         byte[] bytes = new byte[v.remaining()];
@@ -166,7 +288,7 @@ class InBandKeyManagerTest {
     void shouldEncryptRecordHeaders() {
         var kmsService = UnitTestingKmsService.newInstance();
         InMemoryKms kms = kmsService.buildKms(new UnitTestingKmsService.Config());
-        var km = new InBandKeyManager<>(kms, BufferPool.allocating());
+        var km = new InBandKeyManager<>(kms, BufferPool.allocating(), 500_000);
 
         var kekId = kms.generateKey();
 
@@ -194,7 +316,7 @@ class InBandKeyManagerTest {
     void shouldEncryptRecordHeadersForMultipleRecords() throws ExecutionException, InterruptedException, TimeoutException {
         var kmsService = UnitTestingKmsService.newInstance();
         InMemoryKms kms = kmsService.buildKms(new UnitTestingKmsService.Config());
-        var km = new InBandKeyManager<>(kms, BufferPool.allocating());
+        var km = new InBandKeyManager<>(kms, BufferPool.allocating(), 500_000);
 
         var kekId = kms.generateKey();
 
@@ -222,4 +344,28 @@ class InBandKeyManagerTest {
                 new TestingRecord(value2, new Header[]{ new RecordHeader("foo", new byte[]{ 10, 11, 12 }) })), decrypted);
     }
 
+    public TestingDek getSerializedGeneratedEdek(InMemoryKms kms, int i) {
+        DekPair generatedEdek = kms.getGeneratedEdek(i);
+        var edek = generatedEdek.edek();
+        Serde serde = kms.edekSerde();
+        int size = serde.sizeOf(edek);
+        ByteBuffer buffer = ByteBuffer.allocate(size + 2);
+        buffer.putShort((short) size);
+        serde.serialize(edek, buffer);
+        buffer.flip();
+        byte[] bytes = new byte[buffer.remaining()];
+        buffer.get(bytes);
+        return new TestingDek(bytes);
+    }
+
+    @NonNull
+    private static List<TestingDek> extractEdeks(List<TestingRecord> encrypted) {
+        List<TestingDek> deks = encrypted.stream()
+                .map(testingRecord -> Arrays.stream(testingRecord.headers()).filter(header -> header.key().equals("kroxylicious.io/dek")).findFirst().orElseThrow()
+                        .value())
+                .map(TestingDek::new)
+                .toList();
+        return deks;
+    }
+
 }
diff --git a/kroxylicious-filters/kroxylicious-encryption/src/test/java/io/kroxylicious/filter/encryption/inband/KeyContextTest.java b/kroxylicious-filters/kroxylicious-encryption/src/test/java/io/kroxylicious/filter/encryption/inband/KeyContextTest.java
index 0f51103b..f289afa8 100644
--- a/kroxylicious-filters/kroxylicious-encryption/src/test/java/io/kroxylicious/filter/encryption/inband/KeyContextTest.java
+++ b/kroxylicious-filters/kroxylicious-encryption/src/test/java/io/kroxylicious/filter/encryption/inband/KeyContextTest.java
@@ -10,8 +10,10 @@ import java.nio.ByteBuffer;
 import java.nio.charset.StandardCharsets;
 import java.security.SecureRandom;
 import java.util.concurrent.ExecutionException;
+import java.util.concurrent.atomic.AtomicBoolean;
 
 import javax.security.auth.DestroyFailedException;
+import javax.security.auth.Destroyable;
 
 import org.junit.jupiter.api.Test;
 
@@ -44,6 +46,20 @@ class KeyContextTest {
         assertThrows(IllegalArgumentException.class, () -> context.hasAtLeastRemainingEncryptions(-1));
     }
 
+    @Test
+    void testDestroy() {
+        AtomicBoolean called = new AtomicBoolean();
+        KeyContext.destroy(new Destroyable() {
+            @Override
+            public void destroy() throws DestroyFailedException {
+                called.set(true);
+                throw new DestroyFailedException("Eeek!");
+            }
+        });
+
+        assertTrue(called.get());
+    }
+
     @Test
     void testLifecycle() throws ExecutionException, InterruptedException, DestroyFailedException {
         InMemoryKms kms = UnitTestingKmsService.newInstance().buildKms(new UnitTestingKmsService.Config());
@@ -65,6 +81,7 @@ class KeyContextTest {
         assertFalse(output.hasRemaining());
         output.flip();
         bb.flip();
+        context.recordEncryptions(1);
 
         assertTrue(context.hasAtLeastRemainingEncryptions(1));
         assertFalse(context.isExpiredForEncryption(101010));
@@ -77,6 +94,7 @@ class KeyContextTest {
         assertFalse(output.hasRemaining());
         output.flip();
         bb.flip();
+        context.recordEncryptions(1);
 
         assertFalse(context.hasAtLeastRemainingEncryptions(1));
         context.encodedSize(bb.capacity());
@@ -84,7 +102,9 @@ class KeyContextTest {
         var e = assertThrows(ExhaustedDekException.class, () -> context.encode(bb, output));
         assertEquals("No more encryptions", e.getMessage());
 
+        assertFalse(context.isDestroyed());
         // destroy
-        assertThrows(DestroyFailedException.class, context::destroy);
+        context.destroy();
+        assertTrue(context.isDestroyed());
     }
 }
diff --git a/kroxylicious-filters/kroxylicious-encryption/src/test/java/io/kroxylicious/filter/encryption/inband/TestingDek.java b/kroxylicious-filters/kroxylicious-encryption/src/test/java/io/kroxylicious/filter/encryption/inband/TestingDek.java
new file mode 100644
index 00000000..e384e977
--- /dev/null
+++ b/kroxylicious-filters/kroxylicious-encryption/src/test/java/io/kroxylicious/filter/encryption/inband/TestingDek.java
@@ -0,0 +1,29 @@
+/*
+ * 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;
+
+import java.util.Arrays;
+
+public record TestingDek(byte[] serializedEdek) {
+
+    @Override
+    public boolean equals(Object o) {
+        if (this == o) {
+            return true;
+        }
+        if (o == null || getClass() != o.getClass()) {
+            return false;
+        }
+        TestingDek that = (TestingDek) o;
+        return Arrays.equals(serializedEdek, that.serializedEdek);
+    }
+
+    @Override
+    public int hashCode() {
+        return Arrays.hashCode(serializedEdek);
+    }
+}
diff --git a/kroxylicious-kms-provider-kroxylicious-inmemory/src/main/java/io/kroxylicious/kms/provider/kroxylicious/inmemory/IntegrationTestingKmsService.java b/kroxylicious-kms-provider-kroxylicious-inmemory/src/main/java/io/kroxylicious/kms/provider/kroxylicious/inmemory/IntegrationTestingKmsService.java
index 1e311b8a..0828beb1 100644
--- a/kroxylicious-kms-provider-kroxylicious-inmemory/src/main/java/io/kroxylicious/kms/provider/kroxylicious/inmemory/IntegrationTestingKmsService.java
+++ b/kroxylicious-kms-provider-kroxylicious-inmemory/src/main/java/io/kroxylicious/kms/provider/kroxylicious/inmemory/IntegrationTestingKmsService.java
@@ -10,7 +10,6 @@ import java.util.Map;
 import java.util.ServiceLoader;
 import java.util.UUID;
 import java.util.concurrent.ConcurrentHashMap;
-import java.util.concurrent.atomic.AtomicInteger;
 
 import javax.crypto.SecretKey;
 
@@ -65,14 +64,12 @@ public class IntegrationTestingKmsService implements KmsService<IntegrationTesti
     @Override
     public InMemoryKms buildKms(Config options) {
         return KMSES.computeIfAbsent(options.name(), ignored -> {
-            var numGeneratedDeks = new AtomicInteger();
             var keys = new ConcurrentHashMap<UUID, SecretKey>();
             var aliases = new ConcurrentHashMap<String, UUID>();
             return new InMemoryKms(12,
                     128,
                     keys,
-                    aliases,
-                    numGeneratedDeks);
+                    aliases);
         });
     }
 
diff --git a/kroxylicious-kms-provider-kroxylicious-inmemory/src/main/java/io/kroxylicious/kms/provider/kroxylicious/inmemory/UnitTestingKmsService.java b/kroxylicious-kms-provider-kroxylicious-inmemory/src/main/java/io/kroxylicious/kms/provider/kroxylicious/inmemory/UnitTestingKmsService.java
index 149cb601..06ef8965 100644
--- a/kroxylicious-kms-provider-kroxylicious-inmemory/src/main/java/io/kroxylicious/kms/provider/kroxylicious/inmemory/UnitTestingKmsService.java
+++ b/kroxylicious-kms-provider-kroxylicious-inmemory/src/main/java/io/kroxylicious/kms/provider/kroxylicious/inmemory/UnitTestingKmsService.java
@@ -10,7 +10,6 @@ import java.util.Map;
 import java.util.ServiceLoader;
 import java.util.UUID;
 import java.util.concurrent.ConcurrentHashMap;
-import java.util.concurrent.atomic.AtomicInteger;
 
 import javax.crypto.SecretKey;
 
@@ -61,8 +60,6 @@ public class UnitTestingKmsService implements KmsService<UnitTestingKmsService.C
 
     }
 
-    private final AtomicInteger numGeneratedDeks = new AtomicInteger();
-
     private final Map<UUID, SecretKey> keys = new ConcurrentHashMap<>();
     private final Map<String, UUID> aliases = new ConcurrentHashMap<>();
 
@@ -72,8 +69,7 @@ public class UnitTestingKmsService implements KmsService<UnitTestingKmsService.C
         return new InMemoryKms(options.numIvBytes(),
                 options.numAuthBits(),
                 keys,
-                aliases,
-                numGeneratedDeks);
+                aliases);
     }
 
 }