From 0eb76b5b2c20d3ac351ae303f2b6af7c50b97291 Mon Sep 17 00:00:00 2001 From: Gus Brodman Date: Fri, 14 Aug 2026 12:56:34 -0400 Subject: [PATCH] Split QuotaManager into generic and EppServer classes This separates the logic of "how many tokens should a particular user/group have" from "manage the quota given whatever limits, contacting Valkey". This is in preparation for allowing other types of quota management and throttling besides just on the EPP server (e.g. domain-create throttling). We also convert the expirations from seconds to milliseconds (and use a Duration so that this is masked from users). It's better to have the API use a full-fledged Duration object rather than an int, and this allows for finer control over expiration times. Note that we'll probably want to use a sliding window in the future instead of a fixed window, but that's a problem for future us. --- .../registry/eppserver/EppProtocolModule.java | 8 +- .../eppserver/handler/EppServiceHandler.java | 8 +- .../quota/EppServerQuotaManager.java | 91 ++++++++++ .../eppserver/quota/QuotaManager.java | 171 ------------------ .../quota/GenericValkeyQuotaManager.java | 132 ++++++++++++++ .../handler/EppServiceHandlerTest.java | 32 ++-- .../quota/EppServerQuotaManagerTest.java | 126 +++++++++++++ .../eppserver/quota/QuotaManagerTest.java | 159 ---------------- .../quota/GenericValkeyQuotaManagerTest.java | 170 +++++++++++++++++ 9 files changed, 539 insertions(+), 358 deletions(-) create mode 100644 core/src/main/java/google/registry/eppserver/quota/EppServerQuotaManager.java delete mode 100644 core/src/main/java/google/registry/eppserver/quota/QuotaManager.java create mode 100644 core/src/main/java/google/registry/quota/GenericValkeyQuotaManager.java create mode 100644 core/src/test/java/google/registry/eppserver/quota/EppServerQuotaManagerTest.java delete mode 100644 core/src/test/java/google/registry/eppserver/quota/QuotaManagerTest.java create mode 100644 core/src/test/java/google/registry/quota/GenericValkeyQuotaManagerTest.java diff --git a/core/src/main/java/google/registry/eppserver/EppProtocolModule.java b/core/src/main/java/google/registry/eppserver/EppProtocolModule.java index 9529117ec31..c8670c15edb 100644 --- a/core/src/main/java/google/registry/eppserver/EppProtocolModule.java +++ b/core/src/main/java/google/registry/eppserver/EppProtocolModule.java @@ -24,8 +24,9 @@ import google.registry.config.RegistryConfigSettings; import google.registry.eppserver.Protocol.FrontendProtocol; import google.registry.eppserver.handler.EppServiceHandler; -import google.registry.eppserver.quota.QuotaManager; +import google.registry.eppserver.quota.EppServerQuotaManager; import google.registry.networking.handler.SslServerInitializer; +import google.registry.quota.GenericValkeyQuotaManager; import io.netty.channel.ChannelHandler; import io.netty.channel.socket.nio.NioSocketChannel; import io.netty.handler.codec.LengthFieldBasedFrameDecoder; @@ -137,8 +138,9 @@ static SslServerInitializer provideSslServerInitializer( @Provides @Singleton @CommandQuota - static QuotaManager provideCommandQuotaManager( + static EppServerQuotaManager provideCommandQuotaManager( @Config("eppServerQuota") RegistryConfigSettings.Quota quota, Optional jedis) { - return new QuotaManager(quota, jedis.orElse(null), "command"); + return new EppServerQuotaManager( + quota, new GenericValkeyQuotaManager(jedis.orElse(null), "command")); } } diff --git a/core/src/main/java/google/registry/eppserver/handler/EppServiceHandler.java b/core/src/main/java/google/registry/eppserver/handler/EppServiceHandler.java index b4e51915aa6..d3f8f9562c5 100644 --- a/core/src/main/java/google/registry/eppserver/handler/EppServiceHandler.java +++ b/core/src/main/java/google/registry/eppserver/handler/EppServiceHandler.java @@ -26,8 +26,8 @@ import google.registry.config.RegistryConfig.Config; import google.registry.eppserver.EppProtocolModule.CommandQuota; import google.registry.eppserver.metric.FrontendMetrics; +import google.registry.eppserver.quota.EppServerQuotaManager; import google.registry.eppserver.quota.LocalConnectionLimiter; -import google.registry.eppserver.quota.QuotaManager; import google.registry.module.RegistryServlet; import google.registry.request.RequestHandler; import google.registry.util.FakeHttpServletRequest; @@ -76,7 +76,7 @@ public class EppServiceHandler extends SimpleChannelInboundHandler { private final byte[] helloBytes; private final FrontendMetrics metrics; private final LocalConnectionLimiter localConnectionLimiter; - private final QuotaManager commandQuotaManager; + private final EppServerQuotaManager commandQuotaManager; private final Supplier idTokenSupplier; private final String projectId; private final int preLoginReadTimeoutSeconds; @@ -98,7 +98,7 @@ public EppServiceHandler( @Named("hello") byte[] helloBytes, FrontendMetrics metrics, LocalConnectionLimiter localConnectionLimiter, - @CommandQuota QuotaManager commandQuotaManager, + @CommandQuota EppServerQuotaManager commandQuotaManager, @Named("idToken") Supplier idTokenSupplier, @Config("projectId") String projectId, @Config("eppServerPreLoginReadTimeoutSeconds") int preLoginReadTimeoutSeconds) { @@ -227,7 +227,7 @@ private boolean acquireCommandQuota(ChannelHandlerContext ctx) { String throttleId = (authenticatedRegistrarId != null) ? authenticatedRegistrarId : sslClientCertificateHash; if (throttleId != null) { - if (!commandQuotaManager.acquireQuota(new QuotaManager.QuotaRequest(throttleId)).success()) { + if (!commandQuotaManager.acquireQuota(throttleId)) { metrics.registerQuotaRejection("epp_command", throttleId); closeConnection(ctx); return false; diff --git a/core/src/main/java/google/registry/eppserver/quota/EppServerQuotaManager.java b/core/src/main/java/google/registry/eppserver/quota/EppServerQuotaManager.java new file mode 100644 index 00000000000..187d0769ee5 --- /dev/null +++ b/core/src/main/java/google/registry/eppserver/quota/EppServerQuotaManager.java @@ -0,0 +1,91 @@ +// 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.eppserver.quota; + +import com.google.common.collect.ImmutableMap; +import google.registry.config.RegistryConfigSettings.Quota; +import google.registry.config.RegistryConfigSettings.Quota.QuotaGroup; +import google.registry.quota.GenericValkeyQuotaManager; +import java.time.Duration; +import javax.annotation.concurrent.ThreadSafe; + +/** + * Quota management for the EPP server using Redis/Valkey. + * + *

