diff --git a/core/src/main/java/google/registry/flows/domain/token/AllocationTokenFlowUtils.java b/core/src/main/java/google/registry/flows/domain/token/AllocationTokenFlowUtils.java index 702cf6339..af48dc2e1 100644 --- a/core/src/main/java/google/registry/flows/domain/token/AllocationTokenFlowUtils.java +++ b/core/src/main/java/google/registry/flows/domain/token/AllocationTokenFlowUtils.java @@ -277,11 +277,14 @@ public class AllocationTokenFlowUtils { return maybeTokenEntity.get(); } - maybeTokenEntity = AllocationToken.get(VKey.create(AllocationToken.class, token)); - if (maybeTokenEntity.isEmpty()) { - throw new NonexistentAllocationTokenException(); + VKey tokenKey = VKey.create(AllocationToken.class, token); + AllocationToken tokenEntity = + AllocationToken.get(tokenKey).orElseThrow(NonexistentAllocationTokenException::new); + if (tokenEntity.getTokenType().equals(AllocationToken.TokenType.SINGLE_USE)) { + // Reload the token to avoid possible cache race conditions where the token may have already + // been redeemed + tokenEntity = tm().loadByKey(tokenKey); } - AllocationToken tokenEntity = maybeTokenEntity.get(); validateTokenEntity(tokenEntity, registrarId, domainName, now); return tokenEntity; } diff --git a/core/src/main/java/google/registry/model/domain/token/AllocationToken.java b/core/src/main/java/google/registry/model/domain/token/AllocationToken.java index 626647ded..bdb6ee803 100644 --- a/core/src/main/java/google/registry/model/domain/token/AllocationToken.java +++ b/core/src/main/java/google/registry/model/domain/token/AllocationToken.java @@ -52,6 +52,7 @@ import google.registry.model.reporting.HistoryEntry.HistoryEntryId; import google.registry.persistence.VKey; import google.registry.persistence.WithVKey; import google.registry.persistence.converter.AllocationTokenStatusTransitionUserType; +import google.registry.util.NonFinalForTesting; import jakarta.persistence.AttributeOverride; import jakarta.persistence.AttributeOverrides; import jakarta.persistence.Column; @@ -61,6 +62,7 @@ import jakarta.persistence.Enumerated; import jakarta.persistence.Id; import jakarta.persistence.Index; import jakarta.persistence.Table; +import java.time.Duration; import java.time.Instant; import java.util.Map; import java.util.Optional; @@ -374,28 +376,39 @@ public class AllocationToken extends UpdateAutoTimestampEntity implements Builda return ALLOCATION_TOKENS_CACHE.getAll(keys); } - /** A cache that loads the {@link AllocationToken} object for a given AllocationToken VKey. */ - private static final LoadingCache, Optional> - ALLOCATION_TOKENS_CACHE = - CacheUtils.newCacheBuilder(getSingletonCacheRefreshDuration()) - .build( - new CacheLoader<>() { - @Override - public Optional load(VKey key) { - return tm().reTransact(() -> tm().loadByKeyIfPresent(key)); - } + @VisibleForTesting + public static void setCacheForTest(Optional expiry) { + Duration effectiveExpiry = expiry.orElse(getSingletonCacheRefreshDuration()); + ALLOCATION_TOKENS_CACHE = createAllocationTokensCache(effectiveExpiry); + } - @Override - public Map, ? extends Optional> - loadAll(Set> keys) { - return tm().reTransact( - () -> - keys.stream() - .collect( - toImmutableMap( - key -> key, key -> tm().loadByKeyIfPresent(key)))); - } - }); + /** A cache that loads the {@link AllocationToken} object for a given AllocationToken VKey. */ + @NonFinalForTesting + private static LoadingCache, Optional> + ALLOCATION_TOKENS_CACHE = createAllocationTokensCache(getSingletonCacheRefreshDuration()); + + private static LoadingCache, Optional> + createAllocationTokensCache(Duration expiry) { + return CacheUtils.newCacheBuilder(expiry) + .build( + new CacheLoader<>() { + @Override + public Optional load(VKey key) { + return tm().reTransact(() -> tm().loadByKeyIfPresent(key)); + } + + @Override + public Map, ? extends Optional> + loadAll(Set> keys) { + return tm().reTransact( + () -> + keys.stream() + .collect( + toImmutableMap( + key -> key, key -> tm().loadByKeyIfPresent(key)))); + } + }); + } @Override public VKey createVKey() { diff --git a/core/src/test/java/google/registry/flows/domain/token/AllocationTokenFlowUtilsTest.java b/core/src/test/java/google/registry/flows/domain/token/AllocationTokenFlowUtilsTest.java index 6bdcc0617..df8ab722b 100644 --- a/core/src/test/java/google/registry/flows/domain/token/AllocationTokenFlowUtilsTest.java +++ b/core/src/test/java/google/registry/flows/domain/token/AllocationTokenFlowUtilsTest.java @@ -22,6 +22,7 @@ import static google.registry.model.domain.token.AllocationToken.TokenStatus.VAL import static google.registry.model.domain.token.AllocationToken.TokenType.DEFAULT_PROMO; import static google.registry.model.domain.token.AllocationToken.TokenType.SINGLE_USE; import static google.registry.model.domain.token.AllocationToken.TokenType.UNLIMITED_USE; +import static google.registry.persistence.transaction.TransactionManagerFactory.tm; import static google.registry.testing.DatabaseHelper.createTld; import static google.registry.testing.DatabaseHelper.persistResource; import static google.registry.testing.EppExceptionSubject.assertAboutEppExceptions; @@ -45,6 +46,7 @@ import google.registry.flows.custom.DomainPricingCustomLogic; import google.registry.flows.domain.DomainPricingLogic; import google.registry.flows.domain.token.AllocationTokenFlowUtils.AllocationTokenNotInPromotionException; import google.registry.flows.domain.token.AllocationTokenFlowUtils.AllocationTokenNotValidForRegistrarException; +import google.registry.flows.domain.token.AllocationTokenFlowUtils.AlreadyRedeemedAllocationTokenException; import google.registry.flows.domain.token.AllocationTokenFlowUtils.NonexistentAllocationTokenException; import google.registry.model.domain.fee.FeeQueryCommandExtensionItem.CommandName; import google.registry.model.domain.token.AllocationToken; @@ -55,6 +57,7 @@ import google.registry.model.tld.Tld; import google.registry.persistence.transaction.JpaTestExtensions; import google.registry.persistence.transaction.JpaTestExtensions.JpaIntegrationTestExtension; import google.registry.testing.FakeClock; +import google.registry.testing.TestCacheExtension; import java.time.Duration; import java.time.Instant; import java.util.Optional; @@ -71,6 +74,10 @@ class AllocationTokenFlowUtilsTest { final JpaIntegrationTestExtension jpa = new JpaTestExtensions.Builder().withClock(clock).buildIntegrationTestExtension(); + @RegisterExtension + final TestCacheExtension testCacheExtension = + new TestCacheExtension.Builder().withAllocationTokenCache(Duration.ofMinutes(10)).build(); + private final AllocationTokenExtension allocationTokenExtension = mock(AllocationTokenExtension.class); @@ -131,8 +138,13 @@ class AllocationTokenFlowUtilsTest { .build()); when(allocationTokenExtension.getAllocationToken()).thenReturn("tokeN"); assertThat( - AllocationTokenFlowUtils.loadAllocationTokenFromExtension( - "TheRegistrar", "example.tld", clock.now(), Optional.of(allocationTokenExtension))) + tm().transact( + () -> + AllocationTokenFlowUtils.loadAllocationTokenFromExtension( + "TheRegistrar", + "example.tld", + clock.now(), + Optional.of(allocationTokenExtension)))) .hasValue(token); } @@ -148,15 +160,17 @@ class AllocationTokenFlowUtilsTest { .build()); when(allocationTokenExtension.getAllocationToken()).thenReturn("tokeN"); assertThat( - AllocationTokenFlowUtils.loadTokenFromExtensionOrGetDefault( - "TheRegistrar", - clock.now(), - Optional.of(allocationTokenExtension), - tld, - "example.tld", - CommandName.CREATE, - Optional.of(1), - domainPricingLogic)) + tm().transact( + () -> + AllocationTokenFlowUtils.loadTokenFromExtensionOrGetDefault( + "TheRegistrar", + clock.now(), + Optional.of(allocationTokenExtension), + tld, + "example.tld", + CommandName.CREATE, + Optional.of(1), + domainPricingLogic))) .hasValue(token); } @@ -188,15 +202,17 @@ class AllocationTokenFlowUtilsTest { .build()); when(allocationTokenExtension.getAllocationToken()).thenReturn("tokeN"); assertThat( - AllocationTokenFlowUtils.loadTokenFromExtensionOrGetDefault( - "TheRegistrar", - clock.now(), - Optional.of(allocationTokenExtension), - tld, - "example.tld", - CommandName.CREATE, - Optional.of(1), - domainPricingLogic)) + tm().transact( + () -> + AllocationTokenFlowUtils.loadTokenFromExtensionOrGetDefault( + "TheRegistrar", + clock.now(), + Optional.of(allocationTokenExtension), + tld, + "example.tld", + CommandName.CREATE, + Optional.of(1), + domainPricingLogic))) .hasValue(defaultToken); } @@ -268,6 +284,29 @@ class AllocationTokenFlowUtilsTest { assertLoadTokenFromExtensionThrowsException(NonexistentAllocationTokenException.class); } + @Test + void testFailure_loadFromExtension_alreadyRedeemedToken() { + persistResource( + singleUseTokenBuilder().setRedemptionHistoryId(new HistoryEntryId("repoId", 1L)).build()); + assertLoadTokenFromExtensionThrowsException(AlreadyRedeemedAllocationTokenException.class); + } + + @Test + void testFailure_loadFromExtension_singleUseTokenStaleCacheReloadsFromDatabase() { + AllocationToken token = persistResource(singleUseTokenBuilder().build()); + assertThat(AllocationToken.get(token.createVKey())).hasValue(token); + tm().transact( + () -> + tm().put( + token + .asBuilder() + .setRedemptionHistoryId(new HistoryEntryId("repoId", 1L)) + .build())); + // cache still returns the old un-redeemed token + assertThat(AllocationToken.get(token.createVKey()).get().isRedeemed()).isFalse(); + assertLoadTokenFromExtensionThrowsException(AlreadyRedeemedAllocationTokenException.class); + } + @Test void testFailure_tokenInvalidForRegistrar() { persistResource( @@ -341,18 +380,20 @@ class AllocationTokenFlowUtilsTest { // Tokens tied to a domain should throw a catastrophic exception if used for a different domain persistResource(singleUseTokenBuilder().setDomainName("someotherdomain.tld").build()); when(allocationTokenExtension.getAllocationToken()).thenReturn("tokeN"); - assertThrows( - AllocationTokenFlowUtils.AllocationTokenNotValidForDomainException.class, - () -> - AllocationTokenFlowUtils.loadTokenFromExtensionOrGetDefault( - "TheRegistrar", - clock.now(), - Optional.of(allocationTokenExtension), - tld, - "example.tld", - CommandName.CREATE, - Optional.of(1), - domainPricingLogic)); + tm().transact( + () -> + assertThrows( + AllocationTokenFlowUtils.AllocationTokenNotValidForDomainException.class, + () -> + AllocationTokenFlowUtils.loadTokenFromExtensionOrGetDefault( + "TheRegistrar", + clock.now(), + Optional.of(allocationTokenExtension), + tld, + "example.tld", + CommandName.CREATE, + Optional.of(1), + domainPricingLogic))); } @Test @@ -458,17 +499,19 @@ class AllocationTokenFlowUtilsTest { } private void assertLoadTokenFromExtensionThrowsException(Class clazz) { - assertAboutEppExceptions() - .that( - assertThrows( - clazz, - () -> - AllocationTokenFlowUtils.loadAllocationTokenFromExtension( - "TheRegistrar", - "example.tld", - clock.now(), - Optional.of(allocationTokenExtension)))) - .marshalsToXml(); + tm().transact( + () -> + assertAboutEppExceptions() + .that( + assertThrows( + clazz, + () -> + AllocationTokenFlowUtils.loadAllocationTokenFromExtension( + "TheRegistrar", + "example.tld", + tm().getTxTime(), + Optional.of(allocationTokenExtension)))) + .marshalsToXml()); } private AllocationToken.Builder singleUseTokenBuilder() { diff --git a/core/src/test/java/google/registry/testing/TestCacheExtension.java b/core/src/test/java/google/registry/testing/TestCacheExtension.java index 267177cb4..9c623e3fa 100644 --- a/core/src/test/java/google/registry/testing/TestCacheExtension.java +++ b/core/src/test/java/google/registry/testing/TestCacheExtension.java @@ -17,6 +17,7 @@ package google.registry.testing; import com.google.common.collect.ImmutableList; import google.registry.model.EppResource; import google.registry.model.ForeignKeyUtils; +import google.registry.model.domain.token.AllocationToken; import google.registry.model.tld.label.PremiumListDao; import google.registry.model.tmch.ClaimsListDao; import java.time.Duration; @@ -80,6 +81,11 @@ public class TestCacheExtension implements BeforeEachCallback, AfterEachCallback return this; } + public Builder withAllocationTokenCache(Duration expiry) { + cacheHandlers.add(new TestCacheHandler(AllocationToken::setCacheForTest, expiry)); + return this; + } + public TestCacheExtension build() { return new TestCacheExtension(ImmutableList.copyOf(cacheHandlers)); }