Fix bugs in KmsConnectionImpl

A few errors emerged when doing an integration test against the actual API. I've updated the unit tests to reflect the correct behavior.

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

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

@@ -24,6 +24,7 @@ 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 com.google.api.services.cloudkms.v1beta1.model.UpdateCryptoKeyPrimaryVersionRequest;
 import google.registry.config.RegistryConfig.Config;
 import google.registry.keyring.api.KeyringException;
 import java.io.IOException;
@@ -32,11 +33,10 @@ import javax.inject.Inject;
 /** The {@link KmsConnection} which talks to Cloud KMS. */
 class KmsConnectionImpl implements KmsConnection {
 
+  private static final String KMS_LOCATION_FORMAT = "projects/%s/locations/global";
   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;
@@ -68,7 +68,8 @@ class KmsConnectionImpl implements KmsConnection {
         kms.projects()
             .locations()
             .keyRings()
-            .create("global", new KeyRing().setName(fullKeyRingName))
+            .create(getLocationName(projectId), new KeyRing())
+            .setKeyRingId(kmsKeyRingName)
             .execute();
       } else {
         throw jsonException;
@@ -77,39 +78,53 @@ class KmsConnectionImpl implements KmsConnection {
 
     String fullKeyName = getCryptoKeyName(projectId, kmsKeyRingName, cryptoKeyName);
 
+    boolean newCryptoKey = false;
     try {
       kms.projects().locations().keyRings().cryptoKeys().get(fullKeyName).execute();
     } catch (GoogleJsonResponseException jsonException) {
       if (jsonException.getStatusCode() == HttpStatusCodes.STATUS_CODE_NOT_FOUND) {
+        newCryptoKey = true;
         kms.projects()
             .locations()
             .keyRings()
             .cryptoKeys()
-            .create(
+            .create(fullKeyRingName, new CryptoKey().setPurpose("ENCRYPT_DECRYPT"))
-                fullKeyName, new CryptoKey().setName(cryptoKeyName).setPurpose("ENCRYPT_DECRYPT"))
+            .setCryptoKeyId(cryptoKeyName)
             .execute();
       } else {
         throw jsonException;
       }
     }
 
-    CryptoKeyVersion cryptoKeyVersion =
+    // New CryptoKeys start with a CryptoKeyVersion, so we only create a new CryptoKeyVersion and
-        kms.projects()
+    // rotate to it if we're dealing with an existing CryptoKey.
-            .locations()
+    if (!newCryptoKey) {
-            .keyRings()
+      CryptoKeyVersion cryptoKeyVersion =
-            .cryptoKeys()
+          kms.projects()
-            .cryptoKeyVersions()
+              .locations()
-            .create(
+              .keyRings()
-                getCryptoKeyVersionName(projectId, kmsKeyRingName, cryptoKeyName),
+              .cryptoKeys()
-                new CryptoKeyVersion())
+              .cryptoKeyVersions()
-            .execute();
+              .create(fullKeyName, new CryptoKeyVersion())
+              .execute();
+
+      kms.projects()
+          .locations()
+          .keyRings()
+          .cryptoKeys()
+          .updatePrimaryVersion(
+              fullKeyName,
+              new UpdateCryptoKeyPrimaryVersionRequest()
+                  .setCryptoKeyVersionId(getCryptoKeyVersionId(cryptoKeyVersion)))
+          .execute();
+    }
 
     return EncryptResponse.create(
         kms.projects()
             .locations()
             .keyRings()
             .cryptoKeys()
-            .encrypt(cryptoKeyVersion.getName(), new EncryptRequest().encodePlaintext(value))
+            .encrypt(fullKeyName, new EncryptRequest().encodePlaintext(value))
             .execute());
   }
 
@@ -131,17 +146,21 @@ class KmsConnectionImpl implements KmsConnection {
     }
   }
 
-  static String getKeyRingName(String projectId, String kmsKeyRingName) {
+  private static String getLocationName(String projectId) {
+    return String.format(KMS_LOCATION_FORMAT, projectId);
+  }
+
+  private 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) {
+  private static String getCryptoKeyName(
+      String projectId, String kmsKeyRingName, String cryptoKeyName) {
     return String.format(KMS_CRYPTO_KEY_NAME_FORMAT, projectId, kmsKeyRingName, cryptoKeyName);
   }
 
-  static String getCryptoKeyVersionName(
+  private static String getCryptoKeyVersionId(CryptoKeyVersion cryptoKeyVersion) {
-      String projectId, String kmsKeyRingName, String cryptoKeyName) {
+    String cryptoKeyVersionName = cryptoKeyVersion.getName();
-    return String.format(
+    return cryptoKeyVersionName.substring(cryptoKeyVersionName.lastIndexOf('/') + 1);
-        KMS_CRYPTO_KEY_VERSION_NAME_FORMAT, projectId, kmsKeyRingName, cryptoKeyName);
   }
 }

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

@@ -18,6 +18,7 @@ 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.never;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
@@ -32,6 +33,7 @@ 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 com.google.api.services.cloudkms.v1beta1.model.UpdateCryptoKeyPrimaryVersionRequest;
 import java.io.ByteArrayInputStream;
 import org.junit.Before;
 import org.junit.Test;
@@ -54,6 +56,9 @@ public class KmsConnectionImplTest {
   @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.UpdatePrimaryVersion updatePrimaryVersion;
+
   @Mock
   private CloudKMS.Projects.Locations.KeyRings.CryptoKeys.CryptoKeyVersions kmsCryptoKeyVersions;
 
@@ -67,12 +72,15 @@ public class KmsConnectionImplTest {
   @Captor private ArgumentCaptor<KeyRing> keyRing;
   @Captor private ArgumentCaptor<CryptoKey> cryptoKey;
   @Captor private ArgumentCaptor<CryptoKeyVersion> cryptoKeyVersion;
+  @Captor private ArgumentCaptor<String> locationName;
   @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;
 
+  @Captor
+  private ArgumentCaptor<UpdateCryptoKeyPrimaryVersionRequest> updateCryptoKeyPrimaryVersionRequest;
+
   @Before
   public void setUp() throws Exception {
     when(kms.projects()).thenReturn(kmsProjects);
@@ -80,9 +88,11 @@ public class KmsConnectionImplTest {
     when(kmsLocations.keyRings()).thenReturn(kmsKeyRings);
     when(kmsKeyRings.get(anyString())).thenReturn(kmsKeyRingsGet);
     when(kmsKeyRings.create(anyString(), any(KeyRing.class))).thenReturn(kmsKeyRingsCreate);
+    when(kmsKeyRingsCreate.setKeyRingId(anyString())).thenReturn(kmsKeyRingsCreate);
     when(kmsKeyRings.cryptoKeys()).thenReturn(kmsCryptoKeys);
     when(kmsCryptoKeys.get(anyString())).thenReturn(kmsCryptoKeysGet);
     when(kmsCryptoKeys.create(anyString(), any(CryptoKey.class))).thenReturn(kmsCryptoKeysCreate);
+    when(kmsCryptoKeysCreate.setCryptoKeyId(anyString())).thenReturn(kmsCryptoKeysCreate);
     when(kmsCryptoKeys.cryptoKeyVersions()).thenReturn(kmsCryptoKeyVersions);
     when(kmsCryptoKeyVersions.create(anyString(), any(CryptoKeyVersion.class)))
         .thenReturn(kmsCryptoKeyVersionsCreate);
@@ -97,6 +107,9 @@ public class KmsConnectionImplTest {
                 .setCiphertext(KmsTestHelper.DUMMY_ENCRYPTED_VALUE));
     when(kmsCryptoKeys.decrypt(anyString(), any(DecryptRequest.class)))
         .thenReturn(kmsCryptoKeysDecrypt);
+    when(kmsCryptoKeys.updatePrimaryVersion(
+            anyString(), any(UpdateCryptoKeyPrimaryVersionRequest.class)))
+        .thenReturn(updatePrimaryVersion);
   }
 
   @Test
@@ -105,16 +118,17 @@ public class KmsConnectionImplTest {
 
     new KmsConnectionImpl("foo", "bar", kms).encrypt("key", "moo".getBytes(UTF_8));
 
-    verify(kmsKeyRings).create(keyRingName.capture(), keyRing.capture());
+    verify(kmsKeyRings).create(locationName.capture(), keyRing.capture());
-    assertThat(keyRingName.getValue()).isEqualTo("global");
+    assertThat(locationName.getValue()).isEqualTo("projects/foo/locations/global");
-    assertThat(keyRing.getValue())
+    assertThat(keyRing.getValue()).isEqualTo(new KeyRing());
-        .isEqualTo(new KeyRing().setName("projects/foo/locations/global/keyRings/bar"));
+    verify(kmsKeyRingsCreate).setKeyRingId(keyRingName.capture());
+    assertThat(keyRingName.getValue()).isEqualTo("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");
-        "projects/foo/locations/global/keyRings/bar/cryptoKeys/key/cryptoKeyVersions");
   }
 
   @Test
@@ -123,28 +137,44 @@ public class KmsConnectionImplTest {
 
     new KmsConnectionImpl("foo", "bar", kms).encrypt("key", "moo".getBytes(UTF_8));
 
-    verify(kmsCryptoKeys).create(cryptoKeyName.capture(), cryptoKey.capture());
+    verify(kmsCryptoKeys).create(keyRingName.capture(), cryptoKey.capture());
-    assertThat(cryptoKeyName.getValue())
+    assertThat(keyRingName.getValue()).isEqualTo("projects/foo/locations/global/keyRings/bar");
-        .isEqualTo("projects/foo/locations/global/keyRings/bar/cryptoKeys/key");
+    assertThat(cryptoKey.getValue()).isEqualTo(new CryptoKey().setPurpose("ENCRYPT_DECRYPT"));
-    assertThat(cryptoKey.getValue())
+    verify(kmsCryptoKeysCreate).setCryptoKeyId(cryptoKeyName.capture());
-        .isEqualTo(new CryptoKey().setName("key").setPurpose("ENCRYPT_DECRYPT"));
+    assertThat(cryptoKeyName.getValue()).isEqualTo("key");
     verify(kmsCryptoKeysCreate).execute();
+    verify(kmsCryptoKeyVersionsCreate, never()).execute();
+    verify(updatePrimaryVersion, never()).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");
-        "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));
 
+
+    verify(kmsCryptoKeyVersions).create(cryptoKeyName.capture(), cryptoKeyVersion.capture());
+    assertThat(cryptoKeyName.getValue())
+        .isEqualTo("projects/foo/locations/global/keyRings/bar/cryptoKeys/key");
+
+    verify(kmsCryptoKeys)
+        .updatePrimaryVersion(
+            cryptoKeyName.capture(), updateCryptoKeyPrimaryVersionRequest.capture());
+    assertThat(cryptoKeyName.getValue())
+        .isEqualTo("projects/foo/locations/global/keyRings/bar/cryptoKeys/key");
+    assertThat(updateCryptoKeyPrimaryVersionRequest.getValue())
+        .isEqualTo(
+            new UpdateCryptoKeyPrimaryVersionRequest()
+                .setCryptoKeyVersionId(KmsTestHelper.DUMMY_CRYPTO_KEY_VERSION));
+
     verifyEncryptKmsApiCalls(
         "moo",
         "projects/foo/locations/global/keyRings/bar",
-        "projects/foo/locations/global/keyRings/bar/cryptoKeys/key",
+        "projects/foo/locations/global/keyRings/bar/cryptoKeys/key");
-        "projects/foo/locations/global/keyRings/bar/cryptoKeys/key/cryptoKeyVersions");
   }
 
   @Test
@@ -162,10 +192,7 @@ public class KmsConnectionImplTest {
   }
 
   private void verifyEncryptKmsApiCalls(
-      String goldenValue,
+      String goldenValue, String goldenCryptoKeyRingName, String goldenCryptoKeyName)
-      String goldenCryptoKeyRingName,
-      String goldenCryptoKeyName,
-      String goldenCryptoKeyVersionName)
       throws Exception {
     verify(kmsKeyRings).get(keyRingName.capture());
     assertThat(keyRingName.getValue()).isEqualTo(goldenCryptoKeyRingName);
@@ -173,11 +200,8 @@ public class KmsConnectionImplTest {
     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(cryptoKeyName.getValue()).isEqualTo(goldenCryptoKeyName);
     assertThat(encryptRequest.getValue())
         .isEqualTo(new EncryptRequest().encodePlaintext(goldenValue.getBytes(UTF_8)));
   }