Add extensible custom logic to the domain check flow

This also fixes up a hook on the domain create flow custom logic to use a single
parameter, which is the general pattern we want to use going forward. It also
establishes a pattern for custom logic being able to add extensions.

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

java/google/registry/flows/custom/CustomLogicFactory.java

@@ -35,4 +35,9 @@ public class CustomLogicFactory {
       EppInput eppInput, SessionMetadata sessionMetadata) {
     return new DomainCreateFlowCustomLogic(eppInput, sessionMetadata);
   }
+
+  public DomainCheckFlowCustomLogic forDomainCheckFlow(
+      EppInput eppInput, SessionMetadata sessionMetadata) {
+    return new DomainCheckFlowCustomLogic(eppInput, sessionMetadata);
+  }
 }

java/google/registry/flows/custom/CustomLogicModule.java

@@ -28,4 +28,10 @@ public class CustomLogicModule {
       CustomLogicFactory factory, EppInput eppInput, SessionMetadata sessionMetadata) {
     return factory.forDomainCreateFlow(eppInput, sessionMetadata);
   }
+
+  @Provides
+  static DomainCheckFlowCustomLogic provideDomainCheckFlowCustomLogic(
+      CustomLogicFactory factory, EppInput eppInput, SessionMetadata sessionMetadata) {
+    return factory.forDomainCheckFlow(eppInput, sessionMetadata);
+  }
 }

java/google/registry/flows/custom/DomainCheckFlowCustomLogic.java

@@ -0,0 +1,157 @@
+// Copyright 2016 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.flows.custom;
+
+import com.google.auto.value.AutoValue;
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMap;
+import com.google.common.net.InternetDomainName;
+import google.registry.flows.EppException;
+import google.registry.flows.SessionMetadata;
+import google.registry.flows.domain.DomainCheckFlow;
+import google.registry.model.ImmutableObject;
+import google.registry.model.eppinput.EppInput;
+import google.registry.model.eppoutput.CheckData.DomainCheck;
+import google.registry.model.eppoutput.EppResponse.ResponseExtension;
+import org.joda.time.DateTime;
+
+/**
+ * A no-op base class for {@link DomainCheckFlow} custom logic.
+ *
+ * <p>Extend this class and override the hook(s) to perform custom logic.
+ */
+public class DomainCheckFlowCustomLogic extends BaseFlowCustomLogic {
+
+  protected DomainCheckFlowCustomLogic(EppInput eppInput, SessionMetadata sessionMetadata) {
+    super(eppInput, sessionMetadata);
+  }
+
+  /**
+   * A hook that runs before any validation. This is useful to e.g. add allowable extensions.
+   */
+  @SuppressWarnings("unused")
+  public void beforeValidation() throws EppException {
+    // Do nothing.
+  }
+
+  /** A hook that runs at the end of the validation step to perform additional validation. */
+  @SuppressWarnings("unused")
+  public void afterValidation(AfterValidationParameters parameters) throws EppException {
+    // Do nothing.
+  }
+
+  /**
+   * A hook that runs before the response is returned.
+   *
+   * <p>This takes the {@link DomainCheck}s and {@link ResponseExtension}s as input and returns
+   * them, potentially with modifications.
+   */
+  @SuppressWarnings("unused")
+  public BeforeResponseReturnData beforeResponse(BeforeResponseParameters parameters)
+      throws EppException {
+    return BeforeResponseReturnData.newBuilder()
+        .setDomainChecks(parameters.domainChecks())
+        .setResponseExtensions(parameters.responseExtensions())
+        .build();
+  }
+
+  /** A class to encapsulate parameters for a call to {@link #afterValidation}. */
+  @AutoValue
+  public abstract static class AfterValidationParameters extends ImmutableObject {
+
+    public abstract ImmutableMap<String, InternetDomainName> domainNames();
+
+    /**
+     * The time to perform the domain check as of. This defaults to the current time, but can be
+     * overridden in v>=0.12 of the fee extension.
+     */
+    public abstract DateTime asOfDate();
+
+    public static Builder newBuilder() {
+      return new AutoValue_DomainCheckFlowCustomLogic_AfterValidationParameters.Builder();
+    }
+
+    /** Builder for {@link AfterValidationParameters}. */
+    @AutoValue.Builder
+    public abstract static class Builder {
+
+      public abstract Builder setDomainNames(ImmutableMap<String, InternetDomainName> domainNames);
+
+      public abstract Builder setAsOfDate(DateTime asOfDate);
+
+      public abstract AfterValidationParameters build();
+    }
+  }
+
+  /** A class to encapsulate parameters for a call to {@link #beforeResponse}. */
+  @AutoValue
+  public abstract static class BeforeResponseParameters extends ImmutableObject {
+
+    public abstract ImmutableList<DomainCheck> domainChecks();
+
+    public abstract ImmutableList<? extends ResponseExtension> responseExtensions();
+
+    /**
+     * The time to perform the domain check as of. This defaults to the current time, but can be
+     * overridden in v>=0.12 of the fee extension.
+     */
+    public abstract DateTime asOfDate();
+
+    public static Builder newBuilder() {
+      return new AutoValue_DomainCheckFlowCustomLogic_BeforeResponseParameters.Builder();
+    }
+
+    /** Builder for {@link BeforeResponseParameters}. */
+    @AutoValue.Builder
+    public abstract static class Builder {
+
+      public abstract Builder setDomainChecks(ImmutableList<DomainCheck> domainChecks);
+
+      public abstract Builder setResponseExtensions(
+          ImmutableList<? extends ResponseExtension> responseExtensions);
+
+      public abstract Builder setAsOfDate(DateTime asOfDate);
+
+      public abstract BeforeResponseParameters build();
+    }
+  }
+
+  /**
+   * A class to encapsulate parameters for the return values from a call to {@link #beforeResponse}.
+   */
+  @AutoValue
+  public abstract static class BeforeResponseReturnData extends ImmutableObject {
+
+    public abstract ImmutableList<DomainCheck> domainChecks();
+
+    public abstract ImmutableList<? extends ResponseExtension> responseExtensions();
+
+    public static Builder newBuilder() {
+      return new AutoValue_DomainCheckFlowCustomLogic_BeforeResponseReturnData.Builder();
+    }
+
+    /** Builder for {@link BeforeResponseReturnData}. */
+    @AutoValue.Builder
+    public abstract static class Builder {
+
+      public abstract Builder setDomainChecks(ImmutableList<DomainCheck> domainChecks);
+
+      public abstract Builder setResponseExtensions(
+          ImmutableList<? extends ResponseExtension> responseExtensions);
+
+      public abstract BeforeResponseReturnData build();
+    }
+  }
+}

