diff --git a/java/google/registry/proxy/ProxyConfig.java b/java/google/registry/proxy/ProxyConfig.java index e3796eb84..197923644 100644 --- a/java/google/registry/proxy/ProxyConfig.java +++ b/java/google/registry/proxy/ProxyConfig.java @@ -61,6 +61,7 @@ public class ProxyConfig { public int headerLengthBytes; public int readTimeoutSeconds; public String serverHostname; + public Quota quota; } /** Configuration options that apply to WHOIS protocol. */ @@ -70,6 +71,7 @@ public class ProxyConfig { public String relayPath; public int maxMessageLengthBytes; public int readTimeoutSeconds; + public Quota quota; } /** Configuration options that apply to GCP load balancer health check protocol. */ @@ -92,6 +94,21 @@ public class ProxyConfig { public int writeIntervalSeconds; } + /** Configuration options that apply to quota management. */ + public static class Quota { + + /** Quota configuration for a specific set of users. */ + public static class QuotaGroup { + public List userId; + public int tokenAmount; + public int refillSeconds; + } + + public int refreshSeconds; + public QuotaGroup defaultQuota; + public List customQuota; + } + static ProxyConfig getProxyConfig(Environment env) { String defaultYaml = readResourceUtf8(ProxyConfig.class, DEFAULT_CONFIG); String customYaml = diff --git a/java/google/registry/proxy/config/default-config.yaml b/java/google/registry/proxy/config/default-config.yaml index 639d50db3..419c4b4cd 100644 --- a/java/google/registry/proxy/config/default-config.yaml +++ b/java/google/registry/proxy/config/default-config.yaml @@ -82,6 +82,41 @@ epp: # TODO(b/64510444) Remove this after nomulus no longer check sni header. serverHostname: epp.yourdomain.tld + # Quota configuration for EPP + quota: + + # Token database refresh period. Set to 0 to disable refresh. + # + # After the set time period, inactive userIds will be deleted. + refreshSeconds: 0 + + # Default quota for any userId not matched in customQuota. + defaultQuota: + + # List of identifiers, e. g. IP address, certificate hash. + # + # userId for defaultQuota should always be an empty list. Any value + # in the list will be discarded. + # + # There should be no duplicate userIds, either within this list, or + # across quota groups within customQuota. Any duplication will result + # in an error when constructing QuotaConfig. + userId: [] + + # Number of tokens allotted to the matched user. Set to -1 to allow + # infinite quota. + tokenAmount: 100 + + # Token refill period. Set to 0 to disable refill. + # + # After the set time period, the token for the user will be + # reset to tokenAmount. + refillSeconds: 0 + + # List of custom quotas for specific userId. Use the same schema as + # defaultQuota for list entries. + customQuota: [] + whois: port: 43 relayHost: registry-project-id.appspot.com @@ -100,6 +135,36 @@ whois: # idle connection. readTimeoutSeconds: 60 + # Quota configuration for WHOIS + quota: + + # Token database refresh period. Set to 0 to disable refresh. + # + # After the set time period, inactive token buckets will be deleted. + refreshSeconds: 3600 + + # Default quota for any userId not matched in customQuota. + defaultQuota: + + # List of identifiers, e. g. IP address, certificate hash. + # + # userId for defaultQuota should always be an empty list. + userId: [] + + # Number of tokens allotted to the matched user. Set to -1 to allow + # infinite quota. + tokenAmount: 100 + + # Token refill period. Set to 0 to disable refill. + # + # After the set time period, the token for the given user will be + # reset to tokenAmount. + refillSeconds: 600 + + # List of custom quotas for specific userId. Use the same schema as + # defaultQuota for list entries. + customQuota: [] + healthCheck: port: 11111 diff --git a/java/google/registry/proxy/quota/QuotaConfig.java b/java/google/registry/proxy/quota/QuotaConfig.java new file mode 100644 index 000000000..8b3475b44 --- /dev/null +++ b/java/google/registry/proxy/quota/QuotaConfig.java @@ -0,0 +1,64 @@ +// 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.quota; + +import com.google.common.annotations.VisibleForTesting; +import com.google.common.collect.ImmutableMap; +import google.registry.proxy.ProxyConfig.Quota; +import google.registry.proxy.ProxyConfig.Quota.QuotaGroup; +import org.joda.time.Duration; + +/** Value class that stores the quota configuration for a protocol. */ +public class QuotaConfig { + + private final int refreshSeconds; + private final QuotaGroup defaultQuota; + private final ImmutableMap customQuotaMap; + + /** + * Constructs a {@link QuotaConfig} from a {@link Quota}. + * + *

Each {@link QuotaGroup} is keyed to all the {@code userId}s it contains. This allows for + * fast lookup with a {@code userId}. + */ + public QuotaConfig(Quota quota) { + refreshSeconds = quota.refreshSeconds; + defaultQuota = quota.defaultQuota; + ImmutableMap.Builder mapBuilder = new ImmutableMap.Builder<>(); + quota.customQuota.forEach( + quotaGroup -> quotaGroup.userId.forEach(userId -> mapBuilder.put(userId, quotaGroup))); + customQuotaMap = mapBuilder.build(); + } + + @VisibleForTesting + QuotaGroup findQuotaGroup(String userId) { + return customQuotaMap.getOrDefault(userId, defaultQuota); + } + + /** Returns the token amount for the given {@code userId}. */ + public int getTokenAmount(String userId) { + return findQuotaGroup(userId).tokenAmount; + } + + /** Returns the refill period for the given {@code userId}. */ + public Duration getRefillPeriod(String userId) { + return Duration.standardSeconds(findQuotaGroup(userId).refillSeconds); + } + + /** Returns the refresh period for this quota config. */ + public Duration getRefreshPeriod() { + return Duration.standardSeconds(refreshSeconds); + } +} diff --git a/javatests/google/registry/proxy/BUILD b/javatests/google/registry/proxy/BUILD index 328423572..bd2eb6e34 100644 --- a/javatests/google/registry/proxy/BUILD +++ b/javatests/google/registry/proxy/BUILD @@ -10,7 +10,10 @@ load("//java/com/google/testing/builddefs:GenTestRules.bzl", "GenTestRules") java_library( name = "proxy", srcs = glob(["**/*.java"]), - resources = glob(["testdata/*.xml"]), + resources = glob([ + "testdata/*.xml", + "quota/testdata/*.yaml", + ]), deps = [ "//java/google/registry/monitoring/metrics", "//java/google/registry/monitoring/metrics/contrib", @@ -32,6 +35,7 @@ java_library( "@junit", "@org_bouncycastle_bcpkix_jdk15on", "@org_mockito_all", + "@org_yaml_snakeyaml", ], ) diff --git a/javatests/google/registry/proxy/quota/QuotaConfigTest.java b/javatests/google/registry/proxy/quota/QuotaConfigTest.java new file mode 100644 index 000000000..b6ca81bfc --- /dev/null +++ b/javatests/google/registry/proxy/quota/QuotaConfigTest.java @@ -0,0 +1,71 @@ +// 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.quota; + +import static com.google.common.truth.Truth.assertThat; +import static google.registry.testing.JUnitBackports.expectThrows; +import static google.registry.util.ResourceUtils.readResourceUtf8; + +import google.registry.proxy.ProxyConfig.Quota; +import org.joda.time.Duration; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; +import org.yaml.snakeyaml.Yaml; + +/** Unit Tests for {@link QuotaConfig} */ +@RunWith(JUnit4.class) +public class QuotaConfigTest { + + private QuotaConfig quotaConfig; + + private static QuotaConfig loadQuotaConfig(String filename) { + return new QuotaConfig( + new Yaml() + .loadAs(readResourceUtf8(QuotaConfigTest.class, "testdata/" + filename), Quota.class)); + } + + private void validateQuota(String userId, int tokenAmount, int refillSeconds) { + assertThat(quotaConfig.getTokenAmount(userId)).isEqualTo(tokenAmount); + assertThat(quotaConfig.getRefillPeriod(userId)) + .isEqualTo(Duration.standardSeconds(refillSeconds)); + } + + @Test + public void testSuccess_regularConfig() { + quotaConfig = loadQuotaConfig("quota_config_regular.yaml"); + assertThat(quotaConfig.getRefreshPeriod()).isEqualTo(Duration.standardSeconds(3600)); + validateQuota("abc", 10, 60); + validateQuota("987lol", 500, 10); + validateQuota("no_match", 100, 60); + } + + @Test + public void testSuccess_onlyDefault() { + quotaConfig = loadQuotaConfig("quota_config_default.yaml"); + assertThat(quotaConfig.getRefreshPeriod()).isEqualTo(Duration.standardSeconds(3600)); + validateQuota("abc", 100, 60); + validateQuota("987lol", 100, 60); + validateQuota("no_match", 100, 60); + } + + @Test + public void testFailure_duplicateUserId() { + IllegalArgumentException e = + expectThrows( + IllegalArgumentException.class, () -> loadQuotaConfig("quota_config_duplicate.yaml")); + assertThat(e).hasMessageThat().contains("Multiple entries with same key"); + } +} diff --git a/javatests/google/registry/proxy/quota/testdata/quota_config_default.yaml b/javatests/google/registry/proxy/quota/testdata/quota_config_default.yaml new file mode 100644 index 000000000..54d354ca4 --- /dev/null +++ b/javatests/google/registry/proxy/quota/testdata/quota_config_default.yaml @@ -0,0 +1,8 @@ +refreshSeconds: 3600 + +defaultQuota: + userId: [] + tokenAmount: 100 + refillSeconds: 60 + +customQuota: [] diff --git a/javatests/google/registry/proxy/quota/testdata/quota_config_duplicate.yaml b/javatests/google/registry/proxy/quota/testdata/quota_config_duplicate.yaml new file mode 100644 index 000000000..5984b0958 --- /dev/null +++ b/javatests/google/registry/proxy/quota/testdata/quota_config_duplicate.yaml @@ -0,0 +1,14 @@ +refreshSeconds: 3600 + +defaultQuota: + userId: [] + tokenAmount: 100 + refillSeconds: 60 + +customQuota: + - userId: ["abc", "def"] + tokenAmount: 10 + refillSeconds: 60 + - userId: ["xyz123", "def", "luckycat"] + tokenAmount: 500 + refillSeconds: 10 diff --git a/javatests/google/registry/proxy/quota/testdata/quota_config_regular.yaml b/javatests/google/registry/proxy/quota/testdata/quota_config_regular.yaml new file mode 100644 index 000000000..a8c23a627 --- /dev/null +++ b/javatests/google/registry/proxy/quota/testdata/quota_config_regular.yaml @@ -0,0 +1,14 @@ +refreshSeconds: 3600 + +defaultQuota: + userId: [] + tokenAmount: 100 + refillSeconds: 60 + +customQuota: + - userId: ["abc", "def"] + tokenAmount: 10 + refillSeconds: 60 + - userId: ["xyz123", "987lol", "luckycat"] + tokenAmount: 500 + refillSeconds: 10