Make jpaTm for nomulus tool use local credential (#515)

* Make jpaTm for nomulus tool use local credential

* Remove unused methods in RegistryToolEnvironment

* Fix order of annotations

* Remove unused method in PersistenceComponent

* Move the creation of credential to the module

* Move creadential creation to AuthModule

* Add a TODO

core/build.gradle

@@ -197,6 +197,7 @@ dependencies {
   compile deps['com.google.appengine:appengine-remote-api']
   compile deps['com.google.auth:google-auth-library-credentials']
   compile deps['com.google.auth:google-auth-library-oauth2-http']
+  compile deps['com.google.cloud.sql:jdbc-socket-factory-core']
   compile deps['com.google.cloud.sql:postgres-socket-factory']
   compile deps['com.google.code.gson:gson']
   compile deps['com.google.auto.value:auto-value-annotations']

core/src/main/java/google/registry/persistence/PersistenceComponent.java

@@ -19,7 +19,6 @@ import google.registry.config.CredentialModule;
 import google.registry.config.RegistryConfig.ConfigModule;
 import google.registry.keyring.kms.KmsModule;
 import google.registry.persistence.PersistenceModule.AppEngineJpaTm;
-import google.registry.persistence.PersistenceModule.NomulusToolJpaTm;
 import google.registry.persistence.transaction.JpaTransactionManager;
 import google.registry.util.UtilsModule;
 import javax.inject.Singleton;
@@ -39,7 +38,4 @@ public interface PersistenceComponent {
 
   @AppEngineJpaTm
   JpaTransactionManager appEngineJpaTransactionManager();
-
-  @NomulusToolJpaTm
-  JpaTransactionManager nomulusToolJpaTransactionManager();
 }

core/src/main/java/google/registry/persistence/PersistenceModule.java

@@ -22,6 +22,7 @@ import static google.registry.config.RegistryConfig.getHibernateHikariMaximumPoo
 import static google.registry.config.RegistryConfig.getHibernateHikariMinimumIdle;
 import static google.registry.config.RegistryConfig.getHibernateLogSqlQueries;
 
+import com.google.api.client.auth.oauth2.Credential;
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.Maps;
@@ -29,8 +30,10 @@ import dagger.Module;
 import dagger.Provides;
 import google.registry.config.RegistryConfig.Config;
 import google.registry.keyring.kms.KmsKeyring;
+import google.registry.persistence.transaction.CloudSqlCredentialSupplier;
 import google.registry.persistence.transaction.JpaTransactionManager;
 import google.registry.persistence.transaction.JpaTransactionManagerImpl;
+import google.registry.tools.AuthModule.CloudSqlClientCredential;
 import google.registry.util.Clock;
 import java.lang.annotation.Documented;
 import java.util.HashMap;
@@ -118,7 +121,9 @@ public class PersistenceModule {
       @Config("toolsCloudSqlUsername") String username,
       KmsKeyring kmsKeyring,
       @PartialCloudSqlConfigs ImmutableMap<String, String> cloudSqlConfigs,
+      @CloudSqlClientCredential Credential credential,
       Clock clock) {
+    CloudSqlCredentialSupplier.setupCredentialSupplier(credential);
     HashMap<String, String> overrides = Maps.newHashMap(cloudSqlConfigs);
     overrides.put(Environment.USER, username);
     overrides.put(Environment.PASS, kmsKeyring.getToolsCloudSqlPassword());
@@ -158,7 +163,7 @@ public class PersistenceModule {
   /** Dagger qualifier for {@link JpaTransactionManager} used for Nomulus tool. */
   @Qualifier
   @Documented
-  @interface NomulusToolJpaTm {}
+  public @interface NomulusToolJpaTm {}
 
   /** Dagger qualifier for the partial Cloud SQL configs. */
   @Qualifier

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

@@ -0,0 +1,35 @@
+// Copyright 2020 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.persistence.transaction;
+
+import com.google.api.client.auth.oauth2.Credential;
+import com.google.cloud.sql.CredentialFactory;
+
+/** Supplier class to provide {@link Credential} for Cloud SQL library. */
+public class CloudSqlCredentialSupplier implements CredentialFactory {
+  private static Credential credential;
+
+  /** Initialize the supplier with given credential json and scopes. */
+  public static void setupCredentialSupplier(Credential credential) {
+    System.setProperty(
+        CredentialFactory.CREDENTIAL_FACTORY_PROPERTY, CloudSqlCredentialSupplier.class.getName());
+    CloudSqlCredentialSupplier.credential = credential;
+  }
+
+  @Override
+  public Credential create() {
+    return credential;
+  }
+}

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

@@ -16,11 +16,9 @@ package google.registry.persistence.transaction;
 
 import com.google.appengine.api.utils.SystemProperty;
 import com.google.appengine.api.utils.SystemProperty.Environment.Value;
-import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Suppliers;
 import google.registry.model.ofy.DatastoreTransactionManager;
 import google.registry.persistence.DaggerPersistenceComponent;
-import google.registry.tools.RegistryToolEnvironment;
 import google.registry.util.NonFinalForTesting;
 import java.util.function.Supplier;
 
@@ -38,11 +36,10 @@ public class TransactionManagerFactory {
   private TransactionManagerFactory() {}
 
   private static JpaTransactionManager createJpaTransactionManager() {
+    // If we are running a nomulus command, jpaTm will be injected in RegistryCli.java
+    // by calling setJpaTm().
     if (isInAppEngine()) {
       return DaggerPersistenceComponent.create().appEngineJpaTransactionManager();
-    } else if (RegistryToolEnvironment.isInRegistryTool()
-        && RegistryToolEnvironment.isJpaTmEnabled()) {
-      return DaggerPersistenceComponent.create().nomulusToolJpaTransactionManager();
     } else {
       return DummyJpaTransactionManager.create();
     }
@@ -78,8 +75,8 @@ public class TransactionManagerFactory {
     return jpaTm.get();
   }
 
-  @VisibleForTesting
+  /** Sets the return of {@link #jpaTm()} to the given instance of {@link JpaTransactionManager}. */
-  static void setJpaTmForTesting(JpaTransactionManager newJpaTm) {
+  public static void setJpaTm(JpaTransactionManager newJpaTm) {
     jpaTm = Suppliers.ofInstance(newJpaTm);
   }
 }

core/src/main/java/google/registry/tools/AuthModule.java

@@ -20,6 +20,7 @@ import com.google.api.client.auth.oauth2.Credential;
 import com.google.api.client.googleapis.auth.oauth2.GoogleAuthorizationCodeFlow;
 import com.google.api.client.googleapis.auth.oauth2.GoogleClientSecrets;
 import com.google.api.client.googleapis.auth.oauth2.GoogleClientSecrets.Details;
+import com.google.api.client.googleapis.auth.oauth2.GoogleCredential;
 import com.google.api.client.http.javanet.NetHttpTransport;
 import com.google.api.client.json.JsonFactory;
 import com.google.api.client.util.store.AbstractDataStoreFactory;
@@ -42,6 +43,7 @@ import google.registry.util.GoogleCredentialsBundle;
 import java.io.ByteArrayInputStream;
 import java.io.File;
 import java.io.IOException;
+import java.io.UncheckedIOException;
 import java.lang.annotation.Documented;
 import java.lang.annotation.Retention;
 import java.lang.annotation.RetentionPolicy;
@@ -92,6 +94,26 @@ public class AuthModule {
     }
   }
 
+  // TODO(b/138195359): Deprecate this credential once Cloud SQL socket library uses the new auth
+  // library.
+  @Provides
+  @CloudSqlClientCredential
+  public static Credential providesLocalCredentialForCloudSqlClient(
+      @LocalCredentialJson String credentialJson,
+      @Config("localCredentialOauthScopes") ImmutableList<String> credentialScopes) {
+    try {
+      GoogleCredential credential =
+          GoogleCredential.fromStream(new ByteArrayInputStream(credentialJson.getBytes(UTF_8)));
+      if (credential.createScopedRequired()) {
+        credential = credential.createScoped(credentialScopes);
+      }
+      return credential;
+    } catch (IOException e) {
+      throw new UncheckedIOException(
+          "Error occurred while creating a GoogleCredential for Cloud SQL client", e);
+    }
+  }
+
   @Provides
   public static GoogleAuthorizationCodeFlow provideAuthorizationCodeFlow(
       JsonFactory jsonFactory,
@@ -184,6 +206,11 @@ public class AuthModule {
   @Retention(RetentionPolicy.RUNTIME)
   private @interface StoredCredential {}
 
+  /** Dagger qualifier for {@link Credential} used by the Cloud SQL client in the nomulus tool. */
+  @Qualifier
+  @Documented
+  public @interface CloudSqlClientCredential {}
+
   /** Dagger qualifier for the credential qualifier consisting of client and scopes. */
   @Qualifier
   @Documented

core/src/main/java/google/registry/tools/RegistryCli.java

@@ -31,6 +31,7 @@ import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.Iterables;
 import google.registry.config.RegistryConfig;
 import google.registry.model.ofy.ObjectifyService;
+import google.registry.persistence.transaction.TransactionManagerFactory;
 import google.registry.tools.AuthModule.LoginRequiredException;
 import google.registry.tools.params.ParameterFactory;
 import java.io.ByteArrayInputStream;
@@ -237,7 +238,7 @@ final class RegistryCli implements AutoCloseable, CommandRunner {
       // Enable Cloud SQL for command that needs remote API as they will very likely use
       // Cloud SQL after the database migration. Note that the DB password is stored in Datastore
       // and it is already initialized above.
-      RegistryToolEnvironment.enableJpaTm();
+      TransactionManagerFactory.setJpaTm(component.nomulusToolJpaTransactionManager());
     }
 
     command.run();

core/src/main/java/google/registry/tools/RegistryToolComponent.java

@@ -27,6 +27,9 @@ import google.registry.keyring.KeyringModule;
 import google.registry.keyring.api.DummyKeyringModule;
 import google.registry.keyring.api.KeyModule;
 import google.registry.keyring.kms.KmsModule;
+import google.registry.persistence.PersistenceModule;
+import google.registry.persistence.PersistenceModule.NomulusToolJpaTm;
+import google.registry.persistence.transaction.JpaTransactionManager;
 import google.registry.rde.RdeModule;
 import google.registry.request.Modules.DatastoreServiceModule;
 import google.registry.request.Modules.Jackson2Module;
@@ -63,6 +66,7 @@ import javax.inject.Singleton;
       KeyringModule.class,
       KmsModule.class,
       LocalCredentialModule.class,
+      PersistenceModule.class,
       RdeModule.class,
       RequestFactoryModule.class,
       URLFetchServiceModule.class,
@@ -70,7 +74,7 @@ import javax.inject.Singleton;
       UserServiceModule.class,
       UtilsModule.class,
       VoidDnsWriterModule.class,
-      WhoisModule.class,
+      WhoisModule.class
     })
 interface RegistryToolComponent {
   void inject(AckPollMessagesCommand command);
@@ -117,6 +121,9 @@ interface RegistryToolComponent {
   @LocalCredentialJson
   String googleCredentialJson();
 
+  @NomulusToolJpaTm
+  JpaTransactionManager nomulusToolJpaTransactionManager();
+
   @Component.Builder
   interface Builder {
     @BindsInstance

core/src/main/java/google/registry/tools/RegistryToolEnvironment.java

@@ -23,7 +23,6 @@ import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import google.registry.config.RegistryEnvironment;
 import google.registry.config.SystemPropertySetter;
-import google.registry.persistence.transaction.TransactionManagerFactory;
 
 /** Enum of production environments, used for the {@code --environment} flag. */
 public enum RegistryToolEnvironment {
@@ -40,7 +39,6 @@ public enum RegistryToolEnvironment {
 
   private static final ImmutableList<String> FLAGS = ImmutableList.of("-e", "--environment");
   private static RegistryToolEnvironment instance;
-  private static boolean isJpaTmEnabled = false;
   private final RegistryEnvironment actualEnvironment;
   private final ImmutableMap<String, String> extraProperties;
 
@@ -100,26 +98,6 @@ public enum RegistryToolEnvironment {
     }
   }
 
-  /** Returns true if the RegistryToolEnvironment is set up. */
-  public static boolean isInRegistryTool() {
-    return instance != null;
-  }
-
-  /**
-   * Sets the flag to indicate that the running command needs JpaTransactionManager to be enabled.
-   */
-  public static void enableJpaTm() {
-    isJpaTmEnabled = true;
-  }
-
-  /**
-   * Returns true if the JpaTransactionManager is enabled. Note that JpaTm is actually enabled in
-   * {@link TransactionManagerFactory} by reading this flag.
-   */
-  public static boolean isJpaTmEnabled() {
-    return isJpaTmEnabled;
-  }
-
   /** Extracts value from command-line arguments associated with any {@code flags}. */
   private static String getFlagValue(String[] args, Iterable<String> flags) {
     for (String flag : flags) {

core/src/test/java/google/registry/persistence/transaction/JpaTransactionManagerRule.java

@@ -133,12 +133,12 @@ abstract class JpaTransactionManagerRule extends ExternalResource {
             extraEntityClasses);
     JpaTransactionManagerImpl txnManager = new JpaTransactionManagerImpl(emf, clock);
     cachedTm = TransactionManagerFactory.jpaTm();
-    TransactionManagerFactory.setJpaTmForTesting(txnManager);
+    TransactionManagerFactory.setJpaTm(txnManager);
   }
 
   @Override
   public void after() {
-    TransactionManagerFactory.setJpaTmForTesting(cachedTm);
+    TransactionManagerFactory.setJpaTm(cachedTm);
     if (emf != null) {
       emf.close();
       emf = null;

dependencies.gradle

@@ -43,6 +43,7 @@ ext {
       'com.google.auto.value:auto-value-annotations:1.6.3',
       'com.google.auto.value:auto-value:1.6.3',
       'com.google.closure-stylesheets:closure-stylesheets:1.5.0',
+      'com.google.cloud.sql:jdbc-socket-factory-core:1.0.12',
       'com.google.cloud.sql:postgres-socket-factory:1.0.12',
       'com.google.cloud:google-cloud-core:1.59.0',
       'com.google.cloud:google-cloud-storage:1.59.0',