diff --git a/core/src/main/java/google/registry/config/RegistryConfig.java b/core/src/main/java/google/registry/config/RegistryConfig.java index 0affb980f8d..76823c0d9f3 100644 --- a/core/src/main/java/google/registry/config/RegistryConfig.java +++ b/core/src/main/java/google/registry/config/RegistryConfig.java @@ -1172,6 +1172,21 @@ public static Optional provideDomainDropListDriveFolderId( return Optional.ofNullable(config.registryPolicy.domainDropListDriveFolderId); } + /** Returns the duration of the throttling window for domain:create requests. */ + @Provides + @Config("domainCreateThrottleWindowDuration") + public static Duration provideDomainCreateThrottleWindowDuration( + RegistryConfigSettings config) { + return Duration.ofSeconds(config.registryPolicy.domainCreateThrottleWindowDurationSeconds); + } + + /** Returns the number of tokens allowed per throttling window for domain:create requests. */ + @Provides + @Config("domainCreateThrottleWindowTokens") + public static int provideDomainCreateThrottleWindowTokens(RegistryConfigSettings config) { + return config.registryPolicy.domainCreateThrottleWindowTokens; + } + @Singleton @Provides static RegistryConfigSettings provideRegistryConfigSettings() { diff --git a/core/src/main/java/google/registry/config/RegistryConfigSettings.java b/core/src/main/java/google/registry/config/RegistryConfigSettings.java index 68e2e167566..1cdf2540080 100644 --- a/core/src/main/java/google/registry/config/RegistryConfigSettings.java +++ b/core/src/main/java/google/registry/config/RegistryConfigSettings.java @@ -108,6 +108,8 @@ public static class RegistryPolicy { public double sunriseDomainCreateDiscount; public Set tieredPricingPromotionRegistrarIds; public Set noPollMessageOnDeletionRegistrarIds; + public int domainCreateThrottleWindowDurationSeconds; + public int domainCreateThrottleWindowTokens; } public static class DomainExpiryAccessPeriod { diff --git a/core/src/main/java/google/registry/config/files/default-config.yaml b/core/src/main/java/google/registry/config/files/default-config.yaml index 9b903f02238..9370aa6b30a 100644 --- a/core/src/main/java/google/registry/config/files/default-config.yaml +++ b/core/src/main/java/google/registry/config/files/default-config.yaml @@ -198,6 +198,12 @@ registryPolicy: # deletions. noPollMessageOnDeletionRegistrarIds: [] + # Number of seconds for the domain:create throttling window. + domainCreateThrottleWindowDurationSeconds: 10 + + # Number of domain:create requests allowed per throttling window. + domainCreateThrottleWindowTokens: 3 + hibernate: # If set to false, calls to tm().transact() cannot be nested. If set to true, # nested calls to tm().transact() are allowed, as long as they do not specify diff --git a/core/src/main/java/google/registry/flows/FlowRunner.java b/core/src/main/java/google/registry/flows/FlowRunner.java index b93191a8ff2..b1018f23b3b 100644 --- a/core/src/main/java/google/registry/flows/FlowRunner.java +++ b/core/src/main/java/google/registry/flows/FlowRunner.java @@ -23,8 +23,10 @@ import google.registry.flows.FlowModule.RegistrarId; import google.registry.flows.FlowModule.Superuser; import google.registry.flows.FlowModule.Transactional; +import google.registry.flows.quota.FlowQuotaManager; import google.registry.flows.session.LoginFlow; import google.registry.model.eppcommon.Trid; +import google.registry.model.eppinput.EppInput; import google.registry.model.eppoutput.EppOutput; import google.registry.monitoring.whitebox.EppMetric; import google.registry.persistence.PersistenceModule.TransactionIsolationLevel; @@ -56,6 +58,8 @@ public class FlowRunner { @Inject Trid trid; @Inject FlowReporter flowReporter; @Inject JpaTransactionManager jpaTransactionManager; + @Inject EppInput eppInput; + @Inject FlowQuotaManager flowQuotaManager; @Inject FlowRunner() {} @@ -80,6 +84,11 @@ public EppOutput run(final EppMetric.Builder eppMetricBuilder) throws EppExcepti eppMetricBuilder.setCommandNameFromFlow(flowClass.getSimpleName()); final StopwatchLogger stopwatch = new StopwatchLogger(); + // First, acquire quota if necessary + if (!isSuperuser) { + flowQuotaManager.acquireQuota(flowClass, eppInput, registrarId); + } + // We may already be in a transaction, e.g., when invoked by DeleteExpiredDomainsAction. if (!isTransactional || jpaTransactionManager.inTransaction()) { stopwatch.tick("We're in transaction, running the flow now."); diff --git a/core/src/main/java/google/registry/flows/quota/FlowQuotaManager.java b/core/src/main/java/google/registry/flows/quota/FlowQuotaManager.java new file mode 100644 index 00000000000..db3ebe13d37 --- /dev/null +++ b/core/src/main/java/google/registry/flows/quota/FlowQuotaManager.java @@ -0,0 +1,85 @@ +// Copyright 2026 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.quota; + +import com.google.common.base.Ascii; +import com.google.common.flogger.FluentLogger; +import google.registry.flows.EppException; +import google.registry.flows.Flow; +import google.registry.flows.domain.DomainCreateFlow; +import google.registry.model.eppinput.EppInput; +import google.registry.quota.QuotaManager; +import java.time.Duration; +import javax.annotation.concurrent.ThreadSafe; + +/** Quota management for EPP flows using Redis/Valkey. */ +@ThreadSafe +public class FlowQuotaManager { + + private static final FluentLogger logger = FluentLogger.forEnclosingClass(); + + private final QuotaManager quotaManager; + private final int domainCreateThrottleWindowTokens; + private final Duration domainCreateThrottleWindowDuration; + + public FlowQuotaManager( + QuotaManager quotaManager, + int domainCreateThrottleWindowTokens, + Duration domainCreateThrottleWindowDuration) { + this.quotaManager = quotaManager; + this.domainCreateThrottleWindowTokens = domainCreateThrottleWindowTokens; + this.domainCreateThrottleWindowDuration = domainCreateThrottleWindowDuration; + } + + /** Acquires one unit of quota from the quota manager. Throws an exception on failure. */ + public void acquireQuota(Class flowClass, EppInput eppInput, String registrarId) + throws EppException { + // For now at least, we only throttle domain:create requests + if (!flowClass.equals(DomainCreateFlow.class)) { + return; + } + String quotaId = getDomainCreateQuotaId(eppInput, registrarId); + if (!quotaManager.acquireQuota( + quotaId, domainCreateThrottleWindowTokens, domainCreateThrottleWindowDuration)) { + logger.atWarning().log("Failed to acquire domain-create quota for %s", quotaId); + throw new TooManyRequestsException(); + } + } + + private String getDomainCreateQuotaId(EppInput eppInput, String registrarId) + throws MissingTargetIdException { + // Normalize the domain names to lowercase. We're not concerned about total canonicalization or + // verification, as that's caught in the flow itself. + String domainName = + Ascii.toLowerCase( + eppInput.getSingleTargetId().orElseThrow(MissingTargetIdException::new).trim()); + return String.format("%s:%s", registrarId, domainName); + } + + /** Too many requests too quickly. */ + public static class TooManyRequestsException extends EppException.CommandUseErrorException { + public TooManyRequestsException() { + super("Too many requests for this domain"); + } + } + + /** Thrown when a command target identifier (domain name) is missing. */ + public static class MissingTargetIdException + extends EppException.RequiredParameterMissingException { + public MissingTargetIdException() { + super("Required target identifier is missing"); + } + } +} diff --git a/core/src/main/java/google/registry/flows/quota/FlowQuotaModule.java b/core/src/main/java/google/registry/flows/quota/FlowQuotaModule.java new file mode 100644 index 00000000000..a7bc1f03254 --- /dev/null +++ b/core/src/main/java/google/registry/flows/quota/FlowQuotaModule.java @@ -0,0 +1,42 @@ +// Copyright 2026 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.quota; + +import dagger.Module; +import dagger.Provides; +import google.registry.config.RegistryConfig.Config; +import google.registry.quota.NoopQuotaManager; +import google.registry.quota.QuotaManager; +import google.registry.quota.ValkeyQuotaManager; +import jakarta.inject.Singleton; +import java.time.Duration; +import java.util.Optional; +import redis.clients.jedis.UnifiedJedis; + +@Module +public class FlowQuotaModule { + + @Provides + @Singleton + static FlowQuotaManager provideFlowQuotaManager( + Optional jedis, + @Config("domainCreateThrottleWindowDuration") Duration domainCreateThrottleWindowDuration, + @Config("domainCreateThrottleWindowTokens") int domainCreateThrottleWindowTokens) { + QuotaManager quotaManager = + jedis.isPresent() ? new ValkeyQuotaManager(jedis.get(), "flow") : new NoopQuotaManager(); + return new FlowQuotaManager( + quotaManager, domainCreateThrottleWindowTokens, domainCreateThrottleWindowDuration); + } +} diff --git a/core/src/main/java/google/registry/module/RegistryComponent.java b/core/src/main/java/google/registry/module/RegistryComponent.java index cc9806115b7..0e3f041484c 100644 --- a/core/src/main/java/google/registry/module/RegistryComponent.java +++ b/core/src/main/java/google/registry/module/RegistryComponent.java @@ -32,6 +32,7 @@ import google.registry.flows.ServerTridProviderModule; import google.registry.flows.custom.CustomLogicFactoryModule; import google.registry.flows.domain.DomainDeletionTimeCacheModule; +import google.registry.flows.quota.FlowQuotaModule; import google.registry.groups.DirectoryModule; import google.registry.groups.GmailModule; import google.registry.groups.GroupsModule; @@ -68,6 +69,7 @@ DirectoryModule.class, DomainDeletionTimeCacheModule.class, DriveModule.class, + FlowQuotaModule.class, GmailModule.class, GroupsModule.class, GroupssettingsModule.class, diff --git a/core/src/test/java/google/registry/flows/EppTestComponent.java b/core/src/test/java/google/registry/flows/EppTestComponent.java index 8fa0779dab9..85a3ed050b6 100644 --- a/core/src/test/java/google/registry/flows/EppTestComponent.java +++ b/core/src/test/java/google/registry/flows/EppTestComponent.java @@ -27,7 +27,9 @@ import google.registry.flows.custom.TestCustomLogicFactory; import google.registry.flows.domain.DomainDeletionTimeCache; import google.registry.flows.domain.DomainFlowTmchUtils; +import google.registry.flows.quota.FlowQuotaManager; import google.registry.monitoring.whitebox.EppMetric; +import google.registry.quota.NoopQuotaManager; import google.registry.request.Modules.GsonModule; import google.registry.request.RequestScope; import google.registry.request.lock.LockHandler; @@ -40,6 +42,7 @@ import google.registry.util.Clock; import google.registry.util.Sleeper; import jakarta.inject.Singleton; +import java.time.Duration; /** Dagger component for running EPP tests. */ @Singleton @@ -60,6 +63,7 @@ class FakesAndMocksModule { private FakeLockHandler lockHandler; private Sleeper sleeper; private CloudTasksHelper cloudTasksHelper; + private FlowQuotaManager flowQuotaManager; public CloudTasksHelper getCloudTasksHelper() { return cloudTasksHelper; @@ -69,6 +73,10 @@ public EppMetric.Builder getMetricBuilder() { return metricBuilder; } + public FlowQuotaManager getFlowQuotaManager() { + return flowQuotaManager; + } + public static FakesAndMocksModule create(FakeClock clock) { FakesAndMocksModule instance = new FakesAndMocksModule(); CloudTasksHelper cloudTasksHelper = new CloudTasksHelper(clock); @@ -82,6 +90,8 @@ public static FakesAndMocksModule create(FakeClock clock) { instance.metricBuilder = EppMetric.builderForRequest(clock); instance.lockHandler = new FakeLockHandler(true); instance.cloudTasksHelper = cloudTasksHelper; + instance.flowQuotaManager = + new FlowQuotaManager(new NoopQuotaManager(), 3, Duration.ofSeconds(10)); return instance; } @@ -134,6 +144,11 @@ ServerTridProvider provideServerTridProvider() { DomainDeletionTimeCache provideDomainDeletionTimeCache() { return DomainDeletionTimeCache.create(); } + + @Provides + FlowQuotaManager provideFlowQuotaManager() { + return flowQuotaManager; + } } class FakeServerTridProvider implements ServerTridProvider { diff --git a/core/src/test/java/google/registry/flows/FlowRunnerTest.java b/core/src/test/java/google/registry/flows/FlowRunnerTest.java index 4fdc99147ad..cf5efb7ad5a 100644 --- a/core/src/test/java/google/registry/flows/FlowRunnerTest.java +++ b/core/src/test/java/google/registry/flows/FlowRunnerTest.java @@ -17,13 +17,19 @@ import static com.google.common.truth.Truth.assertThat; import static com.google.common.truth.Truth.assertWithMessage; import static google.registry.persistence.transaction.TransactionManagerFactory.tm; +import static google.registry.testing.EppExceptionSubject.assertAboutEppExceptions; import static google.registry.testing.TestDataHelper.loadFile; import static google.registry.testing.TestLogHandlerUtils.findFirstLogMessageByPrefix; import static google.registry.util.DateTimeUtils.START_INSTANT; import static java.nio.charset.StandardCharsets.UTF_8; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyInt; +import static org.mockito.ArgumentMatchers.eq; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.never; import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; import com.google.common.base.Joiner; import com.google.common.base.Splitter; @@ -32,16 +38,23 @@ import com.google.common.net.InetAddresses; import com.google.common.testing.TestLogHandler; import google.registry.flows.certs.CertificateChecker; +import google.registry.flows.domain.DomainCreateFlow; +import google.registry.flows.quota.FlowQuotaManager; +import google.registry.flows.quota.FlowQuotaManager.TooManyRequestsException; import google.registry.model.eppcommon.Trid; +import google.registry.model.eppinput.EppInput; import google.registry.model.eppoutput.EppOutput.ResponseOrGreeting; import google.registry.model.eppoutput.EppResponse; import google.registry.monitoring.whitebox.EppMetric; import google.registry.persistence.PersistenceModule.TransactionIsolationLevel; import google.registry.persistence.transaction.JpaTestExtensions; import google.registry.persistence.transaction.JpaTestExtensions.JpaIntegrationTestExtension; +import google.registry.quota.NoopQuotaManager; +import google.registry.quota.QuotaManager; import google.registry.testing.FakeClock; import google.registry.testing.FakeHttpSession; import google.registry.util.JdkLoggerConfig; +import java.time.Duration; import java.time.Instant; import java.util.List; import java.util.Optional; @@ -112,6 +125,9 @@ void beforeEach() { flowRunner.trid = Trid.create("client-123", "server-456"); flowRunner.flowReporter = mock(FlowReporter.class); flowRunner.jpaTransactionManager = tm(); + flowRunner.eppInput = mock(EppInput.class); + flowRunner.flowQuotaManager = + new FlowQuotaManager(new NoopQuotaManager(), 3, Duration.ofSeconds(10)); } @Test @@ -224,4 +240,31 @@ void testRun_loggingStatement_complexEppInput() throws Exception { String xml = Joiner.on('\n').join(lines.subList(3, lines.size() - 4)); assertThat(xml).isEqualTo(sanitizedDomainCreateXml); } + + @Test + void testRun_quotaExceeded_throwsException() { + flowRunner.flowClass = DomainCreateFlow.class; + when(flowRunner.eppInput.getSingleTargetId()).thenReturn(Optional.of("example.tld")); + QuotaManager mockQuotaManager = mock(QuotaManager.class); + when(mockQuotaManager.acquireQuota( + eq("TheRegistrar:example.tld"), anyInt(), any(Duration.class))) + .thenReturn(false); + flowRunner.flowQuotaManager = new FlowQuotaManager(mockQuotaManager, 3, Duration.ofSeconds(10)); + assertAboutEppExceptions() + .that(assertThrows(TooManyRequestsException.class, () -> flowRunner.run(eppMetricBuilder))) + .marshalsToXml(); + } + + @Test + void testRun_quotaAvailable_succeeds() throws Exception { + flowRunner.flowClass = DomainCreateFlow.class; + when(flowRunner.eppInput.getSingleTargetId()).thenReturn(Optional.of("example.tld")); + QuotaManager mockQuotaManager = mock(QuotaManager.class); + when(mockQuotaManager.acquireQuota( + eq("TheRegistrar:example.tld"), anyInt(), any(Duration.class))) + .thenReturn(true); + flowRunner.flowQuotaManager = new FlowQuotaManager(mockQuotaManager, 3, Duration.ofSeconds(10)); + flowRunner.run(eppMetricBuilder); + verify(mockQuotaManager).acquireQuota("TheRegistrar:example.tld", 3, Duration.ofSeconds(10)); + } } diff --git a/core/src/test/java/google/registry/flows/quota/FlowQuotaManagerTest.java b/core/src/test/java/google/registry/flows/quota/FlowQuotaManagerTest.java new file mode 100644 index 00000000000..67cb54112d2 --- /dev/null +++ b/core/src/test/java/google/registry/flows/quota/FlowQuotaManagerTest.java @@ -0,0 +1,182 @@ +// Copyright 2026 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.quota; + +import static com.google.common.truth.Truth.assertThat; +import static google.registry.testing.EppExceptionSubject.assertAboutEppExceptions; +import static org.junit.jupiter.api.Assertions.assertDoesNotThrow; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +import google.registry.flows.domain.DomainCheckFlow; +import google.registry.flows.domain.DomainCreateFlow; +import google.registry.flows.quota.FlowQuotaManager.MissingTargetIdException; +import google.registry.flows.quota.FlowQuotaManager.TooManyRequestsException; +import google.registry.model.eppinput.EppInput; +import google.registry.quota.ValkeyQuotaManager; +import io.github.ss_bhatt.testcontainers.valkey.ValkeyContainer; +import java.time.Duration; +import java.util.Optional; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.testcontainers.junit.jupiter.Container; +import org.testcontainers.junit.jupiter.Testcontainers; +import redis.clients.jedis.HostAndPort; +import redis.clients.jedis.RedisClient; + +/** Tests for {@link FlowQuotaManager} backed by Valkey. */ +@Testcontainers +class FlowQuotaManagerTest { + + private static final int DEFAULT_TOKENS = 3; + private static final Duration DEFAULT_DURATION = Duration.ofSeconds(10); + + @Container private static final ValkeyContainer valkey = new ValkeyContainer(); + + private RedisClient jedis; + private ValkeyQuotaManager quotaManager; + private final EppInput eppInput = mock(EppInput.class); + + @BeforeEach + void setUp() { + jedis = + RedisClient.builder() + .hostAndPort(new HostAndPort(valkey.getHost(), valkey.getFirstMappedPort())) + .build(); + jedis.flushAll(); + quotaManager = new ValkeyQuotaManager(jedis, "flow"); + } + + @Test + void testAcquireQuota_nonDomainCreateFlow_noOp() { + FlowQuotaManager manager = new FlowQuotaManager(quotaManager, DEFAULT_TOKENS, DEFAULT_DURATION); + assertDoesNotThrow(() -> manager.acquireQuota(DomainCheckFlow.class, eppInput, "TheRegistrar")); + assertThat(jedis.keys("*")).isEmpty(); + } + + @Test + void testAcquireQuota_missingTargetId_throwsMissingTargetIdException() { + FlowQuotaManager manager = new FlowQuotaManager(quotaManager, DEFAULT_TOKENS, DEFAULT_DURATION); + when(eppInput.getSingleTargetId()).thenReturn(Optional.empty()); + MissingTargetIdException thrown = + assertThrows( + MissingTargetIdException.class, + () -> manager.acquireQuota(DomainCreateFlow.class, eppInput, "TheRegistrar")); + assertThat(thrown).hasMessageThat().contains("Required target identifier is missing"); + assertAboutEppExceptions().that(thrown).marshalsToXml(); + } + + @Test + void testAcquireQuota_success() throws Exception { + when(eppInput.getSingleTargetId()).thenReturn(Optional.of("example.tld")); + FlowQuotaManager manager = new FlowQuotaManager(quotaManager, DEFAULT_TOKENS, DEFAULT_DURATION); + + manager.acquireQuota(DomainCreateFlow.class, eppInput, "TheRegistrar"); + assertThat(jedis.get("flow:TheRegistrar:example.tld")).isEqualTo("2"); + + manager.acquireQuota(DomainCreateFlow.class, eppInput, "TheRegistrar"); + assertThat(jedis.get("flow:TheRegistrar:example.tld")).isEqualTo("1"); + } + + @Test + void testAcquireQuota_normalizesDomainNameToLowerCase() throws Exception { + when(eppInput.getSingleTargetId()).thenReturn(Optional.of("EXAMPLE.TLD")); + FlowQuotaManager manager = new FlowQuotaManager(quotaManager, DEFAULT_TOKENS, DEFAULT_DURATION); + + manager.acquireQuota(DomainCreateFlow.class, eppInput, "TheRegistrar"); + assertThat(jedis.get("flow:TheRegistrar:example.tld")).isEqualTo("2"); + } + + @Test + void testAcquireQuota_trimsWhitespace() throws Exception { + when(eppInput.getSingleTargetId()).thenReturn(Optional.of(" example.tld ")); + FlowQuotaManager manager = new FlowQuotaManager(quotaManager, DEFAULT_TOKENS, DEFAULT_DURATION); + + manager.acquireQuota(DomainCreateFlow.class, eppInput, "TheRegistrar"); + assertThat(jedis.get("flow:TheRegistrar:example.tld")).isEqualTo("2"); + } + + @Test + void testAcquireQuota_exceeded_throwsTooManyRequestsException() throws Exception { + when(eppInput.getSingleTargetId()).thenReturn(Optional.of("example.tld")); + FlowQuotaManager manager = new FlowQuotaManager(quotaManager, DEFAULT_TOKENS, DEFAULT_DURATION); + + for (int i = 0; i < 3; i++) { + manager.acquireQuota(DomainCreateFlow.class, eppInput, "TheRegistrar"); + } + + TooManyRequestsException thrown = + assertThrows( + TooManyRequestsException.class, + () -> manager.acquireQuota(DomainCreateFlow.class, eppInput, "TheRegistrar")); + assertThat(thrown).hasMessageThat().contains("Too many requests"); + assertAboutEppExceptions().that(thrown).marshalsToXml(); + } + + @Test + void testAcquireQuota_isolatedByRegistrarId() throws Exception { + when(eppInput.getSingleTargetId()).thenReturn(Optional.of("example.tld")); + FlowQuotaManager manager = new FlowQuotaManager(quotaManager, DEFAULT_TOKENS, DEFAULT_DURATION); + + for (int i = 0; i < 3; i++) { + manager.acquireQuota(DomainCreateFlow.class, eppInput, "RegistrarA"); + } + assertThrows( + TooManyRequestsException.class, + () -> manager.acquireQuota(DomainCreateFlow.class, eppInput, "RegistrarA")); + + // RegistrarB has independent quota + assertDoesNotThrow(() -> manager.acquireQuota(DomainCreateFlow.class, eppInput, "RegistrarB")); + assertThat(jedis.get("flow:RegistrarA:example.tld")).isEqualTo("0"); + assertThat(jedis.get("flow:RegistrarB:example.tld")).isEqualTo("2"); + } + + @Test + void testAcquireQuota_isolatedByDomainName() throws Exception { + FlowQuotaManager manager = new FlowQuotaManager(quotaManager, DEFAULT_TOKENS, DEFAULT_DURATION); + + when(eppInput.getSingleTargetId()).thenReturn(Optional.of("domain1.tld")); + for (int i = 0; i < 3; i++) { + manager.acquireQuota(DomainCreateFlow.class, eppInput, "TheRegistrar"); + } + assertThrows( + TooManyRequestsException.class, + () -> manager.acquireQuota(DomainCreateFlow.class, eppInput, "TheRegistrar")); + + // Different domain has independent quota + when(eppInput.getSingleTargetId()).thenReturn(Optional.of("domain2.tld")); + assertDoesNotThrow( + () -> manager.acquireQuota(DomainCreateFlow.class, eppInput, "TheRegistrar")); + assertThat(jedis.get("flow:TheRegistrar:domain1.tld")).isEqualTo("0"); + assertThat(jedis.get("flow:TheRegistrar:domain2.tld")).isEqualTo("2"); + } + + @Test + void testAcquireQuota_releaseQuotaAfterDuration() throws Exception { + FlowQuotaManager manager = new FlowQuotaManager(quotaManager, 2, Duration.ofMillis(50)); + when(eppInput.getSingleTargetId()).thenReturn(Optional.of("example.tld")); + + manager.acquireQuota(DomainCreateFlow.class, eppInput, "TheRegistrar"); + manager.acquireQuota(DomainCreateFlow.class, eppInput, "TheRegistrar"); + assertThrows( + TooManyRequestsException.class, + () -> manager.acquireQuota(DomainCreateFlow.class, eppInput, "TheRegistrar")); + + Thread.sleep(150); + assertDoesNotThrow( + () -> manager.acquireQuota(DomainCreateFlow.class, eppInput, "TheRegistrar")); + } +} diff --git a/core/src/test/java/google/registry/module/TestRegistryComponent.java b/core/src/test/java/google/registry/module/TestRegistryComponent.java index 2147ce4d3f8..ad7e1a0e2c0 100644 --- a/core/src/test/java/google/registry/module/TestRegistryComponent.java +++ b/core/src/test/java/google/registry/module/TestRegistryComponent.java @@ -28,6 +28,7 @@ import google.registry.flows.ServerTridProviderModule; import google.registry.flows.custom.CustomLogicFactoryModule; import google.registry.flows.domain.DomainDeletionTimeCacheModule; +import google.registry.flows.quota.FlowQuotaModule; import google.registry.groups.GmailModule; import google.registry.groups.GroupsModule; import google.registry.groups.GroupssettingsModule; @@ -60,6 +61,7 @@ CustomLogicFactoryModule.class, DomainDeletionTimeCacheModule.class, DriveModule.class, + FlowQuotaModule.class, GmailModule.class, GroupsModule.class, GroupssettingsModule.class,