Add a class "KmsClientsFactory" to store KmsClients-factories.

We need this because the used mechanism returns global objects which are then modified with "withCredentials" -- in tests this doesn't work.

PiperOrigin-RevId: 539029491

6 files changed, 200 insertions, 12 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 bcb68f667..b0b5f9d44 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
@@ -332,3 +332,11 @@ java_plugin(
         "@maven//:com_google_auto_service_auto_service",
     ],
 )
+
+java_library(
+    name = "kms_clients_factory",
+    srcs = ["KmsClientsFactory.java"],
+    deps = [
+        "@tink_java//src/main/java/com/google/crypto/tink:kms_client",
+    ],
+)
diff --git a/tools/tinkey/src/main/java/com/google/crypto/tink/tinkey/KmsClientsFactory.java b/tools/tinkey/src/main/java/com/google/crypto/tink/tinkey/KmsClientsFactory.java
new file mode 100644
index 000000000..e2a444613
--- /dev/null
+++ b/tools/tinkey/src/main/java/com/google/crypto/tink/tinkey/KmsClientsFactory.java
@@ -0,0 +1,65 @@
+// Copyright 2023 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+//
+////////////////////////////////////////////////////////////////////////////////
+
+package com.google.crypto.tink.tinkey;
+
+import com.google.crypto.tink.KmsClient;
+import java.security.GeneralSecurityException;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.function.Supplier;
+
+/**
+ * Allows creating {@link KmsClient} objects.
+ *
+ * <p>Tink {@link KmsClient} which are registered have the really unfortunate property that they are
+ * modifiable. For Tinkey this should never be a problem, since one always uses the client only
+ * once. However, for testing Tinkey it is a problem. Hence, we avoid the class and simply use this
+ * one here instead.
+ */
+final class KmsClientsFactory {
+
+  private List<Supplier<KmsClient>> factories = new ArrayList<>();
+
+  /** Create a new KmsClientsFactory without any registered factory. */
+  public KmsClientsFactory() {}
+
+  private static final KmsClientsFactory globalInstance = new KmsClientsFactory();
+
+  /** A unique global instance. */
+  public static KmsClientsFactory globalInstance() {
+    return globalInstance;
+  }
+
+  /**
+   * Enumerates all registered factories, creates a new client for each, and returns one if it
+   * supports keyUri.
+   */
+  public KmsClient newClientFor(String keyUri) throws GeneralSecurityException {
+    for (Supplier<KmsClient> factory : factories) {
+      KmsClient client = factory.get();
+      if (client.doesSupport(keyUri)) {
+        return client;
+      }
+    }
+    throw new GeneralSecurityException("Unable to find factory for keyUri: " + keyUri);
+  }
+
+  /** Registers an additional factory. */
+  public void addFactory(Supplier<KmsClient> factory) {
+    factories.add(factory);
+  }
+}
diff --git a/tools/tinkey/src/main/java/com/google/crypto/tink/tinkey/TinkeyTestKmsClient.java b/tools/tinkey/src/main/java/com/google/crypto/tink/tinkey/TinkeyTestKmsClient.java
index 63f215b9d..59909b367 100644
--- a/tools/tinkey/src/main/java/com/google/crypto/tink/tinkey/TinkeyTestKmsClient.java
+++ b/tools/tinkey/src/main/java/com/google/crypto/tink/tinkey/TinkeyTestKmsClient.java
@@ -40,8 +40,19 @@ import java.util.Arrays;
 @AutoService(KmsClient.class)
 public final class TinkeyTestKmsClient implements KmsClient {
 
-  public TinkeyTestKmsClient() {}
+  public TinkeyTestKmsClient() {
+    this(PREFIX);
+  }
+
+  private TinkeyTestKmsClient(String prefix) {
+    this.prefix = prefix;
+  }
+
+  public static KmsClient createForPrefix(String prefix) {
+    return new TinkeyTestKmsClient(prefix);
+  }
 
+  private final String prefix;
   private static final String PREFIX = "tinkey-test-kms-client://";
   private static final String CREDENTIALS_FILE_CONTENTS = "VALID CREDENTIALS";
 
@@ -54,16 +65,16 @@ public final class TinkeyTestKmsClient implements KmsClient {
     Files.write(path, CREDENTIALS_FILE_CONTENTS.getBytes(UTF_8));
   }
 
-  private static String stripPrefix(String str) throws GeneralSecurityException {
-    if (!str.startsWith(PREFIX)) {
+  private static String stripPrefix(String prefix, String str) throws GeneralSecurityException {
+    if (!str.startsWith(prefix)) {
       throw new GeneralSecurityException("Invalid key uri: " + str);
     }
-    return str.substring(PREFIX.length());
+    return str.substring(prefix.length());
   }
 
   @Override
   public boolean doesSupport(String keyUri) {
-    return keyUri.startsWith(PREFIX);
+    return keyUri.startsWith(prefix);
   }
 
   byte[] credentialFileContents = new byte[] {};
@@ -93,7 +104,7 @@ public final class TinkeyTestKmsClient implements KmsClient {
   @Override
   public Aead getAead(String keyUri) throws GeneralSecurityException {
     checkCredentials();
-    String keyset = stripPrefix(keyUri);
+    String keyset = stripPrefix(prefix, keyUri);
     return TinkJsonProtoKeysetFormat.parseKeyset(keyset, InsecureSecretKeyAccess.get())
         .getPrimitive(Aead.class);
   }
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 4b481ff4f..97db3a453 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
@@ -126,3 +126,15 @@ java_test(
         "@tink_java//src/main/java/com/google/crypto/tink/aead:predefined_aead_parameters",
     ],
 )
+
+java_test(
+    name = "KmsClientsFactoryTest",
+    size = "small",
+    srcs = ["KmsClientsFactoryTest.java"],
+    deps = [
+        "@maven//:com_google_truth_truth",
+        "@maven//:junit_junit",
+        "//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",
+    ],
+)
diff --git a/tools/tinkey/src/test/java/com/google/crypto/tink/tinkey/KmsClientsFactoryTest.java b/tools/tinkey/src/test/java/com/google/crypto/tink/tinkey/KmsClientsFactoryTest.java
new file mode 100644
index 000000000..29bbfaa52
--- /dev/null
+++ b/tools/tinkey/src/test/java/com/google/crypto/tink/tinkey/KmsClientsFactoryTest.java
@@ -0,0 +1,66 @@
+// Copyright 2023 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+//
+////////////////////////////////////////////////////////////////////////////////
+
+package com.google.crypto.tink.tinkey;
+
+import static com.google.common.truth.Truth.assertThat;
+import static org.junit.Assert.assertThrows;
+
+import java.security.GeneralSecurityException;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+@RunWith(JUnit4.class)
+public final class KmsClientsFactoryTest {
+  @Test
+  public void testAddAndUse_works() throws Exception {
+    KmsClientsFactory factory = new KmsClientsFactory();
+    factory.addFactory(TinkeyTestKmsClient::new);
+
+    assertThat(factory.newClientFor("tinkey-test-kms-client://some"))
+        .isInstanceOf(TinkeyTestKmsClient.class);
+  }
+
+  @Test
+  public void test_newInstances_differ() throws Exception {
+    KmsClientsFactory factory = new KmsClientsFactory();
+    factory.addFactory(TinkeyTestKmsClient::new);
+
+    assertThat(factory.newClientFor("tinkey-test-kms-client://some"))
+        .isNotEqualTo(factory.newClientFor("tinkey-test-kms-client://some"));
+  }
+
+  @Test
+  public void test_notSupported_throws() throws Exception {
+    KmsClientsFactory factory = new KmsClientsFactory();
+    factory.addFactory(TinkeyTestKmsClient::new);
+
+    assertThrows(
+        GeneralSecurityException.class, () -> factory.newClientFor("not_supported://some"));
+  }
+
+  @Test
+  public void test_multiplePrefixes_works() throws Exception {
+    KmsClientsFactory factory = new KmsClientsFactory();
+    factory.addFactory(() -> TinkeyTestKmsClient.createForPrefix("prefix1:"));
+    factory.addFactory(() -> TinkeyTestKmsClient.createForPrefix("prefix2:"));
+
+    assertThat(factory.newClientFor("prefix1:foo")).isInstanceOf(TinkeyTestKmsClient.class);
+    assertThat(factory.newClientFor("prefix2:bar")).isInstanceOf(TinkeyTestKmsClient.class);
+    assertThrows(GeneralSecurityException.class, () -> factory.newClientFor("prefix3://some"));
+  }
+}
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 cde78affb..9dcf25d4d 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
@@ -18,7 +18,9 @@ package com.google.crypto.tink.tinkey;
 
 import static com.google.common.truth.Truth.assertThat;
 import static java.nio.charset.StandardCharsets.UTF_8;
+import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertThrows;
+import static org.junit.Assert.assertTrue;
 
 import com.google.crypto.tink.Aead;
 import com.google.crypto.tink.KeysetHandle;
@@ -51,9 +53,7 @@ public final class TinkeyTestKmsClientTest {
     KeysetHandle handle = KeysetHandle.generateNew(PredefinedAeadParameters.AES128_GCM);
     String masterKeyUri = TinkeyTestKmsClient.createKeyUri(handle);
     Aead masterKey =
-        KmsClients.getAutoLoaded(masterKeyUri)
-            .withCredentials(credentialPath.toString())
-            .getAead(masterKeyUri);
+        new TinkeyTestKmsClient().withCredentials(credentialPath.toString()).getAead(masterKeyUri);
     Aead manualMasterKey = handle.getPrimitive(Aead.class);
 
     byte[] ciphertext = manualMasterKey.encrypt(new byte[] {}, new byte[] {});
@@ -61,6 +61,15 @@ public final class TinkeyTestKmsClientTest {
   }
 
   @Test
+  public void test_clientAllowsCorrectPrefixes_works() throws Exception {
+    assertTrue(new TinkeyTestKmsClient().doesSupport("tinkey-test-kms-client://"));
+    assertFalse(new TinkeyTestKmsClient().doesSupport("somethingelse://"));
+
+    assertTrue(TinkeyTestKmsClient.createForPrefix("a").doesSupport("a://"));
+    assertFalse(TinkeyTestKmsClient.createForPrefix("a").doesSupport("tinkey-test-kms-client://"));
+  }
+
+  @Test
   public void test_clientCannotBeUsedWithWrongCredentials_throws() throws Exception {
     Path directory = Files.createTempDirectory(/* prefix= */ "");
     Path credentialPath = Paths.get(directory.toString(), "credentials");
@@ -68,8 +77,7 @@ public final class TinkeyTestKmsClientTest {
 
     KeysetHandle handle = KeysetHandle.generateNew(PredefinedAeadParameters.AES128_GCM);
     String masterKeyUri = TinkeyTestKmsClient.createKeyUri(handle);
-    KmsClient client =
-        KmsClients.getAutoLoaded(masterKeyUri).withCredentials(credentialPath.toString());
+    KmsClient client = new TinkeyTestKmsClient().withCredentials(credentialPath.toString());
     assertThrows(GeneralSecurityException.class, () -> client.getAead(masterKeyUri));
   }
 
@@ -77,8 +85,26 @@ public final class TinkeyTestKmsClientTest {
   public void test_clientCannotBeUsedWithoutCallingWithCredential_throws() throws Exception {
     KeysetHandle handle = KeysetHandle.generateNew(PredefinedAeadParameters.AES128_GCM);
     String masterKeyUri = TinkeyTestKmsClient.createKeyUri(handle);
-    KmsClient client = KmsClients.getAutoLoaded(masterKeyUri);
+    KmsClient client = new TinkeyTestKmsClient();
 
     assertThrows(GeneralSecurityException.class, () -> client.getAead(masterKeyUri));
   }
+
+  @Test
+  public void test_differentPrefix_works() throws Exception {
+    Path directory = Files.createTempDirectory(/* prefix= */ "");
+    Path credentialPath = Paths.get(directory.toString(), "credentials");
+    Files.write(credentialPath, "VALID CREDENTIALS".getBytes(UTF_8));
+
+    KeysetHandle handle = KeysetHandle.generateNew(PredefinedAeadParameters.AES128_GCM);
+    String masterKeyUri = TinkeyTestKmsClient.createKeyUri(handle);
+    Aead masterKey =
+        KmsClients.getAutoLoaded(masterKeyUri)
+            .withCredentials(credentialPath.toString())
+            .getAead(masterKeyUri);
+    Aead manualMasterKey = handle.getPrimitive(Aead.class);
+
+    byte[] ciphertext = manualMasterKey.encrypt(new byte[] {}, new byte[] {});
+    assertThat(masterKey.decrypt(ciphertext, new byte[] {})).isEqualTo(new byte[] {});
+  }
 }