Add extensible custom logic to the domain update flow

This also adds beforeValidation hook to DomainCreateCustomFlow.

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

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

@@ -45,4 +45,9 @@ public class CustomLogicFactory {
       EppInput eppInput, SessionMetadata sessionMetadata) {
     return new DomainInfoFlowCustomLogic(eppInput, sessionMetadata);
   }
+
+  public DomainUpdateFlowCustomLogic forDomainUpdateFlow(
+      EppInput eppInput, SessionMetadata sessionMetadata) {
+    return new DomainUpdateFlowCustomLogic(eppInput, sessionMetadata);
+  }
 }

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

@@ -40,4 +40,10 @@ public class CustomLogicModule {
       CustomLogicFactory factory, EppInput eppInput, SessionMetadata sessionMetadata) {
     return factory.forDomainInfoFlow(eppInput, sessionMetadata);
   }
+
+  @Provides
+  static DomainUpdateFlowCustomLogic provideDomainUpdateFlowCustomLogic(
+      CustomLogicFactory factory, EppInput eppInput, SessionMetadata sessionMetadata) {
+    return factory.forDomainUpdateFlow(eppInput, sessionMetadata);
+  }
 }

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

@@ -35,6 +35,12 @@ public class DomainCreateFlowCustomLogic extends BaseFlowCustomLogic {
     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 {

java/google/registry/flows/custom/DomainUpdateFlowCustomLogic.java

@@ -0,0 +1,120 @@
+// 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 google.registry.flows.EppException;
+import google.registry.flows.SessionMetadata;
+import google.registry.flows.domain.DomainUpdateFlow;
+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 {@link DomainUpdateFlow} custom logic.
+ *
+ * <p>Extend this class and override the hook(s) to perform custom logic.
+ */
+public class DomainUpdateFlowCustomLogic extends BaseFlowCustomLogic {
+
+  protected DomainUpdateFlowCustomLogic(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 new entities are persisted, allowing them to be changed.
+   *
+   * <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) throws EppException {
+    return parameters.entityChanges();
+  }
+
+  /** A class to encapsulate parameters for a call to {@link #afterValidation}. */
+  @AutoValue
+  public abstract static class AfterValidationParameters extends ImmutableObject {
+
+    public abstract DomainResource existingDomain();
+
+    public static Builder newBuilder() {
+      return new AutoValue_DomainUpdateFlowCustomLogic_AfterValidationParameters.Builder();
+    }
+
+    /** Builder for {@link AfterValidationParameters}. */
+    @AutoValue.Builder
+    public abstract static class Builder {
+
+      public abstract Builder setExistingDomain(DomainResource existingDomain);
+
+      public abstract AfterValidationParameters build();
+    }
+  }
+
+  /**
+   * A class to encapsulate parameters for a call to {@link #beforeSave}.
+   *
+   * <p>Note that both newDomain and historyEntry are included in entityChanges. They are also
+   * passed separately for convenience, but they are the same instance, and changes to them will
+   * also affect what is persisted from entityChanges.
+   */
+  @AutoValue
+  public abstract static class BeforeSaveParameters extends ImmutableObject {
+
+    public abstract DomainResource existingDomain();
+
+    public abstract DomainResource newDomain();
+
+    public abstract HistoryEntry historyEntry();
+
+
+    public abstract EntityChanges entityChanges();
+
+    public static Builder newBuilder() {
+      return new AutoValue_DomainUpdateFlowCustomLogic_BeforeSaveParameters.Builder();
+    }
+
+    /** Builder for {@link BeforeSaveParameters}. */
+    @AutoValue.Builder
+    public abstract static class Builder {
+
+      public abstract Builder setExistingDomain(DomainResource existingDomain);
+
+      public abstract Builder setNewDomain(DomainResource newDomain);
+
+      public abstract Builder setHistoryEntry(HistoryEntry historyEntry);
+
+      public abstract Builder setEntityChanges(EntityChanges entityChanges);
+
+      public abstract BeforeSaveParameters build();
+    }
+  }
+}

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

@@ -176,6 +176,7 @@ public class DomainCreateFlow implements TransactionalFlow {
         FlagsCreateCommandExtension.class,
         MetadataExtension.class,
         LaunchCreateExtension.class);
+    customLogic.beforeValidation();
     extensionManager.validate();
     validateClientIsLoggedIn(clientId);
     DateTime now = ofy().getTransactionTime();

java/google/registry/flows/domain/DomainUpdateFlow.java

@@ -17,6 +17,7 @@ package google.registry.flows.domain;
 import static com.google.common.base.MoreObjects.firstNonNull;
 import static com.google.common.collect.Sets.symmetricDifference;
 import static com.google.common.collect.Sets.union;
+import static google.registry.flows.FlowUtils.persistEntityChanges;
 import static google.registry.flows.FlowUtils.validateClientIsLoggedIn;
 import static google.registry.flows.ResourceFlowUtils.checkSameValuesNotAddedAndRemoved;
 import static google.registry.flows.ResourceFlowUtils.loadAndVerifyExistence;
@@ -41,7 +42,6 @@ import static google.registry.model.ofy.ObjectifyService.ofy;
 import static google.registry.util.DateTimeUtils.earliestOf;
 
 import com.google.common.base.Optional;
-import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
 import com.googlecode.objectify.Key;
 import google.registry.dns.DnsQueue;
@@ -51,6 +51,10 @@ import google.registry.flows.FlowModule.ClientId;
 import google.registry.flows.FlowModule.Superuser;
 import google.registry.flows.FlowModule.TargetId;
 import google.registry.flows.TransactionalFlow;
+import google.registry.flows.custom.DomainUpdateFlowCustomLogic;
+import google.registry.flows.custom.DomainUpdateFlowCustomLogic.AfterValidationParameters;
+import google.registry.flows.custom.DomainUpdateFlowCustomLogic.BeforeSaveParameters;
+import google.registry.flows.custom.EntityChanges;
 import google.registry.flows.domain.DomainFlowUtils.FeesRequiredForNonFreeUpdateException;
 import google.registry.flows.domain.TldSpecificLogicProxy.EppCommandOperations;
 import google.registry.model.ImmutableObject;
@@ -128,9 +132,8 @@ public final class DomainUpdateFlow implements TransactionalFlow {
    * requires special checking, since you must be able to clear the status off the object with an
    * update.
    */
-  private static final ImmutableSet<StatusValue> UPDATE_DISALLOWED_STATUSES = ImmutableSet.of(
+  private static final ImmutableSet<StatusValue> UPDATE_DISALLOWED_STATUSES =
-      StatusValue.PENDING_DELETE,
+      ImmutableSet.of(StatusValue.PENDING_DELETE, StatusValue.SERVER_UPDATE_PROHIBITED);
-      StatusValue.SERVER_UPDATE_PROHIBITED);
 
   @Inject ResourceCommand resourceCommand;
   @Inject ExtensionManager extensionManager;
@@ -142,6 +145,8 @@ public final class DomainUpdateFlow implements TransactionalFlow {
   @Inject HistoryEntry.Builder historyBuilder;
   @Inject DnsQueue dnsQueue;
   @Inject EppResponse.Builder responseBuilder;
+  @Inject DomainUpdateFlowCustomLogic customLogic;
+
   @Inject DomainUpdateFlow() {}
 
   @Override
@@ -151,12 +156,15 @@ public final class DomainUpdateFlow implements TransactionalFlow {
         FlagsUpdateCommandExtension.class,
         MetadataExtension.class,
         SecDnsUpdateExtension.class);
+    customLogic.beforeValidation();
     extensionManager.validate();
     validateClientIsLoggedIn(clientId);
     DateTime now = ofy().getTransactionTime();
     Update command = cloneAndLinkReferences((Update) resourceCommand, now);
     DomainResource existingDomain = loadAndVerifyExistence(DomainResource.class, targetId, now);
     verifyUpdateAllowed(command, existingDomain, now);
+    customLogic.afterValidation(
+        AfterValidationParameters.newBuilder().setExistingDomain(existingDomain).build());
     HistoryEntry historyEntry = buildHistoryEntry(existingDomain, now);
     DomainResource newDomain = performUpdate(command, existingDomain, now);
     // If the new domain is in the sunrush add grace period and is now publishable to DNS because we
@@ -165,21 +173,30 @@ public final class DomainUpdateFlow implements TransactionalFlow {
       for (GracePeriod gracePeriod : newDomain.getGracePeriods()) {
         if (gracePeriod.isSunrushAddGracePeriod()) {
           newDomain = convertSunrushAddToAdd(newDomain, gracePeriod, historyEntry, now);
-          break;  // There can only be one sunrush add grace period.
+          break; // There can only be one sunrush add grace period.
         }
       }
     }
     validateNewState(newDomain);
     dnsQueue.addDomainRefreshTask(targetId);
     handleExtraFlowLogic(existingDomain, historyEntry, now);
-    ImmutableList.Builder<ImmutableObject> entitiesToSave = new ImmutableList.Builder<>();
+    ImmutableSet.Builder<ImmutableObject> entitiesToSave = new ImmutableSet.Builder<>();
     entitiesToSave.add(newDomain, historyEntry);
     Optional<BillingEvent.OneTime> statusUpdateBillingEvent =
         createBillingEventForStatusUpdates(existingDomain, newDomain, historyEntry, now);
     if (statusUpdateBillingEvent.isPresent()) {
       entitiesToSave.add(statusUpdateBillingEvent.get());
     }
-    ofy().save().entities(entitiesToSave.build());
+    EntityChanges entityChanges =
+        customLogic.beforeSave(
+            BeforeSaveParameters.newBuilder()
+                .setHistoryEntry(historyEntry)
+                .setNewDomain(newDomain)
+                .setExistingDomain(existingDomain)
+                .setEntityChanges(
+                    EntityChanges.newBuilder().setSaves(entitiesToSave.build()).build())
+                .build());
+    persistEntityChanges(entityChanges);
     return responseBuilder.build();
   }