diff --git a/java/google/registry/proxy/EppProtocolModule.java b/java/google/registry/proxy/EppProtocolModule.java index 3a3531e2b..499675642 100644 --- a/java/google/registry/proxy/EppProtocolModule.java +++ b/java/google/registry/proxy/EppProtocolModule.java @@ -26,9 +26,14 @@ import google.registry.proxy.Protocol.BackendProtocol; import google.registry.proxy.Protocol.FrontendProtocol; import google.registry.proxy.handler.EppServiceHandler; import google.registry.proxy.handler.ProxyProtocolHandler; +import google.registry.proxy.handler.QuotaHandler.EppQuotaHandler; import google.registry.proxy.handler.RelayHandler.FullHttpRequestRelayHandler; import google.registry.proxy.handler.SslServerInitializer; import google.registry.proxy.metric.FrontendMetrics; +import google.registry.proxy.quota.QuotaConfig; +import google.registry.proxy.quota.QuotaManager; +import google.registry.proxy.quota.TokenStore; +import google.registry.util.Clock; import google.registry.util.FormattingLogger; import io.netty.channel.ChannelHandler; import io.netty.channel.socket.nio.NioSocketChannel; @@ -37,6 +42,8 @@ import io.netty.handler.codec.LengthFieldPrepender; import io.netty.handler.logging.LoggingHandler; import io.netty.handler.timeout.ReadTimeoutHandler; import java.io.IOException; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.ScheduledExecutorService; import javax.inject.Named; import javax.inject.Provider; import javax.inject.Qualifier; @@ -44,13 +51,13 @@ import javax.inject.Singleton; /** A module that provides the {@link FrontendProtocol} used for epp protocol. */ @Module -class EppProtocolModule { +public class EppProtocolModule { private static final FormattingLogger logger = FormattingLogger.getLoggerForCallerClass(); /** Dagger qualifier to provide epp protocol related handlers and other bindings. */ @Qualifier - @interface EppProtocol {}; + public @interface EppProtocol {}; private static final String PROTOCOL_NAME = "epp"; @@ -79,6 +86,7 @@ class EppProtocolModule { Provider lengthFieldBasedFrameDecoderProvider, Provider lengthFieldPrependerProvider, Provider eppServiceHandlerProvider, + Provider eppQuotaHandlerProvider, Provider loggingHandlerProvider, Provider relayHandlerProvider) { return ImmutableList.of( @@ -88,6 +96,7 @@ class EppProtocolModule { lengthFieldBasedFrameDecoderProvider, lengthFieldPrependerProvider, eppServiceHandlerProvider, + eppQuotaHandlerProvider, loggingHandlerProvider, relayHandlerProvider); } @@ -149,4 +158,19 @@ class EppProtocolModule { helloBytes, metrics); } + + @Provides + @EppProtocol + static TokenStore provideTokenStore( + ProxyConfig config, ScheduledExecutorService refreshExecutor, Clock clock) { + return new TokenStore(new QuotaConfig(config.epp.quota, PROTOCOL_NAME), refreshExecutor, clock); + } + + @Provides + @Singleton + @EppProtocol + static QuotaManager provideQuotaManager( + @EppProtocol TokenStore tokenStore, ExecutorService executorService) { + return new QuotaManager(tokenStore, executorService); + } } diff --git a/java/google/registry/proxy/ProxyModule.java b/java/google/registry/proxy/ProxyModule.java index 3330038a9..70ad5ea5a 100644 --- a/java/google/registry/proxy/ProxyModule.java +++ b/java/google/registry/proxy/ProxyModule.java @@ -48,6 +48,9 @@ import io.netty.handler.ssl.SslProvider; import java.io.IOException; import java.util.Objects; import java.util.Set; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.ScheduledExecutorService; import java.util.logging.ConsoleHandler; import java.util.logging.Handler; import java.util.logging.Level; @@ -255,6 +258,16 @@ public class ProxyModule { return new SystemClock(); } + @Provides + static ExecutorService provideExecutorService() { + return Executors.newWorkStealingPool(); + } + + @Provides + static ScheduledExecutorService provideScheduledExecutorService() { + return Executors.newSingleThreadScheduledExecutor(); + } + @Singleton @Provides ProxyConfig provideProxyConfig(Environment env) { diff --git a/java/google/registry/proxy/WhoisProtocolModule.java b/java/google/registry/proxy/WhoisProtocolModule.java index c453dddf6..6df29cee0 100644 --- a/java/google/registry/proxy/WhoisProtocolModule.java +++ b/java/google/registry/proxy/WhoisProtocolModule.java @@ -22,13 +22,21 @@ import dagger.multibindings.IntoSet; import google.registry.proxy.HttpsRelayProtocolModule.HttpsRelayProtocol; import google.registry.proxy.Protocol.BackendProtocol; import google.registry.proxy.Protocol.FrontendProtocol; +import google.registry.proxy.handler.ProxyProtocolHandler; +import google.registry.proxy.handler.QuotaHandler.WhoisQuotaHandler; import google.registry.proxy.handler.RelayHandler.FullHttpRequestRelayHandler; import google.registry.proxy.handler.WhoisServiceHandler; import google.registry.proxy.metric.FrontendMetrics; +import google.registry.proxy.quota.QuotaConfig; +import google.registry.proxy.quota.QuotaManager; +import google.registry.proxy.quota.TokenStore; +import google.registry.util.Clock; import io.netty.channel.ChannelHandler; import io.netty.handler.codec.LineBasedFrameDecoder; import io.netty.handler.logging.LoggingHandler; import io.netty.handler.timeout.ReadTimeoutHandler; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.ScheduledExecutorService; import javax.inject.Named; import javax.inject.Provider; import javax.inject.Qualifier; @@ -36,11 +44,11 @@ import javax.inject.Singleton; /** A module that provides the {@link FrontendProtocol} used for whois protocol. */ @Module -class WhoisProtocolModule { +public class WhoisProtocolModule { /** Dagger qualifier to provide whois protocol related handlers and other bindings. */ @Qualifier - @interface WhoisProtocol {}; + public @interface WhoisProtocol {}; private static final String PROTOCOL_NAME = "whois"; @@ -63,15 +71,19 @@ class WhoisProtocolModule { @Provides @WhoisProtocol static ImmutableList> provideHandlerProviders( + Provider proxyProtocolHandlerProvider, @WhoisProtocol Provider readTimeoutHandlerProvider, Provider lineBasedFrameDecoderProvider, Provider whoisServiceHandlerProvider, + Provider whoisQuotaHandlerProvider, Provider loggingHandlerProvider, Provider relayHandlerProvider) { return ImmutableList.of( + proxyProtocolHandlerProvider, readTimeoutHandlerProvider, lineBasedFrameDecoderProvider, whoisServiceHandlerProvider, + whoisQuotaHandlerProvider, loggingHandlerProvider, relayHandlerProvider); } @@ -95,4 +107,20 @@ class WhoisProtocolModule { static ReadTimeoutHandler provideReadTimeoutHandler(ProxyConfig config) { return new ReadTimeoutHandler(config.whois.readTimeoutSeconds); } + + @Provides + @WhoisProtocol + static TokenStore provideTokenStore( + ProxyConfig config, ScheduledExecutorService refreshExecutor, Clock clock) { + return new TokenStore( + new QuotaConfig(config.whois.quota, PROTOCOL_NAME), refreshExecutor, clock); + } + + @Provides + @Singleton + @WhoisProtocol + static QuotaManager provideQuotaManager( + @WhoisProtocol TokenStore tokenStore, ExecutorService executorService) { + return new QuotaManager(tokenStore, executorService); + } } diff --git a/java/google/registry/proxy/handler/QuotaHandler.java b/java/google/registry/proxy/handler/QuotaHandler.java new file mode 100644 index 000000000..04477400b --- /dev/null +++ b/java/google/registry/proxy/handler/QuotaHandler.java @@ -0,0 +1,156 @@ +// Copyright 2017 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.proxy.handler; + +import static com.google.common.base.Preconditions.checkNotNull; +import static google.registry.proxy.handler.EppServiceHandler.CLIENT_CERTIFICATE_HASH_KEY; +import static google.registry.proxy.handler.ProxyProtocolHandler.REMOTE_ADDRESS_KEY; + +import google.registry.proxy.EppProtocolModule.EppProtocol; +import google.registry.proxy.WhoisProtocolModule.WhoisProtocol; +import google.registry.proxy.quota.QuotaManager; +import google.registry.proxy.quota.QuotaManager.QuotaRebate; +import google.registry.proxy.quota.QuotaManager.QuotaRequest; +import google.registry.proxy.quota.QuotaManager.QuotaResponse; +import io.netty.channel.ChannelHandlerContext; +import io.netty.channel.ChannelInboundHandlerAdapter; +import java.util.concurrent.Future; +import javax.inject.Inject; + +/** + * Handler that checks quota fulfillment and terminates connection if necessary. + * + *

This handler attempts to acquire quota during the first {@link #channelRead} operation, not + * when connection is established. The reason is that the {@code userId} used for acquiring quota is + * not always available when the connection is just open. + */ +public abstract class QuotaHandler extends ChannelInboundHandlerAdapter { + + protected final QuotaManager quotaManager; + protected QuotaResponse quotaResponse; + + protected QuotaHandler(QuotaManager quotaManager) { + this.quotaManager = quotaManager; + } + + abstract String getUserId(ChannelHandlerContext ctx); + + /** Whether the user id is PII ans should not be logged. IP addresses are considered PII. */ + abstract boolean isUserIdPii(); + + @Override + public void channelRead(ChannelHandlerContext ctx, Object msg) throws Exception { + if (quotaResponse == null) { + String userId = getUserId(ctx); + checkNotNull(userId, "Cannot obtain User ID"); + quotaResponse = quotaManager.acquireQuota(QuotaRequest.create(userId)); + if (!quotaResponse.success()) { + throw new OverQuotaException(isUserIdPii() ? "none" : userId); + } + } + ctx.fireChannelRead(msg); + } + + /** + * Actions to take when the connection terminates. + * + *

Depending on the quota type, the handler either returns the tokens, or does nothing. + */ + @Override + public abstract void channelInactive(ChannelHandlerContext ctx); + + static class OverQuotaException extends Exception { + OverQuotaException(String userId) { + super(String.format("USER ID: %s\nQuota exceeded, terminating connection.", userId)); + } + } + + /** Quota Handler for WHOIS protocol. */ + public static class WhoisQuotaHandler extends QuotaHandler { + + @Inject + WhoisQuotaHandler(@WhoisProtocol QuotaManager quotaManager) { + super(quotaManager); + } + + /** + * Reads user ID from channel attribute {@code REMOTE_ADDRESS_KEY}. + * + *

This attribute is set by {@link ProxyProtocolHandler} when the first frame of message is + * read. + */ + @Override + String getUserId(ChannelHandlerContext ctx) { + return ctx.channel().attr(REMOTE_ADDRESS_KEY).get(); + } + + @Override + boolean isUserIdPii() { + return true; + } + + /** + * Do nothing when connection terminates. + * + *

WHOIS protocol is configured with a QPS type quota, there is no need to return the tokens + * back to the quota store because the quota store will auto-refill tokens based on the QPS. + */ + @Override + public void channelInactive(ChannelHandlerContext ctx) { + ctx.fireChannelInactive(); + } + } + + /** Quota Handler for EPP protocol. */ + public static class EppQuotaHandler extends QuotaHandler { + + @Inject + EppQuotaHandler(@EppProtocol QuotaManager quotaManager) { + super(quotaManager); + } + + /** + * Reads user ID from channel attribute {@code CLIENT_CERTIFICATE_HASH_KEY}. + * + *

This attribute is set by {@link EppServiceHandler} when SSH handshake completes + * successfully. That handler subsequently simulates reading of an EPP HELLO request, in order + * to solicit an EPP GREETING response from the server. The {@link #channelRead} method of this + * handler is called afterward because it is the next handler in the channel pipeline, + * guaranteeing that the {@code CLIENT_CERTIFICATE_HASH_KEY} is always non-null. + */ + @Override + String getUserId(ChannelHandlerContext ctx) { + return ctx.channel().attr(CLIENT_CERTIFICATE_HASH_KEY).get(); + } + + @Override + boolean isUserIdPii() { + return false; + } + + /** + * Returns the leased token back to the token store upon connection termination. + * + *

A connection with concurrent quota needs to do this in order to maintain its quota number + * invariance. + */ + @Override + public void channelInactive(ChannelHandlerContext ctx) { + checkNotNull(quotaResponse, "No quota was leased, return not possible."); + Future unusedFuture = quotaManager.releaseQuota(QuotaRebate.create(quotaResponse)); + ctx.fireChannelInactive(); + } + } +} diff --git a/java/google/registry/proxy/quota/QuotaManager.java b/java/google/registry/proxy/quota/QuotaManager.java index 4cc69e6c5..5515e478a 100644 --- a/java/google/registry/proxy/quota/QuotaManager.java +++ b/java/google/registry/proxy/quota/QuotaManager.java @@ -15,7 +15,6 @@ package google.registry.proxy.quota; import com.google.auto.value.AutoValue; -import google.registry.proxy.quota.QuotaManager.QuotaResponse.Status; import google.registry.proxy.quota.TokenStore.TimestampedInteger; import java.util.concurrent.ExecutorService; import java.util.concurrent.Future; @@ -40,38 +39,36 @@ import org.joda.time.DateTime; @ThreadSafe public class QuotaManager { + /** Value class representing a quota request. */ @AutoValue - abstract static class QuotaRequest { + public abstract static class QuotaRequest { - static QuotaRequest create(String userId) { + public static QuotaRequest create(String userId) { return new AutoValue_QuotaManager_QuotaRequest(userId); } abstract String userId(); } + /** Value class representing a quota response. */ @AutoValue - abstract static class QuotaResponse { - - enum Status { - SUCCESS, - FAILURE, + public abstract static class QuotaResponse { + public static QuotaResponse create( + boolean success, String userId, DateTime grantedTokenRefillTime) { + return new AutoValue_QuotaManager_QuotaResponse(success, userId, grantedTokenRefillTime); } - static QuotaResponse create(Status status, String userId, DateTime grantedTokenRefillTime) { - return new AutoValue_QuotaManager_QuotaResponse(status, userId, grantedTokenRefillTime); - } - - abstract Status status(); + public abstract boolean success(); abstract String userId(); abstract DateTime grantedTokenRefillTime(); } + /** Value class representing a quota rebate. */ @AutoValue - abstract static class QuotaRebate { - static QuotaRebate create(QuotaResponse response) { + public abstract static class QuotaRebate { + public static QuotaRebate create(QuotaResponse response) { return new AutoValue_QuotaManager_QuotaRebate( response.userId(), response.grantedTokenRefillTime()); } @@ -85,25 +82,20 @@ public class QuotaManager { private final ExecutorService backgroundExecutor; - QuotaManager(TokenStore tokenStore, ExecutorService backgroundExecutor) { + public QuotaManager(TokenStore tokenStore, ExecutorService backgroundExecutor) { this.tokenStore = tokenStore; this.backgroundExecutor = backgroundExecutor; tokenStore.scheduleRefresh(); } /** Attempts to acquire requested quota, synchronously. */ - QuotaResponse acquireQuota(QuotaRequest request) { + public QuotaResponse acquireQuota(QuotaRequest request) { TimestampedInteger tokens = tokenStore.take(request.userId()); - Status status = (tokens.value() == 0) ? Status.FAILURE : Status.SUCCESS; - return QuotaResponse.create(status, request.userId(), tokens.timestamp()); + return QuotaResponse.create(tokens.value() != 0, request.userId(), tokens.timestamp()); } - /** - * Returns granted quota to the token store, asynchronously. - * - * @return a {@link Future} representing the asynchronous task to return the quota. - */ - Future releaseQuota(QuotaRebate rebate) { + /** Returns granted quota to the token store, asynchronously. */ + public Future releaseQuota(QuotaRebate rebate) { return backgroundExecutor.submit( () -> tokenStore.put(rebate.userId(), rebate.grantedTokenRefillTime())); } diff --git a/java/google/registry/proxy/quota/TokenStore.java b/java/google/registry/proxy/quota/TokenStore.java index 7d1a87644..017b7d1f7 100644 --- a/java/google/registry/proxy/quota/TokenStore.java +++ b/java/google/registry/proxy/quota/TokenStore.java @@ -79,7 +79,7 @@ public class TokenStore { private final ScheduledExecutorService refreshExecutor; private final Clock clock; - TokenStore(QuotaConfig config, ScheduledExecutorService refreshExecutor, Clock clock) { + public TokenStore(QuotaConfig config, ScheduledExecutorService refreshExecutor, Clock clock) { this.config = config; this.refreshExecutor = refreshExecutor; this.clock = clock; diff --git a/javatests/google/registry/proxy/ProtocolModuleTest.java b/javatests/google/registry/proxy/ProtocolModuleTest.java index 1b1124d37..7f9990516 100644 --- a/javatests/google/registry/proxy/ProtocolModuleTest.java +++ b/javatests/google/registry/proxy/ProtocolModuleTest.java @@ -22,6 +22,7 @@ import com.google.common.base.Supplier; import com.google.common.base.Suppliers; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableSet; +import com.google.common.util.concurrent.MoreExecutors; import dagger.Component; import dagger.Module; import dagger.Provides; @@ -31,6 +32,8 @@ import google.registry.proxy.HttpsRelayProtocolModule.HttpsRelayProtocol; import google.registry.proxy.WhoisProtocolModule.WhoisProtocol; import google.registry.proxy.handler.BackendMetricsHandler; import google.registry.proxy.handler.ProxyProtocolHandler; +import google.registry.proxy.handler.QuotaHandler.EppQuotaHandler; +import google.registry.proxy.handler.QuotaHandler.WhoisQuotaHandler; import google.registry.proxy.handler.RelayHandler.FullHttpRequestRelayHandler; import google.registry.proxy.handler.RelayHandler.FullHttpResponseRelayHandler; import google.registry.proxy.handler.SslClientInitializer; @@ -47,6 +50,9 @@ import io.netty.handler.ssl.util.SelfSignedCertificate; import io.netty.handler.timeout.ReadTimeoutHandler; import java.security.PrivateKey; import java.security.cert.X509Certificate; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.ScheduledExecutorService; import java.util.function.Consumer; import java.util.function.Function; import javax.inject.Named; @@ -95,6 +101,9 @@ public abstract class ProtocolModuleTest { LoggingHandler.class, // Metrics instrumentation is tested separately. BackendMetricsHandler.class, + // Quota management is tested separately. + WhoisQuotaHandler.class, + EppQuotaHandler.class, ReadTimeoutHandler.class); protected EmbeddedChannel channel; @@ -265,5 +274,17 @@ public abstract class ProtocolModuleTest { Clock provideFakeClock() { return fakeClock; } + + @Singleton + @Provides + ExecutorService provideExecutorService() { + return MoreExecutors.newDirectExecutorService(); + } + + @Singleton + @Provides + ScheduledExecutorService provideScheduledExecutorService() { + return Executors.newSingleThreadScheduledExecutor(); + } } } diff --git a/javatests/google/registry/proxy/handler/EppQuotaHandlerTest.java b/javatests/google/registry/proxy/handler/EppQuotaHandlerTest.java new file mode 100644 index 000000000..c558a85a9 --- /dev/null +++ b/javatests/google/registry/proxy/handler/EppQuotaHandlerTest.java @@ -0,0 +1,137 @@ +// Copyright 2017 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.proxy.handler; + +import static com.google.common.truth.Truth.assertThat; +import static google.registry.proxy.handler.EppServiceHandler.CLIENT_CERTIFICATE_HASH_KEY; +import static google.registry.testing.JUnitBackports.expectThrows; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.verifyNoMoreInteractions; +import static org.mockito.Mockito.when; + +import google.registry.proxy.handler.QuotaHandler.EppQuotaHandler; +import google.registry.proxy.handler.QuotaHandler.OverQuotaException; +import google.registry.proxy.quota.QuotaManager; +import google.registry.proxy.quota.QuotaManager.QuotaRebate; +import google.registry.proxy.quota.QuotaManager.QuotaRequest; +import google.registry.proxy.quota.QuotaManager.QuotaResponse; +import io.netty.channel.ChannelFuture; +import io.netty.channel.embedded.EmbeddedChannel; +import org.joda.time.DateTime; +import org.joda.time.DateTimeZone; +import org.joda.time.Duration; +import org.junit.Before; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +/** Unit tests for {@link EppQuotaHandler} */ +@RunWith(JUnit4.class) +public class EppQuotaHandlerTest { + + private final QuotaManager quotaManager = mock(QuotaManager.class); + private final EppQuotaHandler handler = new EppQuotaHandler(quotaManager); + private final EmbeddedChannel channel = new EmbeddedChannel(handler); + private final String clientCertHash = "blah/123!"; + private final DateTime now = DateTime.now(DateTimeZone.UTC); + private final Object message = new Object(); + + @Before + public void setUp() { + channel.attr(CLIENT_CERTIFICATE_HASH_KEY).set(clientCertHash); + } + + @Test + public void testSuccess_quotaGrantedAndReturned() { + when(quotaManager.acquireQuota(QuotaRequest.create(clientCertHash))) + .thenReturn(QuotaResponse.create(true, clientCertHash, now)); + + // First read, acquire quota. + assertThat(channel.writeInbound(message)).isTrue(); + assertThat((Object) channel.readInbound()).isEqualTo(message); + assertThat(channel.isActive()).isTrue(); + verify(quotaManager).acquireQuota(QuotaRequest.create(clientCertHash)); + + // Second read, should not acquire quota again. + Object newMessage = new Object(); + assertThat(channel.writeInbound(newMessage)).isTrue(); + assertThat((Object) channel.readInbound()).isEqualTo(newMessage); + verifyNoMoreInteractions(quotaManager); + + // Channel closed, release quota. + ChannelFuture unusedFuture = channel.close(); + verify(quotaManager) + .releaseQuota(QuotaRebate.create(QuotaResponse.create(true, clientCertHash, now))); + verifyNoMoreInteractions(quotaManager); + } + + @Test + public void testFailure_quotaNotGranted() { + when(quotaManager.acquireQuota(QuotaRequest.create(clientCertHash))) + .thenReturn(QuotaResponse.create(false, clientCertHash, now)); + OverQuotaException e = + expectThrows(OverQuotaException.class, () -> channel.writeInbound(message)); + assertThat(e).hasMessageThat().contains(clientCertHash); + } + + @Test + public void testSuccess_twoChannels_twoUserIds() { + // Set up another user. + final EppQuotaHandler otherHandler = new EppQuotaHandler(quotaManager); + final EmbeddedChannel otherChannel = new EmbeddedChannel(otherHandler); + final String otherClientCertHash = "hola@9x"; + otherChannel.attr(CLIENT_CERTIFICATE_HASH_KEY).set(otherClientCertHash); + final DateTime later = now.plus(Duration.standardSeconds(1)); + + when(quotaManager.acquireQuota(QuotaRequest.create(clientCertHash))) + .thenReturn(QuotaResponse.create(true, clientCertHash, now)); + when(quotaManager.acquireQuota(QuotaRequest.create(otherClientCertHash))) + .thenReturn(QuotaResponse.create(false, otherClientCertHash, later)); + + // Allows the first user. + assertThat(channel.writeInbound(message)).isTrue(); + assertThat((Object) channel.readInbound()).isEqualTo(message); + assertThat(channel.isActive()).isTrue(); + + // Blocks the second user. + OverQuotaException e = + expectThrows(OverQuotaException.class, () -> otherChannel.writeInbound(message)); + assertThat(e).hasMessageThat().contains(otherClientCertHash); + } + + @Test + public void testSuccess_twoChannels_sameUserIds() { + // Set up another channel for the same user. + final EppQuotaHandler otherHandler = new EppQuotaHandler(quotaManager); + final EmbeddedChannel otherChannel = new EmbeddedChannel(otherHandler); + otherChannel.attr(CLIENT_CERTIFICATE_HASH_KEY).set(clientCertHash); + final DateTime later = now.plus(Duration.standardSeconds(1)); + + when(quotaManager.acquireQuota(QuotaRequest.create(clientCertHash))) + .thenReturn(QuotaResponse.create(true, clientCertHash, now)) + .thenReturn(QuotaResponse.create(false, clientCertHash, later)); + + // Allows the first channel. + assertThat(channel.writeInbound(message)).isTrue(); + assertThat((Object) channel.readInbound()).isEqualTo(message); + assertThat(channel.isActive()).isTrue(); + + // Blocks the second channel. + OverQuotaException e = + expectThrows(OverQuotaException.class, () -> otherChannel.writeInbound(message)); + assertThat(e).hasMessageThat().contains(clientCertHash); + } +} diff --git a/javatests/google/registry/proxy/handler/WhoisQuotaHandlerTest.java b/javatests/google/registry/proxy/handler/WhoisQuotaHandlerTest.java new file mode 100644 index 000000000..bb1d33ec8 --- /dev/null +++ b/javatests/google/registry/proxy/handler/WhoisQuotaHandlerTest.java @@ -0,0 +1,146 @@ +// Copyright 2017 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.proxy.handler; + +import static com.google.common.truth.Truth.assertThat; +import static google.registry.proxy.handler.ProxyProtocolHandler.REMOTE_ADDRESS_KEY; +import static google.registry.testing.JUnitBackports.expectThrows; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.verifyNoMoreInteractions; +import static org.mockito.Mockito.when; + +import google.registry.proxy.handler.QuotaHandler.OverQuotaException; +import google.registry.proxy.handler.QuotaHandler.WhoisQuotaHandler; +import google.registry.proxy.quota.QuotaManager; +import google.registry.proxy.quota.QuotaManager.QuotaRequest; +import google.registry.proxy.quota.QuotaManager.QuotaResponse; +import io.netty.channel.ChannelFuture; +import io.netty.channel.embedded.EmbeddedChannel; +import org.joda.time.DateTime; +import org.joda.time.DateTimeZone; +import org.joda.time.Duration; +import org.junit.Before; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +/** Unit tests for {@link WhoisQuotaHandler} */ +@RunWith(JUnit4.class) +public class WhoisQuotaHandlerTest { + + private final QuotaManager quotaManager = mock(QuotaManager.class); + private final WhoisQuotaHandler handler = new WhoisQuotaHandler(quotaManager); + private final EmbeddedChannel channel = new EmbeddedChannel(handler); + private final DateTime now = DateTime.now(DateTimeZone.UTC); + private final String remoteAddress = "127.0.0.1"; + private final Object message = new Object(); + + @Before + public void setUp() { + channel.attr(REMOTE_ADDRESS_KEY).set(remoteAddress); + } + + @Test + public void testSuccess_quotaGranted() { + when(quotaManager.acquireQuota(QuotaRequest.create(remoteAddress))) + .thenReturn(QuotaResponse.create(true, remoteAddress, now)); + + // First read, acquire quota. + assertThat(channel.writeInbound(message)).isTrue(); + assertThat((Object) channel.readInbound()).isEqualTo(message); + assertThat(channel.isActive()).isTrue(); + verify(quotaManager).acquireQuota(QuotaRequest.create(remoteAddress)); + + // Second read, should not acquire quota again. + assertThat(channel.writeInbound(message)).isTrue(); + assertThat((Object) channel.readInbound()).isEqualTo(message); + + // Channel closed, release quota. + ChannelFuture unusedFuture = channel.close(); + verifyNoMoreInteractions(quotaManager); + } + + @Test + public void testFailure_quotaNotGranted() { + when(quotaManager.acquireQuota(QuotaRequest.create(remoteAddress))) + .thenReturn(QuotaResponse.create(false, remoteAddress, now)); + OverQuotaException e = + expectThrows(OverQuotaException.class, () -> channel.writeInbound(message)); + assertThat(e).hasMessageThat().contains("none"); + } + + @Test + public void testSuccess_twoChannels_twoUserIds() { + // Set up another user. + final WhoisQuotaHandler otherHandler = new WhoisQuotaHandler(quotaManager); + final EmbeddedChannel otherChannel = new EmbeddedChannel(otherHandler); + final String otherRemoteAddress = "192.168.0.1"; + otherChannel.attr(REMOTE_ADDRESS_KEY).set(otherRemoteAddress); + final DateTime later = now.plus(Duration.standardSeconds(1)); + + when(quotaManager.acquireQuota(QuotaRequest.create(remoteAddress))) + .thenReturn(QuotaResponse.create(true, remoteAddress, now)); + when(quotaManager.acquireQuota(QuotaRequest.create(otherRemoteAddress))) + .thenReturn(QuotaResponse.create(false, otherRemoteAddress, later)); + + // Allows the first user. + assertThat(channel.writeInbound(message)).isTrue(); + assertThat((Object) channel.readInbound()).isEqualTo(message); + assertThat(channel.isActive()).isTrue(); + + // Blocks the second user. + OverQuotaException e = + expectThrows(OverQuotaException.class, () -> otherChannel.writeInbound(message)); + assertThat(e).hasMessageThat().contains("none"); + } + + @Test + public void testSuccess_oneUser_rateLimited() { + // Set up another channel for the same user. + final WhoisQuotaHandler otherHandler = new WhoisQuotaHandler(quotaManager); + final EmbeddedChannel otherChannel = new EmbeddedChannel(otherHandler); + otherChannel.attr(REMOTE_ADDRESS_KEY).set(remoteAddress); + final DateTime later = now.plus(Duration.standardSeconds(1)); + + // Set up the third channel for the same user + final WhoisQuotaHandler thirdHandler = new WhoisQuotaHandler(quotaManager); + final EmbeddedChannel thirdChannel = new EmbeddedChannel(thirdHandler); + thirdChannel.attr(REMOTE_ADDRESS_KEY).set(remoteAddress); + final DateTime evenLater = now.plus(Duration.standardSeconds(60)); + + when(quotaManager.acquireQuota(QuotaRequest.create(remoteAddress))) + .thenReturn(QuotaResponse.create(true, remoteAddress, now)) + // Throttles the second connection. + .thenReturn(QuotaResponse.create(false, remoteAddress, later)) + // Allows the third connection because token refilled. + .thenReturn(QuotaResponse.create(true, remoteAddress, evenLater)); + + // Allows the first channel. + assertThat(channel.writeInbound(message)).isTrue(); + assertThat((Object) channel.readInbound()).isEqualTo(message); + assertThat(channel.isActive()).isTrue(); + + // Blocks the second channel. + OverQuotaException e = + expectThrows(OverQuotaException.class, () -> otherChannel.writeInbound(message)); + assertThat(e).hasMessageThat().contains("none"); + + // Allows the third channel. + assertThat(thirdChannel.writeInbound(message)).isTrue(); + assertThat((Object) thirdChannel.readInbound()).isEqualTo(message); + assertThat(thirdChannel.isActive()).isTrue(); + } +} diff --git a/javatests/google/registry/proxy/quota/QuotaManagerTest.java b/javatests/google/registry/proxy/quota/QuotaManagerTest.java index a2c3164d7..fe370244e 100644 --- a/javatests/google/registry/proxy/quota/QuotaManagerTest.java +++ b/javatests/google/registry/proxy/quota/QuotaManagerTest.java @@ -15,8 +15,6 @@ package google.registry.proxy.quota; import static com.google.common.truth.Truth.assertThat; -import static google.registry.proxy.quota.QuotaManager.QuotaResponse.Status.FAILURE; -import static google.registry.proxy.quota.QuotaManager.QuotaResponse.Status.SUCCESS; import static org.mockito.Matchers.anyString; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.verify; @@ -56,7 +54,7 @@ public class QuotaManagerTest { request = QuotaRequest.create(USER_ID); response = quotaManager.acquireQuota(request); - assertThat(response.status()).isEqualTo(SUCCESS); + assertThat(response.success()).isTrue(); assertThat(response.userId()).isEqualTo(USER_ID); assertThat(response.grantedTokenRefillTime()).isEqualTo(clock.nowUtc()); } @@ -67,7 +65,7 @@ public class QuotaManagerTest { request = QuotaRequest.create(USER_ID); response = quotaManager.acquireQuota(request); - assertThat(response.status()).isEqualTo(FAILURE); + assertThat(response.success()).isFalse(); assertThat(response.userId()).isEqualTo(USER_ID); assertThat(response.grantedTokenRefillTime()).isEqualTo(clock.nowUtc()); } @@ -75,7 +73,7 @@ public class QuotaManagerTest { @Test public void testSuccess_rebate() throws Exception { DateTime grantedTokenRefillTime = clock.nowUtc(); - response = QuotaResponse.create(SUCCESS, USER_ID, grantedTokenRefillTime); + response = QuotaResponse.create(true, USER_ID, grantedTokenRefillTime); rebate = QuotaRebate.create(response); Future unusedFuture = quotaManager.releaseQuota(rebate); verify(tokenStore).scheduleRefresh();