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)); + } +}