Merge pull request #2954 from confluentinc/fix-shared-keys

Fix sharedKeys cache to account for multiple use of kms key ID
confluentinc · Jan 25, 2024 · 309ab0b · 309ab0b
2 parents 6650405 + 0cb664f
commit 309ab0b
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 4 deletions.
diff --git a/...registry/src/main/java/io/confluent/dekregistry/storage/DefaultDekCacheUpdateHandler.java b/...registry/src/main/java/io/confluent/dekregistry/storage/DefaultDekCacheUpdateHandler.java
@@ -67,7 +67,7 @@ public void handleUpdate(
         if (oldValue instanceof KeyEncryptionKey) {
           KeyEncryptionKey oldKek = (KeyEncryptionKey) oldValue;
           if (oldKek.isShared()) {
-            dekRegistry.getSharedKeys().remove(oldKek.getKmsKeyId());
+            dekRegistry.getSharedKeys().remove(oldKek.getKmsKeyId(), (KeyEncryptionKeyId) key);
             dekRegistry.getMetricsManager().decrementSharedKeyCount(tenant);
           }
         }
@@ -93,7 +93,7 @@ public void handleUpdate(
           dekRegistry.getMetricsManager().incrementSharedKeyCount(tenant);
         } else if (oldKek.isShared() && !kek.isShared()) {
           // Shared -> Not Shared
-          dekRegistry.getSharedKeys().remove(oldKek.getKmsKeyId());
+          dekRegistry.getSharedKeys().remove(oldKek.getKmsKeyId(), (KeyEncryptionKeyId) key);
           dekRegistry.getMetricsManager().decrementSharedKeyCount(tenant);
         }
       }

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
@@ -17,6 +17,9 @@
 
 import com.fasterxml.jackson.core.JsonProcessingException;
 import com.fasterxml.jackson.core.type.TypeReference;
+import com.google.common.collect.Multimaps;
+import com.google.common.collect.SetMultimap;
+import com.google.common.collect.TreeMultimap;
 import com.google.crypto.tink.Aead;
 import com.google.inject.Inject;
 import com.google.inject.Singleton;
@@ -119,7 +122,7 @@ public class DekRegistry implements Closeable {
   private final DekRegistryConfig config;
   // visible for testing
   final Cache<EncryptionKeyId, EncryptionKey> keys;
-  private final Map<String, KeyEncryptionKeyId> sharedKeys = new ConcurrentHashMap<>();
+  private final SetMultimap<String, KeyEncryptionKeyId> sharedKeys;
   private final Map<DekFormat, Cryptor> cryptors;
   private final Map<String, Lock> tenantToLock = new ConcurrentHashMap<>();
   private final AtomicBoolean initialized = new AtomicBoolean();
@@ -137,6 +140,7 @@ public DekRegistry(
       this.config = new DekRegistryConfig(schemaRegistry.config().originalProperties());
       this.keys = createCache(new EncryptionKeyIdSerde(), new EncryptionKeySerde(),
           config.topic(), getCacheUpdateHandler(config));
+      this.sharedKeys = Multimaps.synchronizedSetMultimap(TreeMultimap.create());
       this.cryptors = new ConcurrentHashMap<>();
     } catch (RestConfigException e) {
       throw new IllegalArgumentException("Could not instantiate DekRegistry", e);
@@ -218,7 +222,7 @@ public DekRegistryConfig config() {
     return config;
   }
 
-  protected Map<String, KeyEncryptionKeyId> getSharedKeys() {
+  protected SetMultimap<String, KeyEncryptionKeyId> getSharedKeys() {
     return sharedKeys;
   }