Add DateTime as extensibility param for allocation token logic

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

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

@@ -155,10 +155,13 @@ public final class DomainCheckFlow implements Flow {
     Set<String> existingIds = checkResourcesExist(DomainResource.class, targetIds, now);
     Optional<AllocationTokenExtension> allocationTokenExtension =
         eppInput.getSingleExtension(AllocationTokenExtension.class);
-    ImmutableMap<String, String> tokenCheckResults =
+    ImmutableMap<InternetDomainName, String> tokenCheckResults =
         allocationTokenExtension.isPresent()
             ? allocationTokenFlowUtils.checkDomainsWithToken(
-                targetIds, allocationTokenExtension.get().getAllocationToken(), clientId)
+                ImmutableList.copyOf(domainNames.values()),
+                allocationTokenExtension.get().getAllocationToken(),
+                clientId,
+                now)
             : ImmutableMap.of();
     ImmutableList.Builder<DomainCheck> checks = new ImmutableList.Builder<>();
     for (String targetId : targetIds) {
@@ -182,7 +185,7 @@ public final class DomainCheckFlow implements Flow {
   private Optional<String> getMessageForCheck(
       InternetDomainName domainName,
       Set<String> existingIds,
-      ImmutableMap<String, String> tokenCheckResults,
+      ImmutableMap<InternetDomainName, String> tokenCheckResults,
       DateTime now) {
     if (existingIds.contains(domainName.toString())) {
       return Optional.of("In use");
@@ -204,7 +207,7 @@ public final class DomainCheckFlow implements Flow {
     if (!reservationTypes.isEmpty()) {
       return Optional.of(getTypeOfHighestSeverity(reservationTypes).getMessageForCheck());
     }
-    return Optional.ofNullable(emptyToNull(tokenCheckResults.get(domainName.toString())));
+    return Optional.ofNullable(emptyToNull(tokenCheckResults.get(domainName)));
   }
 
   /** Handle the fee check extension. */

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

@@ -74,6 +74,7 @@ import google.registry.model.billing.BillingEvent.Reason;
 import google.registry.model.billing.BillingEvent.Recurring;
 import google.registry.model.domain.AllocationToken;
 import google.registry.model.domain.DomainApplication;
+import google.registry.model.domain.DomainCommand;
 import google.registry.model.domain.DomainCommand.Create;
 import google.registry.model.domain.DomainResource;
 import google.registry.model.domain.GracePeriod;
@@ -204,7 +205,7 @@ public class DomainCreateFlow implements TransactionalFlow {
     validateClientIsLoggedIn(clientId);
     verifyRegistrarIsActive(clientId);
     DateTime now = ofy().getTransactionTime();
-    Create command = cloneAndLinkReferences((Create) resourceCommand, now);
+    DomainCommand.Create command = cloneAndLinkReferences((Create) resourceCommand, now);
     Period period = command.getPeriod();
     verifyUnitIsYears(period);
     int years = period.getValue();
@@ -261,7 +262,7 @@ public class DomainCreateFlow implements TransactionalFlow {
       }
     }
     Optional<AllocationToken> allocationToken =
-        verifyAllocationTokenIfPresent(domainName, registry, clientId);
+        verifyAllocationTokenIfPresent(command, registry, clientId, now);
     flowCustomLogic.afterValidation(
         DomainCreateFlowCustomLogic.AfterValidationParameters.newBuilder()
             .setDomainName(domainName)
@@ -384,14 +385,14 @@ public class DomainCreateFlow implements TransactionalFlow {
 
   /** Verifies and returns the allocation token if one is specified, otherwise does nothing. */
   private Optional<AllocationToken> verifyAllocationTokenIfPresent(
-      InternetDomainName domainName, Registry registry, String clientId)
+      DomainCommand.Create command, Registry registry, String clientId, DateTime now)
       throws EppException {
     Optional<AllocationTokenExtension> extension =
         eppInput.getSingleExtension(AllocationTokenExtension.class);
     return Optional.ofNullable(
         extension.isPresent()
             ? allocationTokenFlowUtils.verifyToken(
-                domainName, extension.get().getAllocationToken(), registry, clientId)
+                command, extension.get().getAllocationToken(), registry, clientId, now)
             : null);
   }
 

java/google/registry/flows/domain/token/AllocationTokenCustomLogic.java

@@ -18,7 +18,9 @@ import com.google.common.collect.ImmutableMap;
 import com.google.common.net.InternetDomainName;
 import google.registry.flows.EppException;
 import google.registry.model.domain.AllocationToken;
+import google.registry.model.domain.DomainCommand;
 import google.registry.model.registry.Registry;
+import org.joda.time.DateTime;
 
 /**
  * A no-op base class for allocation token custom logic.
@@ -29,15 +31,22 @@ public class AllocationTokenCustomLogic {
 
   /** Performs additional custom logic for verifying a token. */
   public AllocationToken verifyToken(
-      InternetDomainName domainName, AllocationToken token, Registry registry, String clientId)
+      DomainCommand.Create command,
+      AllocationToken token,
+      Registry registry,
+      String clientId,
+      DateTime now)
       throws EppException {
     // Do nothing.
     return token;
   }
 
   /** Performs additional custom logic for performing domain checks using a token. */
-  public ImmutableMap<String, String> checkDomainsWithToken(
-      ImmutableMap<String, String> checkResults, AllocationToken tokenEntity, String clientId) {
+  public ImmutableMap<InternetDomainName, String> checkDomainsWithToken(
+      ImmutableMap<InternetDomainName, String> checkResults,
+      AllocationToken token,
+      String clientId,
+      DateTime now) {
     // Do nothing.
     return checkResults;
   }

java/google/registry/flows/domain/token/AllocationTokenFlowUtils.java

@@ -23,14 +23,15 @@ import google.registry.flows.EppException;
 import google.registry.flows.EppException.AssociationProhibitsOperationException;
 import google.registry.flows.EppException.ParameterValueSyntaxErrorException;
 import google.registry.model.domain.AllocationToken;
+import google.registry.model.domain.DomainCommand;
 import google.registry.model.registry.Registry;
 import google.registry.model.reporting.HistoryEntry;
 import java.util.List;
 import java.util.function.Function;
 import javax.inject.Inject;
+import org.joda.time.DateTime;
 
 /** Utility functions for dealing with {@link AllocationToken}s in domain flows. */
-// TODO: Add a test class.
 public class AllocationTokenFlowUtils {
 
   final AllocationTokenCustomLogic tokenCustomLogic;
@@ -47,7 +48,7 @@ public class AllocationTokenFlowUtils {
    * @throws InvalidAllocationTokenException if the token doesn't exist.
    */
   public AllocationToken verifyToken(
-      InternetDomainName domainName, String token, Registry registry, String clientId)
+      DomainCommand.Create command, String token, Registry registry, String clientId, DateTime now)
       throws EppException {
     AllocationToken tokenEntity = ofy().load().key(Key.create(AllocationToken.class, token)).now();
     if (tokenEntity == null) {
@@ -56,7 +57,7 @@ public class AllocationTokenFlowUtils {
     if (tokenEntity.isRedeemed()) {
       throw new AlreadyRedeemedAllocationTokenException();
     }
-    return tokenCustomLogic.verifyToken(domainName, tokenEntity, registry, clientId);
+    return tokenCustomLogic.verifyToken(command, tokenEntity, registry, clientId, now);
   }
 
   /**
@@ -66,22 +67,26 @@ public class AllocationTokenFlowUtils {
    *     for a a given domain then it does not validate with this allocation token; domains that do
    *     validate have blank messages (i.e. no error).
    */
-  public ImmutableMap<String, String> checkDomainsWithToken(
-      List<String> domainNames, String token, String clientId) {
+  public ImmutableMap<InternetDomainName, String> checkDomainsWithToken(
+      List<InternetDomainName> domainNames, String token, String clientId, DateTime now) {
     AllocationToken tokenEntity = ofy().load().key(Key.create(AllocationToken.class, token)).now();
-    String result;
+    String globalResult;
     if (tokenEntity == null) {
-      result = new InvalidAllocationTokenException().getMessage();
+      globalResult = new InvalidAllocationTokenException().getMessage();
     } else if (tokenEntity.isRedeemed()) {
-      result = AlreadyRedeemedAllocationTokenException.ERROR_MSG_SHORT;
+      globalResult = AlreadyRedeemedAllocationTokenException.ERROR_MSG_SHORT;
     } else {
-      result = "";
+      globalResult = "";
     }
-    ImmutableMap<String, String> checkResults =
+    ImmutableMap<InternetDomainName, String> checkResults =
         domainNames
             .stream()
-            .collect(ImmutableMap.toImmutableMap(Function.identity(), domainName -> result));
-    return tokenCustomLogic.checkDomainsWithToken(checkResults, tokenEntity, clientId);
+            .collect(ImmutableMap.toImmutableMap(Function.identity(), domainName -> globalResult));
+    // Only call custom logic if there wasn't a global allocation token error that applies to all
+    // check results. The custom logic can only add errors, not override existing errors.
+    return globalResult.isEmpty()
+        ? tokenCustomLogic.checkDomainsWithToken(checkResults, tokenEntity, clientId, now)
+        : checkResults;
   }
 
   /** Redeems an {@link AllocationToken}, returning the redeemed copy. */