From 0cb664f184b0f5a597a8195c16a46fa665f7207f Mon Sep 17 00:00:00 2001
From: Robert Yokota <rayokota@gmail.com>
Date: Wed, 24 Jan 2024 17:25:18 -0800
Subject: [PATCH] Fix sharedKeys cache to account for multiple use of kms key
 ID

---
 .../dekregistry/storage/DefaultDekCacheUpdateHandler.java | 4 ++--
 .../io/confluent/dekregistry/storage/DekRegistry.java     | 8 ++++++--
 2 files changed, 8 insertions(+), 4 deletions(-)

diff --git a/dek-registry/src/main/java/io/confluent/dekregistry/storage/DefaultDekCacheUpdateHandler.java b/dek-registry/src/main/java/io/confluent/dekregistry/storage/DefaultDekCacheUpdateHandler.java
index d440645b4a6..ad3538d402c 100644
--- a/dek-registry/src/main/java/io/confluent/dekregistry/storage/DefaultDekCacheUpdateHandler.java
+++ b/dek-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
index 50363cc6b4d..5e7b25f2361 100644
--- 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;
   }