Merge pull request #2889 from confluentinc/refactor-kmsdriver

Minor refactor of KmsDrivers
confluentinc · Dec 8, 2023 · e05ad69 · e05ad69
2 parents d6574a8 + dcecee3
commit e05ad69
Show file tree

Hide file tree

Showing 10 changed files with 29 additions and 119 deletions.
diff --git a/...tion-aws/src/main/java/io/confluent/kafka/schemaregistry/encryption/aws/AwsKmsDriver.java b/...tion-aws/src/main/java/io/confluent/kafka/schemaregistry/encryption/aws/AwsKmsDriver.java
@@ -22,7 +22,6 @@
 import com.amazonaws.auth.PropertiesCredentials;
 import com.amazonaws.services.kms.AWSKMS;
 import com.google.crypto.tink.KmsClient;
-import com.google.crypto.tink.KmsClients;
 import com.google.crypto.tink.integration.awskms.AwsKmsClient;
 import io.confluent.kafka.schemaregistry.encryption.tink.KmsDriver;
 import java.io.ByteArrayInputStream;
@@ -63,16 +62,16 @@ private AWSCredentialsProvider getCredentials(Map<String, ?> configs)
   }
 
   @Override
-  public KmsClient registerKmsClient(Map<String, ?> configs, Optional<String> kekUrl)
+  public KmsClient newKmsClient(Map<String, ?> configs, Optional<String> kekUrl)
       throws GeneralSecurityException {
     AWSKMS testClient = (AWSKMS) getTestClient(configs);
     Optional<AWSCredentialsProvider> creds = testClient != null
         ? Optional.empty()
         : Optional.of(getCredentials(configs));
-    return registerWithAwsKms(kekUrl, creds, testClient);
+    return newKmsClientWithAwsKms(kekUrl, creds, testClient);
   }
 
