1
0
mirror of https://github.com/google/nomulus synced 2026-07-18 05:52:30 +00:00

Reload SINGLE_USE ATs to avoid cache race conditions (#3157)

This commit is contained in:
gbrodman
2026-07-16 23:47:48 -04:00
committed by GitHub
parent 675354ae65
commit 2a04b9be9b
4 changed files with 133 additions and 68 deletions
@@ -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<AllocationToken> 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;
}
@@ -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<VKey<AllocationToken>, Optional<AllocationToken>>
ALLOCATION_TOKENS_CACHE =
CacheUtils.newCacheBuilder(getSingletonCacheRefreshDuration())
.build(
new CacheLoader<>() {
@Override
public Optional<AllocationToken> load(VKey<AllocationToken> key) {
return tm().reTransact(() -> tm().loadByKeyIfPresent(key));
}
@VisibleForTesting
public static void setCacheForTest(Optional<Duration> expiry) {
Duration effectiveExpiry = expiry.orElse(getSingletonCacheRefreshDuration());
ALLOCATION_TOKENS_CACHE = createAllocationTokensCache(effectiveExpiry);
}
@Override
public Map<? extends VKey<AllocationToken>, ? extends Optional<AllocationToken>>
loadAll(Set<? extends VKey<AllocationToken>> 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<VKey<AllocationToken>, Optional<AllocationToken>>
ALLOCATION_TOKENS_CACHE = createAllocationTokensCache(getSingletonCacheRefreshDuration());
private static LoadingCache<VKey<AllocationToken>, Optional<AllocationToken>>
createAllocationTokensCache(Duration expiry) {
return CacheUtils.newCacheBuilder(expiry)
.build(
new CacheLoader<>() {
@Override
public Optional<AllocationToken> load(VKey<AllocationToken> key) {
return tm().reTransact(() -> tm().loadByKeyIfPresent(key));
}
@Override
public Map<? extends VKey<AllocationToken>, ? extends Optional<AllocationToken>>
loadAll(Set<? extends VKey<AllocationToken>> keys) {
return tm().reTransact(
() ->
keys.stream()
.collect(
toImmutableMap(
key -> key, key -> tm().loadByKeyIfPresent(key))));
}
});
}
@Override
public VKey<AllocationToken> createVKey() {
@@ -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<? extends EppException> 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() {
@@ -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));
}