Handles primarily configuration lookup and delegation to the generic quota manager. + */ +@ThreadSafe +public class EppServerQuotaManager { + + private static final Duration DEFAULT_TTL = Duration.ofHours(1); + + private final GenericValkeyQuotaManager quotaManager; + private final QuotaGroup defaultQuota; + private final ImmutableMap customQuotas; + + public EppServerQuotaManager(Quota quota, GenericValkeyQuotaManager quotaManager) { + this.quotaManager = quotaManager; + this.defaultQuota = quota.defaultQuota; + + ImmutableMap.Builder builder = ImmutableMap.builder(); + quota.customQuota.forEach(group -> group.userId.forEach(userId -> builder.put(userId, group))); + this.customQuotas = builder.build(); + } + + /** Attempts to acquire a quota token from Redis. */ + public boolean acquireQuota(String userId) { + QuotaGroup group = customQuotas.getOrDefault(userId, defaultQuota); + + // Unlimited quota check + if (group.tokenAmount < 0) { + return true; + } + + String redisId = getRedisId(group, userId); + return quotaManager.acquireQuota(redisId, group.tokenAmount, getTtl(group)); + } + + /** Refreshes the TTL of an existing quota token. */ + public void refreshQuota(String userId) { + QuotaGroup group = customQuotas.getOrDefault(userId, defaultQuota); + if (group.tokenAmount < 0) { + return; + } + + String redisId = getRedisId(group, userId); + quotaManager.refreshQuota(redisId, getTtl(group)); + } + + /** Returns a token to the pool (used for connection throttling). */ + public void releaseQuota(String userId) { + QuotaGroup group = customQuotas.getOrDefault(userId, defaultQuota); + if (group.tokenAmount < 0) { + return; + } + + String redisId = getRedisId(group, userId); + quotaManager.releaseQuota(redisId, group.tokenAmount); + } + + private String getRedisId(QuotaGroup group, String userId) { + // Use the first ID as the virtual group identity if it's a custom group, + // otherwise isolate each default user by their actual ID. + return (group == defaultQuota || group.userId.isEmpty()) ? userId : group.userId.get(0); + } + + private Duration getTtl(QuotaGroup group) { + return group.refillSeconds > 0 ? Duration.ofSeconds(group.refillSeconds) : DEFAULT_TTL; + } +} diff --git a/core/src/main/java/google/registry/eppserver/quota/QuotaManager.java b/core/src/main/java/google/registry/eppserver/quota/QuotaManager.java deleted file mode 100644 index c493c53988f..00000000000 --- a/core/src/main/java/google/registry/eppserver/quota/QuotaManager.java +++ /dev/null @@ -1,171 +0,0 @@ -// Copyright 2024 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.eppserver.quota; - -import com.google.common.collect.ImmutableMap; -import com.google.common.flogger.FluentLogger; -import google.registry.config.RegistryConfigSettings.Quota; -import google.registry.config.RegistryConfigSettings.Quota.QuotaGroup; -import java.net.URLEncoder; -import java.nio.charset.StandardCharsets; -import javax.annotation.Nullable; -import javax.annotation.concurrent.ThreadSafe; -import redis.clients.jedis.UnifiedJedis; - -/** - * Unified manager for distributed quota enforcement using Redis/Valkey. - * - *

Handles both configuration lookup and atomic Redis operations for connection and command-level - * throttling. - */ -@ThreadSafe -public class QuotaManager { - - private static final FluentLogger logger = FluentLogger.forEnclosingClass(); - private static final int DEFAULT_TTL_SECONDS = 3600; - - /** Lua script to atomically decrement a token bucket with a TTL. */ - private static final String DECR_LUA = - "local current = redis.call('GET', KEYS[1]) " - + "if not current then " - + " redis.call('SET', KEYS[1], ARGV[1] - 1, 'EX', ARGV[2]) " - + " return tonumber(ARGV[1]) - 1 " - + "end " - + "if tonumber(current) <= 0 then " - + " return -1 " - + "end " - + "return redis.call('DECR', KEYS[1])"; - - /** Lua script to atomically increment back a connection token (capped at max). */ - private static final String INCR_LUA = - "local current = redis.call('GET', KEYS[1]) " - + "if current and tonumber(current) < tonumber(ARGV[1]) then " - + " return redis.call('INCR', KEYS[1]) " - + "end " - + "return nil"; - - /** Lua script to refresh the TTL of an existing token bucket. */ - private static final String EXPIRE_LUA = - "if redis.call('EXISTS', KEYS[1]) == 1 then " - + " return redis.call('EXPIRE', KEYS[1], ARGV[1]) " - + "end " - + "return 0"; - - private final UnifiedJedis jedis; - private final String quotaNamespace; - private final QuotaGroup defaultQuota; - private final ImmutableMap customQuotas; - - public QuotaManager(Quota quota, @Nullable UnifiedJedis jedis, String quotaNamespace) { - this.jedis = jedis; - this.quotaNamespace = quotaNamespace; - this.defaultQuota = quota.defaultQuota; - - ImmutableMap.Builder builder = ImmutableMap.builder(); - quota.customQuota.forEach(group -> group.userId.forEach(userId -> builder.put(userId, group))); - this.customQuotas = builder.build(); - } - - public record QuotaRequest(String userId) {} - - public record QuotaResponse(boolean success) {} - - public record QuotaRebate(String userId) {} - - /** Attempts to acquire a quota token from Redis. */ - public QuotaResponse acquireQuota(QuotaRequest request) { - String userId = request.userId(); - QuotaGroup group = customQuotas.getOrDefault(userId, defaultQuota); - - // Unlimited quota check - if (group.tokenAmount < 0) { - return new QuotaResponse(true); - } - - if (jedis == null) { - return new QuotaResponse(true); // Fail open if no Valkey configured - } - - // Use the first ID as the virtual group identity if it's a custom group, - // otherwise isolate each default user by their actual ID. - String redisId = - (group == defaultQuota || group.userId.isEmpty()) ? userId : group.userId.get(0); - String key = String.format("%s:%s", quotaNamespace, redisId); - int ttl = group.refillSeconds > 0 ? group.refillSeconds : DEFAULT_TTL_SECONDS; - - try { - Object result = - jedis.eval(DECR_LUA, 1, key, String.valueOf(group.tokenAmount), String.valueOf(ttl)); - - return new QuotaResponse(((Long) result) >= 0); - } catch (Exception e) { - logger.atSevere().withCause(e).log( - "Valkey error for quota key: %s", URLEncoder.encode(key, StandardCharsets.UTF_8)); - return new QuotaResponse(true); // Fail open - } - } - - /** Refreshes the TTL of an existing quota token. */ - public void refreshQuota(QuotaRequest request) { - if (jedis == null) { - return; - } - - String userId = request.userId(); - QuotaGroup group = customQuotas.getOrDefault(userId, defaultQuota); - if (group.tokenAmount < 0) { - return; - } - - // Use the first ID as the virtual group identity if it's a custom group, - // otherwise isolate each default user by their actual ID. - String redisId = - (group == defaultQuota || group.userId.isEmpty()) ? userId : group.userId.get(0); - String key = String.format("%s:%s", quotaNamespace, redisId); - int ttl = group.refillSeconds > 0 ? group.refillSeconds : DEFAULT_TTL_SECONDS; - - try { - jedis.eval(EXPIRE_LUA, 1, key, String.valueOf(ttl)); - } catch (Exception e) { - logger.atSevere().withCause(e).log( - "Valkey error refreshing quota for: %s", URLEncoder.encode(key, StandardCharsets.UTF_8)); - } - } - - /** Returns a token to the pool (used for connection throttling). */ - public void releaseQuota(QuotaRebate rebate) { - if (jedis == null) { - return; - } - - String userId = rebate.userId(); - QuotaGroup group = customQuotas.getOrDefault(userId, defaultQuota); - if (group.tokenAmount < 0) { - return; - } - - // Use the first ID as the virtual group identity if it's a custom group, - // otherwise isolate each default user by their actual ID. - String redisId = - (group == defaultQuota || group.userId.isEmpty()) ? userId : group.userId.get(0); - String key = String.format("%s:%s", quotaNamespace, redisId); - try { - jedis.eval(INCR_LUA, 1, key, String.valueOf(group.tokenAmount)); - } catch (Exception e) { - logger.atSevere().withCause(e).log( - "Valkey error releasing quota for: %s", URLEncoder.encode(key, StandardCharsets.UTF_8)); - } - } -} diff --git a/core/src/main/java/google/registry/quota/GenericValkeyQuotaManager.java b/core/src/main/java/google/registry/quota/GenericValkeyQuotaManager.java new file mode 100644 index 00000000000..c2f297c633c --- /dev/null +++ b/core/src/main/java/google/registry/quota/GenericValkeyQuotaManager.java @@ -0,0 +1,132 @@ +// 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.quota; + +import static com.google.common.base.Preconditions.checkArgument; +import static java.nio.charset.StandardCharsets.UTF_8; + +import com.google.common.flogger.FluentLogger; +import java.net.URLEncoder; +import java.time.Duration; +import javax.annotation.Nullable; +import javax.annotation.concurrent.ThreadSafe; +import redis.clients.jedis.UnifiedJedis; + +/** Generic quota manager that uses Redis/Valkey as the backing store. */ +@ThreadSafe +public class GenericValkeyQuotaManager { + + private static final FluentLogger logger = FluentLogger.forEnclosingClass(); + + /** + * Lua script to atomically decrement a token bucket with a TTL. + * + *

TODO(b/547996770): maybe use a more complex data structure here (SortedSet?) to manage an + * actual sliding window. Currently this is a fixed window -- the clock "starts" when the first + * request arrives and resets back to 0 entirely once the TTL is hit. + */ + private static final String DECR_LUA = + """ + local current = redis.call('GET', KEYS[1]) + if not current then + redis.call('SET', KEYS[1], ARGV[1] - 1, 'PX', ARGV[2]) + return tonumber(ARGV[1]) - 1 + end + if tonumber(current) <= 0 then + return -1 + end + return redis.call('DECR', KEYS[1]) + """; + + /** Lua script to atomically increment back a connection token (capped at max). */ + private static final String INCR_LUA = + """ + local current = redis.call('GET', KEYS[1]) + if current and tonumber(current) < tonumber(ARGV[1]) then + return redis.call('INCR', KEYS[1]) + end + return nil + """; + + private final UnifiedJedis jedis; + private final String namespace; + + public GenericValkeyQuotaManager(@Nullable UnifiedJedis jedis, String namespace) { + this.jedis = jedis; + this.namespace = namespace; + } + + /** Attempts to acquire a quota token from Valkey. */ + public boolean acquireQuota(String id, int maxTokenAmount, Duration expirationDuration) { + if (jedis == null) { + return true; // Fail open if no Valkey configured + } + checkArgument(expirationDuration.isPositive(), "Duration must be positive"); + checkArgument(maxTokenAmount >= 0, "Max token amount must be non-negative"); + + String key = createValkeyKey(id); + try { + Object result = + jedis.eval( + DECR_LUA, + 1, + key, + String.valueOf(maxTokenAmount), + String.valueOf(expirationDuration.toMillis())); + return (Long) result >= 0; + } catch (Exception e) { + logger.atSevere().withCause(e).log( + "Valkey error for quota key: %s", URLEncoder.encode(key, UTF_8)); + // Fail open + return true; + } + } + + /** Refreshes the TTL of an existing quota token. */ + public void refreshQuota(String id, Duration expirationDuration) { + if (jedis == null) { + return; + } + checkArgument(expirationDuration.isPositive(), "Duration must be positive"); + + String key = createValkeyKey(id); + try { + jedis.pexpire(key, expirationDuration.toMillis()); + } catch (Exception e) { + logger.atSevere().withCause(e).log( + "Valkey error refreshing quota for: %s", URLEncoder.encode(key, UTF_8)); + } + } + + /** Returns a token to the pool (used for connection throttling). */ + public void releaseQuota(String id, int maxTokenAmount) { + if (jedis == null) { + return; + } + checkArgument(maxTokenAmount >= 0, "Max token amount must be non-negative"); + + String key = createValkeyKey(id); + try { + jedis.eval(INCR_LUA, 1, key, String.valueOf(maxTokenAmount)); + } catch (Exception e) { + logger.atSevere().withCause(e).log( + "Valkey error releasing quota for: %s", URLEncoder.encode(key, UTF_8)); + } + } + + private String createValkeyKey(String id) { + return String.format("%s:%s", namespace, id); + } +} diff --git a/core/src/test/java/google/registry/eppserver/handler/EppServiceHandlerTest.java b/core/src/test/java/google/registry/eppserver/handler/EppServiceHandlerTest.java index 520e9e5f64f..689183743dc 100644 --- a/core/src/test/java/google/registry/eppserver/handler/EppServiceHandlerTest.java +++ b/core/src/test/java/google/registry/eppserver/handler/EppServiceHandlerTest.java @@ -30,10 +30,8 @@ import static org.mockito.Mockito.when; import google.registry.eppserver.metric.FrontendMetrics; +import google.registry.eppserver.quota.EppServerQuotaManager; import google.registry.eppserver.quota.LocalConnectionLimiter; -import google.registry.eppserver.quota.QuotaManager; -import google.registry.eppserver.quota.QuotaManager.QuotaRequest; -import google.registry.eppserver.quota.QuotaManager.QuotaResponse; import google.registry.request.RequestHandler; import google.registry.util.FakeHttpServletRequest; import google.registry.util.FakeHttpServletResponse; @@ -68,7 +66,7 @@ class EppServiceHandlerTest { @Mock private FrontendMetrics metrics; @Mock private LocalConnectionLimiter localConnectionLimiter; - @Mock private QuotaManager commandQuotaManager; + @Mock private EppServerQuotaManager commandQuotaManager; @Mock private Supplier idTokenSupplier; @Mock private ChannelHandlerContext ctx; @Mock private Channel channel; @@ -116,9 +114,7 @@ void setUp() { .when(executor) .schedule(any(Runnable.class), anyLong(), any(TimeUnit.class)); - lenient() - .when(commandQuotaManager.acquireQuota(any(QuotaRequest.class))) - .thenReturn(new QuotaResponse(true)); + lenient().when(commandQuotaManager.acquireQuota(any(String.class))).thenReturn(true); } private void setUpSuccessfulHandshake() throws Exception { @@ -149,8 +145,7 @@ void testChannelActive_success() throws Exception { .when(requestHandler) .handleRequest(any(FakeHttpServletRequest.class), any(FakeHttpServletResponse.class)); - when(commandQuotaManager.acquireQuota(any(QuotaRequest.class))) - .thenReturn(new QuotaResponse(true)); + when(commandQuotaManager.acquireQuota(any(String.class))).thenReturn(true); when(idTokenSupplier.get()).thenReturn("fake_id_token"); setUpSuccessfulHandshake(); @@ -192,8 +187,7 @@ void testChannelRead0_registrarQuotaRejected() throws Exception { setUpSuccessfulHandshake(); when(idTokenSupplier.get()).thenReturn("fake_id_token"); - when(commandQuotaManager.acquireQuota(any(QuotaRequest.class))) - .thenReturn(new QuotaResponse(true)); + when(commandQuotaManager.acquireQuota(any(String.class))).thenReturn(true); String eppLoginXml = "RegistrarA"; ByteBuf inFrame = Unpooled.wrappedBuffer(eppLoginXml.getBytes(UTF_8)); @@ -221,8 +215,7 @@ void testChannelRead0_extractsClidAndForwardsRequest() throws Exception { setUpSuccessfulHandshake(); when(idTokenSupplier.get()).thenReturn("fake_id_token"); - when(commandQuotaManager.acquireQuota(any(QuotaRequest.class))) - .thenReturn(new QuotaResponse(true)); + when(commandQuotaManager.acquireQuota(any(String.class))).thenReturn(true); String eppLoginXml = "RegistrarA"; ByteBuf inFrame = Unpooled.wrappedBuffer(eppLoginXml.getBytes(UTF_8)); @@ -248,7 +241,7 @@ void testChannelRead0_extractsClidAndForwardsRequest() throws Exception { // Verify command quota was requested for the cert hash pre-login String certHash = X509Utils.getCertificateHash(certificate); - verify(commandQuotaManager, times(2)).acquireQuota(eq(new QuotaRequest(certHash))); + verify(commandQuotaManager, times(2)).acquireQuota(eq(certHash)); verify(localConnectionLimiter).acquireRegistrar("RegistrarA"); verify(scheduledFuture).cancel(eq(false)); @@ -284,15 +277,14 @@ void testChannelRead0_extractsClidAndForwardsRequest() throws Exception { handler.channelRead0(ctx, inFrame2); // Verify command quota was requested for the authenticated registrar post-login - verify(commandQuotaManager).acquireQuota(eq(new QuotaRequest("RegistrarA"))); + verify(commandQuotaManager).acquireQuota(eq("RegistrarA")); } @Test void testChannelRead0_commandQuotaRejected() throws Exception { setUpSuccessfulHandshake(); - when(commandQuotaManager.acquireQuota(any(QuotaRequest.class))) - .thenReturn(new QuotaResponse(false)); + when(commandQuotaManager.acquireQuota(any(String.class))).thenReturn(false); String eppXml = ""; ByteBuf inFrame = Unpooled.wrappedBuffer(eppXml.getBytes(UTF_8)); @@ -308,8 +300,7 @@ void testChannelRead0_closeSessionHeaderClosesChannel() throws Exception { setUpSuccessfulHandshake(); when(idTokenSupplier.get()).thenReturn("fake_id_token"); - when(commandQuotaManager.acquireQuota(any(QuotaRequest.class))) - .thenReturn(new QuotaResponse(true)); + when(commandQuotaManager.acquireQuota(any(String.class))).thenReturn(true); String eppLogoutXml = ""; ByteBuf inFrame = Unpooled.wrappedBuffer(eppLogoutXml.getBytes(UTF_8)); @@ -348,8 +339,7 @@ void testChannelInactive_postLogin_releasesIpAndRegistrar() throws Exception { setUpSuccessfulHandshake(); when(idTokenSupplier.get()).thenReturn("fake_id_token"); - when(commandQuotaManager.acquireQuota(any(QuotaRequest.class))) - .thenReturn(new QuotaResponse(true)); + when(commandQuotaManager.acquireQuota(any(String.class))).thenReturn(true); when(localConnectionLimiter.acquireRegistrar("RegistrarA")).thenReturn(true); String eppLoginXml = "RegistrarA"; diff --git a/core/src/test/java/google/registry/eppserver/quota/EppServerQuotaManagerTest.java b/core/src/test/java/google/registry/eppserver/quota/EppServerQuotaManagerTest.java new file mode 100644 index 00000000000..fba4c443971 --- /dev/null +++ b/core/src/test/java/google/registry/eppserver/quota/EppServerQuotaManagerTest.java @@ -0,0 +1,126 @@ +// 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.eppserver.quota; + +import static com.google.common.truth.Truth.assertThat; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.verifyNoInteractions; +import static org.mockito.Mockito.when; + +import com.google.common.collect.ImmutableList; +import google.registry.config.RegistryConfigSettings.Quota; +import google.registry.config.RegistryConfigSettings.Quota.QuotaGroup; +import google.registry.quota.GenericValkeyQuotaManager; +import java.time.Duration; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.ExtendWith; +import org.mockito.Mock; +import org.mockito.junit.jupiter.MockitoExtension; + +@ExtendWith(MockitoExtension.class) +class EppServerQuotaManagerTest { + + @Mock private GenericValkeyQuotaManager quotaManager; + + private Quota quotaConfig; + private EppServerQuotaManager manager; + + @BeforeEach + void setUp() { + quotaConfig = new Quota(); + QuotaGroup defaultGroup = new QuotaGroup(); + defaultGroup.tokenAmount = 10; + defaultGroup.refillSeconds = 60; + quotaConfig.defaultQuota = defaultGroup; + + QuotaGroup customGroup = new QuotaGroup(); + customGroup.tokenAmount = 5; + customGroup.refillSeconds = 30; + customGroup.userId = ImmutableList.of("user1"); + quotaConfig.customQuota = ImmutableList.of(customGroup); + + manager = new EppServerQuotaManager(quotaConfig, quotaManager); + } + + @Test + void testAcquireQuota_defaultQuota() { + when(quotaManager.acquireQuota("user2", 10, Duration.ofMinutes(1))).thenReturn(true); + + assertThat(manager.acquireQuota("user2")).isTrue(); + verify(quotaManager).acquireQuota("user2", 10, Duration.ofMinutes(1)); + } + + @Test + void testAcquireQuota_customQuota() { + when(quotaManager.acquireQuota("user1", 5, Duration.ofSeconds(30))).thenReturn(true); + + assertThat(manager.acquireQuota("user1")).isTrue(); + verify(quotaManager).acquireQuota("user1", 5, Duration.ofSeconds(30)); + } + + @Test + void testAcquireQuota_unlimited() { + quotaConfig.defaultQuota.tokenAmount = -1; + manager = new EppServerQuotaManager(quotaConfig, quotaManager); + + assertThat(manager.acquireQuota("user2")).isTrue(); + verifyNoInteractions(quotaManager); + } + + @Test + void testRefreshQuota_success() { + manager.refreshQuota("user2"); + verify(quotaManager).refreshQuota("user2", Duration.ofMinutes(1)); + } + + @Test + void testRefreshQuota_unlimited_noop() { + quotaConfig.defaultQuota.tokenAmount = -1; + manager = new EppServerQuotaManager(quotaConfig, quotaManager); + + manager.refreshQuota("user2"); + verifyNoInteractions(quotaManager); + } + + @Test + void testReleaseQuota_success() { + manager.releaseQuota("user2"); + verify(quotaManager).releaseQuota("user2", 10); + } + + @Test + void testReleaseQuota_unlimited_noop() { + quotaConfig.defaultQuota.tokenAmount = -1; + manager = new EppServerQuotaManager(quotaConfig, quotaManager); + + manager.releaseQuota("user2"); + verifyNoInteractions(quotaManager); + } + + @Test + void testGroupVirtualIdentity_usesFirstIdInList() { + // Modify config so "user1" is accompanied by a virtual group ID "my_group" + quotaConfig.customQuota.get(0).userId = ImmutableList.of("my_group", "user1", "user3"); + manager = new EppServerQuotaManager(quotaConfig, quotaManager); + + when(quotaManager.acquireQuota("my_group", 5, Duration.ofSeconds(30))).thenReturn(true); + + assertThat(manager.acquireQuota("user1")).isTrue(); + assertThat(manager.acquireQuota("user3")).isTrue(); + verify(quotaManager, times(2)).acquireQuota("my_group", 5, Duration.ofSeconds(30)); + } +} diff --git a/core/src/test/java/google/registry/eppserver/quota/QuotaManagerTest.java b/core/src/test/java/google/registry/eppserver/quota/QuotaManagerTest.java deleted file mode 100644 index 0af0e659e4e..00000000000 --- a/core/src/test/java/google/registry/eppserver/quota/QuotaManagerTest.java +++ /dev/null @@ -1,159 +0,0 @@ -// Copyright 2024 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.eppserver.quota; - -import static com.google.common.truth.Truth.assertThat; -import static org.mockito.ArgumentMatchers.anyInt; -import static org.mockito.ArgumentMatchers.anyString; -import static org.mockito.ArgumentMatchers.eq; -import static org.mockito.Mockito.times; -import static org.mockito.Mockito.verify; -import static org.mockito.Mockito.when; - -import com.google.common.collect.ImmutableList; -import google.registry.config.RegistryConfigSettings.Quota; -import google.registry.config.RegistryConfigSettings.Quota.QuotaGroup; -import org.junit.jupiter.api.BeforeEach; -import org.junit.jupiter.api.Test; -import org.junit.jupiter.api.extension.ExtendWith; -import org.mockito.Mock; -import org.mockito.junit.jupiter.MockitoExtension; -import redis.clients.jedis.UnifiedJedis; - -@ExtendWith(MockitoExtension.class) -class QuotaManagerTest { - - @Mock private UnifiedJedis jedis; - - private Quota quotaConfig; - private QuotaManager manager; - - @BeforeEach - void setUp() { - quotaConfig = new Quota(); - QuotaGroup defaultGroup = new QuotaGroup(); - defaultGroup.tokenAmount = 10; - defaultGroup.refillSeconds = 60; - quotaConfig.defaultQuota = defaultGroup; - - QuotaGroup customGroup = new QuotaGroup(); - customGroup.tokenAmount = 5; - customGroup.refillSeconds = 30; - customGroup.userId = ImmutableList.of("user1"); - quotaConfig.customQuota = ImmutableList.of(customGroup); - - manager = new QuotaManager(quotaConfig, jedis, "testQuota"); - } - - @Test - void testAcquireQuota_success() { - when(jedis.eval(anyString(), anyInt(), anyString(), anyString(), anyString())).thenReturn(5L); - - QuotaManager.QuotaResponse response = - manager.acquireQuota(new QuotaManager.QuotaRequest("user2")); - assertThat(response.success()).isTrue(); - verify(jedis).eval(anyString(), eq(1), eq("testQuota:user2"), eq("10"), eq("60")); - } - - @Test - void testAcquireQuota_failure() { - when(jedis.eval(anyString(), anyInt(), anyString(), anyString(), anyString())).thenReturn(-1L); - - QuotaManager.QuotaResponse response = - manager.acquireQuota(new QuotaManager.QuotaRequest("user1")); - assertThat(response.success()).isFalse(); - verify(jedis).eval(anyString(), eq(1), eq("testQuota:user1"), eq("5"), eq("30")); - } - - @Test - void testAcquireQuota_unlimited() { - quotaConfig.defaultQuota.tokenAmount = -1; - manager = new QuotaManager(quotaConfig, jedis, "testQuota"); - - QuotaManager.QuotaResponse response = - manager.acquireQuota(new QuotaManager.QuotaRequest("user2")); - assertThat(response.success()).isTrue(); - } - - @Test - void testAcquireQuota_jedisException_failsOpen() { - when(jedis.eval(anyString(), anyInt(), anyString(), anyString(), anyString())) - .thenThrow(new RuntimeException("Redis error")); - - QuotaManager.QuotaResponse response = - manager.acquireQuota(new QuotaManager.QuotaRequest("user2")); - assertThat(response.success()).isTrue(); - } - - @Test - void testRefreshQuota_success() { - manager.refreshQuota(new QuotaManager.QuotaRequest("user2")); - verify(jedis).eval(anyString(), eq(1), eq("testQuota:user2"), eq("60")); - } - - @Test - void testReleaseQuota_success() { - manager.releaseQuota(new QuotaManager.QuotaRebate("user2")); - verify(jedis).eval(anyString(), eq(1), eq("testQuota:user2"), eq("10")); - } - - @Test - void testGroupVirtualIdentity_usesFirstIdInList() { - // Modify config so "user1" is accompanied by a virtual group ID "my_group" - quotaConfig.customQuota.get(0).userId = ImmutableList.of("my_group", "user1", "user3"); - manager = new QuotaManager(quotaConfig, jedis, "testQuota"); - - when(jedis.eval(anyString(), anyInt(), anyString(), anyString(), anyString())).thenReturn(5L); - - QuotaManager.QuotaResponse response1 = - manager.acquireQuota(new QuotaManager.QuotaRequest("user1")); - QuotaManager.QuotaResponse response2 = - manager.acquireQuota(new QuotaManager.QuotaRequest("user3")); - - assertThat(response1.success()).isTrue(); - assertThat(response2.success()).isTrue(); - // 5 tokens, 30 seconds ttl - verify(jedis, times(2)).eval(anyString(), eq(1), eq("testQuota:my_group"), eq("5"), eq("30")); - } - - @Test - void testGroupVirtualIdentity_exceedsQuota_fails() { - // Modify config so "user1", "user2", "user3" share virtual group ID "my_group" - quotaConfig.customQuota.get(0).userId = ImmutableList.of("my_group", "user1", "user2", "user3"); - manager = new QuotaManager(quotaConfig, jedis, "testQuota"); - - // Simulate Redis returning 1, 0 for successful decrements, and -1 when empty - when(jedis.eval(anyString(), anyInt(), anyString(), anyString(), anyString())) - .thenReturn(1L) - .thenReturn(0L) - .thenReturn(-1L); - - // Act - QuotaManager.QuotaResponse response1 = - manager.acquireQuota(new QuotaManager.QuotaRequest("user1")); - QuotaManager.QuotaResponse response2 = - manager.acquireQuota(new QuotaManager.QuotaRequest("user2")); - QuotaManager.QuotaResponse response3 = - manager.acquireQuota(new QuotaManager.QuotaRequest("user3")); - - // Assert that the third request to the same group fails - assertThat(response1.success()).isTrue(); - assertThat(response2.success()).isTrue(); - assertThat(response3.success()).isFalse(); - - // Verify all 3 requests went to the shared bucket - verify(jedis, times(3)).eval(anyString(), eq(1), eq("testQuota:my_group"), eq("5"), eq("30")); - } -} diff --git a/core/src/test/java/google/registry/quota/GenericValkeyQuotaManagerTest.java b/core/src/test/java/google/registry/quota/GenericValkeyQuotaManagerTest.java new file mode 100644 index 00000000000..a3c222c1a40 --- /dev/null +++ b/core/src/test/java/google/registry/quota/GenericValkeyQuotaManagerTest.java @@ -0,0 +1,170 @@ +// 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.quota; + +import static com.google.common.truth.Truth.assertThat; +import static org.junit.jupiter.api.Assertions.assertDoesNotThrow; + +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; + +@Testcontainers +class GenericValkeyQuotaManagerTest { + + @Container private static final ValkeyContainer valkey = new ValkeyContainer(); + + private RedisClient jedis; + private GenericValkeyQuotaManager quotaManager; + + @BeforeEach + void setUp() { + jedis = + RedisClient.builder() + .hostAndPort(new HostAndPort(valkey.getHost(), valkey.getFirstMappedPort())) + .build(); + jedis.flushAll(); + quotaManager = new GenericValkeyQuotaManager(jedis, "testQuota"); + } + + @Test + void testAcquireQuota_success() { + assertThat(quotaManager.acquireQuota("user1", 5, Duration.ofMinutes(1))).isTrue(); + assertThat(jedis.get("testQuota:user1")).isEqualTo("4"); + assertThat(jedis.ttl("testQuota:user1")).isGreaterThan(0L); + } + + @Test + void testAcquireQuota_exhaustsQuota_thenFails() { + assertThat(quotaManager.acquireQuota("user1", 2, Duration.ofMinutes(1))).isTrue(); + assertThat(jedis.get("testQuota:user1")).isEqualTo("1"); + + assertThat(quotaManager.acquireQuota("user1", 2, Duration.ofMinutes(1))).isTrue(); + assertThat(jedis.get("testQuota:user1")).isEqualTo("0"); + + assertThat(quotaManager.acquireQuota("user1", 2, Duration.ofMinutes(1))).isFalse(); + assertThat(jedis.get("testQuota:user1")).isEqualTo("0"); + } + + @Test + void testAcquireQuota_resetsAfterExpiration() throws Exception { + assertThat(quotaManager.acquireQuota("user1", 1, Duration.ofMillis(50))).isTrue(); + assertThat(quotaManager.acquireQuota("user1", 1, Duration.ofMillis(50))).isFalse(); + + Thread.sleep(150); + + assertThat(quotaManager.acquireQuota("user1", 1, Duration.ofMillis(50))).isTrue(); + } + + @Test + void testAcquireQuota_isolatedByNamespaceAndId() { + GenericValkeyQuotaManager otherQuotaManager = + new GenericValkeyQuotaManager(jedis, "otherQuota"); + + assertThat(quotaManager.acquireQuota("user1", 1, Duration.ofMinutes(1))).isTrue(); + assertThat(quotaManager.acquireQuota("user1", 1, Duration.ofMinutes(1))).isFalse(); + + // user2 in same namespace is independent + assertThat(quotaManager.acquireQuota("user2", 1, Duration.ofMinutes(1))).isTrue(); + + // user1 in other namespace is independent + assertThat(otherQuotaManager.acquireQuota("user1", 1, Duration.ofMinutes(1))).isTrue(); + } + + @Test + void testAcquireQuota_nullJedis_failsOpen() { + GenericValkeyQuotaManager nullJedisManager = new GenericValkeyQuotaManager(null, "testQuota"); + assertThat(nullJedisManager.acquireQuota("user2", 10, Duration.ofMinutes(1))).isTrue(); + } + + @Test + void testAcquireQuota_jedisException_failsOpen() { + jedis.close(); + assertThat(quotaManager.acquireQuota("user2", 10, Duration.ofMinutes(1))).isTrue(); + } + + @Test + void testRefreshQuota_success() { + quotaManager.acquireQuota("user1", 5, Duration.ofSeconds(10)); + + quotaManager.refreshQuota("user1", Duration.ofMinutes(5)); + assertThat(jedis.ttl("testQuota:user1")).isGreaterThan(10L); + } + + @Test + void testRefreshQuota_nonexistentKey_noop() { + quotaManager.refreshQuota("nonexistent", Duration.ofMinutes(5)); + assertThat(jedis.exists("testQuota:nonexistent")).isFalse(); + } + + @Test + void testRefreshQuota_nullJedis_noop() { + GenericValkeyQuotaManager nullJedisManager = new GenericValkeyQuotaManager(null, "testQuota"); + assertDoesNotThrow(() -> nullJedisManager.refreshQuota("user2", Duration.ofMinutes(1))); + } + + @Test + void testRefreshQuota_jedisException_handled() { + jedis.close(); + assertDoesNotThrow(() -> quotaManager.refreshQuota("user2", Duration.ofMinutes(1))); + } + + @Test + void testReleaseQuota_success() { + assertThat(quotaManager.acquireQuota("user1", 1, Duration.ofMinutes(1))).isTrue(); + assertThat(quotaManager.acquireQuota("user1", 1, Duration.ofMinutes(1))).isFalse(); + + quotaManager.releaseQuota("user1", 1); + assertThat(jedis.get("testQuota:user1")).isEqualTo("1"); + + assertThat(quotaManager.acquireQuota("user1", 1, Duration.ofMinutes(1))).isTrue(); + } + + @Test + void testReleaseQuota_cappedAtMax() { + quotaManager.acquireQuota("user1", 3, Duration.ofMinutes(1)); + assertThat(jedis.get("testQuota:user1")).isEqualTo("2"); + + quotaManager.releaseQuota("user1", 3); + assertThat(jedis.get("testQuota:user1")).isEqualTo("3"); + + // Releasing again when already at max should not increment past max + quotaManager.releaseQuota("user1", 3); + assertThat(jedis.get("testQuota:user1")).isEqualTo("3"); + } + + @Test + void testReleaseQuota_nonexistentKey_noop() { + quotaManager.releaseQuota("nonexistent", 5); + assertThat(jedis.exists("testQuota:nonexistent")).isFalse(); + } + + @Test + void testReleaseQuota_nullJedis_noop() { + GenericValkeyQuotaManager nullJedisManager = new GenericValkeyQuotaManager(null, "testQuota"); + assertDoesNotThrow(() -> nullJedisManager.releaseQuota("user2", 10)); + } + + @Test + void testReleaseQuota_jedisException_handled() { + jedis.close(); + assertDoesNotThrow(() -> quotaManager.releaseQuota("user2", 10)); + } +}