-  public static KmsClient registerWithAwsKms(
+  protected static KmsClient newKmsClientWithAwsKms(
       Optional<String> keyUri, Optional<AWSCredentialsProvider> credentials, AWSKMS awsKms)
       throws GeneralSecurityException {
     AwsKmsClient client;
@@ -89,7 +88,6 @@ public static KmsClient registerWithAwsKms(
     if (awsKms != null) {
       setAwsKms(client, awsKms);
     }
-    KmsClients.add(client);
     return client;
   }
 

diff --git a/...zure/src/main/java/io/confluent/kafka/schemaregistry/encryption/azure/AzureKmsClient.java b/...zure/src/main/java/io/confluent/kafka/schemaregistry/encryption/azure/AzureKmsClient.java
@@ -21,20 +21,17 @@
 import com.azure.core.http.HttpPipelineBuilder;
 import com.azure.core.http.policy.ExponentialBackoff;
 import com.azure.core.http.policy.RetryPolicy;
-import com.azure.identity.ClientSecretCredentialBuilder;
 import com.azure.identity.DefaultAzureCredentialBuilder;
 import com.azure.security.keyvault.keys.cryptography.CryptographyClient;
 import com.azure.security.keyvault.keys.cryptography.CryptographyClientBuilder;
 import com.azure.security.keyvault.keys.cryptography.models.EncryptionAlgorithm;
 import com.azure.security.keyvault.keys.implementation.KeyVaultCredentialPolicy;
 import com.google.crypto.tink.Aead;
 import com.google.crypto.tink.KmsClient;
-import com.google.crypto.tink.KmsClients;
 import com.google.crypto.tink.subtle.Validators;
 import java.security.GeneralSecurityException;
 import java.time.Duration;
 import java.util.Locale;
-import java.util.Optional;
 
 /**
  * An implementation of {@link KmsClient} for <a
@@ -171,59 +168,4 @@ public Aead getAead(String uri) throws GeneralSecurityException {
     }
     return new AzureKmsAead(client, this.algorithm);
   }
-
-  /**
-   * Creates and registers a {@link #AzureKmsClient} with the Tink runtime.
-   *
-   * <p></p>Utilize {@link #AzureKmsClient(String)} to create client. loads credentials using
-   * {@link DefaultAzureCredentialBuilder} which expects credentials to provided as environment
-   * variables.
-   *
-   * @throws GeneralSecurityException if keyUrl is missing
-   */
-  public static void register(Optional<String> keyUri) throws GeneralSecurityException {
-    AzureKmsClient client;
-    if (keyUri.isPresent()) {
-      client = new AzureKmsClient(keyUri.get());
-    } else {
-      throw new GeneralSecurityException("key url missing while registering KmsClient");
-    }
-    client.withDefaultCredentials();
-    KmsClients.add(client);
-  }
-
-  /**
-   * Creates and registers a {@link #AzureKmsClient} with the Tink runtime.
-   *
-   * @param keyUri       - azure keyvault key uri
-   * @param tenantId     - keyvault  tenantId
-   * @param clientId     - keyvault  clientId
-   * @param clientSecret - keyvault  clientSecret
-   * @throws GeneralSecurityException if keyUrl is missing
-   */
-  public static void register(Optional<String> keyUri, Optional<String> tenantId,
-      Optional<String> clientId, Optional<String> clientSecret) throws GeneralSecurityException {
-
-    if (!tenantId.isPresent()) {
-      throw new IllegalArgumentException("tenantId must be present");
-    }
-    if (!clientId.isPresent()) {
-      throw new IllegalArgumentException("clientId must be present");
-    }
-    if (!clientSecret.isPresent()) {
-      throw new IllegalArgumentException("clientSecret must be present");
-    }
-
-    AzureKmsClient client;
-    if (keyUri.isPresent()) {
-      client = new AzureKmsClient(keyUri.get());
-    } else {
-      throw new GeneralSecurityException("key url missing while registering KmsClient");
-    }
-    client.withCredentialsProvider(new ClientSecretCredentialBuilder()
-        .clientId(clientId.get())
-        .tenantId(tenantId.get())
-        .clientSecret(clientSecret.get()).build());
-    KmsClients.add(client);
-  }
 }
diff --git a/...zure/src/main/java/io/confluent/kafka/schemaregistry/encryption/azure/AzureKmsDriver.java b/...zure/src/main/java/io/confluent/kafka/schemaregistry/encryption/azure/AzureKmsDriver.java
@@ -21,7 +21,6 @@
 import com.azure.identity.DefaultAzureCredentialBuilder;
 import com.azure.security.keyvault.keys.cryptography.CryptographyClient;
 import com.google.crypto.tink.KmsClient;
-import com.google.crypto.tink.KmsClients;
 import io.confluent.kafka.schemaregistry.encryption.tink.KmsDriver;
 import java.security.GeneralSecurityException;
 import java.util.Map;
@@ -57,16 +56,16 @@ private TokenCredential getCredentials(Map<String, ?> configs) {
   }
 
   @Override
-  public KmsClient registerKmsClient(Map<String, ?> configs, Optional<String> kekUrl)
+  public KmsClient newKmsClient(Map<String, ?> configs, Optional<String> kekUrl)
       throws GeneralSecurityException {
     CryptographyClient testClient = (CryptographyClient) getTestClient(configs);
     Optional<TokenCredential> creds = testClient != null
         ? Optional.empty()
         : Optional.of(getCredentials(configs));
-    return registerWithAzureKms(kekUrl, creds, testClient);
+    return newKmsClientWithAzureKms(kekUrl, creds, testClient);
   }
 
-  public static KmsClient registerWithAzureKms(
+  protected static KmsClient newKmsClientWithAzureKms(
       Optional<String> keyUri, Optional<TokenCredential> credentials,
       CryptographyClient cryptographyClient)
       throws GeneralSecurityException {
@@ -84,7 +83,6 @@ public static KmsClient registerWithAzureKms(
     if (cryptographyClient != null) {
       client.withCryptographyClient(cryptographyClient);
     }
-    KmsClients.add(client);
     return client;
   }
 }

diff --git a/...tion-gcp/src/main/java/io/confluent/kafka/schemaregistry/encryption/gcp/GcpKmsDriver.java b/...tion-gcp/src/main/java/io/confluent/kafka/schemaregistry/encryption/gcp/GcpKmsDriver.java
@@ -19,7 +19,6 @@
 import com.google.api.services.cloudkms.v1.CloudKMS;
 import com.google.auth.oauth2.GoogleCredentials;
 import com.google.crypto.tink.KmsClient;
-import com.google.crypto.tink.KmsClients;
 import com.google.crypto.tink.integration.gcpkms.GcpKmsClient;
 import io.confluent.kafka.schemaregistry.encryption.tink.KmsDriver;
 import java.io.ByteArrayInputStream;
@@ -73,16 +72,16 @@ private GoogleCredentials getCredentials(Map<String, ?> configs)
   }
 
   @Override
-  public KmsClient registerKmsClient(Map<String, ?> configs, Optional<String> kekUrl)
+  public KmsClient newKmsClient(Map<String, ?> configs, Optional<String> kekUrl)
       throws GeneralSecurityException {
     CloudKMS testClient = (CloudKMS) getTestClient(configs);
     Optional<GoogleCredentials> creds = testClient != null
         ? Optional.empty()
         : Optional.ofNullable(getCredentials(configs));
-    return registerWithCloudKms(kekUrl, creds, testClient);
+    return newKmsClientWithCloudKms(kekUrl, creds, testClient);
   }
 
-  public static KmsClient registerWithCloudKms(
+  protected static KmsClient newKmsClientWithCloudKms(
       Optional<String> keyUri, Optional<GoogleCredentials> credentials, CloudKMS cloudKms)
       throws GeneralSecurityException {
     GcpKmsClient client;
@@ -100,7 +99,6 @@ public static KmsClient registerWithCloudKms(
     } else {
       setCloudKms(client, cloudKms);
     }
-    KmsClients.add(client);
     return client;
   }
 

diff --git a/.../src/main/java/io/confluent/kafka/schemaregistry/encryption/hcvault/HcVaultKmsDriver.java b/.../src/main/java/io/confluent/kafka/schemaregistry/encryption/hcvault/HcVaultKmsDriver.java
@@ -17,7 +17,6 @@
 package io.confluent.kafka.schemaregistry.encryption.hcvault;
 
 import com.google.crypto.tink.KmsClient;
-import com.google.crypto.tink.KmsClients;
 import io.confluent.kafka.schemaregistry.encryption.tink.KmsDriver;
 import io.github.jopenlibs.vault.api.Logical;
 import java.security.GeneralSecurityException;
@@ -46,17 +45,17 @@ private String getNamespace(Map<String, ?> configs) {
   }
 
   @Override
-  public KmsClient registerKmsClient(Map<String, ?> configs, Optional<String> kekUrl)
+  public KmsClient newKmsClient(Map<String, ?> configs, Optional<String> kekUrl)
       throws GeneralSecurityException {
     Logical testClient = (Logical) getTestClient(configs);
     Optional<String> creds = testClient != null
         ? Optional.empty()
         : Optional.ofNullable(getToken(configs));
     Optional<String> namespace = Optional.ofNullable(getNamespace(configs));
-    return registerWithHcVaultKms(kekUrl, creds, namespace, testClient);
+    return newKmsClientWithHcVaultKms(kekUrl, creds, namespace, testClient);
   }
 
-  public static KmsClient registerWithHcVaultKms(
+  protected static KmsClient newKmsClientWithHcVaultKms(
       Optional<String> keyUri, Optional<String> credentials,
       Optional<String> namespace, Logical vault)
       throws GeneralSecurityException {
@@ -74,7 +73,6 @@ public static KmsClient registerWithHcVaultKms(
     if (vault != null) {
       client.withVault(vault);
     }
-    KmsClients.add(client);
     return client;
   }
 }

diff --git a/...ption-tink/src/main/java/io/confluent/kafka/schemaregistry/encryption/tink/KmsDriver.java b/...ption-tink/src/main/java/io/confluent/kafka/schemaregistry/encryption/tink/KmsDriver.java
@@ -30,6 +30,9 @@ public interface KmsDriver {
 
   String getKeyUrlPrefix();
 
+  KmsClient newKmsClient(Map<String, ?> configs, Optional<String> kekUrl)
+      throws GeneralSecurityException;
+
   /**
    * @return true if this client does support {@code keyUri}
    */
@@ -41,8 +44,12 @@ default KmsClient getKmsClient(String kekUrl) throws GeneralSecurityException {
     return KmsClients.get(kekUrl);
   }
 
-  KmsClient registerKmsClient(Map<String, ?> configs, Optional<String> kekUrl)
-      throws GeneralSecurityException;
+  default KmsClient registerKmsClient(Map<String, ?> configs, Optional<String> kekUrl)
+      throws GeneralSecurityException {
+    KmsClient client = newKmsClient(configs, kekUrl);
+    KmsClients.add(client);
+    return client;
+  }
 
   default Object getTestClient(Map<String, ?> configs) {
     return configs.get(TEST_CLIENT);

diff --git a/...tion/src/main/java/io/confluent/kafka/schemaregistry/encryption/local/LocalKmsClient.java b/...tion/src/main/java/io/confluent/kafka/schemaregistry/encryption/local/LocalKmsClient.java
@@ -18,7 +18,6 @@
 
 import com.google.crypto.tink.Aead;
 import com.google.crypto.tink.KmsClient;
-import com.google.crypto.tink.KmsClients;
 import com.google.crypto.tink.PrimitiveSet;
 import com.google.crypto.tink.Registry;
 import com.google.crypto.tink.proto.AesGcmKey;
@@ -35,7 +34,6 @@
 import java.security.NoSuchAlgorithmException;
 import java.util.List;
 import java.util.Locale;
-import java.util.Optional;
 import javax.annotation.Nullable;
 
 /**
@@ -57,7 +55,7 @@ public final class LocalKmsClient implements KmsClient {
   private LocalKmsClient() {
   }
 
-  private LocalKmsClient(String uri, String secret, List<String> oldSecrets)
+  public LocalKmsClient(String uri, String secret, List<String> oldSecrets)
       throws GeneralSecurityException {
     if (!uri.toLowerCase(Locale.US).startsWith(PREFIX)) {
       throw new IllegalArgumentException("key URI must start with " + PREFIX);
@@ -146,18 +144,4 @@ public Aead getAead(String uri) throws GeneralSecurityException {
 
     return aead;
   }
-
-  /**
-   * Creates and registers a {@link #LocalKmsClient} with the Tink runtime.
-   *
-   * <p>If {@code keyUri} is present, it is the only key that the new client will support.
-   * Otherwise
-   * the new client supports all local KMS keys.
-   */
-  public static KmsClient register(Optional<String> keyUri, String secret, List<String> oldSecrets)
-      throws GeneralSecurityException {
-    KmsClient client = new LocalKmsClient(keyUri.orElse(PREFIX), secret, oldSecrets);
-    KmsClients.add(client);
-    return client;
-  }
 }
diff --git a/...tion/src/main/java/io/confluent/kafka/schemaregistry/encryption/local/LocalKmsDriver.java b/...tion/src/main/java/io/confluent/kafka/schemaregistry/encryption/local/LocalKmsDriver.java
@@ -65,9 +65,10 @@ private List<String> getOldSecrets(Map<String, ?> configs) {
   }
 
   @Override
-  public KmsClient registerKmsClient(Map<String, ?> configs, Optional<String> kekUrl)
+  public KmsClient newKmsClient(Map<String, ?> configs, Optional<String> kekUrl)
       throws GeneralSecurityException {
-    return LocalKmsClient.register(kekUrl, getSecret(configs), getOldSecrets(configs));
+    return new LocalKmsClient(
+        kekUrl.orElse(LocalKmsClient.PREFIX), getSecret(configs), getOldSecrets(configs));
   }
 }
 
diff --git a/dek-registry/src/test/java/io/confluent/dekregistry/testutil/TestKmsClient.java b/dek-registry/src/test/java/io/confluent/dekregistry/testutil/TestKmsClient.java
@@ -18,7 +18,6 @@
 
 import com.google.crypto.tink.Aead;
 import com.google.crypto.tink.KmsClient;
-import com.google.crypto.tink.KmsClients;
 import com.google.crypto.tink.PrimitiveSet;
 import com.google.crypto.tink.Registry;
 import com.google.crypto.tink.proto.AesGcmKey;
@@ -34,7 +33,6 @@
 import java.security.MessageDigest;
 import java.security.NoSuchAlgorithmException;
 import java.util.Locale;
-import java.util.Optional;
 import javax.annotation.Nullable;
 
 /**
@@ -53,7 +51,7 @@ public final class TestKmsClient implements KmsClient {
   private TestKmsClient() {
   }
 
-  private TestKmsClient(String uri, String secret)
+  public TestKmsClient(String uri, String secret)
       throws GeneralSecurityException {
     if (!uri.toLowerCase(Locale.US).startsWith(PREFIX)) {
       throw new IllegalArgumentException("key URI must start with " + PREFIX);
@@ -139,18 +137,4 @@ public Aead getAead(String uri) throws GeneralSecurityException {
 
     return aead;
   }
-
-  /**
-   * Creates and registers a {@link #TestKmsClient} with the Tink runtime.
-   *
-   * <p>If {@code keyUri} is present, it is the only key that the new client will support.
-   * Otherwise
-   * the new client supports all local KMS keys.
-   */
-  public static KmsClient register(Optional<String> keyUri, String secret)
-      throws GeneralSecurityException {
-    KmsClient client = new TestKmsClient(keyUri.orElse(PREFIX), secret);
-    KmsClients.add(client);
-    return client;
-  }
-}
+}
diff --git a/dek-registry/src/test/java/io/confluent/dekregistry/testutil/TestKmsDriver.java b/dek-registry/src/test/java/io/confluent/dekregistry/testutil/TestKmsDriver.java
@@ -35,9 +35,9 @@ public String getKeyUrlPrefix() {
   }
 
   @Override
-  public KmsClient registerKmsClient(Map<String, ?> configs, Optional<String> kekUrl)
+  public KmsClient newKmsClient(Map<String, ?> configs, Optional<String> kekUrl)
       throws GeneralSecurityException {
-    return TestKmsClient.register(kekUrl, SECRET);
+    return new TestKmsClient(kekUrl.orElse(TestKmsClient.PREFIX), SECRET);
   }
 }