java/google/registry/flows/custom/DomainCreateFlowCustomLogic.java

@@ -18,15 +18,16 @@ import com.google.auto.value.AutoValue;
 import com.google.common.net.InternetDomainName;
 import google.registry.flows.EppException;
 import google.registry.flows.SessionMetadata;
+import google.registry.flows.domain.DomainCreateFlow;
 import google.registry.model.ImmutableObject;
 import google.registry.model.domain.DomainResource;
 import google.registry.model.eppinput.EppInput;
 import google.registry.model.reporting.HistoryEntry;
 
 /**
- * A no-op base class for domain create flow custom logic.
+ * A no-op base class for {@link DomainCreateFlow} custom logic.
  *
- * <p>Extend this class and override the hooks to perform custom logic.
+ * <p>Extend this class and override the hook(s) to perform custom logic.
  */
 public class DomainCreateFlowCustomLogic extends BaseFlowCustomLogic {
 
@@ -41,17 +42,16 @@ public class DomainCreateFlowCustomLogic extends BaseFlowCustomLogic {
   }
 
   /**
-   * A hook that runs before new entities are persisted.
+   * A hook that runs before new entities are persisted, allowing them to be changed.
    *
-   * <p>This takes the new entities as input and returns the actual entities to save. It is
-   * important to be careful when changing the flow behavior for existing entities, because the core
-   * logic across many different flows expects the existence of these entities and many of the
-   * fields on them.
+   * <p>It returns the actual entity changes that should be persisted to Datastore. It is important
+   * to be careful when changing the flow behavior for existing entities, because the core logic
+   * across many different flows expects the existence of these entities and many of the fields on
+   * them.
    */
   @SuppressWarnings("unused")
-  public EntityChanges beforeSave(BeforeSaveParameters parameters, EntityChanges entityChanges)
-      throws EppException {
-    return entityChanges;
+  public EntityChanges beforeSave(BeforeSaveParameters parameters) throws EppException {
+    return parameters.entityChanges();
   }
 
   /** A class to encapsulate parameters for a call to {@link #afterValidation}. */
@@ -59,6 +59,7 @@ public class DomainCreateFlowCustomLogic extends BaseFlowCustomLogic {
   public abstract static class AfterValidationParameters extends ImmutableObject {
 
     public abstract InternetDomainName domainName();
+
     public abstract int years();
 
     public static Builder newBuilder() {
@@ -68,8 +69,11 @@ public class DomainCreateFlowCustomLogic extends BaseFlowCustomLogic {
     /** Builder for {@link AfterValidationParameters}. */
     @AutoValue.Builder
     public abstract static class Builder {
+
       public abstract Builder setDomainName(InternetDomainName domainName);
+
       public abstract Builder setYears(int years);
+
       public abstract AfterValidationParameters build();
     }
   }
@@ -79,7 +83,11 @@ public class DomainCreateFlowCustomLogic extends BaseFlowCustomLogic {
   public abstract static class BeforeSaveParameters extends ImmutableObject {
 
     public abstract DomainResource newDomain();
+
     public abstract HistoryEntry historyEntry();
+
+    public abstract EntityChanges entityChanges();
+
     public abstract int years();
 
     public static Builder newBuilder() {
@@ -89,9 +97,15 @@ public class DomainCreateFlowCustomLogic extends BaseFlowCustomLogic {
     /** Builder for {@link BeforeSaveParameters}. */
     @AutoValue.Builder
     public abstract static class Builder {
+
       public abstract Builder setNewDomain(DomainResource newDomain);
+
       public abstract Builder setHistoryEntry(HistoryEntry historyEntry);
+
+      public abstract Builder setEntityChanges(EntityChanges entityChanges);
+
       public abstract Builder setYears(int years);
+
       public abstract BeforeSaveParameters build();
     }
   }

java/google/registry/flows/domain/DomainCheckFlow.java

@@ -41,6 +41,9 @@ import google.registry.flows.ExtensionManager;
 import google.registry.flows.Flow;
 import google.registry.flows.FlowModule.ClientId;
 import google.registry.flows.FlowModule.Superuser;
+import google.registry.flows.custom.DomainCheckFlowCustomLogic;
+import google.registry.flows.custom.DomainCheckFlowCustomLogic.BeforeResponseParameters;
+import google.registry.flows.custom.DomainCheckFlowCustomLogic.BeforeResponseReturnData;
 import google.registry.model.domain.DomainApplication;
 import google.registry.model.domain.DomainCommand.Check;
 import google.registry.model.domain.DomainResource;
@@ -106,11 +109,13 @@ public final class DomainCheckFlow implements Flow {
   @Inject @Superuser boolean isSuperuser;
   @Inject Clock clock;
   @Inject EppResponse.Builder responseBuilder;
+  @Inject DomainCheckFlowCustomLogic customLogic;
   @Inject DomainCheckFlow() {}
 
   @Override
   public EppResponse run() throws EppException {
     extensionManager.register(FeeCheckCommandExtension.class, LaunchCheckExtension.class);
+    customLogic.beforeValidation();
     extensionManager.validate();
     validateClientIsLoggedIn(clientId);
     List<String> targetIds = ((Check) resourceCommand).getTargetIds();
@@ -133,15 +138,27 @@ public final class DomainCheckFlow implements Flow {
       }
     }
     ImmutableMap<String, InternetDomainName> domainNames = domains.build();
+    customLogic.afterValidation(
+        DomainCheckFlowCustomLogic.AfterValidationParameters.newBuilder()
+            .setDomainNames(domainNames)
+            .setAsOfDate(now)
+            .build());
     Set<String> existingIds = checkResourcesExist(DomainResource.class, targetIds, now);
     ImmutableList.Builder<DomainCheck> checks = new ImmutableList.Builder<>();
     for (String targetId : targetIds) {
       String message = getMessageForCheck(domainNames.get(targetId), existingIds, now);
       checks.add(DomainCheck.create(message == null, targetId, message));
     }
+    BeforeResponseReturnData responseData =
+        customLogic.beforeResponse(
+            BeforeResponseParameters.newBuilder()
+                .setDomainChecks(checks.build())
+                .setResponseExtensions(getResponseExtensions(domainNames, now))
+                .setAsOfDate(now)
+                .build());
     return responseBuilder
-        .setResData(DomainCheckData.create(checks.build()))
-        .setExtensions(getResponseExtensions(domainNames, now))
+        .setResData(DomainCheckData.create(responseData.domainChecks()))
+        .setExtensions(responseData.responseExtensions())
         .build();
   }
 
@@ -170,7 +187,6 @@ public final class DomainCheckFlow implements Flow {
     return reservationType.getMessageForCheck();
   }
 
-
   /** Handle the fee check extension. */
   private ImmutableList<? extends ResponseExtension> getResponseExtensions(
       ImmutableMap<String, InternetDomainName> domainNames, DateTime now) throws EppException {

java/google/registry/flows/domain/DomainCreateFlow.java

@@ -309,9 +309,10 @@ public class DomainCreateFlow implements TransactionalFlow {
             DomainCreateFlowCustomLogic.BeforeSaveParameters.newBuilder()
                 .setNewDomain(newDomain)
                 .setHistoryEntry(historyEntry)
+                .setEntityChanges(
+                    EntityChanges.newBuilder().setSaves(entitiesToSave.build()).build())
                 .setYears(years)
-                .build(),
-            EntityChanges.newBuilder().setSaves(entitiesToSave.build()).build());
+                .build());
     persistEntityChanges(entityChanges);
 
     return responseBuilder

java/google/registry/flows/domain/RegistryExtraFlowLogic.java

@@ -27,6 +27,7 @@ import org.joda.time.DateTime;
  * Interface for classes which provide extra registry logic for things like TLD-specific rules and
  * discounts.
  */
+@Deprecated
 public interface RegistryExtraFlowLogic {
 
   /** Gets the flags to be returned for application info commands. */