Avoid KmsClients.getAutoLoaded in Tinkey and explicitly register factories to create new clients.

PiperOrigin-RevId: 539614473

9 files changed, 35 insertions, 24 deletions

diff --git a/tools/tinkey/src/main/java/com/google/crypto/tink/tinkey/BUILD.bazel b/tools/tinkey/src/main/java/com/google/crypto/tink/tinkey/BUILD.bazel
index b0b5f9d44..eaf4db8e3 100644
--- a/tools/tinkey/src/main/java/com/google/crypto/tink/tinkey/BUILD.bazel
+++ b/tools/tinkey/src/main/java/com/google/crypto/tink/tinkey/BUILD.bazel
@@ -37,6 +37,7 @@ java_library(
     name = "tinkey",
     srcs = ["Tinkey.java"],
     deps = [
+        ":kms_clients_factory",
         ":tinkey_commands",
         "@maven//:args4j_args4j",
         "@tink_java//src/main/java/com/google/crypto/tink/daead:deterministic_aead_config",
@@ -47,6 +48,8 @@ java_library(
         "@tink_java//src/main/java/com/google/crypto/tink/prf:prf_config",
         "@tink_java//src/main/java/com/google/crypto/tink/signature:signature_config",
         "@tink_java//src/main/java/com/google/crypto/tink/streamingaead:streaming_aead_config",
+        "@tink_java//src/main/java/com/google/crypto/tink/integration/awskms:aws_kms_client",
+        "@tink_java//src/main/java/com/google/crypto/tink/integration/gcpkms:gcp_kms_client",
     ],
 )
 
@@ -108,12 +111,8 @@ java_library(
 java_library(
     name = "tinkey_util",
     srcs = ["TinkeyUtil.java"],
-    runtime_deps = [
-        # Tinkey automatically loads these KMS clients at runtime.
-        "@tink_java//src/main/java/com/google/crypto/tink/integration/awskms:aws_kms_client",
-        "@tink_java//src/main/java/com/google/crypto/tink/integration/gcpkms:gcp_kms_client",
-    ],
     deps = [
+        ":kms_clients_factory",
         "@com_google_protobuf//:protobuf_javalite",
         "@tink_java//proto:tink_java_proto",
         "@tink_java//src/main/java/com/google/crypto/tink:aead",
@@ -125,7 +124,6 @@ java_library(
         "@tink_java//src/main/java/com/google/crypto/tink:key_template",
         "@tink_java//src/main/java/com/google/crypto/tink:keyset_reader",
         "@tink_java//src/main/java/com/google/crypto/tink:keyset_writer",
-        "@tink_java//src/main/java/com/google/crypto/tink:kms_clients",
         "@tink_java//src/main/java/com/google/crypto/tink:registry_cluster",
     ],
 )
diff --git a/tools/tinkey/src/main/java/com/google/crypto/tink/tinkey/Tinkey.java b/tools/tinkey/src/main/java/com/google/crypto/tink/tinkey/Tinkey.java
index d4588c29e..f3fc79403 100644
--- a/tools/tinkey/src/main/java/com/google/crypto/tink/tinkey/Tinkey.java
+++ b/tools/tinkey/src/main/java/com/google/crypto/tink/tinkey/Tinkey.java
@@ -18,6 +18,8 @@ package com.google.crypto.tink.tinkey;
 
 import com.google.crypto.tink.daead.DeterministicAeadConfig;
 import com.google.crypto.tink.hybrid.HybridConfig;
+import com.google.crypto.tink.integration.awskms.AwsKmsClient;
+import com.google.crypto.tink.integration.gcpkms.GcpKmsClient;
 import com.google.crypto.tink.jwt.JwtMacConfig;
 import com.google.crypto.tink.jwt.JwtSignatureConfig;
 import com.google.crypto.tink.keyderivation.KeyDerivationConfig;
@@ -41,6 +43,10 @@ public final class Tinkey {
     JwtMacConfig.register();
     KeyDerivationConfig.register();
     // place holder for Internal Prps. DO NOT EDIT.
+
+    KmsClientsFactory.globalInstance().addFactory(AwsKmsClient::new);
+    KmsClientsFactory.globalInstance().addFactory(GcpKmsClient::new);
+
     TinkeyCommands commands = new TinkeyCommands();
     CmdLineParser parser = new CmdLineParser(commands);
 
diff --git a/tools/tinkey/src/main/java/com/google/crypto/tink/tinkey/TinkeyUtil.java b/tools/tinkey/src/main/java/com/google/crypto/tink/tinkey/TinkeyUtil.java
index 27f708dc2..ccce483c6 100644
--- a/tools/tinkey/src/main/java/com/google/crypto/tink/tinkey/TinkeyUtil.java
+++ b/tools/tinkey/src/main/java/com/google/crypto/tink/tinkey/TinkeyUtil.java
@@ -27,7 +27,6 @@ import com.google.crypto.tink.KeysetHandle;
 import com.google.crypto.tink.KeysetManager;
 import com.google.crypto.tink.KeysetReader;
 import com.google.crypto.tink.KeysetWriter;
-import com.google.crypto.tink.KmsClients;
 import com.google.crypto.tink.proto.OutputPrefixType;
 import com.google.protobuf.ByteString;
 import java.io.ByteArrayInputStream;
@@ -216,7 +215,8 @@ final class TinkeyUtil {
     KeysetWriter writer = createKeysetWriter(outputStream, outFormat);
     if (masterKeyUri != null) {
       Aead masterKey =
-          KmsClients.getAutoLoaded(masterKeyUri)
+          KmsClientsFactory.globalInstance()
+              .newClientFor(masterKeyUri)
               .withCredentials(credentialPath)
               .getAead(masterKeyUri);
       handle.write(writer, masterKey);
@@ -235,7 +235,8 @@ final class TinkeyUtil {
     KeysetReader reader = createKeysetReader(inputStream, inFormat);
     if (masterKeyUri != null) {
       Aead masterKey =
-          KmsClients.getAutoLoaded(masterKeyUri)
+          KmsClientsFactory.globalInstance()
+              .newClientFor(masterKeyUri)
               .withCredentials(credentialPath)
               .getAead(masterKeyUri);
       return KeysetHandle.read(reader, masterKey);
diff --git a/tools/tinkey/src/test/java/com/google/crypto/tink/tinkey/AddKeyCommandTest.java b/tools/tinkey/src/test/java/com/google/crypto/tink/tinkey/AddKeyCommandTest.java
index c4c07a845..2c6c9d1e0 100644
--- a/tools/tinkey/src/test/java/com/google/crypto/tink/tinkey/AddKeyCommandTest.java
+++ b/tools/tinkey/src/test/java/com/google/crypto/tink/tinkey/AddKeyCommandTest.java
@@ -44,6 +44,7 @@ public class AddKeyCommandTest {
   public static void setUp() throws Exception {
     AeadConfig.register();
     MacConfig.register();
+    KmsClientsFactory.globalInstance().addFactory(TinkeyTestKmsClient::new);
   }
 
   @Test
diff --git a/tools/tinkey/src/test/java/com/google/crypto/tink/tinkey/BUILD.bazel b/tools/tinkey/src/test/java/com/google/crypto/tink/tinkey/BUILD.bazel
index 97db3a453..d44998d59 100644
--- a/tools/tinkey/src/test/java/com/google/crypto/tink/tinkey/BUILD.bazel
+++ b/tools/tinkey/src/test/java/com/google/crypto/tink/tinkey/BUILD.bazel
@@ -12,6 +12,7 @@ java_test(
     ],
     deps = [
         "//tinkey/src/main/java/com/google/crypto/tink/tinkey",
+        "//tinkey/src/main/java/com/google/crypto/tink/tinkey:kms_clients_factory",
         "//tinkey/src/main/java/com/google/crypto/tink/tinkey:tinkey_util",
         "@maven//:com_google_truth_truth",
         "@maven//:junit_junit",
@@ -19,7 +20,6 @@ java_test(
         "@tink_java//src/main/java/com/google/crypto/tink:cleartext_keyset_handle",
         "@tink_java//src/main/java/com/google/crypto/tink:insecure_secret_key_access",
         "@tink_java//src/main/java/com/google/crypto/tink:key",
-        "@tink_java//src/main/java/com/google/crypto/tink:kms_clients",
         "@tink_java//src/main/java/com/google/crypto/tink:private_key",
         "@tink_java//src/main/java/com/google/crypto/tink:registry_cluster",
         "@tink_java//src/main/java/com/google/crypto/tink:tink_json_proto_keyset_format",
@@ -43,11 +43,11 @@ java_test(
     ],
     deps = [
         "//tinkey/src/main/java/com/google/crypto/tink/tinkey",
+        "//tinkey/src/main/java/com/google/crypto/tink/tinkey:kms_clients_factory",
         "@maven//:com_google_truth_truth",
         "@maven//:junit_junit",
         "@tink_java//src/main/java/com/google/crypto/tink:aead",
         "@tink_java//src/main/java/com/google/crypto/tink:insecure_secret_key_access",
-        "@tink_java//src/main/java/com/google/crypto/tink:kms_clients",
         "@tink_java//src/main/java/com/google/crypto/tink:registry_cluster",
         "@tink_java//src/main/java/com/google/crypto/tink:tink_json_proto_keyset_format",
         "@tink_java//src/main/java/com/google/crypto/tink:tink_proto_keyset_format",
@@ -66,6 +66,7 @@ java_test(
     ],
     deps = [
         "//tinkey/src/main/java/com/google/crypto/tink/tinkey",
+        "//tinkey/src/main/java/com/google/crypto/tink/tinkey:kms_clients_factory",
         "//tinkey/src/main/java/com/google/crypto/tink/tinkey:tinkey_test_kms_client",
         "@maven//:com_google_truth_truth",
         "@maven//:junit_junit",
@@ -93,11 +94,11 @@ java_test(
     ],
     deps = [
         "//tinkey/src/main/java/com/google/crypto/tink/tinkey",
+        "//tinkey/src/main/java/com/google/crypto/tink/tinkey:kms_clients_factory",
         "@maven//:com_google_truth_truth",
         "@maven//:junit_junit",
         "@tink_java//src/main/java/com/google/crypto/tink:aead",
         "@tink_java//src/main/java/com/google/crypto/tink:insecure_secret_key_access",
-        "@tink_java//src/main/java/com/google/crypto/tink:kms_clients",
         "@tink_java//src/main/java/com/google/crypto/tink:registry_cluster",
         "@tink_java//src/main/java/com/google/crypto/tink:tink_json_proto_keyset_format",
         "@tink_java//src/main/java/com/google/crypto/tink:tink_proto_keyset_format",
@@ -115,12 +116,12 @@ java_test(
         "//tinkey/src/main/java/com/google/crypto/tink/tinkey:tinkey_test_kms_client",
     ],
     deps = [
+        "//tinkey/src/main/java/com/google/crypto/tink/tinkey:kms_clients_factory",
         "//tinkey/src/main/java/com/google/crypto/tink/tinkey:tinkey_test_kms_client",
         "@maven//:com_google_truth_truth",
         "@maven//:junit_junit",
         "@tink_java//src/main/java/com/google/crypto/tink:aead",
         "@tink_java//src/main/java/com/google/crypto/tink:kms_client",
-        "@tink_java//src/main/java/com/google/crypto/tink:kms_clients",
         "@tink_java//src/main/java/com/google/crypto/tink:registry_cluster",
         "@tink_java//src/main/java/com/google/crypto/tink/aead:aead_config",
         "@tink_java//src/main/java/com/google/crypto/tink/aead:predefined_aead_parameters",
diff --git a/tools/tinkey/src/test/java/com/google/crypto/tink/tinkey/CreateKeysetCommandTest.java b/tools/tinkey/src/test/java/com/google/crypto/tink/tinkey/CreateKeysetCommandTest.java
index 30a15ee82..9b83b8f5d 100644
--- a/tools/tinkey/src/test/java/com/google/crypto/tink/tinkey/CreateKeysetCommandTest.java
+++ b/tools/tinkey/src/test/java/com/google/crypto/tink/tinkey/CreateKeysetCommandTest.java
@@ -22,7 +22,6 @@ import static java.nio.charset.StandardCharsets.UTF_8;
 import com.google.crypto.tink.Aead;
 import com.google.crypto.tink.InsecureSecretKeyAccess;
 import com.google.crypto.tink.KeysetHandle;
-import com.google.crypto.tink.KmsClients;
 import com.google.crypto.tink.TinkJsonProtoKeysetFormat;
 import com.google.crypto.tink.TinkProtoKeysetFormat;
 import com.google.crypto.tink.mac.PredefinedMacParameters;
@@ -113,7 +112,8 @@ public class CreateKeysetCommandTest {
     Tinkey.main(commandLine.split(" "));
 
     Aead masterKeyAead =
-        KmsClients.getAutoLoaded(TestUtil.GCP_KMS_TEST_KEY_URI)
+        KmsClientsFactory.globalInstance()
+            .newClientFor(TestUtil.GCP_KMS_TEST_KEY_URI)
             .withCredentials(TestUtil.SERVICE_ACCOUNT_FILE)
             .getAead(TestUtil.GCP_KMS_TEST_KEY_URI);
 
@@ -141,7 +141,8 @@ public class CreateKeysetCommandTest {
     Tinkey.main(commandLine.split(" "));
 
     Aead masterKeyAead =
-        KmsClients.getAutoLoaded(TestUtil.GCP_KMS_TEST_KEY_URI)
+        KmsClientsFactory.globalInstance()
+            .newClientFor(TestUtil.GCP_KMS_TEST_KEY_URI)
             .withCredentials(TestUtil.SERVICE_ACCOUNT_FILE)
             .getAead(TestUtil.GCP_KMS_TEST_KEY_URI);
 
diff --git a/tools/tinkey/src/test/java/com/google/crypto/tink/tinkey/CreatePublicKeysetCommandTest.java b/tools/tinkey/src/test/java/com/google/crypto/tink/tinkey/CreatePublicKeysetCommandTest.java
index a21eaf79b..fb4d910e6 100644
--- a/tools/tinkey/src/test/java/com/google/crypto/tink/tinkey/CreatePublicKeysetCommandTest.java
+++ b/tools/tinkey/src/test/java/com/google/crypto/tink/tinkey/CreatePublicKeysetCommandTest.java
@@ -24,7 +24,6 @@ import com.google.crypto.tink.Aead;
 import com.google.crypto.tink.InsecureSecretKeyAccess;
 import com.google.crypto.tink.Key;
 import com.google.crypto.tink.KeysetHandle;
-import com.google.crypto.tink.KmsClients;
 import com.google.crypto.tink.PrivateKey;
 import com.google.crypto.tink.TinkJsonProtoKeysetFormat;
 import com.google.crypto.tink.TinkProtoKeysetFormat;
@@ -130,7 +129,8 @@ public class CreatePublicKeysetCommandTest {
         KeysetHandle.generateNew(Ed25519Parameters.create(Ed25519Parameters.Variant.TINK));
 
     Aead masterKeyAead =
-        KmsClients.getAutoLoaded(TestUtil.GCP_KMS_TEST_KEY_URI)
+        KmsClientsFactory.globalInstance()
+            .newClientFor(TestUtil.GCP_KMS_TEST_KEY_URI)
             .withCredentials(TestUtil.SERVICE_ACCOUNT_FILE)
             .getAead(TestUtil.GCP_KMS_TEST_KEY_URI);
     String serializedKeyset =
@@ -170,7 +170,8 @@ public class CreatePublicKeysetCommandTest {
         KeysetHandle.generateNew(Ed25519Parameters.create(Ed25519Parameters.Variant.TINK));
 
     Aead masterKeyAead =
-        KmsClients.getAutoLoaded(TestUtil.GCP_KMS_TEST_KEY_URI)
+        KmsClientsFactory.globalInstance()
+            .newClientFor(TestUtil.GCP_KMS_TEST_KEY_URI)
             .withCredentials(TestUtil.SERVICE_ACCOUNT_FILE)
             .getAead(TestUtil.GCP_KMS_TEST_KEY_URI);
     byte[] serializedKeyset =
diff --git a/tools/tinkey/src/test/java/com/google/crypto/tink/tinkey/RotateKeysetCommandTest.java b/tools/tinkey/src/test/java/com/google/crypto/tink/tinkey/RotateKeysetCommandTest.java
index cfce43db7..8c5eaeebc 100644
--- a/tools/tinkey/src/test/java/com/google/crypto/tink/tinkey/RotateKeysetCommandTest.java
+++ b/tools/tinkey/src/test/java/com/google/crypto/tink/tinkey/RotateKeysetCommandTest.java
@@ -23,7 +23,6 @@ import static org.junit.Assert.assertThrows;
 import com.google.crypto.tink.Aead;
 import com.google.crypto.tink.InsecureSecretKeyAccess;
 import com.google.crypto.tink.KeysetHandle;
-import com.google.crypto.tink.KmsClients;
 import com.google.crypto.tink.TinkJsonProtoKeysetFormat;
 import com.google.crypto.tink.TinkProtoKeysetFormat;
 import com.google.crypto.tink.mac.MacConfig;
@@ -126,7 +125,8 @@ public class RotateKeysetCommandTest {
     Path outputFile = Paths.get(path.toString(), "output");
 
     Aead masterKeyAead =
-        KmsClients.getAutoLoaded(TestUtil.GCP_KMS_TEST_KEY_URI)
+        KmsClientsFactory.globalInstance()
+            .newClientFor(TestUtil.GCP_KMS_TEST_KEY_URI)
             .withCredentials(TestUtil.SERVICE_ACCOUNT_FILE)
             .getAead(TestUtil.GCP_KMS_TEST_KEY_URI);
 
@@ -174,7 +174,8 @@ public class RotateKeysetCommandTest {
     Path outputFile = Paths.get(path.toString(), "output");
 
     Aead masterKeyAead =
-        KmsClients.getAutoLoaded(TestUtil.GCP_KMS_TEST_KEY_URI)
+        KmsClientsFactory.globalInstance()
+            .newClientFor(TestUtil.GCP_KMS_TEST_KEY_URI)
             .withCredentials(TestUtil.SERVICE_ACCOUNT_FILE)
             .getAead(TestUtil.GCP_KMS_TEST_KEY_URI);
 
diff --git a/tools/tinkey/src/test/java/com/google/crypto/tink/tinkey/TinkeyTestKmsClientTest.java b/tools/tinkey/src/test/java/com/google/crypto/tink/tinkey/TinkeyTestKmsClientTest.java
index 9dcf25d4d..67b0f98e1 100644
--- a/tools/tinkey/src/test/java/com/google/crypto/tink/tinkey/TinkeyTestKmsClientTest.java
+++ b/tools/tinkey/src/test/java/com/google/crypto/tink/tinkey/TinkeyTestKmsClientTest.java
@@ -25,7 +25,6 @@ import static org.junit.Assert.assertTrue;
 import com.google.crypto.tink.Aead;
 import com.google.crypto.tink.KeysetHandle;
 import com.google.crypto.tink.KmsClient;
-import com.google.crypto.tink.KmsClients;
 import com.google.crypto.tink.aead.AeadConfig;
 import com.google.crypto.tink.aead.PredefinedAeadParameters;
 import java.nio.file.Files;
@@ -42,6 +41,7 @@ public final class TinkeyTestKmsClientTest {
   @BeforeClass
   public static void setUp() throws Exception {
     AeadConfig.register();
+    KmsClientsFactory.globalInstance().addFactory(TinkeyTestKmsClient::new);
   }
 
   @Test
@@ -99,7 +99,8 @@ public final class TinkeyTestKmsClientTest {
     KeysetHandle handle = KeysetHandle.generateNew(PredefinedAeadParameters.AES128_GCM);
     String masterKeyUri = TinkeyTestKmsClient.createKeyUri(handle);
     Aead masterKey =
-        KmsClients.getAutoLoaded(masterKeyUri)
+        KmsClientsFactory.globalInstance()
+            .newClientFor(masterKeyUri)
             .withCredentials(credentialPath.toString())
             .getAead(masterKeyUri);
     Aead manualMasterKey = handle.getPrimitive(Aead.class);