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..d0f99f459f5 --- /dev/null +++ b/core/src/main/java/google/registry/flows/quota/FlowQuotaManager.java @@ -0,0 +1,81 @@ +// 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.collect.ImmutableList; +import com.google.common.collect.ImmutableMap; +import com.google.common.flogger.FluentLogger; +import google.registry.flows.EppException; +import google.registry.flows.Flow; +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. + * + *
This is provided the base {@link QuotaManager} as well as a list of {@link
+ * FlowQuotaParameters}. The latter is used to define per-flow throttling characteristics so that we
+ * can check quota before the flow class itself is instantiated.
+ */
+@ThreadSafe
+public class FlowQuotaManager {
+
+ private static final FluentLogger logger = FluentLogger.forEnclosingClass();
+
+ private final QuotaManager quotaManager;
+ private final ImmutableMap Normally, we'd wish to define these in the flow class itself by using some interface that
+ * might be named something like "ThrottlingFlow". However, the flow class itself is not
+ * instantiated until a transaction has been opened (if the flow is a transactional flow). This is
+ * suboptimal, given that one of the primary purposes of throttling is to reduce database load. As a
+ * result, we define the throttling parameters outside the flow class itself so the throttling /
+ * quota can be checked before the class is actually instantiated.
+ */
+public interface FlowQuotaParameters {
+
+ /** The flow that is being throttled. */
+ Class extends Flow> getFlowClass();
+
+ /**
+ * The ID of the quota, e.g. a domain-name-registrar-ID combo.
+ *
+ * Note: the registrar ID may be empty if we're not authenticated yet, like a LoginFlow.
+ */
+ String getQuotaId(EppInput eppInput, String registrarId);
+
+ /** The maximum number of tokens / quota for each ID for this flow. */
+ int getMaxQuotaAllowed();
+
+ /** The fixed window over which to throttle requests with the same ID. */
+ Duration getWindowDuration();
+}
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..f90a2e1314b 100644
--- a/core/src/test/java/google/registry/flows/EppTestComponent.java
+++ b/core/src/test/java/google/registry/flows/EppTestComponent.java
@@ -14,6 +14,7 @@
package google.registry.flows;
+import com.google.common.collect.ImmutableList;
import dagger.Component;
import dagger.Module;
import dagger.Provides;
@@ -27,7 +28,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;
@@ -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 =
+ FlowQuotaManager.create(new NoopQuotaManager(), ImmutableList.of());
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..00ee9906352 100644
--- a/core/src/test/java/google/registry/flows/FlowRunnerTest.java
+++ b/core/src/test/java/google/registry/flows/FlowRunnerTest.java
@@ -21,27 +21,40 @@
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;
+import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableSet;
import com.google.common.collect.ImmutableSortedMap;
import com.google.common.net.InetAddresses;
import com.google.common.testing.TestLogHandler;
import google.registry.flows.certs.CertificateChecker;
+import google.registry.flows.quota.FlowQuotaManager;
+import google.registry.flows.quota.FlowQuotaManager.TooManyRequestsException;
+import google.registry.flows.quota.FlowQuotaParameters;
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 =
+ FlowQuotaManager.create(new NoopQuotaManager(), ImmutableList.of());
}
@Test
@@ -224,4 +240,51 @@ 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() {
+ QuotaManager mockQuotaManager = mock(QuotaManager.class);
+ when(mockQuotaManager.acquireQuota(
+ eq("TheRegistrar:test-quota-id"), anyInt(), any(Duration.class)))
+ .thenReturn(false);
+ flowRunner.flowQuotaManager =
+ FlowQuotaManager.create(mockQuotaManager, ImmutableList.of(new TestFlowQuotaParameters()));
+
+ assertThrows(TooManyRequestsException.class, () -> flowRunner.run(eppMetricBuilder));
+ }
+
+ @Test
+ void testRun_quotaAvailable_succeeds() throws Exception {
+ QuotaManager mockQuotaManager = mock(QuotaManager.class);
+ when(mockQuotaManager.acquireQuota(
+ eq("TheRegistrar:test-quota-id"), anyInt(), any(Duration.class)))
+ .thenReturn(true);
+ flowRunner.flowQuotaManager =
+ FlowQuotaManager.create(mockQuotaManager, ImmutableList.of(new TestFlowQuotaParameters()));
+
+ flowRunner.run(eppMetricBuilder);
+ verify(mockQuotaManager).acquireQuota("TheRegistrar:test-quota-id", 10, Duration.ofMinutes(1));
+ }
+
+ private static class TestFlowQuotaParameters implements FlowQuotaParameters {
+ @Override
+ public Class extends Flow> getFlowClass() {
+ return TestCommandFlow.class;
+ }
+
+ @Override
+ public String getQuotaId(EppInput eppInput, String registrarId) {
+ return registrarId + ":test-quota-id";
+ }
+
+ @Override
+ public int getMaxQuotaAllowed() {
+ return 10;
+ }
+
+ @Override
+ public Duration getWindowDuration() {
+ return Duration.ofMinutes(1);
+ }
+ }
}
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..31b7f5cbe2d
--- /dev/null
+++ b/core/src/test/java/google/registry/flows/quota/FlowQuotaManagerTest.java
@@ -0,0 +1,180 @@
+// 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 com.google.common.collect.ImmutableList;
+import google.registry.flows.Flow;
+import google.registry.flows.quota.FlowQuotaManager.TooManyRequestsException;
+import google.registry.model.eppinput.EppInput;
+import google.registry.model.eppoutput.EppOutput.ResponseOrGreeting;
+import google.registry.quota.ValkeyQuotaManager;
+import io.github.ss_bhatt.testcontainers.valkey.ValkeyContainer;
+import java.time.Duration;
+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 {
+
+ @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_noConfiguredParameters_noOp() {
+ FlowQuotaManager manager = FlowQuotaManager.create(quotaManager, ImmutableList.of());
+ assertDoesNotThrow(() -> manager.acquireQuota(FlowA.class, eppInput, "TheRegistrar"));
+ assertThat(jedis.keys("*")).isEmpty();
+ }
+
+ @Test
+ void testAcquireQuota_unconfiguredFlow_noOp() {
+ FlowQuotaParameters paramA =
+ createParameters(FlowA.class, "quota-a", 10, Duration.ofMinutes(1));
+ FlowQuotaManager manager = FlowQuotaManager.create(quotaManager, ImmutableList.of(paramA));
+ assertDoesNotThrow(() -> manager.acquireQuota(FlowB.class, eppInput, "TheRegistrar"));
+ assertThat(jedis.keys("*")).isEmpty();
+ }
+
+ @Test
+ void testAcquireQuota_success() throws Exception {
+ FlowQuotaParameters paramA = createParameters(FlowA.class, "quota-a", 2, Duration.ofMinutes(1));
+ FlowQuotaManager manager = FlowQuotaManager.create(quotaManager, ImmutableList.of(paramA));
+
+ manager.acquireQuota(FlowA.class, eppInput, "TheRegistrar");
+ assertThat(jedis.get("flow:TheRegistrar:quota-a")).isEqualTo("1");
+
+ manager.acquireQuota(FlowA.class, eppInput, "TheRegistrar");
+ assertThat(jedis.get("flow:TheRegistrar:quota-a")).isEqualTo("0");
+ }
+
+ @Test
+ void testAcquireQuota_exceeded_throwsTooManyRequestsException() throws Exception {
+ FlowQuotaParameters paramA = createParameters(FlowA.class, "quota-a", 2, Duration.ofMinutes(1));
+ FlowQuotaManager manager = FlowQuotaManager.create(quotaManager, ImmutableList.of(paramA));
+
+ manager.acquireQuota(FlowA.class, eppInput, "TheRegistrar");
+ manager.acquireQuota(FlowA.class, eppInput, "TheRegistrar");
+
+ TooManyRequestsException thrown =
+ assertThrows(
+ TooManyRequestsException.class,
+ () -> manager.acquireQuota(FlowA.class, eppInput, "TheRegistrar"));
+ assertThat(thrown).hasMessageThat().contains("Too many requests");
+ assertAboutEppExceptions().that(thrown).marshalsToXml();
+ }
+
+ @Test
+ void testAcquireQuota_releaseQuotaAfterDuration() throws Exception {
+ FlowQuotaParameters paramA = createParameters(FlowA.class, "quota-a", 2, Duration.ofMillis(10));
+ FlowQuotaManager manager = FlowQuotaManager.create(quotaManager, ImmutableList.of(paramA));
+
+ manager.acquireQuota(FlowA.class, eppInput, "TheRegistrar");
+ Thread.sleep(30);
+ assertDoesNotThrow(() -> manager.acquireQuota(FlowA.class, eppInput, "TheRegistrar"));
+ }
+
+ @Test
+ void testAcquireQuota_isolatedByRegistrarId() throws Exception {
+ FlowQuotaParameters paramA = createParameters(FlowA.class, "quota-a", 1, Duration.ofMinutes(1));
+ FlowQuotaManager manager = FlowQuotaManager.create(quotaManager, ImmutableList.of(paramA));
+
+ manager.acquireQuota(FlowA.class, eppInput, "RegistrarA");
+ assertThrows(
+ TooManyRequestsException.class,
+ () -> manager.acquireQuota(FlowA.class, eppInput, "RegistrarA"));
+
+ // RegistrarB has independent quota
+ assertDoesNotThrow(() -> manager.acquireQuota(FlowA.class, eppInput, "RegistrarB"));
+ assertThat(jedis.get("flow:RegistrarA:quota-a")).isEqualTo("0");
+ assertThat(jedis.get("flow:RegistrarB:quota-a")).isEqualTo("0");
+ }
+
+ @Test
+ void testAcquireQuota_multipleFlowParameters() throws Exception {
+ FlowQuotaParameters paramA = createParameters(FlowA.class, "quota-a", 2, Duration.ofMinutes(1));
+ FlowQuotaParameters paramB = createParameters(FlowB.class, "quota-b", 5, Duration.ofMinutes(5));
+ FlowQuotaManager manager =
+ FlowQuotaManager.create(quotaManager, ImmutableList.of(paramA, paramB));
+
+ manager.acquireQuota(FlowA.class, eppInput, "TheRegistrar");
+ manager.acquireQuota(FlowB.class, eppInput, "TheRegistrar");
+
+ assertThat(jedis.get("flow:TheRegistrar:quota-a")).isEqualTo("1");
+ assertThat(jedis.get("flow:TheRegistrar:quota-b")).isEqualTo("4");
+ }
+
+ static class FlowA implements Flow {
+ @Override
+ public ResponseOrGreeting run() {
+ return null;
+ }
+ }
+
+ static class FlowB implements Flow {
+ @Override
+ public ResponseOrGreeting run() {
+ return null;
+ }
+ }
+
+ private static FlowQuotaParameters createParameters(
+ Class extends Flow> flowClass, String quotaPrefix, int maxQuota, Duration window) {
+ return new FlowQuotaParameters() {
+ @Override
+ public Class extends Flow> getFlowClass() {
+ return flowClass;
+ }
+
+ @Override
+ public String getQuotaId(EppInput input, String registrarId) {
+ return registrarId + ":" + quotaPrefix;
+ }
+
+ @Override
+ public int getMaxQuotaAllowed() {
+ return maxQuota;
+ }
+
+ @Override
+ public Duration getWindowDuration() {
+ return window;
+ }
+ };
+ }
+}
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,