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 String kmsKeyRingName;
-  private final String projectId;
+  private final KmsConnection kmsConnection;
 
   @Inject
-  KmsKeyring(
-      @Config("cloudKmsProjectId") String projectId,
-      @Config("cloudKmsKeyRing") String kmsKeyringName,
-      CloudKMS kms) {
-    this.projectId = projectId;
-    this.kmsKeyRingName = kmsKeyringName;
-    this.kms = kms;
+  KmsKeyring(KmsConnection kmsConnection) {
+    this.kmsConnection = kmsConnection;
   }
 
   @Override
@@ -212,32 +196,10 @@ public class KmsKeyring implements Keyring {
     String encryptedData = ofy().load().key(secret.getLatestRevision()).now().getEncryptedValue();
 
     try {
-      return kms.projects()
-          .locations()
-          .keyRings()
-          .cryptoKeys()
-          .decrypt(
-              getCryptoKeyName(projectId, kmsKeyRingName, secret.getName()),
-              new DecryptRequest().setCiphertext(encryptedData))
-          .execute()
-          .decodePlaintext();
+      return kmsConnection.decrypt(secret.getName(), encryptedData);
     } 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 String projectId;
-  private final String kmsKeyRingName;
-  private final CloudKMS kms;
-
+  private final KmsConnection kmsConnection;
   private final HashMap<String, byte[]> secretValues;
 
   @Inject
-  public KmsUpdater(
-      @Config("cloudKmsProjectId") String projectId,
-      @Config("cloudKmsKeyRing") String kmsKeyRingName,
-      CloudKMS kms) {
-    this.projectId = projectId;
-    this.kmsKeyRingName = kmsKeyRingName;
-    this.kms = kms;
+  public KmsUpdater(KmsConnection kmsConnection) {
+    this.kmsConnection = kmsConnection;
 
     // 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 fullKeyName = getCryptoKeyName(projectId, kmsKeyRingName, keyName);
-
-      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());
+      String secretName = entry.getKey();
+      encryptedValues.put(secretName, kmsConnection.encrypt(secretName, entry.getValue()));
     }
     return encryptedValues.build();
   }
@@ -262,8 +187,8 @@ public final class KmsUpdater {
 
                   KmsSecretRevision secretRevision =
                       new KmsSecretRevision.Builder()
-                          .setEncryptedValue(revisionData.getCiphertext())
-                          .setKmsCryptoKeyVersionName(revisionData.getName())
+                          .setEncryptedValue(revisionData.ciphertext())
+                          .setKmsCryptoKeyVersionName(revisionData.cryptoKeyVersionName())
                           .setParent(secretName)
                           .build();
                   ofy()