Convert Kms* classes to use SQL when appropriate (#1043)

* Convert Kms* classes to use SQL when appropriate

core/src/main/java/google/registry/keyring/kms/KmsKeyring.java

@@ -19,6 +19,7 @@ import static com.google.common.base.CaseFormat.UPPER_UNDERSCORE;
 import static com.google.common.base.Preconditions.checkState;
 import static google.registry.model.common.EntityGroupRoot.getCrossTldKey;
 import static google.registry.model.ofy.ObjectifyService.ofy;
+import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 
 import com.googlecode.objectify.Key;
 import google.registry.config.RegistryConfig.Config;
@@ -26,7 +27,10 @@ import google.registry.keyring.api.KeySerializer;
 import google.registry.keyring.api.Keyring;
 import google.registry.keyring.api.KeyringException;
 import google.registry.model.server.KmsSecret;
+import google.registry.model.server.KmsSecretRevision;
+import google.registry.model.server.KmsSecretRevisionSqlDao;
 import java.io.IOException;
+import java.util.Optional;
 import javax.inject.Inject;
 import org.bouncycastle.openpgp.PGPException;
 import org.bouncycastle.openpgp.PGPKeyPair;
@@ -201,13 +205,21 @@ public class KmsKeyring implements Keyring {
   }
 
   private byte[] getDecryptedData(String keyName) {
-    KmsSecret secret =
+    String encryptedData;
-        ofy().load().key(Key.create(getCrossTldKey(), KmsSecret.class, keyName)).now();
+    if (tm().isOfy()) {
-    checkState(secret != null, "Requested secret '%s' does not exist.", keyName);
+      KmsSecret secret =
-    String encryptedData = ofy().load().key(secret.getLatestRevision()).now().getEncryptedValue();
+          ofy().load().key(Key.create(getCrossTldKey(), KmsSecret.class, keyName)).now();
+      checkState(secret != null, "Requested secret '%s' does not exist.", keyName);
+      encryptedData = ofy().load().key(secret.getLatestRevision()).now().getEncryptedValue();
+    } else {
+      Optional<KmsSecretRevision> revision =
+          tm().transact(() -> KmsSecretRevisionSqlDao.getLatestRevision(keyName));
+      checkState(revision.isPresent(), "Requested secret '%s' does not exist.", keyName);
+      encryptedData = revision.get().getEncryptedValue();
+    }
 
     try {
-      return kmsConnection.decrypt(secret.getName(), encryptedData);
+      return kmsConnection.decrypt(keyName, encryptedData);
     } catch (Exception e) {
       throw new KeyringException(
           String.format("CloudKMS decrypt operation failed for secret %s", keyName), e);

core/src/main/java/google/registry/keyring/kms/KmsUpdater.java

@@ -34,7 +34,6 @@ import static google.registry.keyring.kms.KmsKeyring.StringKeyLabel.RDE_SSH_CLIE
 import static google.registry.keyring.kms.KmsKeyring.StringKeyLabel.RDE_SSH_CLIENT_PUBLIC_STRING;
 import static google.registry.keyring.kms.KmsKeyring.StringKeyLabel.SAFE_BROWSING_API_KEY;
 import static google.registry.keyring.kms.KmsKeyring.StringKeyLabel.TOOLS_CLOUD_SQL_PASSWORD_STRING;
-import static google.registry.model.ofy.ObjectifyService.ofy;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.util.PreconditionsUtils.checkArgumentNotNull;
 
@@ -194,8 +193,7 @@ public final class KmsUpdater {
    */
   private static void persistEncryptedValues(
       final ImmutableMap<String, EncryptResponse> encryptedValues) {
-    tm()
+    tm().transact(
-        .transact(
             () -> {
               for (Map.Entry<String, EncryptResponse> entry : encryptedValues.entrySet()) {
                 String secretName = entry.getKey();
@@ -207,7 +205,7 @@ public final class KmsUpdater {
                         .setKmsCryptoKeyVersionName(revisionData.cryptoKeyVersionName())
                         .setParent(secretName)
                         .build();
-                ofy().save().entities(secretRevision, KmsSecret.create(secretName, secretRevision));
+                tm().putAll(secretRevision, KmsSecret.create(secretName, secretRevision));
               }
             });
   }

core/src/main/java/google/registry/persistence/transaction/JpaTransactionManagerImpl.java

@@ -36,6 +36,7 @@ import google.registry.model.index.ForeignKeyIndex.ForeignKeyContactIndex;
 import google.registry.model.index.ForeignKeyIndex.ForeignKeyDomainIndex;
 import google.registry.model.index.ForeignKeyIndex.ForeignKeyHostIndex;
 import google.registry.model.ofy.DatastoreTransactionManager;
+import google.registry.model.server.KmsSecret;
 import google.registry.persistence.JpaRetries;
 import google.registry.persistence.VKey;
 import google.registry.util.Clock;
@@ -73,7 +74,8 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
           EppResourceIndex.class,
           ForeignKeyContactIndex.class,
           ForeignKeyDomainIndex.class,
-          ForeignKeyHostIndex.class);
+          ForeignKeyHostIndex.class,
+          KmsSecret.class);
 
   // EntityManagerFactory is thread safe.
   private final EntityManagerFactory emf;

core/src/test/java/google/registry/keyring/kms/KmsKeyringTest.java

@@ -15,22 +15,23 @@
 package google.registry.keyring.kms;
 
 import static com.google.common.truth.Truth.assertThat;
-import static google.registry.testing.DatabaseHelper.persistResources;
+import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 
-import com.google.common.collect.ImmutableList;
 import google.registry.keyring.api.KeySerializer;
 import google.registry.model.server.KmsSecret;
 import google.registry.model.server.KmsSecretRevision;
 import google.registry.testing.AppEngineExtension;
 import google.registry.testing.BouncyCastleProviderExtension;
+import google.registry.testing.DualDatabaseTest;
+import google.registry.testing.TestOfyAndSql;
 import org.bouncycastle.openpgp.PGPKeyPair;
 import org.bouncycastle.openpgp.PGPPrivateKey;
 import org.bouncycastle.openpgp.PGPPublicKey;
 import org.junit.jupiter.api.BeforeEach;
-import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.RegisterExtension;
 
 /** Unit tests for {@link KmsKeyring}. */
+@DualDatabaseTest
 class KmsKeyringTest {
 
   @RegisterExtension
@@ -47,21 +48,21 @@ class KmsKeyringTest {
     keyring = new KmsKeyring(new FakeKmsConnection());
   }
 
-  @Test
+  @TestOfyAndSql
   void test_getCloudSqlPassword() {
     saveCleartextSecret("cloud-sql-password-string");
     String cloudSqlPassword = keyring.getCloudSqlPassword();
     assertThat(cloudSqlPassword).isEqualTo("cloud-sql-password-stringmoo");
   }
 
-  @Test
+  @TestOfyAndSql
   void test_getToolsCloudSqlPassword() {
     saveCleartextSecret("tools-cloud-sql-password-string");
     String toolsCloudSqlPassword = keyring.getToolsCloudSqlPassword();
     assertThat(toolsCloudSqlPassword).isEqualTo("tools-cloud-sql-password-stringmoo");
   }
 
-  @Test
+  @TestOfyAndSql
   void test_getRdeSigningKey() throws Exception {
     saveKeyPairSecret("rde-signing-public", "rde-signing-private");
     PGPKeyPair rdeSigningKey = keyring.getRdeSigningKey();
@@ -69,7 +70,7 @@ class KmsKeyringTest {
         .isEqualTo(KeySerializer.serializeKeyPair(KmsTestHelper.getKeyPair()));
   }
 
-  @Test
+  @TestOfyAndSql
   void test_getRdeStagingEncryptionKey() throws Exception {
     savePublicKeySecret("rde-staging-public");
     PGPPublicKey rdeStagingEncryptionKey = keyring.getRdeStagingEncryptionKey();
@@ -77,7 +78,7 @@ class KmsKeyringTest {
         .isEqualTo(KmsTestHelper.getPublicKey().getFingerprint());
   }
 
-  @Test
+  @TestOfyAndSql
   void test_getRdeStagingDecryptionKey() throws Exception {
     savePrivateKeySecret("rde-staging-private");
     savePublicKeySecret("rde-staging-public");
@@ -90,7 +91,7 @@ class KmsKeyringTest {
         .isEqualTo(KeySerializer.serializeKeyPair(KmsTestHelper.getKeyPair()));
   }
 
-  @Test
+  @TestOfyAndSql
   void test_getRdeReceiverKey() throws Exception {
     savePublicKeySecret("rde-receiver-public");
     PGPPublicKey rdeReceiverKey = keyring.getRdeReceiverKey();
@@ -98,7 +99,7 @@ class KmsKeyringTest {
         .isEqualTo(KmsTestHelper.getPublicKey().getFingerprint());
   }
 
-  @Test
+  @TestOfyAndSql
   void test_getBrdaSigningKey() throws Exception {
     saveKeyPairSecret("brda-signing-public", "brda-signing-private");
     PGPKeyPair brdaSigningKey = keyring.getBrdaSigningKey();
@@ -106,7 +107,7 @@ class KmsKeyringTest {
         .isEqualTo(KeySerializer.serializeKeyPair(KmsTestHelper.getKeyPair()));
   }
 
-  @Test
+  @TestOfyAndSql
   void test_getBrdaReceiverKey() throws Exception {
     savePublicKeySecret("brda-receiver-public");
     PGPPublicKey brdaReceiverKey = keyring.getBrdaReceiverKey();
@@ -114,49 +115,49 @@ class KmsKeyringTest {
         .isEqualTo(KmsTestHelper.getPublicKey().getFingerprint());
   }
 
-  @Test
+  @TestOfyAndSql
   void test_getRdeSshClientPublicKey() {
     saveCleartextSecret("rde-ssh-client-public-string");
     String rdeSshClientPublicKey = keyring.getRdeSshClientPublicKey();
     assertThat(rdeSshClientPublicKey).isEqualTo("rde-ssh-client-public-stringmoo");
   }
 
-  @Test
+  @TestOfyAndSql
   void test_getRdeSshClientPrivateKey() {
     saveCleartextSecret("rde-ssh-client-private-string");
     String rdeSshClientPrivateKey = keyring.getRdeSshClientPrivateKey();
     assertThat(rdeSshClientPrivateKey).isEqualTo("rde-ssh-client-private-stringmoo");
   }
 
-  @Test
+  @TestOfyAndSql
   void test_getIcannReportingPassword() {
     saveCleartextSecret("icann-reporting-password-string");
     String icannReportingPassword = keyring.getIcannReportingPassword();
     assertThat(icannReportingPassword).isEqualTo("icann-reporting-password-stringmoo");
   }
 
-  @Test
+  @TestOfyAndSql
   void test_getMarksdbDnlLoginAndPassword() {
     saveCleartextSecret("marksdb-dnl-login-string");
     String marksdbDnlLoginAndPassword = keyring.getMarksdbDnlLoginAndPassword();
     assertThat(marksdbDnlLoginAndPassword).isEqualTo("marksdb-dnl-login-stringmoo");
   }
 
-  @Test
+  @TestOfyAndSql
   void test_getMarksdbLordnPassword() {
     saveCleartextSecret("marksdb-lordn-password-string");
     String marksdbLordnPassword = keyring.getMarksdbLordnPassword();
     assertThat(marksdbLordnPassword).isEqualTo("marksdb-lordn-password-stringmoo");
   }
 
-  @Test
+  @TestOfyAndSql
   void test_getMarksdbSmdrlLoginAndPassword() {
     saveCleartextSecret("marksdb-smdrl-login-string");
     String marksdbSmdrlLoginAndPassword = keyring.getMarksdbSmdrlLoginAndPassword();
     assertThat(marksdbSmdrlLoginAndPassword).isEqualTo("marksdb-smdrl-login-stringmoo");
   }
 
-  @Test
+  @TestOfyAndSql
   void test_getJsonCredential() {
     saveCleartextSecret("json-credential-string");
     String jsonCredential = keyring.getJsonCredential();
@@ -173,7 +174,7 @@ class KmsKeyringTest {
             .setParent(secretName)
             .build();
     KmsSecret secret = KmsSecret.create(secretName, secretRevision);
-    persistResources(ImmutableList.of(secretRevision, secret));
+    tm().transact(() -> tm().putAll(secretRevision, secret));
   }
 
   private static void saveCleartextSecret(String secretName) {

core/src/test/java/google/registry/keyring/kms/KmsUpdaterTest.java

@@ -17,21 +17,25 @@ package google.registry.keyring.kms;
 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 google.registry.persistence.transaction.TransactionManagerFactory.tm;
 
 import com.googlecode.objectify.Key;
 import google.registry.keyring.api.KeySerializer;
 import google.registry.model.server.KmsSecret;
 import google.registry.model.server.KmsSecretRevision;
+import google.registry.model.server.KmsSecretRevisionSqlDao;
 import google.registry.testing.AppEngineExtension;
 import google.registry.testing.BouncyCastleProviderExtension;
+import google.registry.testing.DualDatabaseTest;
+import google.registry.testing.TestOfyAndSql;
 import java.io.IOException;
 import org.bouncycastle.openpgp.PGPKeyPair;
 import org.bouncycastle.openpgp.PGPPublicKey;
 import org.junit.jupiter.api.BeforeEach;
-import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.RegisterExtension;
 
 /** Unit tests for {@link KmsUpdater} */
+@DualDatabaseTest
 public class KmsUpdaterTest {
 
   @RegisterExtension
@@ -48,7 +52,7 @@ public class KmsUpdaterTest {
     updater = new KmsUpdater(new FakeKmsConnection());
   }
 
-  @Test
+  @TestOfyAndSql
   void test_setMultipleSecrets() {
     updater
         .setMarksdbDnlLoginAndPassword("value1")
@@ -68,7 +72,7 @@ public class KmsUpdaterTest {
         "json-credential-string", "json-credential-string/foo", getCiphertext("value3"));
   }
 
-  @Test
+  @TestOfyAndSql
   void test_setBrdaReceiverKey() throws Exception {
     updater.setBrdaReceiverPublicKey(KmsTestHelper.getPublicKey()).update();
 
@@ -78,7 +82,7 @@ public class KmsUpdaterTest {
         getCiphertext(KmsTestHelper.getPublicKey()));
   }
 
-  @Test
+  @TestOfyAndSql
   void test_setBrdaSigningKey() throws Exception {
     updater.setBrdaSigningKey(KmsTestHelper.getKeyPair()).update();
 
@@ -92,7 +96,7 @@ public class KmsUpdaterTest {
         getCiphertext(KmsTestHelper.getPublicKey()));
   }
 
-  @Test
+  @TestOfyAndSql
   void test_setCloudSqlPassword() {
     updater.setCloudSqlPassword("value1").update();
 
@@ -100,7 +104,7 @@ public class KmsUpdaterTest {
         "cloud-sql-password-string", "cloud-sql-password-string/foo", getCiphertext("value1"));
   }
 
-  @Test
+  @TestOfyAndSql
   void test_setToolsCloudSqlPassword() {
     updater.setToolsCloudSqlPassword("value1").update();
 
@@ -110,7 +114,7 @@ public class KmsUpdaterTest {
         getCiphertext("value1"));
   }
 
-  @Test
+  @TestOfyAndSql
   void test_setIcannReportingPassword() {
     updater.setIcannReportingPassword("value1").update();
 
@@ -120,7 +124,7 @@ public class KmsUpdaterTest {
         getCiphertext("value1"));
   }
 
-  @Test
+  @TestOfyAndSql
   void test_setJsonCredential() {
     updater.setJsonCredential("value1").update();
 
@@ -128,7 +132,7 @@ public class KmsUpdaterTest {
         "json-credential-string", "json-credential-string/foo", getCiphertext("value1"));
   }
 
-  @Test
+  @TestOfyAndSql
   void test_setMarksdbDnlLoginAndPassword() {
     updater.setMarksdbDnlLoginAndPassword("value1").update();
 
@@ -136,7 +140,7 @@ public class KmsUpdaterTest {
         "marksdb-dnl-login-string", "marksdb-dnl-login-string/foo", getCiphertext("value1"));
   }
 
-  @Test
+  @TestOfyAndSql
   void test_setMarksdbLordnPassword() {
     updater.setMarksdbLordnPassword("value1").update();
 
@@ -146,7 +150,7 @@ public class KmsUpdaterTest {
         getCiphertext("value1"));
   }
 
-  @Test
+  @TestOfyAndSql
   void test_setMarksdbSmdrlLoginAndPassword() {
     updater.setMarksdbSmdrlLoginAndPassword("value1").update();
 
@@ -154,7 +158,7 @@ public class KmsUpdaterTest {
         "marksdb-smdrl-login-string", "marksdb-smdrl-login-string/foo", getCiphertext("value1"));
   }
 
-  @Test
+  @TestOfyAndSql
   void test_setRdeReceiverKey() throws Exception {
     updater.setRdeReceiverPublicKey(KmsTestHelper.getPublicKey()).update();
 
@@ -165,7 +169,7 @@ public class KmsUpdaterTest {
             KeySerializer.serializePublicKey(KmsTestHelper.getPublicKey())));
   }
 
-  @Test
+  @TestOfyAndSql
   void test_setRdeSigningKey() throws Exception {
     updater.setRdeSigningKey(KmsTestHelper.getKeyPair()).update();
 
@@ -179,7 +183,7 @@ public class KmsUpdaterTest {
         getCiphertext(KmsTestHelper.getPublicKey()));
   }
 
-  @Test
+  @TestOfyAndSql
   void test_setRdeSshClientPrivateKey() {
     updater.setRdeSshClientPrivateKey("value1").update();
 
@@ -189,7 +193,7 @@ public class KmsUpdaterTest {
         getCiphertext("value1"));
   }
 
-  @Test
+  @TestOfyAndSql
   void test_setRdeSshClientPublicKey() {
     updater.setRdeSshClientPublicKey("value1").update();
 
@@ -199,7 +203,7 @@ public class KmsUpdaterTest {
         getCiphertext("value1"));
   }
 
-  @Test
+  @TestOfyAndSql
   void test_setRdeStagingKey() throws Exception {
     updater.setRdeStagingKey(KmsTestHelper.getKeyPair()).update();
 
@@ -213,13 +217,18 @@ public class KmsUpdaterTest {
         getCiphertext(KmsTestHelper.getPublicKey()));
   }
 
-
   private static void verifySecretAndSecretRevisionWritten(
       String secretName, String expectedCryptoKeyVersionName, String expectedEncryptedValue) {
-    KmsSecret secret =
+    KmsSecretRevision secretRevision;
-        ofy().load().key(Key.create(getCrossTldKey(), KmsSecret.class, secretName)).now();
+    if (tm().isOfy()) {
-    assertThat(secret).isNotNull();
+      KmsSecret secret =
-    KmsSecretRevision secretRevision = ofy().load().key(secret.getLatestRevision()).now();
+          ofy().load().key(Key.create(getCrossTldKey(), KmsSecret.class, secretName)).now();
+      assertThat(secret).isNotNull();
+      secretRevision = ofy().load().key(secret.getLatestRevision()).now();
+    } else {
+      secretRevision =
+          tm().transact(() -> KmsSecretRevisionSqlDao.getLatestRevision(secretName).get());
+    }
     assertThat(secretRevision.getKmsCryptoKeyVersionName()).isEqualTo(expectedCryptoKeyVersionName);
     assertThat(secretRevision.getEncryptedValue()).isEqualTo(expectedEncryptedValue);
   }