Abstract KMS code with KmsConnection and create a fake KmsConnection

This simplifies the tests for KmsKeyring and KmsUpdater.

This is a followup to []

-------------
Created by MOE: https://github.com/google/moe
MOE_MIGRATED_REVID=148496758

java/google/registry/keyring/kms/BUILD

@@ -15,6 +15,7 @@ java_library(
         "//third_party/java/objectify:objectify-v4_1",
         "@com_google_api_client",
         "@com_google_apis_google_api_services_cloudkms",
+        "@com_google_auto_value",
         "@com_google_code_findbugs_jsr305",
         "@com_google_dagger",
         "@com_google_guava",

java/google/registry/keyring/kms/EncryptResponse.java

@@ -0,0 +1,41 @@
+// Copyright 2017 The Nomulus Authors. All Rights Reserved.
+//
+// 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 google.registry.keyring.kms;
+
+import com.google.auto.value.AutoValue;
+import com.google.common.annotations.VisibleForTesting;
+
+/**
+ * A value type class containing a Cloud KMS encrypted and encoded ciphertext, and the name of the
+ * CryptoKeyVersion used to encrypt it.
+ */
+@AutoValue
+abstract class EncryptResponse {
+
+  static EncryptResponse create(
+      com.google.api.services.cloudkms.v1beta1.model.EncryptResponse cloudKmsEncryptResponse) {
+    return new AutoValue_EncryptResponse(
+        cloudKmsEncryptResponse.getCiphertext(), cloudKmsEncryptResponse.getName());
+  }
+
+  @VisibleForTesting
+  static EncryptResponse create(String ciphertext, String cryptoKeyVersionName) {
+    return new AutoValue_EncryptResponse(ciphertext, cryptoKeyVersionName);
+  }
+
+  abstract String ciphertext();
+
+  abstract String cryptoKeyVersionName();
+}

java/google/registry/keyring/kms/KmsConnection.java

@@ -0,0 +1,42 @@
+// Copyright 2017 The Nomulus Authors. All Rights Reserved.
+//
+// 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 google.registry.keyring.kms;
+
+import java.io.IOException;
+
+/** An abstraction to simplify Cloud KMS operations. */
+interface KmsConnection {
+
+  /**
+   * The maximum allowable secret size, as set by Cloud KMS.
+   *
+   * @see <a
+   *     href="https://cloud.google.com/kms/docs/reference/rest/v1beta1/projects.locations.keyRings.cryptoKeys/encrypt#request-body">projects.locations.keyRings.cryptoKeys.encrypt</a>
+   */
+  int MAX_SECRET_SIZE_BYTES = 64 * 1024;
+
+  /**
+   * Encrypts a plaintext with CryptoKey {@code cryptoKeyName} on KeyRing {@code keyRingName}.
+   *
+   * <p>The latest CryptoKeyVersion is used to encrypt the value. The value must not be larger than
+   * {@code MAX_SECRET_SIZE_BYTES}.
+   *
+   * <p>If no applicable CryptoKey or CryptoKeyVersion exist, they will be created.
+   */
+  EncryptResponse encrypt(String cryptoKeyName, byte[] plaintext) throws IOException;
+
+  /** Decrypts a Cloud KMS encrypted and encoded value with CryptoKey {@code cryptoKeyName}. */
+  byte[] decrypt(String cryptoKeyName, String encodedCiphertext) throws IOException;
+}

java/google/registry/keyring/kms/KmsConnectionImpl.java

@@ -0,0 +1,147 @@
+// Copyright 2017 The Nomulus Authors. All Rights Reserved.
+//
+// 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 google.registry.keyring.kms;
+
+import static com.google.common.base.Preconditions.checkArgument;
+
+import com.google.api.client.googleapis.json.GoogleJsonResponseException;
+import com.google.api.client.http.HttpStatusCodes;
+import com.google.api.services.cloudkms.v1beta1.CloudKMS;
+import com.google.api.services.cloudkms.v1beta1.model.CryptoKey;
+import com.google.api.services.cloudkms.v1beta1.model.CryptoKeyVersion;
+import com.google.api.services.cloudkms.v1beta1.model.DecryptRequest;
+import com.google.api.services.cloudkms.v1beta1.model.EncryptRequest;
+import com.google.api.services.cloudkms.v1beta1.model.KeyRing;
+import google.registry.config.RegistryConfig.Config;
+import google.registry.keyring.api.KeyringException;
+import java.io.IOException;
+import javax.inject.Inject;
+
+/** The {@link KmsConnection} which talks to Cloud KMS. */
+class KmsConnectionImpl implements KmsConnection {
+
+  private static final String KMS_KEYRING_NAME_FORMAT = "projects/%s/locations/global/keyRings/%s";
+  private static final String KMS_CRYPTO_KEY_NAME_FORMAT =
+      "projects/%s/locations/global/keyRings/%s/cryptoKeys/%s";
+  private static final String KMS_CRYPTO_KEY_VERSION_NAME_FORMAT =
+      "projects/%s/locations/global/keyRings/%s/cryptoKeys/%s/cryptoKeyVersions";
+
+  private final CloudKMS kms;
+  private final String kmsKeyRingName;
+  private final String projectId;
+
+  @Inject
+  KmsConnectionImpl(
+      @Config("cloudKmsProjectId") String projectId,
+      @Config("cloudKmsKeyRing") String kmsKeyringName,
+      CloudKMS kms) {
+    this.projectId = projectId;
+    this.kmsKeyRingName = kmsKeyringName;
+    this.kms = kms;
+  }
+
+  @Override
+  public EncryptResponse encrypt(String cryptoKeyName, byte[] value) throws IOException {
+    checkArgument(
+        value.length <= MAX_SECRET_SIZE_BYTES,
+        "Value to encrypt was larger than %s bytes",
+        MAX_SECRET_SIZE_BYTES);
+    String fullKeyRingName = getKeyRingName(projectId, kmsKeyRingName);
+    try {
+      kms.projects().locations().keyRings().get(fullKeyRingName).execute();
+    } catch (GoogleJsonResponseException jsonException) {
+      if (jsonException.getStatusCode() == HttpStatusCodes.STATUS_CODE_NOT_FOUND) {
+        // Create the KeyRing in the "global" namespace. Encryption keys will be accessible from all
+        // GCP regions.
+        kms.projects()
+            .locations()
+            .keyRings()
+            .create("global", new KeyRing().setName(fullKeyRingName))
+            .execute();
+      } else {
+        throw jsonException;
+      }
+    }
+
+    String fullKeyName = getCryptoKeyName(projectId, kmsKeyRingName, cryptoKeyName);
+
+    try {
+      kms.projects().locations().keyRings().cryptoKeys().get(fullKeyName).execute();
+    } catch (GoogleJsonResponseException jsonException) {
+      if (jsonException.getStatusCode() == HttpStatusCodes.STATUS_CODE_NOT_FOUND) {
+        kms.projects()
+            .locations()
+            .keyRings()
+            .cryptoKeys()
+            .create(
+                fullKeyName, new CryptoKey().setName(cryptoKeyName).setPurpose("ENCRYPT_DECRYPT"))
+            .execute();
+      } else {
+        throw jsonException;
+      }
+    }
+
+    CryptoKeyVersion cryptoKeyVersion =
+        kms.projects()
+            .locations()
+            .keyRings()
+            .cryptoKeys()
+            .cryptoKeyVersions()
+            .create(
+                getCryptoKeyVersionName(projectId, kmsKeyRingName, cryptoKeyName),
+                new CryptoKeyVersion())
+            .execute();
+
+    return EncryptResponse.create(
+        kms.projects()
+            .locations()
+            .keyRings()
+            .cryptoKeys()
+            .encrypt(cryptoKeyVersion.getName(), new EncryptRequest().encodePlaintext(value))
+            .execute());
+  }
+
+  @Override
+  public byte[] decrypt(String cryptoKeyName, String encodedCiphertext) {
+    try {
+      return kms.projects()
+          .locations()
+          .keyRings()
+          .cryptoKeys()
+          .decrypt(
+              getCryptoKeyName(projectId, kmsKeyRingName, cryptoKeyName),
+              new DecryptRequest().setCiphertext(encodedCiphertext))
+          .execute()
+          .decodePlaintext();
+    } catch (IOException e) {
+      throw new KeyringException(
+          String.format("CloudKMS decrypt operation failed for secret %s", cryptoKeyName), e);
+    }
+  }
+
+  static String getKeyRingName(String projectId, String kmsKeyRingName) {
+    return String.format(KMS_KEYRING_NAME_FORMAT, projectId, kmsKeyRingName);
+  }
+
+  static String getCryptoKeyName(String projectId, String kmsKeyRingName, String cryptoKeyName) {
+    return String.format(KMS_CRYPTO_KEY_NAME_FORMAT, projectId, kmsKeyRingName, cryptoKeyName);
+  }
+
+  static String getCryptoKeyVersionName(
+      String projectId, String kmsKeyRingName, String cryptoKeyName) {
+    return String.format(
+        KMS_CRYPTO_KEY_VERSION_NAME_FORMAT, projectId, kmsKeyRingName, cryptoKeyName);
+  }
+}

java/google/registry/keyring/kms/KmsKeyring.java

@@ -20,10 +20,7 @@ import static google.registry.model.common.EntityGroupRoot.getCrossTldKey;
 import static google.registry.model.ofy.ObjectifyService.ofy;
 import static java.nio.charset.StandardCharsets.UTF_8;
 
-import com.google.api.services.cloudkms.v1beta1.CloudKMS;
-import com.google.api.services.cloudkms.v1beta1.model.DecryptRequest;
 import com.googlecode.objectify.Key;
-import google.registry.config.RegistryConfig.Config;
 import google.registry.keyring.api.Keyring;
 import google.registry.keyring.api.KeyringException;
 import google.registry.keyring.api.PgpHelper;
@@ -52,12 +49,6 @@ import org.bouncycastle.openpgp.operator.bc.BcPGPDigestCalculatorProvider;
  */
 public class KmsKeyring implements Keyring {
 
-  private static final String KMS_KEYRING_NAME_FORMAT = "projects/%s/locations/global/keyRings/%s";
-  private static final String KMS_CRYPTO_KEY_NAME_FORMAT =
-      "projects/%s/locations/global/keyRings/%s/cryptoKeys/%s";
-  private static final String KMS_CRYPTO_KEY_VERSION_NAME_FORMAT =
-      "projects/%s/locations/global/keyRings/%s/cryptoKeys/%s/cryptoKeyVersions";
-
   static final String BRAINTREE_PRIVATE_KEY_NAME = "braintree-private-key";
   static final String BRDA_RECEIVER_PUBLIC_NAME = "brda-receiver-public";
   static final String BRDA_SIGNING_PRIVATE_NAME = "brda-signing-private";
@@ -75,18 +66,11 @@ public class KmsKeyring implements Keyring {
   static final String RDE_STAGING_PRIVATE_NAME = "rde-staging-private";
   static final String RDE_STAGING_PUBLIC_NAME = "rde-staging-public";
 
-  private final CloudKMS kms;
+  private final KmsConnection kmsConnection;
-  private final String kmsKeyRingName;
-  private final String projectId;
 
   @Inject
-  KmsKeyring(
+  KmsKeyring(KmsConnection kmsConnection) {
-      @Config("cloudKmsProjectId") String projectId,
+    this.kmsConnection = kmsConnection;
-      @Config("cloudKmsKeyRing") String kmsKeyringName,
-      CloudKMS kms) {
-    this.projectId = projectId;
-    this.kmsKeyRingName = kmsKeyringName;
-    this.kms = kms;
   }
 
   @Override
@@ -212,32 +196,10 @@ public class KmsKeyring implements Keyring {
     String encryptedData = ofy().load().key(secret.getLatestRevision()).now().getEncryptedValue();
 
     try {
-      return kms.projects()
+      return kmsConnection.decrypt(secret.getName(), encryptedData);
-          .locations()
-          .keyRings()
-          .cryptoKeys()
-          .decrypt(
-              getCryptoKeyName(projectId, kmsKeyRingName, secret.getName()),
-              new DecryptRequest().setCiphertext(encryptedData))
-          .execute()
-          .decodePlaintext();
     } catch (IOException e) {
       throw new KeyringException(
           String.format("CloudKMS decrypt operation failed for secret %s", keyName), e);
     }
   }
-
-  static String getKeyRingName(String projectId, String kmsKeyRingName) {
-    return String.format(KMS_KEYRING_NAME_FORMAT, projectId, kmsKeyRingName);
-  }
-
-  static String getCryptoKeyName(String projectId, String kmsKeyRingName, String cryptoKeyName) {
-    return String.format(KMS_CRYPTO_KEY_NAME_FORMAT, projectId, kmsKeyRingName, cryptoKeyName);
-  }
-
-  static String getCryptoKeyVersionName(
-      String projectId, String kmsKeyRingName, String cryptoKeyName) {
-    return String.format(
-        KMS_CRYPTO_KEY_VERSION_NAME_FORMAT, projectId, kmsKeyRingName, cryptoKeyName);
-  }
 }

java/google/registry/keyring/kms/KmsModule.java

@@ -39,4 +39,9 @@ public final class KmsModule {
         .setApplicationName(projectId)
         .build();
   }
+
+  @Provides
+  static KmsConnection provideKmsAdapter(KmsConnectionImpl kmsAdapter) {
+    return kmsAdapter;
+  }
 }

java/google/registry/keyring/kms/KmsUpdater.java

@@ -32,23 +32,12 @@ import static google.registry.keyring.kms.KmsKeyring.RDE_SSH_CLIENT_PRIVATE_NAME
 import static google.registry.keyring.kms.KmsKeyring.RDE_SSH_CLIENT_PUBLIC_NAME;
 import static google.registry.keyring.kms.KmsKeyring.RDE_STAGING_PRIVATE_NAME;
 import static google.registry.keyring.kms.KmsKeyring.RDE_STAGING_PUBLIC_NAME;
-import static google.registry.keyring.kms.KmsKeyring.getCryptoKeyName;
-import static google.registry.keyring.kms.KmsKeyring.getCryptoKeyVersionName;
-import static google.registry.keyring.kms.KmsKeyring.getKeyRingName;
 import static google.registry.model.ofy.ObjectifyService.ofy;
 import static google.registry.util.PreconditionsUtils.checkArgumentNotNull;
 import static java.nio.charset.StandardCharsets.UTF_8;
 
-import com.google.api.client.googleapis.json.GoogleJsonResponseException;
-import com.google.api.services.cloudkms.v1beta1.CloudKMS;
-import com.google.api.services.cloudkms.v1beta1.model.CryptoKey;
-import com.google.api.services.cloudkms.v1beta1.model.CryptoKeyVersion;
-import com.google.api.services.cloudkms.v1beta1.model.EncryptRequest;
-import com.google.api.services.cloudkms.v1beta1.model.EncryptResponse;
-import com.google.api.services.cloudkms.v1beta1.model.KeyRing;
 import com.google.common.collect.ImmutableMap;
 import com.googlecode.objectify.VoidWork;
-import google.registry.config.RegistryConfig.Config;
 import google.registry.model.server.KmsSecret;
 import google.registry.model.server.KmsSecretRevision;
 import java.io.IOException;
@@ -65,22 +54,12 @@ import org.bouncycastle.openpgp.bc.BcPGPSecretKeyRing;
  */
 public final class KmsUpdater {
 
-  private static final int RESOURCE_NOT_FOUND = 404;
+  private final KmsConnection kmsConnection;
-
-  private final String projectId;
-  private final String kmsKeyRingName;
-  private final CloudKMS kms;
-
   private final HashMap<String, byte[]> secretValues;
 
   @Inject
-  public KmsUpdater(
+  public KmsUpdater(KmsConnection kmsConnection) {
-      @Config("cloudKmsProjectId") String projectId,
+    this.kmsConnection = kmsConnection;
-      @Config("cloudKmsKeyRing") String kmsKeyRingName,
-      CloudKMS kms) {
-    this.projectId = projectId;
-    this.kmsKeyRingName = kmsKeyRingName;
-    this.kms = kms;
 
     // Use LinkedHashMap to preserve insertion order on update() to simplify testing and debugging
     this.secretValues = new LinkedHashMap<>();
@@ -180,64 +159,10 @@ public final class KmsUpdater {
    */
   private ImmutableMap<String, EncryptResponse> encryptValues(Map<String, byte[]> keyValues)
       throws IOException {
-    String fullKeyRingName = getKeyRingName(projectId, kmsKeyRingName);
-    try {
-      kms.projects().locations().keyRings().get(fullKeyRingName).execute();
-    } catch (GoogleJsonResponseException jsonException) {
-      if (jsonException.getStatusCode() == RESOURCE_NOT_FOUND) {
-        // Create the KeyRing in the "global" namespace. Encryption keys will be accessible from all
-        // GCP regions.
-        kms.projects()
-            .locations()
-            .keyRings()
-            .create("global", new KeyRing().setName(fullKeyRingName))
-            .execute();
-      } else {
-        throw jsonException;
-      }
-    }
-
     ImmutableMap.Builder<String, EncryptResponse> encryptedValues = new ImmutableMap.Builder<>();
     for (Map.Entry<String, byte[]> entry : keyValues.entrySet()) {
-      String keyName = entry.getKey();
+      String secretName = entry.getKey();
-      String fullKeyName = getCryptoKeyName(projectId, kmsKeyRingName, keyName);
+      encryptedValues.put(secretName, kmsConnection.encrypt(secretName, entry.getValue()));
-
-      try {
-        kms.projects().locations().keyRings().cryptoKeys().get(fullKeyName).execute();
-      } catch (GoogleJsonResponseException jsonException) {
-        if (jsonException.getStatusCode() == RESOURCE_NOT_FOUND) {
-          kms.projects()
-              .locations()
-              .keyRings()
-              .cryptoKeys()
-              .create(fullKeyName, new CryptoKey().setName(keyName).setPurpose("ENCRYPT_DECRYPT"))
-              .execute();
-        } else {
-          throw jsonException;
-        }
-      }
-
-      CryptoKeyVersion cryptoKeyVersion =
-          kms.projects()
-              .locations()
-              .keyRings()
-              .cryptoKeys()
-              .cryptoKeyVersions()
-              .create(
-                  getCryptoKeyVersionName(projectId, kmsKeyRingName, keyName),
-                  new CryptoKeyVersion())
-              .execute();
-
-      encryptedValues.put(
-          keyName,
-          kms.projects()
-              .locations()
-              .keyRings()
-              .cryptoKeys()
-              .encrypt(
-                  cryptoKeyVersion.getName(),
-                  new EncryptRequest().encodePlaintext(entry.getValue()))
-              .execute());
     }
     return encryptedValues.build();
   }
@@ -262,8 +187,8 @@ public final class KmsUpdater {
 
                   KmsSecretRevision secretRevision =
                       new KmsSecretRevision.Builder()
-                          .setEncryptedValue(revisionData.getCiphertext())
+                          .setEncryptedValue(revisionData.ciphertext())
-                          .setKmsCryptoKeyVersionName(revisionData.getName())
+                          .setKmsCryptoKeyVersionName(revisionData.cryptoKeyVersionName())
                           .setParent(secretName)
                           .build();
                   ofy()

javatests/google/registry/keyring/kms/FakeKmsConnection.java

@@ -0,0 +1,46 @@
+// Copyright 2017 The Nomulus Authors. All Rights Reserved.
+//
+// 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 google.registry.keyring.kms;
+
+import com.google.common.io.BaseEncoding;
+import java.io.IOException;
+import org.bouncycastle.util.Arrays;
+
+class FakeKmsConnection implements KmsConnection {
+
+  FakeKmsConnection() {}
+
+  /**
+   * Returns a dummy {@link EncryptResponse}.
+   *
+   * <p>The "encrypted value" in the response is the provided value reversed and then base64-encoded
+   * and the name of the cryptoKeyVersion is {@code cryptoKeyName + "/foo"}.
+   */
+  @Override
+  public EncryptResponse encrypt(String cryptoKeyName, byte[] plaintext) throws IOException {
+    return EncryptResponse.create(
+        BaseEncoding.base64().encode(Arrays.reverse(plaintext)), cryptoKeyName + "/foo");
+  }
+
+  /**
+   * Returns a "decrypted" plaintext.
+   *
+   * <p>The plaintext is the encodedCiphertext base64-decoded and then reversed.
+   */
+  @Override
+  public byte[] decrypt(String cryptoKeyName, String encodedCiphertext) throws IOException {
+    return Arrays.reverse(BaseEncoding.base64().decode(encodedCiphertext));
+  }
+}

javatests/google/registry/keyring/kms/KmsConnectionImplTest.java

@@ -0,0 +1,194 @@
+// Copyright 2017 The Nomulus Authors. All Rights Reserved.
+//
+// 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 google.registry.keyring.kms;
+
+import static com.google.common.truth.Truth.assertThat;
+import static java.nio.charset.StandardCharsets.UTF_8;
+import static org.mockito.Matchers.any;
+import static org.mockito.Matchers.anyString;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+import com.google.api.client.googleapis.json.GoogleJsonResponseException;
+import com.google.api.client.http.HttpResponse;
+import com.google.api.client.http.HttpResponseException;
+import com.google.api.services.cloudkms.v1beta1.CloudKMS;
+import com.google.api.services.cloudkms.v1beta1.model.CryptoKey;
+import com.google.api.services.cloudkms.v1beta1.model.CryptoKeyVersion;
+import com.google.api.services.cloudkms.v1beta1.model.DecryptRequest;
+import com.google.api.services.cloudkms.v1beta1.model.DecryptResponse;
+import com.google.api.services.cloudkms.v1beta1.model.EncryptRequest;
+import com.google.api.services.cloudkms.v1beta1.model.EncryptResponse;
+import com.google.api.services.cloudkms.v1beta1.model.KeyRing;
+import java.io.ByteArrayInputStream;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.ArgumentCaptor;
+import org.mockito.Captor;
+import org.mockito.Mock;
+import org.mockito.runners.MockitoJUnitRunner;
+
+@RunWith(MockitoJUnitRunner.class)
+public class KmsConnectionImplTest {
+
+  @Mock private CloudKMS kms;
+  @Mock private CloudKMS.Projects kmsProjects;
+  @Mock private CloudKMS.Projects.Locations kmsLocations;
+  @Mock private CloudKMS.Projects.Locations.KeyRings kmsKeyRings;
+  @Mock private CloudKMS.Projects.Locations.KeyRings.Get kmsKeyRingsGet;
+  @Mock private CloudKMS.Projects.Locations.KeyRings.Create kmsKeyRingsCreate;
+  @Mock private CloudKMS.Projects.Locations.KeyRings.CryptoKeys kmsCryptoKeys;
+  @Mock private CloudKMS.Projects.Locations.KeyRings.CryptoKeys.Get kmsCryptoKeysGet;
+  @Mock private CloudKMS.Projects.Locations.KeyRings.CryptoKeys.Create kmsCryptoKeysCreate;
+
+  @Mock
+  private CloudKMS.Projects.Locations.KeyRings.CryptoKeys.CryptoKeyVersions kmsCryptoKeyVersions;
+
+  @Mock
+  private CloudKMS.Projects.Locations.KeyRings.CryptoKeys.CryptoKeyVersions.Create
+      kmsCryptoKeyVersionsCreate;
+
+  @Mock private CloudKMS.Projects.Locations.KeyRings.CryptoKeys.Encrypt kmsCryptoKeysEncrypt;
+  @Mock private CloudKMS.Projects.Locations.KeyRings.CryptoKeys.Decrypt kmsCryptoKeysDecrypt;
+
+  @Captor private ArgumentCaptor<KeyRing> keyRing;
+  @Captor private ArgumentCaptor<CryptoKey> cryptoKey;
+  @Captor private ArgumentCaptor<CryptoKeyVersion> cryptoKeyVersion;
+  @Captor private ArgumentCaptor<String> keyRingName;
+  @Captor private ArgumentCaptor<String> cryptoKeyName;
+  @Captor private ArgumentCaptor<String> cryptoKeyVersionName;
+  @Captor private ArgumentCaptor<EncryptRequest> encryptRequest;
+  @Captor private ArgumentCaptor<DecryptRequest> decryptRequest;
+
+  @Before
+  public void setUp() throws Exception {
+    when(kms.projects()).thenReturn(kmsProjects);
+    when(kmsProjects.locations()).thenReturn(kmsLocations);
+    when(kmsLocations.keyRings()).thenReturn(kmsKeyRings);
+    when(kmsKeyRings.get(anyString())).thenReturn(kmsKeyRingsGet);
+    when(kmsKeyRings.create(anyString(), any(KeyRing.class))).thenReturn(kmsKeyRingsCreate);
+    when(kmsKeyRings.cryptoKeys()).thenReturn(kmsCryptoKeys);
+    when(kmsCryptoKeys.get(anyString())).thenReturn(kmsCryptoKeysGet);
+    when(kmsCryptoKeys.create(anyString(), any(CryptoKey.class))).thenReturn(kmsCryptoKeysCreate);
+    when(kmsCryptoKeys.cryptoKeyVersions()).thenReturn(kmsCryptoKeyVersions);
+    when(kmsCryptoKeyVersions.create(anyString(), any(CryptoKeyVersion.class)))
+        .thenReturn(kmsCryptoKeyVersionsCreate);
+    when(kmsCryptoKeyVersionsCreate.execute())
+        .thenReturn(new CryptoKeyVersion().setName(KmsTestHelper.DUMMY_CRYPTO_KEY_VERSION));
+    when(kmsCryptoKeys.encrypt(anyString(), any(EncryptRequest.class)))
+        .thenReturn(kmsCryptoKeysEncrypt);
+    when(kmsCryptoKeysEncrypt.execute())
+        .thenReturn(
+            new EncryptResponse()
+                .setName(KmsTestHelper.DUMMY_CRYPTO_KEY_VERSION)
+                .setCiphertext(KmsTestHelper.DUMMY_ENCRYPTED_VALUE));
+    when(kmsCryptoKeys.decrypt(anyString(), any(DecryptRequest.class)))
+        .thenReturn(kmsCryptoKeysDecrypt);
+  }
+
+  @Test
+  public void test_encrypt_createsKeyRingIfNotFound() throws Exception {
+    when(kmsKeyRingsGet.execute()).thenThrow(createNotFoundException());
+
+    new KmsConnectionImpl("foo", "bar", kms).encrypt("key", "moo".getBytes(UTF_8));
+
+    verify(kmsKeyRings).create(keyRingName.capture(), keyRing.capture());
+    assertThat(keyRingName.getValue()).isEqualTo("global");
+    assertThat(keyRing.getValue())
+        .isEqualTo(new KeyRing().setName("projects/foo/locations/global/keyRings/bar"));
+    verify(kmsKeyRingsCreate).execute();
+    verifyEncryptKmsApiCalls(
+        "moo",
+        "projects/foo/locations/global/keyRings/bar",
+        "projects/foo/locations/global/keyRings/bar/cryptoKeys/key",
+        "projects/foo/locations/global/keyRings/bar/cryptoKeys/key/cryptoKeyVersions");
+  }
+
+  @Test
+  public void test_encrypt_newCryptoKey() throws Exception {
+    when(kmsCryptoKeysGet.execute()).thenThrow(createNotFoundException());
+
+    new KmsConnectionImpl("foo", "bar", kms).encrypt("key", "moo".getBytes(UTF_8));
+
+    verify(kmsCryptoKeys).create(cryptoKeyName.capture(), cryptoKey.capture());
+    assertThat(cryptoKeyName.getValue())
+        .isEqualTo("projects/foo/locations/global/keyRings/bar/cryptoKeys/key");
+    assertThat(cryptoKey.getValue())
+        .isEqualTo(new CryptoKey().setName("key").setPurpose("ENCRYPT_DECRYPT"));
+    verify(kmsCryptoKeysCreate).execute();
+    verifyEncryptKmsApiCalls(
+        "moo",
+        "projects/foo/locations/global/keyRings/bar",
+        "projects/foo/locations/global/keyRings/bar/cryptoKeys/key",
+        "projects/foo/locations/global/keyRings/bar/cryptoKeys/key/cryptoKeyVersions");
+  }
+
+  @Test
+  public void test_encrypt() throws Exception {
+    new KmsConnectionImpl("foo", "bar", kms).encrypt("key", "moo".getBytes(UTF_8));
+
+    verifyEncryptKmsApiCalls(
+        "moo",
+        "projects/foo/locations/global/keyRings/bar",
+        "projects/foo/locations/global/keyRings/bar/cryptoKeys/key",
+        "projects/foo/locations/global/keyRings/bar/cryptoKeys/key/cryptoKeyVersions");
+  }
+
+  @Test
+  public void test_decrypt() throws Exception {
+    when(kmsCryptoKeysDecrypt.execute())
+        .thenReturn(new DecryptResponse().encodePlaintext("moo".getBytes(UTF_8)));
+
+    byte[] plaintext = new KmsConnectionImpl("foo", "bar", kms).decrypt("key", "blah");
+
+    verify(kmsCryptoKeys).decrypt(cryptoKeyName.capture(), decryptRequest.capture());
+    assertThat(cryptoKeyName.getValue())
+        .isEqualTo("projects/foo/locations/global/keyRings/bar/cryptoKeys/key");
+    assertThat(decryptRequest.getValue()).isEqualTo(new DecryptRequest().setCiphertext("blah"));
+    assertThat(plaintext).isEqualTo("moo".getBytes(UTF_8));
+  }
+
+  private void verifyEncryptKmsApiCalls(
+      String goldenValue,
+      String goldenCryptoKeyRingName,
+      String goldenCryptoKeyName,
+      String goldenCryptoKeyVersionName)
+      throws Exception {
+    verify(kmsKeyRings).get(keyRingName.capture());
+    assertThat(keyRingName.getValue()).isEqualTo(goldenCryptoKeyRingName);
+
+    verify(kmsCryptoKeys).get(cryptoKeyName.capture());
+    assertThat(cryptoKeyName.getValue()).isEqualTo(goldenCryptoKeyName);
+
+    verify(kmsCryptoKeyVersions).create(cryptoKeyVersionName.capture(), cryptoKeyVersion.capture());
+    assertThat(cryptoKeyVersionName.getValue()).isEqualTo(goldenCryptoKeyVersionName);
+
+    verify(kmsCryptoKeys).encrypt(cryptoKeyName.capture(), encryptRequest.capture());
+    assertThat(cryptoKeyName.getValue()).isEqualTo(KmsTestHelper.DUMMY_CRYPTO_KEY_VERSION);
+    assertThat(encryptRequest.getValue())
+        .isEqualTo(new EncryptRequest().encodePlaintext(goldenValue.getBytes(UTF_8)));
+  }
+
+  private static GoogleJsonResponseException createNotFoundException() throws Exception {
+    ByteArrayInputStream inputStream = new ByteArrayInputStream("".getBytes(UTF_8));
+    HttpResponse response = GoogleJsonResponseExceptionHelper.createHttpResponse(404, inputStream);
+    HttpResponseException.Builder httpResponseExceptionBuilder =
+        new HttpResponseException.Builder(response);
+    httpResponseExceptionBuilder.setStatusCode(404);
+    httpResponseExceptionBuilder.setStatusMessage("NOT_FOUND");
+    return new GoogleJsonResponseException(httpResponseExceptionBuilder, null);
+  }
+}

javatests/google/registry/keyring/kms/KmsKeyringTest.java

@@ -17,20 +17,13 @@ package google.registry.keyring.kms;
 import static com.google.common.truth.Truth.assertThat;
 import static google.registry.testing.DatastoreHelper.persistResources;
 import static java.nio.charset.StandardCharsets.UTF_8;
-import static org.mockito.Matchers.any;
-import static org.mockito.Matchers.anyString;
-import static org.mockito.Mockito.times;
-import static org.mockito.Mockito.verify;
-import static org.mockito.Mockito.when;
 
-import com.google.api.services.cloudkms.v1beta1.CloudKMS;
-import com.google.api.services.cloudkms.v1beta1.model.DecryptRequest;
-import com.google.api.services.cloudkms.v1beta1.model.DecryptResponse;
 import com.google.common.collect.ImmutableList;
 import google.registry.model.server.KmsSecret;
 import google.registry.model.server.KmsSecretRevision;
 import google.registry.model.server.KmsSecretRevision.Builder;
 import google.registry.testing.AppEngineRule;
+import java.io.IOException;
 import org.bouncycastle.openpgp.PGPKeyPair;
 import org.bouncycastle.openpgp.PGPPrivateKey;
 import org.bouncycastle.openpgp.PGPPublicKey;
@@ -38,307 +31,181 @@ import org.junit.Before;
 import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
-import org.mockito.ArgumentCaptor;
+import org.junit.runners.JUnit4;
-import org.mockito.Captor;
-import org.mockito.Mock;
-import org.mockito.runners.MockitoJUnitRunner;
 
-@RunWith(MockitoJUnitRunner.class)
+@RunWith(JUnit4.class)
 public class KmsKeyringTest {
 
   @Rule public final AppEngineRule appEngine = AppEngineRule.builder().withDatastore().build();
 
-  @Mock private CloudKMS kms;
-  @Mock private CloudKMS.Projects kmsProjects;
-  @Mock private CloudKMS.Projects.Locations kmsLocations;
-  @Mock private CloudKMS.Projects.Locations.KeyRings kmsKeyRings;
-  @Mock private CloudKMS.Projects.Locations.KeyRings.CryptoKeys kmsCryptoKeys;
-  @Mock private CloudKMS.Projects.Locations.KeyRings.CryptoKeys.Decrypt kmsCryptoKeysDecrypt;
-
-  @Captor private ArgumentCaptor<String> cryptoKeyName;
-  @Captor private ArgumentCaptor<DecryptRequest> decryptRequest;
-
   private KmsKeyring keyring;
 
   @Before
-  public void setUp() throws Exception {
+  public void setUp() {
-    keyring = new KmsKeyring("foo", "bar", kms);
+    keyring = new KmsKeyring(new FakeKmsConnection());
-
-    when(kms.projects()).thenReturn(kmsProjects);
-    when(kmsProjects.locations()).thenReturn(kmsLocations);
-    when(kmsLocations.keyRings()).thenReturn(kmsKeyRings);
-    when(kmsKeyRings.cryptoKeys()).thenReturn(kmsCryptoKeys);
-    when(kmsCryptoKeys.decrypt(anyString(), any(DecryptRequest.class)))
-        .thenReturn(kmsCryptoKeysDecrypt);
   }
 
   @Test
   public void test_getRdeSigningKey() throws Exception {
-    persistSecret("rde-signing-private");
+    saveKeyPairSecret("rde-signing-public", "rde-signing-private");
-    persistSecret("rde-signing-public");
-    when(kmsCryptoKeysDecrypt.execute())
-        .thenReturn(
-            new DecryptResponse()
-                .encodePlaintext(KmsTestHelper.getPublicKeyring().getPublicKey().getEncoded()))
-        .thenReturn(
-            new DecryptResponse().encodePlaintext(KmsTestHelper.getPrivateKeyring().getEncoded()));
 
     PGPKeyPair rdeSigningKey = keyring.getRdeSigningKey();
 
-    verify(kmsCryptoKeys, times(2)).decrypt(cryptoKeyName.capture(), decryptRequest.capture());
-    assertThat(cryptoKeyName.getAllValues())
-        .isEqualTo(
-            ImmutableList.of(
-                "projects/foo/locations/global/keyRings/bar/cryptoKeys/rde-signing-public",
-                "projects/foo/locations/global/keyRings/bar/cryptoKeys/rde-signing-private"));
-    assertThat(decryptRequest.getAllValues())
-        .isEqualTo(
-            ImmutableList.of(
-                new DecryptRequest().setCiphertext(KmsTestHelper.DUMMY_ENCRYPTED_VALUE),
-                new DecryptRequest().setCiphertext(KmsTestHelper.DUMMY_ENCRYPTED_VALUE)));
     assertThat(rdeSigningKey.getKeyID())
         .isEqualTo(KmsTestHelper.getPublicKeyring().getPublicKey().getKeyID());
   }
 
   @Test
   public void test_getRdeStagingEncryptionKey() throws Exception {
-    persistSecret("rde-staging-public");
+    savePublicKeySecret("rde-staging-public");
-    when(kmsCryptoKeysDecrypt.execute())
-        .thenReturn(
-            new DecryptResponse()
-                .encodePlaintext(KmsTestHelper.getPublicKeyring().getPublicKey().getEncoded()));
 
     PGPPublicKey rdeStagingEncryptionKey = keyring.getRdeStagingEncryptionKey();
 
-    verify(kmsCryptoKeys).decrypt(cryptoKeyName.capture(), decryptRequest.capture());
-    assertThat(cryptoKeyName.getValue())
-        .isEqualTo("projects/foo/locations/global/keyRings/bar/cryptoKeys/rde-staging-public");
-    assertThat(decryptRequest.getValue())
-        .isEqualTo(new DecryptRequest().setCiphertext(KmsTestHelper.DUMMY_ENCRYPTED_VALUE));
     assertThat(rdeStagingEncryptionKey.getFingerprint())
         .isEqualTo(KmsTestHelper.getPublicKeyring().getPublicKey().getFingerprint());
   }
 
   @Test
   public void test_getRdeStagingDecryptionKey() throws Exception {
-    persistSecret("rde-staging-private");
+    savePrivateKeySecret("rde-staging-private");
-    when(kmsCryptoKeysDecrypt.execute())
-        .thenReturn(
-            new DecryptResponse().encodePlaintext(KmsTestHelper.getPrivateKeyring().getEncoded()));
 
     PGPPrivateKey rdeStagingDecryptionKey = keyring.getRdeStagingDecryptionKey();
 
-    verify(kmsCryptoKeys).decrypt(cryptoKeyName.capture(), decryptRequest.capture());
-    assertThat(cryptoKeyName.getValue())
-        .isEqualTo("projects/foo/locations/global/keyRings/bar/cryptoKeys/rde-staging-private");
-    assertThat(decryptRequest.getValue())
-        .isEqualTo(new DecryptRequest().setCiphertext(KmsTestHelper.DUMMY_ENCRYPTED_VALUE));
     assertThat(rdeStagingDecryptionKey.getKeyID())
         .isEqualTo(KmsTestHelper.getPrivateKeyring().getSecretKey().getKeyID());
   }
 
   @Test
   public void test_getRdeReceiverKey() throws Exception {
-    persistSecret("rde-receiver-public");
+    savePublicKeySecret("rde-receiver-public");
-    when(kmsCryptoKeysDecrypt.execute())
-        .thenReturn(
-            new DecryptResponse().encodePlaintext(KmsTestHelper.getPublicKeyring().getEncoded()));
 
     PGPPublicKey rdeReceiverKey = keyring.getRdeReceiverKey();
 
-    verify(kmsCryptoKeys).decrypt(cryptoKeyName.capture(), decryptRequest.capture());
-    assertThat(cryptoKeyName.getValue())
-        .isEqualTo("projects/foo/locations/global/keyRings/bar/cryptoKeys/rde-receiver-public");
-    assertThat(decryptRequest.getValue())
-        .isEqualTo(new DecryptRequest().setCiphertext(KmsTestHelper.DUMMY_ENCRYPTED_VALUE));
     assertThat(rdeReceiverKey.getFingerprint())
         .isEqualTo(KmsTestHelper.getPublicKeyring().getPublicKey().getFingerprint());
   }
 
   @Test
   public void test_getBrdaSigningKey() throws Exception {
-    persistSecret("brda-signing-public");
+    saveKeyPairSecret("brda-signing-public", "brda-signing-private");
-    persistSecret("brda-signing-private");
-    when(kmsCryptoKeysDecrypt.execute())
-        .thenReturn(
-            new DecryptResponse()
-                .encodePlaintext(KmsTestHelper.getPublicKeyring().getPublicKey().getEncoded()))
-        .thenReturn(
-            new DecryptResponse().encodePlaintext(KmsTestHelper.getPrivateKeyring().getEncoded()));
 
     PGPKeyPair brdaSigningKey = keyring.getBrdaSigningKey();
 
-    verify(kmsCryptoKeys, times(2)).decrypt(cryptoKeyName.capture(), decryptRequest.capture());
-    assertThat(cryptoKeyName.getAllValues())
-        .isEqualTo(
-            ImmutableList.of(
-                "projects/foo/locations/global/keyRings/bar/cryptoKeys/brda-signing-public",
-                "projects/foo/locations/global/keyRings/bar/cryptoKeys/brda-signing-private"));
-    assertThat(decryptRequest.getAllValues())
-        .isEqualTo(
-            ImmutableList.of(
-                new DecryptRequest().setCiphertext(KmsTestHelper.DUMMY_ENCRYPTED_VALUE),
-                new DecryptRequest().setCiphertext(KmsTestHelper.DUMMY_ENCRYPTED_VALUE)));
     assertThat(brdaSigningKey.getKeyID())
         .isEqualTo(KmsTestHelper.getPrivateKeyring().getPublicKey().getKeyID());
   }
 
   @Test
   public void test_getBrdaReceiverKey() throws Exception {
-    persistSecret("brda-receiver-public");
+    savePublicKeySecret("brda-receiver-public");
-    when(kmsCryptoKeysDecrypt.execute())
-        .thenReturn(
-            new DecryptResponse()
-                .encodePlaintext(KmsTestHelper.getPublicKeyring().getPublicKey().getEncoded()));
 
     PGPPublicKey brdaReceiverKey = keyring.getBrdaReceiverKey();
 
-    verify(kmsCryptoKeys).decrypt(cryptoKeyName.capture(), decryptRequest.capture());
-    assertThat(cryptoKeyName.getValue())
-        .isEqualTo("projects/foo/locations/global/keyRings/bar/cryptoKeys/brda-receiver-public");
-    assertThat(decryptRequest.getValue())
-        .isEqualTo(new DecryptRequest().setCiphertext(KmsTestHelper.DUMMY_ENCRYPTED_VALUE));
     assertThat(brdaReceiverKey.getFingerprint())
         .isEqualTo(KmsTestHelper.getPublicKeyring().getPublicKey().getFingerprint());
   }
 
   @Test
   public void test_getRdeSshClientPublicKey() throws Exception {
-    persistSecret("rde-ssh-client-public");
+    saveCleartextSecret("rde-ssh-client-public");
-    when(kmsCryptoKeysDecrypt.execute())
-        .thenReturn(new DecryptResponse().encodePlaintext(KmsTestHelper.DUMMY_KEY.getBytes(UTF_8)));
 
     String rdeSshClientPublicKey = keyring.getRdeSshClientPublicKey();
 
-    verify(kmsCryptoKeys).decrypt(cryptoKeyName.capture(), decryptRequest.capture());
+    assertThat(rdeSshClientPublicKey).isEqualTo("rde-ssh-client-publicmoo");
-    assertThat(cryptoKeyName.getValue())
-        .isEqualTo("projects/foo/locations/global/keyRings/bar/cryptoKeys/rde-ssh-client-public");
-    assertThat(decryptRequest.getValue())
-        .isEqualTo(new DecryptRequest().setCiphertext(KmsTestHelper.DUMMY_ENCRYPTED_VALUE));
-    assertThat(rdeSshClientPublicKey).isEqualTo(KmsTestHelper.DUMMY_KEY);
   }
 
   @Test
   public void test_getRdeSshClientPrivateKey() throws Exception {
-    persistSecret("rde-ssh-client-private");
+    saveCleartextSecret("rde-ssh-client-private");
-    when(kmsCryptoKeysDecrypt.execute())
-        .thenReturn(new DecryptResponse().encodePlaintext(KmsTestHelper.DUMMY_KEY.getBytes(UTF_8)));
 
     String rdeSshClientPrivateKey = keyring.getRdeSshClientPrivateKey();
 
-    verify(kmsCryptoKeys).decrypt(cryptoKeyName.capture(), decryptRequest.capture());
+    assertThat(rdeSshClientPrivateKey).isEqualTo("rde-ssh-client-privatemoo");
-    assertThat(cryptoKeyName.getValue())
-        .isEqualTo("projects/foo/locations/global/keyRings/bar/cryptoKeys/rde-ssh-client-private");
-    assertThat(decryptRequest.getValue())
-        .isEqualTo(new DecryptRequest().setCiphertext(KmsTestHelper.DUMMY_ENCRYPTED_VALUE));
-    assertThat(rdeSshClientPrivateKey).isEqualTo(KmsTestHelper.DUMMY_KEY);
   }
 
   @Test
   public void test_getIcannReportingPassword() throws Exception {
-    persistSecret("icann-reporting-password");
+    saveCleartextSecret("icann-reporting-password");
-    when(kmsCryptoKeysDecrypt.execute())
-        .thenReturn(new DecryptResponse().encodePlaintext("icann123".getBytes(UTF_8)));
 
     String icannReportingPassword = keyring.getIcannReportingPassword();
 
-    verify(kmsCryptoKeys).decrypt(cryptoKeyName.capture(), decryptRequest.capture());
+    assertThat(icannReportingPassword).isEqualTo("icann-reporting-passwordmoo");
-    assertThat(cryptoKeyName.getValue())
-        .isEqualTo(
-            "projects/foo/locations/global/keyRings/bar/cryptoKeys/icann-reporting-password");
-    assertThat(decryptRequest.getValue())
-        .isEqualTo(new DecryptRequest().setCiphertext(KmsTestHelper.DUMMY_ENCRYPTED_VALUE));
-    assertThat(icannReportingPassword).isEqualTo("icann123");
   }
 
   @Test
   public void test_getMarksdbDnlLogin() throws Exception {
-    persistSecret("marksdb-dnl-login");
+    saveCleartextSecret("marksdb-dnl-login");
-    when(kmsCryptoKeysDecrypt.execute())
-        .thenReturn(new DecryptResponse().encodePlaintext(KmsTestHelper.DUMMY_KEY.getBytes(UTF_8)));
 
     String marksdbDnlLogin = keyring.getMarksdbDnlLogin();
 
-    verify(kmsCryptoKeys).decrypt(cryptoKeyName.capture(), decryptRequest.capture());
+    assertThat(marksdbDnlLogin).isEqualTo("marksdb-dnl-loginmoo");
-    assertThat(cryptoKeyName.getValue())
-        .isEqualTo("projects/foo/locations/global/keyRings/bar/cryptoKeys/marksdb-dnl-login");
-    assertThat(decryptRequest.getValue())
-        .isEqualTo(new DecryptRequest().setCiphertext(KmsTestHelper.DUMMY_ENCRYPTED_VALUE));
-    assertThat(marksdbDnlLogin).isEqualTo(KmsTestHelper.DUMMY_KEY);
   }
 
   @Test
   public void test_getMarksdbLordnPassword() throws Exception {
-    persistSecret("marksdb-lordn-password");
+    saveCleartextSecret("marksdb-lordn-password");
-    when(kmsCryptoKeysDecrypt.execute())
-        .thenReturn(new DecryptResponse().encodePlaintext(KmsTestHelper.DUMMY_KEY.getBytes(UTF_8)));
 
     String marksdbLordnPassword = keyring.getMarksdbLordnPassword();
 
-    verify(kmsCryptoKeys).decrypt(cryptoKeyName.capture(), decryptRequest.capture());
+    assertThat(marksdbLordnPassword).isEqualTo("marksdb-lordn-passwordmoo");
-    assertThat(cryptoKeyName.getValue())
-        .isEqualTo("projects/foo/locations/global/keyRings/bar/cryptoKeys/marksdb-lordn-password");
-    assertThat(decryptRequest.getValue())
-        .isEqualTo(new DecryptRequest().setCiphertext(KmsTestHelper.DUMMY_ENCRYPTED_VALUE));
-    assertThat(marksdbLordnPassword).isEqualTo(KmsTestHelper.DUMMY_KEY);
   }
 
   @Test
   public void test_getMarksdbSmdrlLogin() throws Exception {
-    persistSecret("marksdb-smdrl-login");
+    saveCleartextSecret("marksdb-smdrl-login");
-    when(kmsCryptoKeysDecrypt.execute())
-        .thenReturn(new DecryptResponse().encodePlaintext(KmsTestHelper.DUMMY_KEY.getBytes(UTF_8)));
 
     String marksdbSmdrlLogin = keyring.getMarksdbSmdrlLogin();
 
-    verify(kmsCryptoKeys).decrypt(cryptoKeyName.capture(), decryptRequest.capture());
+    assertThat(marksdbSmdrlLogin).isEqualTo("marksdb-smdrl-loginmoo");
-    assertThat(cryptoKeyName.getValue())
+
-        .isEqualTo("projects/foo/locations/global/keyRings/bar/cryptoKeys/marksdb-smdrl-login");
-    assertThat(decryptRequest.getValue())
-        .isEqualTo(new DecryptRequest().setCiphertext(KmsTestHelper.DUMMY_ENCRYPTED_VALUE));
-    assertThat(marksdbSmdrlLogin).isEqualTo(KmsTestHelper.DUMMY_KEY);
   }
 
   @Test
   public void test_getJsonCredential() throws Exception {
-    persistSecret("json-credential");
+    saveCleartextSecret("json-credential");
-    when(kmsCryptoKeysDecrypt.execute())
-        .thenReturn(new DecryptResponse().encodePlaintext(KmsTestHelper.DUMMY_KEY.getBytes(UTF_8)));
 
     String jsonCredential = keyring.getJsonCredential();
 
-    verify(kmsCryptoKeys).decrypt(cryptoKeyName.capture(), decryptRequest.capture());
+    assertThat(jsonCredential).isEqualTo("json-credentialmoo");
-    assertThat(cryptoKeyName.getValue())
-        .isEqualTo("projects/foo/locations/global/keyRings/bar/cryptoKeys/json-credential");
-    assertThat(decryptRequest.getValue())
-        .isEqualTo(new DecryptRequest().setCiphertext(KmsTestHelper.DUMMY_ENCRYPTED_VALUE));
-    assertThat(jsonCredential).isEqualTo(KmsTestHelper.DUMMY_KEY);
   }
 
   @Test
   public void test_getBraintreePrivateKey() throws Exception {
-    persistSecret("braintree-private-key");
+    saveCleartextSecret("braintree-private-key");
-    when(kmsCryptoKeysDecrypt.execute())
-        .thenReturn(new DecryptResponse().encodePlaintext(KmsTestHelper.DUMMY_KEY.getBytes(UTF_8)));
 
     String braintreePrivateKey = keyring.getBraintreePrivateKey();
 
-    verify(kmsCryptoKeys).decrypt(cryptoKeyName.capture(), decryptRequest.capture());
+    assertThat(braintreePrivateKey).isEqualTo("braintree-private-keymoo");
-    assertThat(cryptoKeyName.getValue())
-        .isEqualTo("projects/foo/locations/global/keyRings/bar/cryptoKeys/braintree-private-key");
-    assertThat(decryptRequest.getValue())
-        .isEqualTo(new DecryptRequest().setCiphertext(KmsTestHelper.DUMMY_ENCRYPTED_VALUE));
-    assertThat(braintreePrivateKey).isEqualTo(KmsTestHelper.DUMMY_KEY);
   }
 
-  private static void persistSecret(String secretName) {
+  private static void persistSecret(String secretName, byte[] secretValue) throws IOException {
+    KmsConnection kmsConnection = new FakeKmsConnection();
+
     KmsSecretRevision secretRevision =
         new Builder()
-            .setEncryptedValue(KmsTestHelper.DUMMY_ENCRYPTED_VALUE)
+            .setEncryptedValue(kmsConnection.encrypt(secretName, secretValue).ciphertext())
             .setKmsCryptoKeyVersionName(KmsTestHelper.DUMMY_CRYPTO_KEY_VERSION)
             .setParent(secretName)
             .build();
     KmsSecret secret = KmsSecret.create(secretName, secretRevision);
     persistResources(ImmutableList.of(secretRevision, secret));
   }
+
+  private static void saveCleartextSecret(String secretName) throws Exception {
+    persistSecret(secretName, (secretName + "moo").getBytes(UTF_8));
+  }
+
+  private static void savePublicKeySecret(String publicKeyName) throws Exception {
+    persistSecret(publicKeyName, KmsTestHelper.getPublicKeyring().getPublicKey().getEncoded());
+  }
+
+  private static void savePrivateKeySecret(String privateKeyName) throws Exception {
+    persistSecret(privateKeyName, KmsTestHelper.getPrivateKeyring().getEncoded());
+  }
+
+  private static void saveKeyPairSecret(String publicKeyName, String privateKeyName)
+      throws Exception {
+    savePublicKeySecret(publicKeyName);
+    savePrivateKeySecret(privateKeyName);
+  }
 }

javatests/google/registry/keyring/kms/KmsTestHelper.java

@@ -25,7 +25,6 @@ import org.bouncycastle.openpgp.bc.BcPGPSecretKeyRing;
 /** Stores dummy values for test use in {@link KmsUpdaterTest} and {@link KmsKeyringTest}. */
 final class KmsTestHelper {
 
-  static final String DUMMY_KEY = "the quick brown fox";
   static final String DUMMY_CRYPTO_KEY_VERSION = "cheeseburger";
   static final String DUMMY_ENCRYPTED_VALUE = "meow";
 

javatests/google/registry/keyring/kms/KmsUpdaterTest.java

@@ -18,123 +18,28 @@ import static com.google.common.truth.Truth.assertThat;
 import static google.registry.model.common.EntityGroupRoot.getCrossTldKey;
 import static google.registry.model.ofy.ObjectifyService.ofy;
 import static java.nio.charset.StandardCharsets.UTF_8;
-import static org.mockito.Matchers.anyString;
-import static org.mockito.Mockito.any;
-import static org.mockito.Mockito.times;
-import static org.mockito.Mockito.verify;
-import static org.mockito.Mockito.when;
 
-import com.google.api.client.googleapis.json.GoogleJsonResponseException;
-import com.google.api.client.http.HttpResponse;
-import com.google.api.client.http.HttpResponseException;
-import com.google.api.services.cloudkms.v1beta1.CloudKMS;
-import com.google.api.services.cloudkms.v1beta1.model.CryptoKey;
-import com.google.api.services.cloudkms.v1beta1.model.CryptoKeyVersion;
-import com.google.api.services.cloudkms.v1beta1.model.EncryptRequest;
-import com.google.api.services.cloudkms.v1beta1.model.EncryptResponse;
-import com.google.api.services.cloudkms.v1beta1.model.KeyRing;
-import com.google.common.collect.ImmutableList;
 import com.googlecode.objectify.Key;
 import google.registry.model.server.KmsSecret;
 import google.registry.model.server.KmsSecretRevision;
 import google.registry.testing.AppEngineRule;
-import java.io.ByteArrayInputStream;
+import java.io.IOException;
 import org.junit.Before;
 import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
-import org.mockito.ArgumentCaptor;
+import org.junit.runners.JUnit4;
-import org.mockito.Captor;
-import org.mockito.Mock;
-import org.mockito.runners.MockitoJUnitRunner;
 
-@RunWith(MockitoJUnitRunner.class)
+@RunWith(JUnit4.class)
 public class KmsUpdaterTest {
 
   @Rule public final AppEngineRule appEngine = AppEngineRule.builder().withDatastore().build();
 
-  @Mock private CloudKMS kms;
-  @Mock private CloudKMS.Projects kmsProjects;
-  @Mock private CloudKMS.Projects.Locations kmsLocations;
-  @Mock private CloudKMS.Projects.Locations.KeyRings kmsKeyRings;
-  @Mock private CloudKMS.Projects.Locations.KeyRings.Get kmsKeyRingsGet;
-  @Mock private CloudKMS.Projects.Locations.KeyRings.Create kmsKeyRingsCreate;
-  @Mock private CloudKMS.Projects.Locations.KeyRings.CryptoKeys kmsCryptoKeys;
-  @Mock private CloudKMS.Projects.Locations.KeyRings.CryptoKeys.Get kmsCryptoKeysGet;
-  @Mock private CloudKMS.Projects.Locations.KeyRings.CryptoKeys.Create kmsCryptoKeysCreate;
-
-  @Mock
-  private CloudKMS.Projects.Locations.KeyRings.CryptoKeys.CryptoKeyVersions kmsCryptoKeyVersions;
-
-  @Mock
-  private CloudKMS.Projects.Locations.KeyRings.CryptoKeys.CryptoKeyVersions.Create
-      kmsCryptoKeyVersionsCreate;
-
-  @Mock private CloudKMS.Projects.Locations.KeyRings.CryptoKeys.Encrypt kmsCryptoKeysEncrypt;
-
-  @Captor private ArgumentCaptor<KeyRing> keyRing;
-  @Captor private ArgumentCaptor<CryptoKey> cryptoKey;
-  @Captor private ArgumentCaptor<CryptoKeyVersion> cryptoKeyVersion;
-  @Captor private ArgumentCaptor<String> keyRingName;
-  @Captor private ArgumentCaptor<String> cryptoKeyName;
-  @Captor private ArgumentCaptor<String> cryptoKeyVersionName;
-  @Captor private ArgumentCaptor<EncryptRequest> encryptRequest;
-
   private KmsUpdater updater;
 
   @Before
-  public void setUp() throws Exception {
+  public void setUp() {
-    when(kms.projects()).thenReturn(kmsProjects);
+    updater = new KmsUpdater(new FakeKmsConnection());
-    when(kmsProjects.locations()).thenReturn(kmsLocations);
-    when(kmsLocations.keyRings()).thenReturn(kmsKeyRings);
-    when(kmsKeyRings.get(anyString())).thenReturn(kmsKeyRingsGet);
-    when(kmsKeyRings.create(anyString(), any(KeyRing.class))).thenReturn(kmsKeyRingsCreate);
-    when(kmsKeyRings.cryptoKeys()).thenReturn(kmsCryptoKeys);
-    when(kmsCryptoKeys.get(anyString())).thenReturn(kmsCryptoKeysGet);
-    when(kmsCryptoKeys.create(anyString(), any(CryptoKey.class))).thenReturn(kmsCryptoKeysCreate);
-    when(kmsCryptoKeys.cryptoKeyVersions()).thenReturn(kmsCryptoKeyVersions);
-    when(kmsCryptoKeyVersions.create(anyString(), any(CryptoKeyVersion.class)))
-        .thenReturn(kmsCryptoKeyVersionsCreate);
-    when(kmsCryptoKeyVersionsCreate.execute())
-        .thenReturn(new CryptoKeyVersion().setName(KmsTestHelper.DUMMY_CRYPTO_KEY_VERSION));
-    when(kmsCryptoKeys.encrypt(anyString(), any(EncryptRequest.class)))
-        .thenReturn(kmsCryptoKeysEncrypt);
-    when(kmsCryptoKeysEncrypt.execute())
-        .thenReturn(
-            new EncryptResponse()
-                .setName(KmsTestHelper.DUMMY_CRYPTO_KEY_VERSION)
-                .setCiphertext(KmsTestHelper.DUMMY_ENCRYPTED_VALUE));
-
-    updater = new KmsUpdater("foo", "bar", kms);
-  }
-
-  @Test
-  public void test_close_createsNewKeyRing_ifNotFound() throws Exception {
-    when(kmsKeyRingsGet.execute()).thenThrow(createNotFoundException());
-
-    updater.setBraintreePrivateKey(KmsTestHelper.DUMMY_KEY);
-    updater.update();
-
-    verify(kmsKeyRings).create(keyRingName.capture(), keyRing.capture());
-    assertThat(keyRingName.getValue()).isEqualTo("global");
-    assertThat(keyRing.getValue())
-        .isEqualTo(new KeyRing().setName("projects/foo/locations/global/keyRings/bar"));
-    verify(kmsKeyRingsCreate).execute();
-  }
-
-  @Test
-  public void test_close_createsNewCryptoKey_ifNotFound() throws Exception {
-    when(kmsCryptoKeysGet.execute()).thenThrow(createNotFoundException());
-
-    updater.setBraintreePrivateKey(KmsTestHelper.DUMMY_KEY);
-    updater.update();
-
-    verify(kmsCryptoKeys).create(cryptoKeyName.capture(), cryptoKey.capture());
-    assertThat(cryptoKeyName.getValue())
-        .isEqualTo("projects/foo/locations/global/keyRings/bar/cryptoKeys/braintree-private-key");
-    assertThat(cryptoKey.getValue())
-        .isEqualTo(new CryptoKey().setName("braintree-private-key").setPurpose("ENCRYPT_DECRYPT"));
-    verify(kmsCryptoKeysCreate).execute();
   }
 
   @Test
@@ -142,336 +47,159 @@ public class KmsUpdaterTest {
     updater
         .setBraintreePrivateKey("value1")
         .setIcannReportingPassword("value2")
-        .setJsonCredential("value3");
+        .setJsonCredential("value3")
-    updater.update();
+        .update();
 
-    verify(kmsCryptoKeys, times(3)).get(cryptoKeyName.capture());
+    verifySecretAndSecretRevisionWritten(
-    assertThat(cryptoKeyName.getAllValues())
+        "braintree-private-key", "braintree-private-key/foo", getCiphertext("value1"));
-        .isEqualTo(
+    verifySecretAndSecretRevisionWritten(
-            ImmutableList.of(
+        "icann-reporting-password", "icann-reporting-password/foo", getCiphertext("value2"));
-                "projects/foo/locations/" + "global/keyRings/bar/cryptoKeys/braintree-private-key",
+    verifySecretAndSecretRevisionWritten(
-                "projects/foo/locations/"
+        "json-credential", "json-credential/foo", getCiphertext("value3"));
-                    + "global/keyRings/bar/cryptoKeys/icann-reporting-password",
-                "projects/foo/locations/" + "global/keyRings/bar/cryptoKeys/json-credential"));
-
-    verify(kmsCryptoKeyVersions, times(3))
-        .create(cryptoKeyVersionName.capture(), cryptoKeyVersion.capture());
-    assertThat(cryptoKeyVersionName.getAllValues())
-        .isEqualTo(
-            ImmutableList.of(
-                "projects/foo/locations/"
-                    + "global/keyRings/bar/cryptoKeys/braintree-private-key/cryptoKeyVersions",
-                "projects/foo/locations/"
-                    + "global/keyRings/bar/cryptoKeys/icann-reporting-password/cryptoKeyVersions",
-                "projects/foo/locations/"
-                    + "global/keyRings/bar/cryptoKeys/json-credential/cryptoKeyVersions"));
-
-    verify(kmsCryptoKeys, times(3)).encrypt(cryptoKeyName.capture(), encryptRequest.capture());
-    assertThat(cryptoKeyName.getValue()).isEqualTo(KmsTestHelper.DUMMY_CRYPTO_KEY_VERSION);
-    assertThat(encryptRequest.getAllValues())
-        .isEqualTo(
-            ImmutableList.of(
-                new EncryptRequest().encodePlaintext("value1".getBytes(UTF_8)),
-                new EncryptRequest().encodePlaintext("value2".getBytes(UTF_8)),
-                new EncryptRequest().encodePlaintext("value3".getBytes(UTF_8))));
-
-    KmsSecret firstSecret = loadSecret("braintree-private-key");
-    assertThat(firstSecret).isNotNull();
-    assertThat(firstSecret.getLatestRevision()).isNotNull();
-
-    KmsSecret secondSecret = loadSecret("icann-reporting-password");
-    assertThat(secondSecret).isNotNull();
-    assertThat(secondSecret.getLatestRevision()).isNotNull();
-
-    KmsSecret thirdSecret = loadSecret("icann-reporting-password");
-    assertThat(thirdSecret).isNotNull();
-    assertThat(thirdSecret.getLatestRevision()).isNotNull();
   }
 
   @Test
   public void test_setBraintreePrivateKey() throws Exception {
-    updater.setBraintreePrivateKey(KmsTestHelper.DUMMY_KEY);
+    updater.setBraintreePrivateKey("value1").update();
-    updater.update();
 
-    verifyKmsApiCallsAndDatastoreWrites(
+    verifySecretAndSecretRevisionWritten(
-        "braintree-private-key",
+        "braintree-private-key", "braintree-private-key/foo", getCiphertext("value1"));
-        KmsTestHelper.DUMMY_KEY.getBytes(UTF_8),
-        "projects/foo/locations/global/keyRings/bar",
-        "projects/foo/locations/global/keyRings/bar/cryptoKeys/braintree-private-key",
-        "projects/foo/locations/"
-            + "global/keyRings/bar/cryptoKeys/braintree-private-key/cryptoKeyVersions");
   }
 
   @Test
   public void test_setBrdaReceiverKey() throws Exception {
-    updater.setBrdaReceiverPublicKey(KmsTestHelper.getPublicKeyring().getPublicKey());
+    updater.setBrdaReceiverPublicKey(KmsTestHelper.getPublicKeyring().getPublicKey()).update();
-    updater.update();
 
-    verifyKmsApiCallsAndDatastoreWrites(
+    verifySecretAndSecretRevisionWritten(
         "brda-receiver-public",
-        KmsTestHelper.getPublicKeyring().getPublicKey().getEncoded(),
+        "brda-receiver-public/foo",
-        "projects/foo/locations/global/keyRings/bar",
+        getCiphertext(KmsTestHelper.getPublicKeyring().getPublicKey().getEncoded()));
-        "projects/foo/locations/global/keyRings/bar/cryptoKeys/brda-receiver-public",
-        "projects/foo/locations/"
-            + "global/keyRings/bar/cryptoKeys/brda-receiver-public/cryptoKeyVersions");
   }
 
   @Test
   public void test_setBrdaSigningKey() throws Exception {
-    updater.setBrdaSigningKey(KmsTestHelper.getPrivateKeyring());
+    updater.setBrdaSigningKey(KmsTestHelper.getPrivateKeyring()).update();
-    updater.update();
 
-    verifyKmsApiCallsAndDatastoreWrites(
+    verifySecretAndSecretRevisionWritten(
         "brda-signing-private",
-        KmsTestHelper.getPrivateKeyring().getEncoded(),
+        "brda-signing-private/foo",
-        "projects/foo/locations/global/keyRings/bar",
+        getCiphertext(KmsTestHelper.getPrivateKeyring().getEncoded()));
-        "projects/foo/locations/global/keyRings/bar/cryptoKeys/brda-signing-private",
+    verifySecretAndSecretRevisionWritten(
-        "projects/foo/locations/"
-            + "global/keyRings/bar/cryptoKeys/brda-signing-private/cryptoKeyVersions",
         "brda-signing-public",
-        KmsTestHelper.getPublicKeyring().getPublicKey().getEncoded(),
+        "brda-signing-public/foo",
-        "projects/foo/locations/global/keyRings/bar/cryptoKeys/brda-signing-public",
+        getCiphertext(KmsTestHelper.getPrivateKeyring().getPublicKey().getEncoded()));
-        "projects/foo/locations/"
-            + "global/keyRings/bar/cryptoKeys/brda-signing-public/cryptoKeyVersions");
   }
 
   @Test
   public void test_setIcannReportingPassword() throws Exception {
-    updater.setIcannReportingPassword(KmsTestHelper.DUMMY_KEY);
+    updater.setIcannReportingPassword("value1").update();
-    updater.update();
 
-    verifyKmsApiCallsAndDatastoreWrites(
+    verifySecretAndSecretRevisionWritten(
-        "icann-reporting-password",
+        "icann-reporting-password", "icann-reporting-password/foo", getCiphertext("value1"));
-        KmsTestHelper.DUMMY_KEY.getBytes(UTF_8),
-        "projects/foo/locations/global/keyRings/bar",
-        "projects/foo/locations/global/keyRings/bar/cryptoKeys/icann-reporting-password",
-        "projects/foo/locations/"
-            + "global/keyRings/bar/cryptoKeys/icann-reporting-password/cryptoKeyVersions");
   }
 
   @Test
   public void test_setJsonCredential() throws Exception {
-    updater.setJsonCredential(KmsTestHelper.DUMMY_KEY);
+    updater.setJsonCredential("value1").update();
-    updater.update();
 
-    verifyKmsApiCallsAndDatastoreWrites(
+    verifySecretAndSecretRevisionWritten(
-        "json-credential",
+        "json-credential", "json-credential/foo", getCiphertext("value1"));
-        KmsTestHelper.DUMMY_KEY.getBytes(UTF_8),
-        "projects/foo/locations/global/keyRings/bar",
-        "projects/foo/locations/global/keyRings/bar/cryptoKeys/json-credential",
-        "projects/foo/locations/"
-            + "global/keyRings/bar/cryptoKeys/json-credential/cryptoKeyVersions");
   }
 
   @Test
   public void test_setMarksdbDnlLogin() throws Exception {
-    updater.setMarksdbDnlLogin(KmsTestHelper.DUMMY_KEY);
+    updater.setMarksdbDnlLogin("value1").update();
-    updater.update();
 
-    verifyKmsApiCallsAndDatastoreWrites(
+    verifySecretAndSecretRevisionWritten(
-        "marksdb-dnl-login",
+        "marksdb-dnl-login", "marksdb-dnl-login/foo", getCiphertext("value1"));
-        KmsTestHelper.DUMMY_KEY.getBytes(UTF_8),
-        "projects/foo/locations/global/keyRings/bar",
-        "projects/foo/locations/global/keyRings/bar/cryptoKeys/marksdb-dnl-login",
-        "projects/foo/locations/"
-            + "global/keyRings/bar/cryptoKeys/marksdb-dnl-login/cryptoKeyVersions");
   }
 
   @Test
   public void test_setMarksdbLordnPassword() throws Exception {
-    updater.setMarksdbLordnPassword(KmsTestHelper.DUMMY_KEY);
+    updater.setMarksdbLordnPassword("value1").update();
-    updater.update();
 
-    verifyKmsApiCallsAndDatastoreWrites(
+    verifySecretAndSecretRevisionWritten(
-        "marksdb-lordn-password",
+        "marksdb-lordn-password", "marksdb-lordn-password/foo", getCiphertext("value1"));
-        KmsTestHelper.DUMMY_KEY.getBytes(UTF_8),
-        "projects/foo/locations/global/keyRings/bar",
-        "projects/foo/locations/global/keyRings/bar/cryptoKeys/marksdb-lordn-password",
-        "projects/foo/locations/"
-            + "global/keyRings/bar/cryptoKeys/marksdb-lordn-password/cryptoKeyVersions");
   }
 
   @Test
   public void test_setMarksdbSmdrlLogin() throws Exception {
-    updater.setMarksdbSmdrlLogin(KmsTestHelper.DUMMY_KEY);
+    updater.setMarksdbSmdrlLogin("value1").update();
-    updater.update();
 
-    verifyKmsApiCallsAndDatastoreWrites(
+    verifySecretAndSecretRevisionWritten(
-        "marksdb-smdrl-login",
+        "marksdb-smdrl-login", "marksdb-smdrl-login/foo", getCiphertext("value1"));
-        KmsTestHelper.DUMMY_KEY.getBytes(UTF_8),
-        "projects/foo/locations/global/keyRings/bar",
-        "projects/foo/locations/global/keyRings/bar/cryptoKeys/marksdb-smdrl-login",
-        "projects/foo/locations/"
-            + "global/keyRings/bar/cryptoKeys/marksdb-smdrl-login/cryptoKeyVersions");
   }
 
   @Test
   public void test_setRdeReceiverKey() throws Exception {
-    updater.setRdeReceiverPublicKey(KmsTestHelper.getPublicKeyring().getPublicKey());
+    updater.setRdeReceiverPublicKey(KmsTestHelper.getPublicKeyring().getPublicKey()).update();
-    updater.update();
 
-    verifyKmsApiCallsAndDatastoreWrites(
+    verifySecretAndSecretRevisionWritten(
         "rde-receiver-public",
-        KmsTestHelper.getPublicKeyring().getPublicKey().getEncoded(),
+        "rde-receiver-public/foo",
-        "projects/foo/locations/global/keyRings/bar",
+        getCiphertext(KmsTestHelper.getPublicKeyring().getPublicKey().getEncoded()));
-        "projects/foo/locations/global/keyRings/bar/cryptoKeys/rde-receiver-public",
-        "projects/foo/locations/"
-            + "global/keyRings/bar/cryptoKeys/rde-receiver-public/cryptoKeyVersions");
   }
 
   @Test
   public void test_setRdeSigningKey() throws Exception {
-    updater.setRdeSigningKey(KmsTestHelper.getPrivateKeyring());
+    updater.setRdeSigningKey(KmsTestHelper.getPrivateKeyring()).update();
-    updater.update();
 
-    verifyKmsApiCallsAndDatastoreWrites(
+    verifySecretAndSecretRevisionWritten(
         "rde-signing-private",
-        KmsTestHelper.getPrivateKeyring().getEncoded(),
+        "rde-signing-private/foo",
-        "projects/foo/locations/global/keyRings/bar",
+        getCiphertext(KmsTestHelper.getPrivateKeyring().getEncoded()));
-        "projects/foo/locations/global/keyRings/bar/cryptoKeys/rde-signing-private",
+    verifySecretAndSecretRevisionWritten(
-        "projects/foo/locations/"
-            + "global/keyRings/bar/cryptoKeys/rde-signing-private/cryptoKeyVersions",
         "rde-signing-public",
-        KmsTestHelper.getPublicKeyring().getPublicKey().getEncoded(),
+        "rde-signing-public/foo",
-        "projects/foo/locations/global/keyRings/bar/cryptoKeys/rde-signing-public",
+        getCiphertext(KmsTestHelper.getPrivateKeyring().getPublicKey().getEncoded()));
-        "projects/foo/locations/"
-            + "global/keyRings/bar/cryptoKeys/rde-signing-public/cryptoKeyVersions");
   }
 
   @Test
   public void test_setRdeSshClientPrivateKey() throws Exception {
-    updater.setRdeSshClientPrivateKey(KmsTestHelper.DUMMY_KEY);
+    updater.setRdeSshClientPrivateKey("value1").update();
-    updater.update();
 
-    verifyKmsApiCallsAndDatastoreWrites(
+    verifySecretAndSecretRevisionWritten(
-        "rde-ssh-client-private",
+        "rde-ssh-client-private", "rde-ssh-client-private/foo", getCiphertext("value1"));
-        KmsTestHelper.DUMMY_KEY.getBytes(UTF_8),
-        "projects/foo/locations/global/keyRings/bar",
-        "projects/foo/locations/global/keyRings/bar/cryptoKeys/rde-ssh-client-private",
-        "projects/foo/locations/"
-            + "global/keyRings/bar/cryptoKeys/rde-ssh-client-private/cryptoKeyVersions");
   }
 
   @Test
   public void test_setRdeSshClientPublicKey() throws Exception {
-    updater.setRdeSshClientPublicKey(KmsTestHelper.DUMMY_KEY);
+    updater.setRdeSshClientPublicKey("value1").update();
-    updater.update();
 
-    verifyKmsApiCallsAndDatastoreWrites(
+    verifySecretAndSecretRevisionWritten(
-        "rde-ssh-client-public",
+        "rde-ssh-client-public", "rde-ssh-client-public/foo", getCiphertext("value1"));
-        KmsTestHelper.DUMMY_KEY.getBytes(UTF_8),
-        "projects/foo/locations/global/keyRings/bar",
-        "projects/foo/locations/global/keyRings/bar/cryptoKeys/rde-ssh-client-public",
-        "projects/foo/locations/"
-            + "global/keyRings/bar/cryptoKeys/rde-ssh-client-public/cryptoKeyVersions");
   }
 
   @Test
   public void test_setRdeStagingKey() throws Exception {
-    updater.setRdeStagingKey(KmsTestHelper.getPrivateKeyring());
+    updater.setRdeStagingKey(KmsTestHelper.getPrivateKeyring()).update();
-    updater.update();
 
-    verifyKmsApiCallsAndDatastoreWrites(
+    verifySecretAndSecretRevisionWritten(
         "rde-staging-private",
-        KmsTestHelper.getPrivateKeyring().getEncoded(),
+        "rde-staging-private/foo",
-        "projects/foo/locations/global/keyRings/bar",
+        getCiphertext(KmsTestHelper.getPrivateKeyring().getEncoded()));
-        "projects/foo/locations/global/keyRings/bar/cryptoKeys/rde-staging-private",
+    verifySecretAndSecretRevisionWritten(
-        "projects/foo/locations/"
-            + "global/keyRings/bar/cryptoKeys/rde-staging-private/cryptoKeyVersions",
         "rde-staging-public",
-        KmsTestHelper.getPublicKeyring().getPublicKey().getEncoded(),
+        "rde-staging-public/foo",
-        "projects/foo/locations/global/keyRings/bar/cryptoKeys/rde-staging-public",
+        getCiphertext(KmsTestHelper.getPrivateKeyring().getPublicKey().getEncoded()));
-        "projects/foo/locations/"
-            + "global/keyRings/bar/cryptoKeys/rde-staging-public/cryptoKeyVersions");
   }
 
-  private void verifyKmsApiCallsAndDatastoreWrites(
-      String secretName,
-      byte[] goldenValue,
-      String goldenCryptoKeyRingName,
-      String goldenCryptoKeyName,
-      String goldenCryptoKeyVersionName)
-      throws Exception {
-    verify(kmsKeyRings).get(keyRingName.capture());
-    assertThat(keyRingName.getValue()).isEqualTo(goldenCryptoKeyRingName);
 
-    verify(kmsCryptoKeys).get(cryptoKeyName.capture());
+  private static void verifySecretAndSecretRevisionWritten(
-    assertThat(cryptoKeyName.getValue()).isEqualTo(goldenCryptoKeyName);
+      String secretName, String expectedCryptoKeyVersionName, String expectedEncryptedValue) {
-
+    KmsSecret secret =
-    verify(kmsCryptoKeyVersions).create(cryptoKeyVersionName.capture(), cryptoKeyVersion.capture());
+        ofy().load().key(Key.create(getCrossTldKey(), KmsSecret.class, secretName)).now();
-    assertThat(cryptoKeyVersionName.getValue()).isEqualTo(goldenCryptoKeyVersionName);
+    assertThat(secret).isNotNull();
-
-    verify(kmsCryptoKeys).encrypt(cryptoKeyName.capture(), encryptRequest.capture());
-    assertThat(cryptoKeyName.getValue()).isEqualTo(KmsTestHelper.DUMMY_CRYPTO_KEY_VERSION);
-    assertThat(encryptRequest.getValue())
-        .isEqualTo(new EncryptRequest().encodePlaintext(goldenValue));
-
-    KmsSecret secret = loadSecret(secretName);
     KmsSecretRevision secretRevision = ofy().load().key(secret.getLatestRevision()).now();
-    assertThat(secretRevision.getKmsCryptoKeyVersionName())
+    assertThat(secretRevision.getKmsCryptoKeyVersionName()).isEqualTo(expectedCryptoKeyVersionName);
-        .isEqualTo(KmsTestHelper.DUMMY_CRYPTO_KEY_VERSION);
+    assertThat(secretRevision.getEncryptedValue()).isEqualTo(expectedEncryptedValue);
-    assertThat(secretRevision.getEncryptedValue()).isEqualTo(KmsTestHelper.DUMMY_ENCRYPTED_VALUE);
   }
 
-  /** Variant of {@code verifyKmsApiCallsAndDatastoreWrites} for key pairs. */
+  private static String getCiphertext(byte[] plaintext) throws IOException {
-  private void verifyKmsApiCallsAndDatastoreWrites(
+    return new FakeKmsConnection().encrypt("blah", plaintext).ciphertext();
-      String firstSecretName,
-      byte[] firstGoldenValue,
-      String goldenCryptoKeyRingName,
-      String firstGoldenCryptoKeyName,
-      String firstGoldenCryptoKeyVersionName,
-      String secondSecretName,
-      byte[] secondGoldenValue,
-      String secondGoldenCryptoKeyName,
-      String secondGoldenCryptoKeyVersionName)
-      throws Exception {
-    verify(kmsKeyRings, times(1)).get(keyRingName.capture());
-    assertThat(keyRingName.getValue()).isEqualTo(goldenCryptoKeyRingName);
-
-    verify(kmsCryptoKeys, times(2)).get(cryptoKeyName.capture());
-    assertThat(cryptoKeyName.getAllValues())
-        .isEqualTo(ImmutableList.of(firstGoldenCryptoKeyName, secondGoldenCryptoKeyName));
-
-    verify(kmsCryptoKeyVersions, times(2))
-        .create(cryptoKeyVersionName.capture(), cryptoKeyVersion.capture());
-    assertThat(cryptoKeyVersionName.getAllValues())
-        .isEqualTo(
-            ImmutableList.of(firstGoldenCryptoKeyVersionName, secondGoldenCryptoKeyVersionName));
-
-    verify(kmsCryptoKeys, times(2)).encrypt(cryptoKeyName.capture(), encryptRequest.capture());
-    assertThat(cryptoKeyName.getValue()).isEqualTo(KmsTestHelper.DUMMY_CRYPTO_KEY_VERSION);
-    assertThat(encryptRequest.getAllValues())
-        .isEqualTo(
-            ImmutableList.of(
-                new EncryptRequest().encodePlaintext(firstGoldenValue),
-                new EncryptRequest().encodePlaintext(secondGoldenValue)));
-
-    KmsSecret secret = loadSecret(firstSecretName);
-    KmsSecretRevision secretRevision = ofy().load().key(secret.getLatestRevision()).now();
-    assertThat(secretRevision.getKmsCryptoKeyVersionName())
-        .isEqualTo(KmsTestHelper.DUMMY_CRYPTO_KEY_VERSION);
-    assertThat(secretRevision.getEncryptedValue()).isEqualTo(KmsTestHelper.DUMMY_ENCRYPTED_VALUE);
-
-    KmsSecret secondSecret = loadSecret(secondSecretName);
-    KmsSecretRevision secondSecretRevision =
-        ofy().load().key(secondSecret.getLatestRevision()).now();
-    assertThat(secondSecretRevision.getKmsCryptoKeyVersionName())
-        .isEqualTo(KmsTestHelper.DUMMY_CRYPTO_KEY_VERSION);
-    assertThat(secondSecretRevision.getEncryptedValue())
-        .isEqualTo(KmsTestHelper.DUMMY_ENCRYPTED_VALUE);
   }
 
-  private static GoogleJsonResponseException createNotFoundException() throws Exception {
+  private static String getCiphertext(String plaintext) throws IOException {
-    ByteArrayInputStream inputStream = new ByteArrayInputStream("".getBytes(UTF_8));
+    return getCiphertext(plaintext.getBytes(UTF_8));
-    HttpResponse response = GoogleJsonResponseExceptionHelper.createHttpResponse(404, inputStream);
-    HttpResponseException.Builder httpResponseExceptionBuilder =
-        new HttpResponseException.Builder(response);
-    httpResponseExceptionBuilder.setStatusCode(404);
-    httpResponseExceptionBuilder.setStatusMessage("NOT_FOUND");
-    return new GoogleJsonResponseException(httpResponseExceptionBuilder, null);
-  }
-
-  private static KmsSecret loadSecret(String secret) {
-    return ofy().load().key(Key.create(getCrossTldKey(), KmsSecret.class, secret)).now();
   }
 }