Modify DomainCreateFlow to check for an applicable defaultPromoToken (#1904)

* Modify DomainCreateFlow to check for an applicable defaultPromoToken

* Add handling for deleted tokens

* Change cache to allocation token cache

* Abstract away cache methods

* Use AllocationToken.getAll in create flow

* Filter out empty tokens

core/src/main/java/google/registry/flows/domain/DomainCreateFlow.java

@@ -15,6 +15,8 @@
 package google.registry.flows.domain;
 
 import static com.google.common.base.Preconditions.checkArgument;
+import static com.google.common.base.Preconditions.checkState;
+import static com.google.common.collect.ImmutableList.toImmutableList;
 import static com.google.common.collect.ImmutableSet.toImmutableSet;
 import static google.registry.flows.FlowUtils.persistEntityChanges;
 import static google.registry.flows.FlowUtils.validateRegistrarIsLoggedIn;
@@ -52,6 +54,7 @@ import static google.registry.model.tld.Registry.TldState.QUIET_PERIOD;
 import static google.registry.model.tld.Registry.TldState.START_DATE_SUNRISE;
 import static google.registry.model.tld.label.ReservationType.NAME_COLLISION;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
+import static google.registry.util.CollectionUtils.isNullOrEmpty;
 import static google.registry.util.DateTimeUtils.END_OF_TIME;
 import static google.registry.util.DateTimeUtils.leapSafeAddYears;
 
@@ -61,6 +64,7 @@ import com.google.common.collect.ImmutableSet;
 import com.google.common.net.InternetDomainName;
 import google.registry.dns.DnsQueue;
 import google.registry.flows.EppException;
+import google.registry.flows.EppException.AssociationProhibitsOperationException;
 import google.registry.flows.EppException.CommandUseErrorException;
 import google.registry.flows.EppException.ParameterValuePolicyErrorException;
 import google.registry.flows.ExtensionManager;
@@ -117,7 +121,9 @@ import google.registry.model.tld.Registry.TldType;
 import google.registry.model.tld.label.ReservationType;
 import google.registry.model.tmch.ClaimsList;
 import google.registry.model.tmch.ClaimsListDao;
+import google.registry.persistence.VKey;
 import google.registry.tmch.LordnTaskUtils;
+import java.util.Map;
 import java.util.Optional;
 import javax.annotation.Nullable;
 import javax.inject.Inject;
@@ -270,6 +276,9 @@ public final class DomainCreateFlow implements TransactionalFlow {
             registrarId,
             now,
             eppInput.getSingleExtension(AllocationTokenExtension.class));
+    if (!allocationToken.isPresent() && !registry.getDefaultPromoTokens().isEmpty()) {
+      allocationToken = checkForDefaultToken(registry, command);
+    }
     boolean isAnchorTenant =
         isAnchorTenant(
             domainName, allocationToken, eppInput.getSingleExtension(MetadataExtension.class));
@@ -434,6 +443,36 @@ public final class DomainCreateFlow implements TransactionalFlow {
         .build();
   }
 
+  private Optional<AllocationToken> checkForDefaultToken(
+      Registry registry, DomainCommand.Create command) throws EppException {
+    Map<VKey<AllocationToken>, Optional<AllocationToken>> tokens =
+        AllocationToken.getAll(registry.getDefaultPromoTokens());
+    ImmutableList<Optional<AllocationToken>> tokenList =
+        registry.getDefaultPromoTokens().stream()
+            .map(tokens::get)
+            .filter(Optional::isPresent)
+            .collect(toImmutableList());
+    checkState(
+        !isNullOrEmpty(tokenList),
+        "Failure while loading default TLD promotions from the database");
+    // Check if any of the tokens are valid for this domain registration
+    for (Optional<AllocationToken> token : tokenList) {
+      try {
+        AllocationTokenFlowUtils.validateToken(
+            InternetDomainName.from(command.getDomainName()),
+            token.get(),
+            registrarId,
+            tm().getTransactionTime());
+      } catch (AssociationProhibitsOperationException e) {
+        // Allocation token was not valid for this registration, continue to check the next token in
+        // the list
+        continue;
+      }
+      // Only use the first valid token in the list
+      return token;
+    }
+    return Optional.empty();
+  }
   /**
    * Verifies that signed marks are only sent during sunrise.
    *

core/src/main/java/google/registry/flows/domain/token/AllocationTokenFlowUtils.java

@@ -108,7 +108,7 @@ public class AllocationTokenFlowUtils {
    *
    * @throws EppException if the token is invalid in any way
    */
-  private static void validateToken(
+  public static void validateToken(
       InternetDomainName domainName, AllocationToken token, String registrarId, DateTime now)
       throws EppException {
 

core/src/main/java/google/registry/model/domain/token/AllocationToken.java

@@ -16,16 +16,22 @@ package google.registry.model.domain.token;
 
 import static com.google.common.base.Preconditions.checkArgument;
 import static com.google.common.base.Preconditions.checkState;
+import static com.google.common.collect.ImmutableMap.toImmutableMap;
+import static google.registry.config.RegistryConfig.getSingletonCacheRefreshDuration;
 import static google.registry.model.domain.token.AllocationToken.TokenStatus.CANCELLED;
 import static google.registry.model.domain.token.AllocationToken.TokenStatus.ENDED;
 import static google.registry.model.domain.token.AllocationToken.TokenStatus.NOT_STARTED;
 import static google.registry.model.domain.token.AllocationToken.TokenStatus.VALID;
+import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.util.CollectionUtils.forceEmptyToNull;
 import static google.registry.util.CollectionUtils.nullToEmptyImmutableCopy;
 import static google.registry.util.PreconditionsUtils.checkArgumentNotNull;
 
+import com.github.benmanes.caffeine.cache.CacheLoader;
+import com.github.benmanes.caffeine.cache.LoadingCache;
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Strings;
+import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableMultimap;
 import com.google.common.collect.ImmutableSet;
@@ -34,6 +40,7 @@ import com.google.common.collect.Range;
 import google.registry.flows.EppException;
 import google.registry.flows.domain.DomainFlowUtils;
 import google.registry.model.Buildable;
+import google.registry.model.CacheUtils;
 import google.registry.model.CreateAutoTimestamp;
 import google.registry.model.UpdateAutoTimestampEntity;
 import google.registry.model.billing.BillingEvent.RenewalPriceBehavior;
@@ -41,6 +48,7 @@ import google.registry.model.common.TimedTransitionProperty;
 import google.registry.model.reporting.HistoryEntry.HistoryEntryId;
 import google.registry.persistence.VKey;
 import google.registry.persistence.WithVKey;
+import java.util.Map;
 import java.util.Optional;
 import java.util.Set;
 import javax.annotation.Nullable;
@@ -267,6 +275,39 @@ public class AllocationToken extends UpdateAutoTimestampEntity implements Builda
     return STATIC_TOKEN_BEHAVIORS.getOrDefault(token, TokenBehavior.DEFAULT);
   }
 
+  public static Optional<AllocationToken> get(VKey<AllocationToken> key) {
+    return ALLOCATION_TOKENS_CACHE.get(key);
+  }
+
+  public static Map<VKey<AllocationToken>, Optional<AllocationToken>> getAll(
+      ImmutableList<VKey<AllocationToken>> keys) {
+    return ALLOCATION_TOKENS_CACHE.getAll(keys);
+  }
+
+  /** A cache that loads the {@link AllocationToken} object for a given AllocationToken VKey. */
+  private static final LoadingCache<VKey<AllocationToken>, Optional<AllocationToken>>
+      ALLOCATION_TOKENS_CACHE =
+          CacheUtils.newCacheBuilder(getSingletonCacheRefreshDuration())
+              .build(
+                  new CacheLoader<VKey<AllocationToken>, Optional<AllocationToken>>() {
+                    @Override
+                    public Optional<AllocationToken> load(VKey<AllocationToken> key) {
+                      return tm().transact(() -> tm().loadByKeyIfPresent(key));
+                    }
+
+                    @Override
+                    public Map<VKey<AllocationToken>, Optional<AllocationToken>> loadAll(
+                        Iterable<? extends VKey<AllocationToken>> keys) {
+                      ImmutableSet<VKey<AllocationToken>> keySet = ImmutableSet.copyOf(keys);
+                      return tm().transact(
+                              () ->
+                                  keySet.stream()
+                                      .collect(
+                                          toImmutableMap(
+                                              key -> key, key -> tm().loadByKeyIfPresent(key))));
+                    }
+                  });
+
   @Override
   public VKey<AllocationToken> createVKey() {
     if (!AllocationToken.TokenBehavior.DEFAULT.equals(getTokenBehavior())) {

core/src/test/java/google/registry/flows/domain/DomainCreateFlowTest.java

@@ -25,6 +25,7 @@ import static google.registry.model.billing.BillingEvent.RenewalPriceBehavior.DE
 import static google.registry.model.billing.BillingEvent.RenewalPriceBehavior.NONPREMIUM;
 import static google.registry.model.billing.BillingEvent.RenewalPriceBehavior.SPECIFIED;
 import static google.registry.model.domain.fee.Fee.FEE_EXTENSION_URIS;
+import static google.registry.model.domain.token.AllocationToken.TokenType.DEFAULT_PROMO;
 import static google.registry.model.domain.token.AllocationToken.TokenType.PACKAGE;
 import static google.registry.model.domain.token.AllocationToken.TokenType.SINGLE_USE;
 import static google.registry.model.domain.token.AllocationToken.TokenType.UNLIMITED_USE;
@@ -1612,6 +1613,239 @@ class DomainCreateFlowTest extends ResourceFlowTestCase<DomainCreateFlow, Domain
         .marshalsToXml();
   }
 
+  @Test
+  void testSuccess_usesDefaultToken() throws Exception {
+    persistContactsAndHosts();
+    AllocationToken defaultToken1 =
+        persistResource(
+            new AllocationToken.Builder()
+                .setToken("aaaaa")
+                .setTokenType(DEFAULT_PROMO)
+                .setAllowedRegistrarIds(ImmutableSet.of("NewRegistrar"))
+                .setAllowedTlds(ImmutableSet.of("tld"))
+                .build());
+    AllocationToken defaultToken2 =
+        persistResource(
+            new AllocationToken.Builder()
+                .setToken("bbbbb")
+                .setTokenType(DEFAULT_PROMO)
+                .setAllowedRegistrarIds(ImmutableSet.of("TheRegistrar"))
+                .setAllowedTlds(ImmutableSet.of("tld"))
+                .build());
+    persistResource(
+        Registry.get("tld")
+            .asBuilder()
+            .setDefaultPromoTokens(
+                ImmutableList.of(defaultToken1.createVKey(), defaultToken2.createVKey()))
+            .build());
+    runTest_defaultToken("bbbbb");
+  }
+
+  @Test
+  void testSuccess_doesNotUseDefaultTokenWhenTokenPassedIn() throws Exception {
+    persistContactsAndHosts();
+    persistResource(
+        new AllocationToken.Builder()
+            .setToken("abc123")
+            .setTokenType(UNLIMITED_USE)
+            .setDiscountFraction(0.5)
+            .build());
+    AllocationToken defaultToken1 =
+        persistResource(
+            new AllocationToken.Builder()
+                .setToken("aaaaa")
+                .setTokenType(DEFAULT_PROMO)
+                .setAllowedRegistrarIds(ImmutableSet.of("NewRegistrar"))
+                .setAllowedTlds(ImmutableSet.of("tld"))
+                .build());
+    AllocationToken defaultToken2 =
+        persistResource(
+            new AllocationToken.Builder()
+                .setToken("bbbbb")
+                .setTokenType(DEFAULT_PROMO)
+                .setAllowedRegistrarIds(ImmutableSet.of("TheRegistrar"))
+                .setAllowedTlds(ImmutableSet.of("tld"))
+                .build());
+    persistResource(
+        Registry.get("tld")
+            .asBuilder()
+            .setDefaultPromoTokens(
+                ImmutableList.of(defaultToken1.createVKey(), defaultToken2.createVKey()))
+            .build());
+    clock.advanceOneMilli();
+    setEppInput(
+        "domain_create_allocationtoken.xml",
+        ImmutableMap.of("DOMAIN", "example.tld", "YEARS", "2"));
+    runFlowAssertResponse(
+        loadFile("domain_create_response.xml", ImmutableMap.of("DOMAIN", "example.tld")));
+    BillingEvent.OneTime billingEvent =
+        Iterables.getOnlyElement(tm().transact(() -> tm().loadAllOf(BillingEvent.OneTime.class)));
+    assertThat(billingEvent.getTargetId()).isEqualTo("example.tld");
+    assertThat(billingEvent.getCost()).isEqualTo(Money.of(USD, BigDecimal.valueOf(19.5)));
+    assertThat(billingEvent.getAllocationToken().get().getKey()).isEqualTo("abc123");
+  }
+
+  @Test
+  void testSuccess_noValidDefaultToken() throws Exception {
+    persistContactsAndHosts();
+    AllocationToken defaultToken1 =
+        persistResource(
+            new AllocationToken.Builder()
+                .setToken("aaaaa")
+                .setTokenType(DEFAULT_PROMO)
+                .setAllowedRegistrarIds(ImmutableSet.of("NewRegistrar"))
+                .setAllowedTlds(ImmutableSet.of("tld"))
+                .build());
+    AllocationToken defaultToken2 =
+        persistResource(
+            new AllocationToken.Builder()
+                .setToken("bbbbb")
+                .setTokenType(DEFAULT_PROMO)
+                .setAllowedRegistrarIds(ImmutableSet.of("OtherRegistrar"))
+                .setAllowedTlds(ImmutableSet.of("tld"))
+                .build());
+    persistResource(
+        Registry.get("tld")
+            .asBuilder()
+            .setDefaultPromoTokens(
+                ImmutableList.of(defaultToken1.createVKey(), defaultToken2.createVKey()))
+            .build());
+    doSuccessfulTest();
+  }
+
+  void testSuccess_onlyUseFirstValidDefaultToken() throws Exception {
+    persistContactsAndHosts();
+    AllocationToken defaultToken1 =
+        persistResource(
+            new AllocationToken.Builder()
+                .setToken("aaaaa")
+                .setTokenType(DEFAULT_PROMO)
+                .setAllowedRegistrarIds(ImmutableSet.of("NewRegistrar"))
+                .setAllowedTlds(ImmutableSet.of("tld"))
+                .build());
+    AllocationToken defaultToken2 =
+        persistResource(
+            new AllocationToken.Builder()
+                .setToken("bbbbb")
+                .setTokenType(DEFAULT_PROMO)
+                .setAllowedRegistrarIds(ImmutableSet.of("NewRegistrar"))
+                .setAllowedTlds(ImmutableSet.of("tld"))
+                .build());
+    persistResource(
+        Registry.get("tld")
+            .asBuilder()
+            .setDefaultPromoTokens(
+                ImmutableList.of(defaultToken1.createVKey(), defaultToken2.createVKey()))
+            .build());
+    runTest_defaultToken("aaaaa");
+  }
+
+  void testSuccess_registryHasDeletedDefaultToken() throws Exception {
+    persistContactsAndHosts();
+    AllocationToken defaultToken1 =
+        persistResource(
+            new AllocationToken.Builder()
+                .setToken("aaaaa")
+                .setTokenType(DEFAULT_PROMO)
+                .setAllowedRegistrarIds(ImmutableSet.of("NewRegistrar"))
+                .setAllowedTlds(ImmutableSet.of("tld"))
+                .build());
+    AllocationToken defaultToken2 =
+        persistResource(
+            new AllocationToken.Builder()
+                .setToken("bbbbb")
+                .setTokenType(DEFAULT_PROMO)
+                .setAllowedRegistrarIds(ImmutableSet.of("TheRegistrar"))
+                .setAllowedTlds(ImmutableSet.of("tld"))
+                .build());
+    persistResource(
+        Registry.get("tld")
+            .asBuilder()
+            .setDefaultPromoTokens(
+                ImmutableList.of(defaultToken1.createVKey(), defaultToken2.createVKey()))
+            .build());
+    DatabaseHelper.deleteResource(defaultToken1);
+    runTest_defaultToken("bbbbb");
+  }
+
+  @Test
+  void testSuccess_defaultTokenAppliesCorrectPrice() throws Exception {
+    persistContactsAndHosts();
+    AllocationToken defaultToken1 =
+        persistResource(
+            new AllocationToken.Builder()
+                .setToken("aaaaa")
+                .setTokenType(DEFAULT_PROMO)
+                .setAllowedRegistrarIds(ImmutableSet.of("NewRegistrar"))
+                .setAllowedTlds(ImmutableSet.of("tld"))
+                .build());
+    AllocationToken defaultToken2 =
+        persistResource(
+            new AllocationToken.Builder()
+                .setToken("bbbbb")
+                .setTokenType(DEFAULT_PROMO)
+                .setDiscountFraction(0.5)
+                .setAllowedRegistrarIds(ImmutableSet.of("TheRegistrar"))
+                .setAllowedTlds(ImmutableSet.of("tld"))
+                .build());
+    persistResource(
+        Registry.get("tld")
+            .asBuilder()
+            .setDefaultPromoTokens(
+                ImmutableList.of(defaultToken1.createVKey(), defaultToken2.createVKey()))
+            .build());
+    BillingEvent.OneTime billingEvent = runTest_defaultToken("bbbbb");
+    assertThat(billingEvent.getCost()).isEqualTo(Money.of(USD, BigDecimal.valueOf(19.5)));
+  }
+
+  @Test
+  void testSuccess_skipsOverMissingDefaultToken() throws Exception {
+    persistContactsAndHosts();
+    AllocationToken defaultToken1 =
+        persistResource(
+            new AllocationToken.Builder()
+                .setToken("aaaaa")
+                .setTokenType(DEFAULT_PROMO)
+                .setAllowedRegistrarIds(ImmutableSet.of("NewRegistrar"))
+                .setAllowedTlds(ImmutableSet.of("tld"))
+                .build());
+    AllocationToken defaultToken2 =
+        persistResource(
+            new AllocationToken.Builder()
+                .setToken("bbbbb")
+                .setTokenType(DEFAULT_PROMO)
+                .setDiscountFraction(0.5)
+                .setAllowedRegistrarIds(ImmutableSet.of("TheRegistrar"))
+                .setAllowedTlds(ImmutableSet.of("tld"))
+                .build());
+    persistResource(
+        Registry.get("tld")
+            .asBuilder()
+            .setDefaultPromoTokens(
+                ImmutableList.of(defaultToken1.createVKey(), defaultToken2.createVKey()))
+            .build());
+    DatabaseHelper.deleteResource(defaultToken1);
+    BillingEvent.OneTime billingEvent = runTest_defaultToken("bbbbb");
+    assertThat(billingEvent.getCost()).isEqualTo(Money.of(USD, BigDecimal.valueOf(19.5)));
+  }
+
+  BillingEvent.OneTime runTest_defaultToken(String token) throws Exception {
+    setEppInput("domain_create.xml", ImmutableMap.of("DOMAIN", "example.tld"));
+    runFlowAssertResponse(
+        loadFile(
+            "domain_create_response_wildcard.xml",
+            new ImmutableMap.Builder<String, String>()
+                .put("DOMAIN", "example.tld")
+                .put("CRDATE", "1999-04-03T22:00:00.0Z")
+                .put("EXDATE", "2001-04-03T22:00:00.0Z")
+                .build()));
+    BillingEvent.OneTime billingEvent =
+        Iterables.getOnlyElement(tm().transact(() -> tm().loadAllOf(BillingEvent.OneTime.class)));
+    assertThat(billingEvent.getTargetId()).isEqualTo("example.tld");
+    assertThat(billingEvent.getAllocationToken().get().getKey()).isEqualTo(token);
+    return billingEvent;
+  }
+
   @Test
   void testSuccess_superuserReserved() throws Exception {
     setEppInput("domain_create_reserved.xml");