Repository: kroxylicious/kroxylicious (java)

Issue:
Add support for AES_GCM 256bit keys
**Is your feature request related to a problem? Please describe.**
Our documentation says we support 256 bit keys but the implementation only offers 128 (as its the only version the JDK mandates)

**Describe the solution you'd like**
Make `AEC_GCM_256` available

**Describe alternatives you've considered**
A clear and concise description of any alternative solutions or features you've considered.

**Additional context**
Add any other context or screenshots about the feature request here.


Pull request:
Change AES GCM cipher to require a 256bit key
### Type of change

- Bugfix

### Description

Closes #1044

The original design and recommendations were for a 256bit key but our CipherSpec declared it was 128bit. Then it's algorithm, AES/GCM/NoPadding, didn't enforce the length so it happily used a 256bit key from the KMS so there was a mismatch between the name and what was enforce.

This change aligns the CipherSpec naming on 256 and updates the algorithm to enforce the key size.

### 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 Sonarcloud 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 eafd24bc..8574063c 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -4,6 +4,7 @@ Please enumerate **all user-facing** changes using format `<githib issue/pr numb
 
 ## 0.5.0
 
+* [#1050](https://github.com/kroxylicious/kroxylicious/pull/1050): Change AES GCM cipher to require a 256bit key
 * [#1049](https://github.com/kroxylicious/kroxylicious/pull/1049): Add deprecated EnvelopeEncryption filter to ease migration to RecordEncryption filter
 * [#1043](https://github.com/kroxylicious/kroxylicious/pull/1043): Rename EnvelopeEncryption filter to RecordEncryption
 * [#1029](https://github.com/kroxylicious/kroxylicious/pull/1029): Upgrade to Kafka 3.7.0
diff --git a/kroxylicious-filters/kroxylicious-record-encryption/src/main/java/io/kroxylicious/filter/encryption/dek/CipherSpec.java b/kroxylicious-filters/kroxylicious-record-encryption/src/main/java/io/kroxylicious/filter/encryption/dek/CipherSpec.java
index 40953218..b9d1b09d 100644
--- a/kroxylicious-filters/kroxylicious-record-encryption/src/main/java/io/kroxylicious/filter/encryption/dek/CipherSpec.java
+++ b/kroxylicious-filters/kroxylicious-record-encryption/src/main/java/io/kroxylicious/filter/encryption/dek/CipherSpec.java
@@ -29,11 +29,11 @@ import javax.crypto.spec.IvParameterSpec;
 public enum CipherSpec {
 
     /**
-     * AES/GCM with 128-bit key, 96-bit IV and 128-bit tag.
+     * AES/GCM with 256-bit key, 96-bit IV and 128-bit tag.
      * @see <a href="https://www.ietf.org/rfc/rfc5116.txt">RFC-5116</a>
      */
-    AES_128_GCM_128((byte) 0,
-            "AES/GCM/NoPadding",
+    AES_256_GCM_128((byte) 0,
+            "AES_256/GCM/NoPadding",
             1L << 32 // 2^32
     ) {
 
@@ -128,7 +128,7 @@ public enum CipherSpec {
     public static CipherSpec fromPersistentId(int persistentId) {
         switch (persistentId) {
             case 0:
-                return CipherSpec.AES_128_GCM_128;
+                return CipherSpec.AES_256_GCM_128;
             case 1:
                 return CipherSpec.CHACHA20_POLY1305;
             default:
diff --git a/kroxylicious-filters/kroxylicious-record-encryption/src/main/java/io/kroxylicious/filter/encryption/inband/EncryptionDekCache.java b/kroxylicious-filters/kroxylicious-record-encryption/src/main/java/io/kroxylicious/filter/encryption/inband/EncryptionDekCache.java
index 2e32562c..2bd15e96 100644
--- a/kroxylicious-filters/kroxylicious-record-encryption/src/main/java/io/kroxylicious/filter/encryption/inband/EncryptionDekCache.java
+++ b/kroxylicious-filters/kroxylicious-record-encryption/src/main/java/io/kroxylicious/filter/encryption/inband/EncryptionDekCache.java
@@ -41,7 +41,7 @@ public class EncryptionDekCache<K, E> {
     private record CacheKey<K>(K kek, CipherSpec cipherSpec) {}
 
     private static <K> CacheKey<K> cacheKey(EncryptionScheme<K> encryptionScheme) {
-        return new CacheKey<>(encryptionScheme.kekId(), CipherSpec.AES_128_GCM_128);
+        return new CacheKey<>(encryptionScheme.kekId(), CipherSpec.AES_256_GCM_128);
     }
 
     private final DekManager<K, E> dekManager;


Test patch:
diff --git a/kroxylicious-filters/kroxylicious-record-encryption/src/test/java/io/kroxylicious/filter/encryption/inband/FixedDekKmsService.java b/kroxylicious-filters/kroxylicious-record-encryption/src/test/java/io/kroxylicious/filter/encryption/inband/FixedDekKmsService.java
new file mode 100644
index 00000000..214c296f
--- /dev/null
+++ b/kroxylicious-filters/kroxylicious-record-encryption/src/test/java/io/kroxylicious/filter/encryption/inband/FixedDekKmsService.java
@@ -0,0 +1,103 @@
+/*
+ * 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.nio.ByteBuffer;
+import java.security.NoSuchAlgorithmException;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionStage;
+
+import javax.crypto.KeyGenerator;
+import javax.crypto.SecretKey;
+
+import io.kroxylicious.kms.service.DekPair;
+import io.kroxylicious.kms.service.Kms;
+import io.kroxylicious.kms.service.KmsService;
+import io.kroxylicious.kms.service.Serde;
+
+import edu.umd.cs.findbugs.annotations.NonNull;
+
+public class FixedDekKmsService implements KmsService<FixedDekKmsService.Config, ByteBuffer, ByteBuffer> {
+
+    private final SecretKey dek;
+    private final ByteBuffer edek;
+    private static final ByteBuffer KEK_ID = ByteBuffer.wrap(new byte[]{ 1, 2, 3 });
+    private final FixedEdekKms fixedEdekKms = new FixedEdekKms();
+
+    FixedDekKmsService(int keysize) {
+        try {
+            KeyGenerator generator = KeyGenerator.getInstance("AES");
+            generator.init(keysize);
+            dek = generator.generateKey();
+            edek = ByteBuffer.wrap(dek.getEncoded()); // not concerned with wrapping/unwrapping
+        }
+        catch (NoSuchAlgorithmException e) {
+            throw new RuntimeException(e);
+        }
+    }
+
+    @NonNull
+    @Override
+    public Kms<ByteBuffer, ByteBuffer> buildKms(Config options) {
+        return fixedEdekKms;
+    }
+
+    public ByteBuffer getKekId() {
+        return KEK_ID;
+    }
+
+    public Serde<ByteBuffer> edekSerde() {
+        return fixedEdekKms.edekSerde();
+    }
+
+    record Config() {}
+
+    private class FixedEdekKms implements Kms<ByteBuffer, ByteBuffer> {
+
+        @NonNull
+        @Override
+        public CompletionStage<DekPair<ByteBuffer>> generateDekPair(@NonNull ByteBuffer kekRef) {
+            return CompletableFuture.completedFuture(new DekPair<>(edek, dek));
+        }
+
+        @NonNull
+        @Override
+        public CompletionStage<SecretKey> decryptEdek(@NonNull ByteBuffer edek) {
+            return CompletableFuture.completedFuture(dek);
+        }
+
+        @NonNull
+        @Override
+        public Serde<ByteBuffer> edekSerde() {
+            return new Serde<>() {
+                @Override
+                public int sizeOf(ByteBuffer object) {
+                    return object.remaining();
+                }
+
+                @Override
+                public void serialize(ByteBuffer object, @NonNull ByteBuffer buffer) {
+                    var p0 = object.position();
+                    buffer.put(object);
+                    object.position(p0);
+                }
+
+                @Override
+                public ByteBuffer deserialize(@NonNull ByteBuffer buffer) {
+                    throw new UnsupportedOperationException();
+                }
+            };
+        }
+
+        @NonNull
+        @Override
+        public CompletionStage<ByteBuffer> resolveAlias(@NonNull String alias) {
+            return CompletableFuture.completedFuture(KEK_ID);
+        }
+    }
+
+}
diff --git a/kroxylicious-filters/kroxylicious-record-encryption/src/test/java/io/kroxylicious/filter/encryption/inband/RecordEncryptorTest.java b/kroxylicious-filters/kroxylicious-record-encryption/src/test/java/io/kroxylicious/filter/encryption/inband/RecordEncryptorTest.java
index 677e2f30..bf7bbd4f 100644
--- a/kroxylicious-filters/kroxylicious-record-encryption/src/test/java/io/kroxylicious/filter/encryption/inband/RecordEncryptorTest.java
+++ b/kroxylicious-filters/kroxylicious-record-encryption/src/test/java/io/kroxylicious/filter/encryption/inband/RecordEncryptorTest.java
@@ -6,17 +6,17 @@
 
 package io.kroxylicious.filter.encryption.inband;
 
+import java.io.Closeable;
 import java.nio.ByteBuffer;
-import java.security.NoSuchAlgorithmException;
 import java.util.Set;
-
-import javax.crypto.KeyGenerator;
-import javax.crypto.SecretKey;
+import java.util.concurrent.CompletionStage;
+import java.util.concurrent.TimeUnit;
 
 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.record.RecordBatch;
+import org.junit.jupiter.api.AfterAll;
 import org.junit.jupiter.api.BeforeAll;
 import org.junit.jupiter.api.Disabled;
 import org.junit.jupiter.api.Test;
@@ -26,76 +26,82 @@ import io.kroxylicious.filter.encryption.EncryptionVersion;
 import io.kroxylicious.filter.encryption.RecordField;
 import io.kroxylicious.filter.encryption.dek.CipherSpec;
 import io.kroxylicious.filter.encryption.dek.Dek;
+import io.kroxylicious.filter.encryption.dek.DekException;
+import io.kroxylicious.filter.encryption.dek.DekManager;
 import io.kroxylicious.filter.encryption.records.RecordTransform;
 import io.kroxylicious.kms.service.Serde;
 import io.kroxylicious.test.assertj.KafkaAssertions;
 import io.kroxylicious.test.record.RecordTestUtils;
 
-import edu.umd.cs.findbugs.annotations.NonNull;
-
-import static org.mockito.ArgumentMatchers.anyInt;
-import static org.mockito.Mockito.doReturn;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
 import static org.mockito.Mockito.mock;
 
 class RecordEncryptorTest {
 
-    private static Dek.Decryptor DECRYPTOR;
-
-    private static Dek.Encryptor ENCRYPTOR;
-    private static ByteBuffer EDEK;
-    private static Serde<ByteBuffer> EDEK_SERDE;
+    private static TestComponents components;
+    private static FixedDekKmsService fixedDekKmsService;
+
+    record TestComponents(
+                          Dek.Encryptor encryptor,
+                          Dek.Decryptor decryptor,
+                          Serde<ByteBuffer> edekSerde)
+            implements Closeable {
+        public void close() {
+            encryptor.close();
+            decryptor.close();
+        }
+    }
 
     @BeforeAll
-    public static void initKeyContext() throws NoSuchAlgorithmException, ReflectiveOperationException {
-        var generator = KeyGenerator.getInstance("AES");
-        var key = generator.generateKey();
-        EDEK = ByteBuffer.wrap(key.getEncoded()); // it doesn't matter for this test that it's not encrypted
-
-        var dek = mock(Dek.class);
-        doReturn(EDEK).when(dek).edek();
-        var encryptorConstructor = Dek.Encryptor.class.getDeclaredConstructor(Dek.class, CipherSpec.class, SecretKey.class, Integer.TYPE);
-        encryptorConstructor.setAccessible(true);
-        ENCRYPTOR = encryptorConstructor.newInstance(dek, CipherSpec.AES_128_GCM_128, key, 1_000_000);
-        doReturn(ENCRYPTOR).when(dek).encryptor(anyInt());
-
-        var decryptorConstructor = Dek.Decryptor.class.getDeclaredConstructor(Dek.class, CipherSpec.class, SecretKey.class);
-        decryptorConstructor.setAccessible(true);
-        DECRYPTOR = decryptorConstructor.newInstance(dek, CipherSpec.AES_128_GCM_128, key);
-        doReturn(DECRYPTOR).when(dek).decryptor();
-
-        EDEK_SERDE = new Serde<ByteBuffer>() {
-            @Override
-            public int sizeOf(ByteBuffer object) {
-                return object.remaining();
-            }
-
-            @Override
-            public void serialize(
-                                  ByteBuffer object,
-                                  @NonNull ByteBuffer buffer) {
-                var p0 = object.position();
-                buffer.put(object);
-                object.position(p0);
-            }
-
-            @Override
-            public ByteBuffer deserialize(@NonNull ByteBuffer buffer) {
-                throw new UnsupportedOperationException();
-            }
-        };
+    public static void initKeyContext() {
+        components = setup(256);
+    }
+
+    @AfterAll
+    public static void teardown() {
+        components.close();
+    }
+
+    static TestComponents setup(int keysize) {
+        try {
+            fixedDekKmsService = new FixedDekKmsService(keysize);
+            DekManager<ByteBuffer, ByteBuffer> manager = new DekManager<>(fixedDekKmsService, new FixedDekKmsService.Config(), 10000);
+            CompletionStage<Dek<ByteBuffer>> dekCompletionStage = manager.generateDek(fixedDekKmsService.getKekId(), CipherSpec.AES_256_GCM_128);
+            Dek<ByteBuffer> dek = dekCompletionStage.toCompletableFuture().get(0, TimeUnit.SECONDS);
+            var encryptor = dek.encryptor(1000);
+            var decryptor = dek.decryptor();
+            return new TestComponents(encryptor, decryptor, fixedDekKmsService.edekSerde());
+        }
+        catch (Exception e) {
+            throw new RuntimeException(e);
+        }
+    }
+
+    @Test
+    void aes256KeyMustBe256bits() {
+        try (TestComponents componentsWith128BitDek = setup(128)) {
+            Set<RecordField> fields = Set.of(RecordField.RECORD_VALUE);
+            long offset = 55L;
+            long timestamp = System.currentTimeMillis();
+            var key = "hello";
+            var value = "world";
+
+            assertThatThrownBy(() -> encryptSingleRecord(componentsWith128BitDek, fields, offset, timestamp, key, value)).isInstanceOf(DekException.class)
+                    .hasMessageContaining("The key must be 32 bytes");
+        }
     }
 
-    private Record encryptSingleRecord(Set<RecordField> fields, long offset, long timestamp, String key, String value, Header... headers) {
+    private Record encryptSingleRecord(TestComponents components, Set<RecordField> fields, long offset, long timestamp, String key, String value, Header... headers) {
         var re = new RecordEncryptor(
                 "topic", 0,
                 EncryptionVersion.V1,
                 new EncryptionScheme<>("key", fields),
-                EDEK_SERDE,
+                components.edekSerde,
                 ByteBuffer.allocate(100));
 
         Record record = RecordTestUtils.record(RecordBatch.MAGIC_VALUE_V2, offset, timestamp, key, value, headers);
 
-        return transformRecord(re, ENCRYPTOR, record);
+        return transformRecord(re, components.encryptor, record);
     }
 
     private <S> Record transformRecord(RecordTransform<S> recordTransform, S state, Record record) {
@@ -120,7 +126,7 @@ class RecordEncryptorTest {
         var value = "world";
 
         // When
-        var t = encryptSingleRecord(fields, offset, timestamp, key, value);
+        var t = encryptSingleRecord(components, fields, offset, timestamp, key, value);
 
         // Then
         KafkaAssertions.assertThat(t)
@@ -134,7 +140,7 @@ class RecordEncryptorTest {
 
         // And when
         var rd = new RecordDecryptor("topic", 0);
-        var rt = transformRecord(rd, new DecryptState(EncryptionVersion.V1).withDecryptor(DECRYPTOR), t);
+        var rt = transformRecord(rd, new DecryptState(EncryptionVersion.V1).withDecryptor(components.decryptor), t);
 
         // Then
         KafkaAssertions.assertThat(rt)
@@ -158,7 +164,7 @@ class RecordEncryptorTest {
         var header = new RecordHeader("bob", null);
 
         // When
-        var t = encryptSingleRecord(fields, offset, timestamp, key, value, header);
+        var t = encryptSingleRecord(components, fields, offset, timestamp, key, value, header);
 
         // Then
         KafkaAssertions.assertThat(t)
@@ -173,7 +179,7 @@ class RecordEncryptorTest {
         // And when
         var rd = new RecordDecryptor("topic", 0); // index -> new DecryptState(t, EncryptionVersion.V1, DECRYPTOR)
 
-        var rt = transformRecord(rd, new DecryptState(EncryptionVersion.V1).withDecryptor(DECRYPTOR), t);
+        var rt = transformRecord(rd, new DecryptState(EncryptionVersion.V1).withDecryptor(components.decryptor), t);
 
         // Then
         KafkaAssertions.assertThat(rt)
@@ -196,7 +202,7 @@ class RecordEncryptorTest {
         var value = (String) null;
 
         // When
-        var t = encryptSingleRecord(fields, offset, timestamp, key, value);
+        var t = encryptSingleRecord(components, fields, offset, timestamp, key, value);
 
         // Then
         KafkaAssertions.assertThat(t)
@@ -230,7 +236,7 @@ class RecordEncryptorTest {
         var header = new RecordHeader("bob", null);
 
         // When
-        var t = encryptSingleRecord(fields, offset, timestamp, key, value, header);
+        var t = encryptSingleRecord(components, fields, offset, timestamp, key, value, header);
 
         // Then
         KafkaAssertions.assertThat(t)
@@ -244,7 +250,7 @@ class RecordEncryptorTest {
 
         // And when
         var rd = new RecordDecryptor("topic", 0);
-        var rt = transformRecord(rd, new DecryptState<>(EncryptionVersion.V1).withDecryptor(DECRYPTOR), t);
+        var rt = transformRecord(rd, new DecryptState<>(EncryptionVersion.V1).withDecryptor(components.decryptor), t);
 
         // Then
         KafkaAssertions.assertThat(rt)
@@ -269,7 +275,7 @@ class RecordEncryptorTest {
         var header = new RecordHeader("bob", null);
 
         // When
-        var t = encryptSingleRecord(fields, offset, timestamp, key, value, header);
+        var t = encryptSingleRecord(components, fields, offset, timestamp, key, value, header);
 
         // Then
         KafkaAssertions.assertThat(t)
@@ -296,7 +302,7 @@ class RecordEncryptorTest {
         var header = new RecordHeader("bob", null);
 
         // When
-        var t = encryptSingleRecord(fields, offset, timestamp, key, value, header);
+        var t = encryptSingleRecord(components, fields, offset, timestamp, key, value, header);
 
         // Then
         KafkaAssertions.assertThat(t)
@@ -323,7 +329,7 @@ class RecordEncryptorTest {
         var header = new RecordHeader("bob", null);
 
         // When
-        var t = encryptSingleRecord(fields, offset, timestamp, key, value, header);
+        var t = encryptSingleRecord(components, fields, offset, timestamp, key, value, header);
 
         // Then
         KafkaAssertions.assertThat(t)