Refactor getAead method so it can be overridden

confluentinc · Dec 12, 2023 · 1493691 · 1493691
1 parent 1ae8d40
commit 1493691
Showing 1 changed file with 6 additions and 2 deletions.
diff --git a/dek-registry/src/main/java/io/confluent/dekregistry/storage/DekRegistry.java b/dek-registry/src/main/java/io/confluent/dekregistry/storage/DekRegistry.java
@@ -576,7 +576,7 @@ public DataEncryptionKey createDek(String kekName, CreateDekRequest request)
   protected DataEncryptionKey generateEncryptedDek(KeyEncryptionKey kek, DataEncryptionKey key)
       throws DekGenerationException {
     try {
-      Aead aead = kek.toKekEntity().toAead(config.originals());
+      Aead aead = getAead(kek);
       // Generate new dek
       byte[] rawDek = getCryptor(key.getAlgorithm()).generateKey();
       byte[] encryptedDek = aead.encrypt(rawDek, EMPTY_AAD);
@@ -595,7 +595,7 @@ protected DataEncryptionKey generateRawDek(KeyEncryptionKey kek, DataEncryptionK
       throws DekGenerationException {
     try {
       // Decrypt dek
-      Aead aead = kek.toKekEntity().toAead(config.originals());
+      Aead aead = getAead(kek);
       byte[] encryptedDek = Base64.getDecoder().decode(
           key.getEncryptedKeyMaterial().getBytes(StandardCharsets.UTF_8));
       byte[] rawDek = aead.decrypt(encryptedDek, EMPTY_AAD);
@@ -613,6 +613,10 @@ protected DataEncryptionKey generateRawDek(KeyEncryptionKey kek, DataEncryptionK
     }
   }
 
+  protected Aead getAead(KeyEncryptionKey kek) throws GeneralSecurityException {
+    return kek.toKekEntity().toAead(config.originals());
+  }
+
   public Kek putKekOrForward(String name, UpdateKekRequest request,
       Map<String, String> headerProperties) throws SchemaRegistryException {
     String tenant = schemaRegistry.tenant();