From 0bdb81e7d300070d2cca08c11023fb0102eb4b56 Mon Sep 17 00:00:00 2001
From: Robert Yokota <rayokota@gmail.com>
Date: Thu, 11 Jan 2024 15:29:53 -0800
Subject: [PATCH] DGS-9732 Omit decrypted DEK when registering DEK in IMPORT
 mode

---
 .../dekregistry/storage/DekRegistry.java      |  6 ++++-
 .../dekregistry/web/rest/RestApiTest.java     | 26 +++++++++++++++----
 2 files changed, 26 insertions(+), 6 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
index 3628a4dc783..651db3ca6f2 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
@@ -49,6 +49,7 @@
 import io.confluent.kafka.schemaregistry.exceptions.SchemaRegistryRequestForwardingException;
 import io.confluent.kafka.schemaregistry.exceptions.UnknownLeaderException;
 import io.confluent.kafka.schemaregistry.storage.KafkaSchemaRegistry;
+import io.confluent.kafka.schemaregistry.storage.Mode;
 import io.confluent.kafka.schemaregistry.storage.SchemaRegistry;
 import io.confluent.kafka.schemaregistry.utils.JacksonMapper;
 import io.confluent.rest.RestConfigException;
@@ -568,7 +569,10 @@ public DataEncryptionKey createDek(String kekName, CreateDekRequest request)
     // Retrieve key with ts set
     key = (DataEncryptionKey) keys.get(keyId);
     if (kek.isShared()) {
-      key = generateRawDek(kek, key);
+      Mode mode = schemaRegistry.getModeInScope(request.getSubject());
+      if (mode != Mode.IMPORT) {
+        key = generateRawDek(kek, key);
+      }
     }
     return key;
   }
diff --git a/dek-registry/src/test/java/io/confluent/dekregistry/web/rest/RestApiTest.java b/dek-registry/src/test/java/io/confluent/dekregistry/web/rest/RestApiTest.java
index be9438c0fb8..55ba7038dc7 100644
--- a/dek-registry/src/test/java/io/confluent/dekregistry/web/rest/RestApiTest.java
+++ b/dek-registry/src/test/java/io/confluent/dekregistry/web/rest/RestApiTest.java
@@ -17,6 +17,7 @@
 import static io.confluent.dekregistry.storage.DekRegistry.X_FORWARD_HEADER;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertNull;
 import static org.junit.Assert.fail;
 
 import com.google.common.collect.ImmutableList;
@@ -43,6 +44,7 @@
 import io.confluent.kafka.schemaregistry.encryption.tink.DekFormat;
 import io.confluent.kafka.schemaregistry.rest.SchemaRegistryConfig;
 import io.confluent.kafka.schemaregistry.storage.KafkaSchemaRegistry;
+import io.confluent.kafka.schemaregistry.storage.Mode;
 import io.confluent.kafka.schemaregistry.storage.RuleSetHandler;
 import java.nio.charset.StandardCharsets;
 import java.util.Base64;
@@ -109,7 +111,7 @@ public io.confluent.kafka.schemaregistry.storage.RuleSet transform(RuleSet ruleS
   public void testBasic() throws Exception {
     Map<String, String> headers = new HashMap<>();
     headers.put("Content-Type", Versions.SCHEMA_REGISTRY_V1_JSON_WEIGHTED);
-    testBasic(headers);
+    testBasic(headers, false);
   }
 
   @Test
@@ -117,10 +119,17 @@ public void testForwarding() throws Exception {
     Map<String, String> headers = new HashMap<>();
     headers.put("Content-Type", Versions.SCHEMA_REGISTRY_V1_JSON_WEIGHTED);
     headers.put(X_FORWARD_HEADER, "false");
-    testBasic(headers);
+    testBasic(headers, true);
   }
 
-  private void testBasic(Map<String, String> headers) throws Exception {
+  @Test
+  public void testBasicImport() throws Exception {
+    Map<String, String> headers = new HashMap<>();
+    headers.put("Content-Type", Versions.SCHEMA_REGISTRY_V1_JSON_WEIGHTED);
+    testBasic(headers, true);
+  }
+
+  private void testBasic(Map<String, String> headers, boolean isImport) throws Exception {
     String kekName = "kek1";
     String kmsType = "test-kms";
     String kmsKeyId = "myid";
@@ -130,6 +139,10 @@ private void testBasic(Map<String, String> headers) throws Exception {
     DekFormat algorithm = DekFormat.AES256_GCM;
     Kek kek = new Kek(kekName, kmsType, kmsKeyId, null, null, false, null, null);
 
+    if (isImport) {
+      client.setMode("IMPORT");
+    }
+
     // Create kek
     Kek newKek = client.createKek(headers, kekName, kmsType, kmsKeyId, null, null, false, false);
     assertEquals(kek, newKek);
@@ -264,7 +277,11 @@ private void testBasic(Map<String, String> headers) throws Exception {
     // Create dek w/o key material, receive both encrypted and decrypted key material
     newDek = client.createDek(headers, kekName, subject2, null, algorithm, null, false);
     assertNotNull(newDek.getEncryptedKeyMaterial());
-    assertNotNull(newDek.getKeyMaterial());
+    if (isImport) {
+      assertNull(newDek.getKeyMaterial());
+    } else {
+      assertNotNull(newDek.getKeyMaterial());
+    }
     assertNotNull(newDek.getTimestamp());
 
     // Create versioned dek
@@ -388,7 +405,6 @@ private void testBasic(Map<String, String> headers) throws Exception {
       assertEquals(DekRegistryErrors.KEY_NOT_FOUND_ERROR_CODE, e.getErrorCode());
     }
   }
-
   @Test
   public void testUnknownKmsType() throws Exception {
     String kekName = "kek1";